曙光DCU高效训练指南：从conda环境到sbatch脚本的全流程优化
在深度学习研究领域，高效利用计算资源是每个开发者必须掌握的技能。曙光DCU作为国产高性能加速卡，为AI训练提供了强大的算力支持。然而，许多用户在集群环境中进行大规模训练时，常常面临环境配置混乱、资源分配不合理、任务管理低效等问题。本文将分享一套经过实战检验的工作流程，帮助您在曙光DCU集群上建立标准化、自动化的训练任务管理体系。
1. 构建可复现的conda虚拟环境
环境隔离是保证实验可复现性的第一步。在曙光DCU集群上，我们推荐使用conda进行环境管理，它不仅能够隔离不同项目所需的依赖，还能确保环境配置的一致性。
1.1 conda环境初始化
首先需要在集群上安装miniconda。由于曙光集群通常采用AMD架构，建议从官方渠道获取适配的conda版本：
# 下载并安装miniconda wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda3安装完成后，需要将conda加入环境变量：
echo ". $HOME/miniconda3/etc/profile.d/conda.sh" >> ~/.bashrc source ~/.bashrc1.2 创建专用训练环境
针对不同的训练任务，建议创建独立的环境。例如，为PyTorch项目创建环境：
conda create -n pytorch_train python=3.8 -y conda activate pytorch_train注意：曙光DCU需要安装特定版本的PyTorch，不能直接使用官方版本。安装命令示例：
pip install /public/software/apps/DeepLearning/whl/dtk-22.10/pytorch/torch-1.10.0a0+git36449ea-cp38-cp38-linux_x86_64.whl
环境创建完成后，建议将依赖包列表导出保存：
conda env export > environment.yml pip freeze > requirements.txt这样可以在其他节点或未来重建环境时快速恢复：
conda env create -f environment.yml pip install -r requirements.txt2. 编写高效的sbatch脚本
sbatch是曙光集群上最常用的作业提交方式，通过脚本可以一次性完成资源申请、环境配置和任务执行。
2.1 基础sbatch脚本结构
一个典型的训练脚本应包含以下部分：
#!/bin/bash #SBATCH -J my_training_job # 作业名称 #SBATCH -p hebhdnormal # 指定队列 #SBATCH -N 1 # 节点数 #SBATCH --gres=dcu:4 # 每节点DCU数量 #SBATCH -o %j.out # 标准输出文件 #SBATCH -e %j.err # 标准错误文件 # 加载必要模块 module rm compiler/rocm/3.3 module load compiler/rocm/dtk-22.10 # 激活conda环境 source $HOME/miniconda3/etc/profile.d/conda.sh conda activate pytorch_train # 运行训练脚本 python train.py --batch-size 256 --epochs 1002.2 资源申请策略优化
合理申请资源可以显著提高作业调度效率：
| 资源类型 | 推荐配置 | 适用场景 |
|---|---|---|
| DCU数量 | 4-8卡 | 中等规模模型 |
| 节点数 | 1-2节点 | 数据并行训练 |
| 内存 | 64-128G | 大型数据集 |
| 时间限制 | 24-72h | 长期训练任务 |
对于超参数搜索等需要多次运行的任务，可以使用作业数组：
#SBATCH --array=1-10%3 # 同时运行3个任务，共10个参数组合 python train.py --lr 0.$SLURM_ARRAY_TASK_ID3. 任务监控与日志管理
良好的日志系统是长期实验的基础保障。
3.1 实时监控工具
曙光集群提供了多种监控方式：
hy-smi：查看DCU使用情况squeue -u $USER：查看个人作业状态sacct：查看作业历史记录
建议在训练代码中加入定期保存检查点和日志的功能：
import torch from torch.utils.tensorboard import SummaryWriter writer = SummaryWriter() for epoch in range(epochs): # 训练代码... writer.add_scalar('Loss/train', loss.item(), epoch) if epoch % 10 == 0: torch.save(model.state_dict(), f'checkpoint_{epoch}.pth')3.2 日志分析技巧
结合grep等工具可以快速定位关键信息：
# 查找错误信息 grep -i error job_12345.err # 统计训练速度 grep 'iter/s' job_12345.out | awk '{sum+=$NF} END {print sum/NR}' # 提取验证准确率 grep 'val_acc' job_12345.out | awk '{print $NF}' > acc.txt对于长期运行的作业，可以设置邮件通知：
#SBATCH --mail-type=FAIL,END #SBATCH --mail-user=your@email.com4. 高级优化技巧
4.1 环境缓存加速
在多个节点上重复创建相同环境会浪费大量时间。可以先将环境打包：
conda pack -n pytorch_train -o pytorch_train.tar.gz然后在脚本中解压使用：
# 在作业脚本中添加 mkdir -p $HOME/conda_envs tar -xzf pytorch_train.tar.gz -C $HOME/conda_envs source $HOME/conda_envs/bin/activate4.2 数据预处理优化
大规模训练中，数据加载常常成为瓶颈。建议：
- 使用内存映射文件格式（如HDF5）
- 预处理数据并存储在本地SSD
- 采用多进程数据加载：
from torch.utils.data import DataLoader loader = DataLoader(dataset, batch_size=256, num_workers=4, pin_memory=True)4.3 混合精度训练
DCU支持混合精度训练，可显著减少显存占用并提升速度：
from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() for inputs, targets in loader: with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()5. 常见问题解决方案
在实际使用中，我们积累了一些典型问题的应对策略：
问题1：作业排队时间过长
- 尝试申请非高峰时段资源
- 适当减少单次申请的DCU数量
- 使用
--qos=low降低QoS等级
问题2：显存不足
# 在PyTorch中减少显存占用技巧 torch.backends.cudnn.benchmark = True # 启用cudnn自动优化 torch.cuda.empty_cache() # 定期清空缓存问题3：训练意外中断
建议在代码中加入断点续训功能：
if os.path.exists('latest_checkpoint.pth'): state = torch.load('latest_checkpoint.pth') model.load_state_dict(state['model']) optimizer.load_state_dict(state['optimizer']) start_epoch = state['epoch']问题4：多卡训练效率低
确保正确设置数据并行：
model = torch.nn.DataParallel(model) # 或者更精细控制的分布式训练 torch.distributed.init_process_group(backend='nccl') model = torch.nn.parallel.DistributedDataParallel(model)经过多个项目的实践验证，这套工作流程能够将曙光DCU集群的利用率提升30%以上，同时大幅减少环境配置和任务管理的时间成本。关键在于建立标准化的操作规范，并通过脚本实现自动化管理。
