在DCU上把PyTorch跑起来，我踩过的5个坑和一份踩坑指南

如果你也在国产加速卡上折腾过PyTorch，一定经历过torch.cuda.is_available()返回False时的心跳骤停。这篇文章不讲虚的，从环境部署到分布式训练，把关键点串一遍。

---

1. PyTorch的底子：它凭什么成了主流

PyTorch最早是Facebook从Torch7（Lua语言那版）改过来的，核心思路很简单——用Python包一层，保持易用性，底层靠C++/CUDA扛性能。目前GitHub上有97k+ Star。

跟它同期竞争的框架不少，但活到今天且还活跃的没几个了：

框架	Star数	维护方	现状
PyTorch	97.3k	Meta	动态图起家，2.x后通过torch.compile补齐图优化短板
TensorFlow	194k	Google	2.x默认Eager模式，静态图靠tf.function，部署生态强
Caffe	34.8k	BVLC	推理性能不错，但社区基本凉了
MxNet	20.8k	DMLC/Amazon	已基本停止发展
PaddlePaddle	23.6k	百度	面向工业+国产生态，动静混合模式

PyTorch能活成今天这样，核心就三块东西，缺一不可：

说白了：Tensor负责干活，Autograd负责自动算梯度，NN负责把网络搭起来。三层配合，写模型就像搭积木。

---

2. DCU上的PyTorch：安装不是pip install torch就完事的

在国产加速卡（DCU）上装PyTorch，跟NVIDIA GPU上最大的区别是：你不能直接用官方PyTorch。官方PyTorch绑的是CUDA后端，DCU走的是另一套路线——基于ROCm生态的HIP接口。

安装有四种路子：

方式	适用场景	要点
pip安装	日常开发首选	从DAS源或下载whl包安装，配置镜像源加速
Anaconda	多环境隔离	创建独立环境，注意venv路径
源码编译	定制化需求	需设置环境变量、安装依赖、运行编译脚本
Docker	快速复现 / CI	加载镜像后直接跑，环境零污染

2.1 pip安装：最常用的方式

公网源地址：http://pypi.sourcefind.cn:666/source/packages，在DAS PyPI上选择对应软件栈版本的仓库和torch包，直接复制安装命令即可。

# 示例（版本号已脱敏处理）
pip install torch==x.x.x+das \
  -i http://pypi.sourcefind.cn:666/release/dtk-xxx-rc1/+simple/ \
  --trusted-host pypi.sourcefind.cn

⚠️ 坑1：不同DTK版本的包不通用，选错版本装完会直接报符号未定义。务必确认das后缀和DTK版本匹配。

2.2 环境验证：三行命令确认是否就绪

# 1. 加载DTK环境（物理机和部分容器需手动执行）
source /opt/dtk/env.sh

# 2. 查看加速卡信息
rocminfo | grep gfx

# 3. 查看卡状态（类似nvidia-smi）
hy-smi

然后用Python验证PyTorch是否认到卡：

import torch

print(torch.cuda.is_available())        # 应为 True
print(torch.__version__)                # 如 x.x.x+das
print(torch.version.__hcu_version__)    # 带das后缀的HCU版本号

⚠️ 坑2：source /opt/dtk/env.sh这一步太容易忘。容器里有时自动加载了，物理机必手动执行。忘了的话hy-smi能用但PyTorch认不到卡——因为动态库路径没设对。

2.3 三方库支持：这些常用库都能跑

DCU生态对主流PyTorch三方库的移植情况相当不错：

项目	软件包	支持状态
计算机视觉	TorchVision	ROCm官方已支持
3D视觉	pytorch3d	已完成移植
混精计算	APEX	部分移植（不含BNP模块），ROCm官方已支持
分布式训练优化	DeepSpeed	已完成移植
推理优化	torch-ni	已完成移植
音频处理	torchaudio	已完成移植
检测/分割/关键点	MMCV全家桶	已完成移植，ROCm官方已支持
目标检测	Detectron2	已完成移植
分布式MoE	fastmoe	已完成移植

---

3. torch.compile：一行代码让模型快一截

PyTorch 2.x最大的亮点就是torch.compile。它的思路是：你照常写Eager模式的代码，它在底层帮你做JIT编译优化，不动你一行业务逻辑。

使用方式极其简单：

# 定义模型
model = MyModel().cuda()

# 一行代码包裹
compiled_model = torch.compile(model)

# 后续训练/推理逻辑完全不变
output = compiled_model(input)

常用参数一览：

参数	作用	推荐场景
fullgraph=True	强制生成完整计算图	模型结构固定，追求极致性能
dynamic=True	支持动态shape输入	输入尺寸不固定的NLP任务
mode="max-autotune"	自动搜索最优kernel配置	推理部署场景，编译慢但跑得快
mode="reduce-overhead"	减少CPU调度开销	小模型高频推理

⚠️ 坑3：mode="max-autotune"第一次编译极慢（有时候几分钟），但编译完的性能提升对得住等待。如果只是快速验证功能，先用默认模式。

---

4. 分布式通信：后端怎么选，mpirun怎么配

PyTorch分布式训练的核心通信模块是torch.distributed。三种后端各有利弊：

一句话总结选型：GPU分布式训练用RCCL，CPU用Gloo，超算环境用MPI。

初始化代码长这样：

import torch.distributed as dist

dist.init_process_group(
    backend='nccl',            # 后端：nccl / gloo / mpi
    init_method='tcp://...',   # 初始化方式：tcp / file / env
    world_size=4,              # 总进程数
    rank=0,                    # 当前进程序号
    timeout=timedelta(minutes=30)  # 超时，nccl默认10分钟，gloo默认30分钟
)

4.1 mpirun多节点启动

先用Slurm申请资源（关键参数已脱敏处理）：

#!/bin/bash
#SBATCH -p <队列名>
#SBATCH -N 2                   # 节点数
#SBATCH --ntasks-per-node=8    # 每节点任务数
#SBATCH --cpus-per-task=16     # 每任务CPU核数
#SBATCH --gres=dcu:8           # 每节点加速卡数
#SBATCH --exclusive            # 独占节点保证性能

申请到节点后启动分布式任务：

mpirun --allow-run-as-root \
  -np 16 \
  -H <node1>:8,<node2>:8 \
  -x ROCM_PATH \
  -x NCCL_SOCKET_IFNAME=ib0 \
  -x NCCL_IB_HCA=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1 \
  -x NCCL_NET_GDR_LEVEL=4 \
  -x NCCL_MAX_NCHANNELS=16 \
  -x LD_LIBRARY_PATH \
  /path/to/your/program -b 4096 -e 1G -f 2 -g 1 -d half

关键参数说明：

参数	含义
-np 16	启动16个进程
-H node1:8,node2:8	每节点8个进程
-x	将环境变量传递给远程进程
NCCL_SOCKET_IFNAME=ib0	指定InfiniBand网卡
NCCL_NET_GDR_LEVEL=4	启用GPUDirect RDMA

⚠️ 坑4：多节点通信时网卡指定错了一个字符就跑不起来。ib0还是ibp0还是别的，先用ibstat确认实际网卡名。另外NCCL_IB_HCA的端口号格式（mlx5_0:1）在不同机器上可能不同，别直接抄。

---

5. 调优三板斧：numa绑定、DataLoader、Profiler

5.1 基础优化项

**第一步：CPU拓扑绑定。**用lscpu看NUMA节点分布，然后：

numactl --cpunodebind=<节点ID> --membind=<节点ID> <任务启动命令>

跨NUMA节点访问内存延迟高很多，不绑的话大模型训练会有明显的吞吐波动。

第二步：DataLoader参数调优。

DataLoader(
    dataset,
    batch_size=64,
    num_workers=4,           # 数据加载线程数
    pin_memory=True,          # 预分配锁页内存
    prefetch_factor=4,        # 每个worker预加载batch数
    persistent_workers=True   # worker常驻，避免反复fork
)

第三步：开启cuDNN benchmark。

torch.backends.cudnn.benchmark = True

⚠️ 坑5：cudnn.benchmark=True在输入shape固定时效果明显，但如果每个batch的输入尺寸都在变（比如NLP里padding长度飘忽），开了反而会拖慢——因为每次都要重新搜索最优算法。

5.2 pdb调试：分布式场景下的救命工具

PyTorch分布式出问题时，print大法经常不够用。pdb（Python Debugger）是最轻量的调试方案：

# 方式1：代码内设置断点
import pdb; pdb.set_trace()
# 或直接用Python 3.7+的breakpoint()
breakpoint()

# 方式2：命令行启动
python -m pdb code.py

常用命令速查：

分类	命令	作用
执行控制	n (next)	单步执行
	s (step)	进入函数内部
	c (continue)	执行到下一个断点
	r (return)	执行到当前函数返回
	q (quit)	退出调试
信息查看	p / pp	打印变量值
	whatis var	查看变量类型
	w (where)	当前调用栈
	a (argument)	打印函数参数
断点管理	b <行号>	设置断点
	cl <文件:行号>	清除断点
	l (list)	显示上下文代码

5.3 torch.profiler：找到性能瓶颈的最终手段

典型用法：

import torch.profiler as profiler
from torch.profiler import ProfilerActivity, schedule

prof_schedule = schedule(
    wait=5,        # 等5步预热
    warmup=2,      # 预热2步
    active=3,      # 采集3步
    repeat=1       # 重复1轮
)

with profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,        # 记录每个Tensor形状
    schedule=prof_schedule,
    profile_memory=True,       # 追踪显存分配/释放
    with_stack=True            # 记录Python调用栈
) as prof:
    # 你的训练/推理代码
    for step, data in enumerate(dataloader):
        output = model(data)
        loss.backward()
        optimizer.step()
        prof.step()

# 输出统计表
print(prof.key_averages().table(
    sort_by="self_cuda_time_total",
    row_limit=10
))

# 导出Chrome Trace文件
prof.export_chrome_trace("trace.json")
# 把trace.json拖到 https://ui.perfetto.dev/ 可视化分析

输出示例（数据已脱敏）：

Name                     Self CUDA%  CUDA total  # of Calls
vectorized_elementwise   12.55%      xxx us      8
MemcpyDeviceToHost        9.51%      xxx us      2
reduce_kernel             8.37%      xxx us      2
unrolled_elementwise      6.46%      xxx us      4

从输出能快速定位哪些算子最耗时——MemcpyDeviceToHost占比高说明CPU-GPU数据传输是瓶颈，reduce_kernel高说明规约操作需要优化。

---

6. 总结

本文把在DCU上用PyTorch从头装到调优的关键节点串了一遍：

1. 安装：走DAS源，注意DTK版本匹配，别用官方PyTorch
2. 环境验证：source /opt/dtk/env.sh别忘，torch.cuda.is_available()确认
3. torch.compile：一行代码换性能，但max-autotune第一次编译要有耐心
4. 分布式：GPU训练无脑选RCCL，多节点注意网卡名和IB配置
5. 调优：NUMA绑定 + DataLoader参数 + profiler定位瓶颈，三步走

完整源码和更多细节可以参考公开源：http://pypi.sourcefind.cn:666/source/packages

---