在曙光超算上跑PyTorch：一份给AI研究员的DCU与GPU双卡实战避坑指南
当AI研究员第一次接触曙光超算平台时，面对国产DCU和英伟达GPU两种异构计算资源，往往会陷入配置混乱、性能波动的困境。本文将从实战角度出发，分享如何高效利用这两种硬件资源运行PyTorch代码，特别针对环境配置、作业提交中的典型问题进行深度解析。
1. 环境准备：双卡配置的差异化处理
1.1 DCU环境搭建的关键细节
国产DCU（Deep Computing Unit）在曙光超算上的支持有其特殊性。与常见的CUDA环境不同，DCU基于ROCm架构，需要特别注意驱动版本与PyTorch版本的严格匹配：
# 加载特定版本的ROCm编译器 module switch compiler/dtk/22.04.1 # 安装定制版PyTorch pip install /public/software/apps/DeepLearning/whl/dtk-22.04.2/torch-1.10.0a0_gitc7f69d6_dtk22.04.2-cp37-cp37m-manylinux2014_x86_64.whl注意：DCU环境必须配置以下关键环境变量，否则会导致库链接失败：
export LD_LIBRARY_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/lib:$LD_LIBRARY_PATH
1.2 GPU环境的快速配置
相比之下，英伟达GPU的环境配置更为标准化：
# 加载CUDA工具包 module load apps/cuda/10.2 # 通过conda安装官方PyTorch conda install pytorch torchvision cudatoolkit=10.2 -c pytorch两种硬件的环境差异对比如下：
| 配置项 | DCU环境 | GPU环境 |
|---|---|---|
| 驱动架构 | ROCm | CUDA |
| PyTorch安装 | 需使用平台提供的定制whl包 | 可直接从官方渠道安装 |
| 编译器版本 | 必须严格匹配dtk版本 | 相对宽松的版本兼容性 |
| 环境变量 | 需手动配置LD_LIBRARY_PATH | 通常自动配置 |
2. 作业提交：双卡系统的实战策略
2.1 资源申请与节点分配
在曙光超算上，DCU和GPU使用不同的资源队列和申请参数：
# 申请DCU节点（2块DCU卡） salloc -p dcu_queue -N 1 --gres=dcu:2 # 申请GPU节点（1块Tesla V100） salloc -p gpu_queue -N 1 --gres=gpu:1提示：使用
whichpartition命令可查看当前可用队列，不同超算中心的队列命名规则可能不同。
2.2 作业脚本编写技巧
一个完整的作业脚本应包含以下关键部分：
#!/bin/bash #SBATCH -J pytorch_job #SBATCH -p dcu_queue #SBATCH -N 1 #SBATCH --gres=dcu:2 # 环境初始化 module switch compiler/dtk/22.04.1 source ~/pytorch_env.sh source activate my_env # 运行PyTorch训练脚本 python train.py --batch-size 256 --device dcu对于GPU作业，需要将dcu替换为gpu，并加载对应的CUDA模块。
3. 性能调优：应对DCU的不稳定性
3.1 常见性能问题分析
多位研究员反馈，DCU在运行相同代码时可能出现：
- 训练时间波动达30%-50%
- 相同超参下模型收敛性不一致
- 显存管理不如GPU稳定
3.2 实测有效的优化手段
通过大量实验验证，以下方法可提升DCU稳定性：
批大小调整：DCU对大批次支持较好，建议从256开始尝试
环境隔离：为每个项目创建独立的conda环境
显存监控：添加定期显存清理逻辑
import torch def clean_cache(): torch.dcu.empty_cache()混合精度训练：
from torch.cuda.amp import GradScaler scaler = GradScaler() with torch.autocast(device_type='dcu'): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()
4. 结果验证与故障排查
4.1 硬件识别检查
在提交作业前，建议先交互式验证硬件识别状态：
import torch print(f"PyTorch版本: {torch.__version__}") print(f"DCU可用: {torch.dcu.is_available()}") print(f"GPU可用: {torch.cuda.is_available()}") print(f"可用DCU数量: {torch.dcu.device_count()}")4.2 常见错误解决方案
| 错误现象 | 可能原因 | 解决方案 |
|---|---|---|
| ImportError: librocblas.so | 环境变量未正确配置 | 检查LD_LIBRARY_PATH包含ROCm库路径 |
| DCU out of memory | 显存碎片积累 | 减小批大小或增加显存清理频率 |
| 性能大幅波动 | 计算单元调度不均 | 尝试绑定计算核心：numactl --cpunodebind=0 |
| 作业卡在PD状态 | 队列资源不足 | 换用非高峰时段或调整资源请求量 |
在实际项目中，我通常会准备两套环境配置脚本，分别针对DCU和GPU进行优化。当DCU表现不稳定时，可以快速切换到GPU环境进行结果验证。这种双轨策略既能利用国产算力，又能确保关键实验的可靠性。
