在曙光超算上跑PyTorch:从DCU到GPU,一份给AI新手的保姆级环境配置与作业提交指南

在曙光超算上跑PyTorch:从DCU到GPU,一份给AI新手的保姆级环境配置与作业提交指南 在曙光超算上跑PyTorch从DCU到GPU的实战指南第一次登录曙光超算平台时面对陌生的命令行界面和复杂的作业调度系统大多数AI新手都会感到手足无措。本文将带你从零开始逐步完成在曙光超算上配置PyTorch环境、提交训练任务的全过程特别针对国产DCU加速卡与NVIDIA GPU的不同操作要点进行对比说明。1. 环境准备选择适合的硬件平台曙光超算平台通常提供两种计算加速硬件国产DCU加速卡和NVIDIA GPU。选择哪种硬件取决于你的具体需求DCU加速卡国产自主研发适合支持国产化技术的项目NVIDIA GPU生态成熟社区支持完善适合需要特定CUDA功能的研究1.1 登录与资源查看首先通过SSH登录超算平台ssh your_usernamelogin.node.address登录后查看可用计算分区whichpartition常见输出示例分区名称节点类型每节点DCU数每节点GPU数dcuDCU节点40gpuGPU节点022. DCU环境配置详解2.1 创建Python虚拟环境建议为每个项目创建独立的conda环境module load apps/anaconda3 conda create -n pytorch_dcu python3.7 conda activate pytorch_dcu2.2 安装PyTorch for DCUDCU需要安装特定版本的PyTorch通常由超算中心提供预编译的whl包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 pip install /public/software/apps/DeepLearning/whl/dtk-22.04.2/torchvision-0.10.0a0_dtk22.04.2_e17f5ea-cp37-cp37m-manylinux2014_x86_64.whl2.3 环境变量配置创建环境变量配置文件~/pytorch_env.shexport LD_LIBRARY_PATH/public/software/apps/DeepLearning/PyTorch_Lib/lib:/public/software/apps/DeepLearning/PyTorch_Lib/lmdb-0.9.24-build/lib:/public/software/apps/DeepLearning/PyTorch_Lib/opencv-2.4.13.6-build/lib:/public/software/apps/DeepLearning/PyTorch_Lib/openblas-0.3.7-build/lib:$LD_LIBRARY_PATH加载环境变量source ~/pytorch_env.sh3. GPU环境配置流程3.1 创建Python虚拟环境module load apps/anaconda3 conda create -n pytorch_gpu python3.8 conda activate pytorch_gpu3.2 安装PyTorch for GPU查询可用的CUDA版本module avail cuda加载合适的CUDA模块module load apps/cuda/11.3通过pip安装PyTorchpip install torch1.12.1cu113 torchvision0.13.1cu113 torchaudio0.12.1 --extra-index-url https://download.pytorch.org/whl/cu1134. 作业提交与管理4.1 编写Slurm作业脚本DCU作业脚本示例dcu_job.sh#!/bin/bash #SBATCH -J pytorch_demo # 作业名称 #SBATCH -p dcu # 使用dcu分区 #SBATCH -N 1 # 使用1个节点 #SBATCH --gresdcu:2 # 每节点使用2块DCU卡 #SBATCH -o %j.out # 输出文件 #SBATCH -e %j.err # 错误文件 # 加载必要模块 module switch compiler/dtk/22.04.1 source ~/pytorch_env.sh source activate pytorch_dcu # 运行Python脚本 python train.py --batch-size 64 --epochs 10GPU作业脚本示例gpu_job.sh#!/bin/bash #SBATCH -J pytorch_demo # 作业名称 #SBATCH -p gpu # 使用gpu分区 #SBATCH -N 1 # 使用1个节点 #SBATCH --gresgpu:2 # 每节点使用2块GPU卡 #SBATCH -o %j.out # 错误文件 #SBATCH -e %j.err # 输出文件 # 加载必要模块 module load apps/cuda/11.3 source activate pytorch_gpu # 运行Python脚本 python train.py --batch-size 64 --epochs 104.2 提交与监控作业提交作业sbatch dcu_job.sh # 对于DCU作业 sbatch gpu_job.sh # 对于GPU作业查看作业状态squeue -u $USER常见作业状态说明状态代码含义应采取的行动PD排队中等待或调整资源请求R运行中可监控输出CG正在结束等待完成CD已完成查看结果F失败检查错误日志并重新提交终止作业scancel jobid实时查看输出tail -f jobid.out5. 性能调优与问题排查5.1 DCU与GPU性能对比根据实际测试经验两种硬件在典型深度学习任务中的表现差异任务类型DCU相对性能GPU相对性能稳定性CNN图像分类70-80%基准100%中等RNN文本生成60-70%基准100%较低Transformer训练75-85%基准100%中等数值计算90-95%基准100%高5.2 常见问题解决方案问题1PyTorch无法识别加速设备解决方案DCU环境import torch print(torch.cuda.is_available()) # 应该返回True如果返回False检查是否正确加载了dtk模块环境变量是否配置正确是否安装了正确版本的PyTorchGPU环境 确保CUDA版本与PyTorch版本匹配可通过以下命令验证import torch print(torch.cuda.is_available()) print(torch.version.cuda)问题2作业长时间排队优化策略调整资源请求减少GPU/DCU数量或缩短预计运行时间选择非高峰期提交作业使用更高优先级队列如有问题3内存不足错误解决方法减小batch size使用梯度累积技术优化模型内存占用6. 高级技巧与最佳实践6.1 多卡训练配置对于DCUimport torch import torch.distributed as dist dist.init_process_group(hccl) model torch.nn.parallel.DistributedDataParallel(model)对于GPUimport torch import torch.distributed as dist dist.init_process_group(nccl) model torch.nn.parallel.DistributedDataParallel(model)6.2 混合精度训练DCU示例from torch.cuda.amp import autocast, GradScaler scaler GradScaler() with autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()6.3 资源监控在作业脚本中添加资源监控# 在作业脚本中添加 module load apps/htop htop -d 10 htop.log 7. 实际案例图像分类任务全流程7.1 数据准备建议将数据集放在共享存储空间mkdir -p /share/username/datasets/cifar107.2 训练脚本示例train.pyimport torch import torchvision from torch import nn, optim # 设备配置 device torch.device(cuda if torch.cuda.is_available() else cpu) # 数据加载 transform torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) trainset torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) trainloader torch.utils.data.DataLoader(trainset, batch_size64, shuffleTrue) # 模型定义 model torchvision.models.resnet18(num_classes10).to(device) criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.01, momentum0.9) # 训练循环 for epoch in range(10): for i, (inputs, labels) in enumerate(trainloader): inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() if i % 100 0: print(fEpoch [{epoch1}/10], Step [{i1}/{len(trainloader)}], Loss: {loss.item():.4f})7.3 结果分析与保存训练完成后保存模型和日志torch.save(model.state_dict(), cifar10_resnet18.pth)在本地分析结果# 将结果文件下载到本地 scp usernamelogin.node.address:~/jobid.out .