1. Jetson平台编译Pytorch的必要性

最近在Jetson Xavier NX上部署MMDeploy时，遇到了一个棘手的问题——torch.distributed.ReduceOp组件缺失。这个经历让我意识到，在边缘计算设备上使用预编译的Pytorch wheel文件可能会遇到功能缺失的情况。经过一番折腾，最终通过源码编译解决了问题，今天就把完整过程分享给大家。

为什么要在Jetson上源码编译Pytorch？主要有三个原因：

1. 预编译版本功能不全：官方提供的wheel文件为了减小体积，往往会裁剪掉一些功能模块，分布式训练相关的torch.distributed就是重灾区
2. 硬件适配需求：Jetson平台的ARM架构与常规x86服务器不同，需要针对性地优化编译参数
3. 版本控制需求：当需要特定版本的Pytorch时，源码编译是最可靠的方式

实测发现，从Pytorch 1.11版本开始，官方就不再提供完整功能的预编译版本了。如果你在Python中执行import torch; print(hasattr(torch.distributed, 'ReduceOp'))返回False，那就说明你遇到了和我同样的问题。

2. 环境准备与版本确认

2.1 检查JetPack版本

在开始编译前，首先要确认你的Jetson设备环境。推荐使用jtop工具查看系统信息：

sudo -H pip install jetson-stats
sudo jtop

在我的设备上显示的信息如下：

• JetPack版本：5.1.2
• CUDA版本：11.4
• L4T版本：35.3.1

这些信息至关重要，因为它们决定了你能使用的Pytorch版本。我曾经尝试用不匹配的版本组合，结果编译过程各种报错，白白浪费了好几个小时。

2.2 选择正确的Pytorch版本

访问Nvidia官方论坛可以找到JetPack与Pytorch的版本对应关系。以我的环境为例：

• JetPack 5.1.2 → Pytorch 2.1.0
• CUDA 11.4 → cuDNN 8.6

千万不要贪图新版本，我曾经强行编译Pytorch 2.2.0，结果发现很多算子不支持Jetson的ARM架构。官方维护的版本虽然旧一些，但稳定性有保障。

3. 源码获取与准备工作

3.1 下载Pytorch源码

这里有个大坑需要注意——Pytorch的源码仓库包含了大量第三方子模块。如果直接下载zip包，这些子模块会是空的，导致编译失败。正确做法是使用git递归克隆：

git clone --recursive --branch v2.1.0 https://github.com/pytorch/pytorch
cd pytorch

如果网络不稳定导致克隆失败，可以尝试先克隆主仓库，再手动下载子模块：

git clone --branch v2.1.0 https://github.com/pytorch/pytorch
cd pytorch
git submodule update --init --recursive

3.2 配置编译环境

我强烈建议使用conda来管理Python环境。但要注意，在Jetson上应该使用Archiconda而不是Anaconda：

wget https://github.com/Archiconda/build-tools/releases/download/0.2.3/Archiconda3-0.2.3-Linux-aarch64.sh
bash Archiconda3-0.2.3-Linux-aarch64.sh

创建专用环境并安装基础依赖：

conda create -n torch-build python=3.8 -y
conda activate torch-build
sudo apt-get install build-essential cmake libopenblas-dev libopenmpi-dev
pip install -r requirements.txt

4. 关键编译配置与参数

4.1 必须设置的编译选项

在pytorch目录下创建编译配置文件至关重要。以下是我的配置经验：

export USE_NCCL=1
export USE_DISTRIBUTED=1  # 这个必须设为1！
export USE_CUDA=1
export USE_QNNPACK=0
export USE_PYTORCH_QNNPACK=0
export TORCH_CUDA_ARCH_LIST="5.3;6.2;7.2"  # 根据你的Jetson型号调整
export PYTORCH_BUILD_VERSION=2.1.0
export PYTORCH_BUILD_NUMBER=1

特别提醒：USE_DISTRIBUTED=1这个选项如果漏掉，即使编译成功也会缺少分布式训练功能。我就因为这个疏忽重新编译了一次，多花了两个小时。

4.2 解决常见编译错误

在Jetson上编译时可能会遇到这些问题：

1. 内存不足：建议增加swap空间

   sudo fallocate -l 8G /swapfile
   sudo chmod 600 /swapfile
   sudo mkswap /swapfile
   sudo swapon /swapfile
   

2. 第三方库缺失：确保安装了所有依赖

   sudo apt-get install libjpeg-dev libopenmpi-dev openmpi-bin openmpi-common
   

3. 编译卡死：尝试限制并行编译线程数

   export MAX_JOBS=2
   

5. 编译与安装过程

5.1 执行编译命令

清理之前的编译结果后开始正式编译：

rm -rf build
python setup.py clean
python setup.py bdist_wheel

编译时间取决于你的Jetson型号：

• Jetson Nano：约4-5小时
• Jetson Xavier NX：约2-3小时
• Jetson AGX Orin：约1小时

建议在tmux或screen会话中运行编译，防止SSH断开导致中断。

5.2 安装编译结果

编译完成后，在dist目录下会生成wheel文件：

pip install dist/torch-2.1.0-cp38-cp38-linux_aarch64.whl

验证安装是否成功：

import torch
print(torch.__version__)  # 应该显示2.1.0
print(hasattr(torch.distributed, 'ReduceOp'))  # 应该返回True

5.3 安装对应版本的torchvision

为了完整使用Pytorch功能，还需要编译安装匹配版本的torchvision：

git clone https://github.com/pytorch/vision torchvision
cd torchvision
git checkout tags/v0.16.0 -b v0.16.0
export BUILD_VERSION=0.16.0
pip install -e .

6. 分布式功能验证与性能测试

6.1 基础功能验证

编写一个简单的测试脚本验证分布式功能：

import torch
import torch.distributed as dist

dist.init_process_group(backend='nccl')
tensor = torch.ones(1).cuda()
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
print(tensor)  # 应该输出tensor([2.], device='cuda:0')

如果运行时报错，可能是NCCL配置有问题。尝试设置环境变量：

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0

6.2 多卡训练测试

在支持多卡的Jetson设备上（如AGX Xavier），可以测试真正的分布式训练：

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
model = nn.Linear(10, 10).cuda()
ddp_model = DDP(model, device_ids=[torch.cuda.current_device()])
optimizer = optim.SGD(ddp_model.parameters(), lr=0.01)

# 训练循环...

7. 实际项目中的应用技巧

在MMDeploy等实际项目中使用时，有几个注意事项：

1. 确保所有节点使用相同版本的Pytorch
2. 分布式训练时batch size要合理设置，避免Jetson内存溢出
3. 监控GPU温度，必要时启用风扇控制：

sudo sh -c 'echo 255 > /sys/devices/pwm-fan/target_pwm'

对于长期运行的训练任务，建议使用nohup配合tmux：

tmux new -s training
nohup python train.py --distributed > log.txt 2>&1 &

如果在使用过程中遇到奇怪的问题，可以尝试清除Pytorch缓存：

rm -rf ~/.cache/torch