在曙光超算上高效运行PyTorch任务的Slurm实战指南1. 超算环境初探与准备工作曙光超算作为国内领先的高性能计算平台为深度学习研究提供了强大的算力支持。与个人工作站不同超算环境采用Slurm作业管理系统进行资源调度这对刚接触超算的研究者来说需要一定的适应过程。我们先从基础环境配置开始逐步构建完整的PyTorch运行环境。关键准备工作清单确保拥有有效的超算账号和SSH访问权限了解基础Linux命令如ls、cd、vim等准备至少2GB的临时存储空间用于环境安装确认可用的计算节点GPU型号如NVIDIA V100/A100注意超算登录节点仅用于文件管理和作业提交禁止直接运行计算密集型任务否则可能导致账号被限制。超算环境通常采用模块化管理系统常用软件需要通过module命令加载。对于PyTorch环境建议优先使用Conda创建独立环境# 创建Python3.8环境的示例 conda create -n pytorch_env python3.8 -y conda activate pytorch_env2. PyTorch环境精准配置策略在超算环境中安装PyTorch需要特别注意版本兼容性包括CUDA驱动、cuDNN等底层依赖。根据曙光超算的典型配置我们推荐以下安装方案版本匹配对照表组件推荐版本验证命令PyTorch1.13.0python -c import torch; print(torch.__version__)CUDA11.6nvcc --versioncuDNN8.3.2conda listPython3.8.xpython --version使用conda安装PyTorch的黄金命令组合conda install pytorch1.13.0 torchvision0.14.0 torchaudio0.13.0 \ pytorch-cuda11.6 -c pytorch -c nvidia安装完成后建议运行以下验证脚本确认GPU可用性import torch print(fPyTorch版本: {torch.__version__}) print(fCUDA可用: {torch.cuda.is_available()}) print(fGPU数量: {torch.cuda.device_count()}) print(f当前GPU: {torch.cuda.current_device()}) print(fGPU名称: {torch.cuda.get_device_name(0)})3. Slurm作业管理核心技巧Slurm是超算资源调度的核心工具掌握其使用技巧能大幅提升工作效率。下面我们分解关键操作场景3.1 资源监控与查询实时掌握系统状态是合理提交作业的前提# 查看可用分区和节点状态 sinfo -o %20P %5D %14F %6t %N # 查看个人作业队列 squeue -u $USER -o %.8i %.12j %.4t %.10M %.6D %.4C %.10m %R状态代码速查表代码含义应对措施PD排队中检查资源请求是否合理R运行中监控资源使用情况CG完成中等待系统释放资源F失败检查错误日志TO超时调整时间限制3.2 作业提交实战模板针对PyTorch任务特点我们设计了一个智能化的Slurm脚本模板#!/bin/bash #SBATCH --job-namept_train # 作业名称 #SBATCH --partitiongpu # 指定GPU分区 #SBATCH --nodes1 # 节点数 #SBATCH --ntasks-per-node1 # 每节点任务数 #SBATCH --cpus-per-task8 # 每任务CPU核心 #SBATCH --gresgpu:2 # 每节点GPU数量 #SBATCH --time24:00:00 # 运行时限 #SBATCH --output%j.out # 标准输出路径 #SBATCH --error%j.err # 错误输出路径 # 环境初始化 module purge module load cuda/11.6 source activate pytorch_env # 主程序执行 python -u train.py \ --batch-size 256 \ --epochs 100 \ --lr 0.1 \ --gpus $CUDA_VISIBLE_DEVICES提示使用#SBATCH --mail-typeFAIL --mail-useryouremail可在作业失败时接收邮件通知。4. 模型加载与数据管理进阶方案超算环境通常限制外网访问这给预训练模型加载带来挑战。我们提供两种可靠解决方案4.1 离线模型预加载方案在本地下载所需模型文件如HuggingFace或TorchVision模型通过SCP上传到超算scp model.pth usernamecluster:/public/home/username/.cache/torch/hub/checkpoints/在代码中指定本地路径model torch.hub.load_state_dict_from_local( /public/home/username/.cache/torch/hub/checkpoints/model.pth)4.2 代理访问技术方案部分超算提供内部代理服务可通过以下方式临时启用外网访问# 设置HTTP代理具体参数需咨询超算管理员 export http_proxyhttp://proxy.cluster:8080 export https_proxyhttp://proxy.cluster:8080 # 测试网络连接 curl -I https://pytorch.org常见问题排查清单遇到SSL证书错误添加--insecure参数或更新CA证书连接超时检查代理地址和端口是否正确认证失败确认账号权限和密码有效性5. 高效调试与性能优化在超算环境中调试PyTorch任务需要特殊技巧5.1 交互式调试模式# 申请交互式会话2小时限制示例 salloc -N 1 -p debug --gresgpu:1 -t 2:00:00 # 进入计算节点 ssh $SLURM_NODELIST # 启动Python调试 python -m pdb train.py5.2 性能监控工具实时监控GPU使用情况# 在计算节点执行 nvidia-smi -l 1 # 每秒刷新GPU状态 # 查看进程级资源使用 htop -u $USER性能优化检查表确保数据加载使用num_workers4*cpus_per_task验证pin_memoryTrue是否带来加速检查混合精度训练torch.cuda.amp是否启用监控GPU利用率是否持续高于80%6. 实战案例图像分类全流程以一个完整的ImageNet分类任务为例展示端到端的实现过程数据准备阶段# 创建软链接避免数据拷贝 ln -s /dataset/ImageNet/train ./data/train ln -s /dataset/ImageNet/val ./data/val优化后的Slurm脚本#!/bin/bash #SBATCH --nodes2 #SBATCH --ntasks-per-node4 #SBATCH --gresgpu:4 #SBATCH --cpus-per-task8 # 数据并行启动命令 python -m torch.distributed.launch \ --nproc_per_node4 \ --nnodes$SLURM_NNODES \ --node_rank$SLURM_NODEID \ --master_addr$MASTER_ADDR \ train_distributed.py结果自动收集# 训练完成后自动打包结果 import tarfile with tarfile.open(results.tar.gz, w:gz) as tar: tar.add(checkpoints) tar.add(logs)7. 高级技巧与经验分享在实际使用中这些技巧往往能事半功倍资源利用最大化策略使用--exclusive参数独占节点避免干扰合理设置--mem参数防止内存溢出通过--dependencyafterok:JOBID创建任务链实用命令组合# 一键取消所有作业 squeue -u $USER -h -o %i | xargs -n 1 scancel # 查看历史作业统计 sacct -S $(date -d 7 days ago %Y-%m-%d) \ --formatJobID,JobName,Partition,AllocCPUS,Elapsed,State环境迁移技巧# 导出conda环境 conda env export environment.yml # 在新节点快速重建 conda env create -f environment.yml在超算上运行PyTorch任务时最常遇到的坑是环境变量继承问题。特别是在使用sbatch提交作业时建议在脚本开头显式执行module purge清除继承的模块环境然后按需加载特定版本的CUDA和cuDNN。另一个经验是对于长时间运行的训练任务务必添加模型定期保存逻辑并利用torch.save的_use_new_zipfile_serializationTrue参数确保大模型可靠存储。
在曙光超算上跑PyTorch？这份保姆级Slurm避坑指南请收好（含完整脚本模板）
在曙光超算上高效运行PyTorch任务的Slurm实战指南1. 超算环境初探与准备工作曙光超算作为国内领先的高性能计算平台为深度学习研究提供了强大的算力支持。与个人工作站不同超算环境采用Slurm作业管理系统进行资源调度这对刚接触超算的研究者来说需要一定的适应过程。我们先从基础环境配置开始逐步构建完整的PyTorch运行环境。关键准备工作清单确保拥有有效的超算账号和SSH访问权限了解基础Linux命令如ls、cd、vim等准备至少2GB的临时存储空间用于环境安装确认可用的计算节点GPU型号如NVIDIA V100/A100注意超算登录节点仅用于文件管理和作业提交禁止直接运行计算密集型任务否则可能导致账号被限制。超算环境通常采用模块化管理系统常用软件需要通过module命令加载。对于PyTorch环境建议优先使用Conda创建独立环境# 创建Python3.8环境的示例 conda create -n pytorch_env python3.8 -y conda activate pytorch_env2. PyTorch环境精准配置策略在超算环境中安装PyTorch需要特别注意版本兼容性包括CUDA驱动、cuDNN等底层依赖。根据曙光超算的典型配置我们推荐以下安装方案版本匹配对照表组件推荐版本验证命令PyTorch1.13.0python -c import torch; print(torch.__version__)CUDA11.6nvcc --versioncuDNN8.3.2conda listPython3.8.xpython --version使用conda安装PyTorch的黄金命令组合conda install pytorch1.13.0 torchvision0.14.0 torchaudio0.13.0 \ pytorch-cuda11.6 -c pytorch -c nvidia安装完成后建议运行以下验证脚本确认GPU可用性import torch print(fPyTorch版本: {torch.__version__}) print(fCUDA可用: {torch.cuda.is_available()}) print(fGPU数量: {torch.cuda.device_count()}) print(f当前GPU: {torch.cuda.current_device()}) print(fGPU名称: {torch.cuda.get_device_name(0)})3. Slurm作业管理核心技巧Slurm是超算资源调度的核心工具掌握其使用技巧能大幅提升工作效率。下面我们分解关键操作场景3.1 资源监控与查询实时掌握系统状态是合理提交作业的前提# 查看可用分区和节点状态 sinfo -o %20P %5D %14F %6t %N # 查看个人作业队列 squeue -u $USER -o %.8i %.12j %.4t %.10M %.6D %.4C %.10m %R状态代码速查表代码含义应对措施PD排队中检查资源请求是否合理R运行中监控资源使用情况CG完成中等待系统释放资源F失败检查错误日志TO超时调整时间限制3.2 作业提交实战模板针对PyTorch任务特点我们设计了一个智能化的Slurm脚本模板#!/bin/bash #SBATCH --job-namept_train # 作业名称 #SBATCH --partitiongpu # 指定GPU分区 #SBATCH --nodes1 # 节点数 #SBATCH --ntasks-per-node1 # 每节点任务数 #SBATCH --cpus-per-task8 # 每任务CPU核心 #SBATCH --gresgpu:2 # 每节点GPU数量 #SBATCH --time24:00:00 # 运行时限 #SBATCH --output%j.out # 标准输出路径 #SBATCH --error%j.err # 错误输出路径 # 环境初始化 module purge module load cuda/11.6 source activate pytorch_env # 主程序执行 python -u train.py \ --batch-size 256 \ --epochs 100 \ --lr 0.1 \ --gpus $CUDA_VISIBLE_DEVICES提示使用#SBATCH --mail-typeFAIL --mail-useryouremail可在作业失败时接收邮件通知。4. 模型加载与数据管理进阶方案超算环境通常限制外网访问这给预训练模型加载带来挑战。我们提供两种可靠解决方案4.1 离线模型预加载方案在本地下载所需模型文件如HuggingFace或TorchVision模型通过SCP上传到超算scp model.pth usernamecluster:/public/home/username/.cache/torch/hub/checkpoints/在代码中指定本地路径model torch.hub.load_state_dict_from_local( /public/home/username/.cache/torch/hub/checkpoints/model.pth)4.2 代理访问技术方案部分超算提供内部代理服务可通过以下方式临时启用外网访问# 设置HTTP代理具体参数需咨询超算管理员 export http_proxyhttp://proxy.cluster:8080 export https_proxyhttp://proxy.cluster:8080 # 测试网络连接 curl -I https://pytorch.org常见问题排查清单遇到SSL证书错误添加--insecure参数或更新CA证书连接超时检查代理地址和端口是否正确认证失败确认账号权限和密码有效性5. 高效调试与性能优化在超算环境中调试PyTorch任务需要特殊技巧5.1 交互式调试模式# 申请交互式会话2小时限制示例 salloc -N 1 -p debug --gresgpu:1 -t 2:00:00 # 进入计算节点 ssh $SLURM_NODELIST # 启动Python调试 python -m pdb train.py5.2 性能监控工具实时监控GPU使用情况# 在计算节点执行 nvidia-smi -l 1 # 每秒刷新GPU状态 # 查看进程级资源使用 htop -u $USER性能优化检查表确保数据加载使用num_workers4*cpus_per_task验证pin_memoryTrue是否带来加速检查混合精度训练torch.cuda.amp是否启用监控GPU利用率是否持续高于80%6. 实战案例图像分类全流程以一个完整的ImageNet分类任务为例展示端到端的实现过程数据准备阶段# 创建软链接避免数据拷贝 ln -s /dataset/ImageNet/train ./data/train ln -s /dataset/ImageNet/val ./data/val优化后的Slurm脚本#!/bin/bash #SBATCH --nodes2 #SBATCH --ntasks-per-node4 #SBATCH --gresgpu:4 #SBATCH --cpus-per-task8 # 数据并行启动命令 python -m torch.distributed.launch \ --nproc_per_node4 \ --nnodes$SLURM_NNODES \ --node_rank$SLURM_NODEID \ --master_addr$MASTER_ADDR \ train_distributed.py结果自动收集# 训练完成后自动打包结果 import tarfile with tarfile.open(results.tar.gz, w:gz) as tar: tar.add(checkpoints) tar.add(logs)7. 高级技巧与经验分享在实际使用中这些技巧往往能事半功倍资源利用最大化策略使用--exclusive参数独占节点避免干扰合理设置--mem参数防止内存溢出通过--dependencyafterok:JOBID创建任务链实用命令组合# 一键取消所有作业 squeue -u $USER -h -o %i | xargs -n 1 scancel # 查看历史作业统计 sacct -S $(date -d 7 days ago %Y-%m-%d) \ --formatJobID,JobName,Partition,AllocCPUS,Elapsed,State环境迁移技巧# 导出conda环境 conda env export environment.yml # 在新节点快速重建 conda env create -f environment.yml在超算上运行PyTorch任务时最常遇到的坑是环境变量继承问题。特别是在使用sbatch提交作业时建议在脚本开头显式执行module purge清除继承的模块环境然后按需加载特定版本的CUDA和cuDNN。另一个经验是对于长时间运行的训练任务务必添加模型定期保存逻辑并利用torch.save的_use_new_zipfile_serializationTrue参数确保大模型可靠存储。
