从零到一：手把手教你用Slurm在超算集群上跑通第一个PyTorch模型（含Anaconda环境配置）

从零到一手把手教你用Slurm在超算集群上跑通第一个PyTorch模型含Anaconda环境配置当你第一次拿到超算账号时面对陌生的命令行界面和复杂的作业调度系统可能会感到无从下手。本文将带你一步步完成从登录超算集群到成功运行PyTorch模型的全过程即使你没有任何超算使用经验。1. 准备工作连接超算集群在开始之前你需要准备以下工具和信息已申请的超算账号通常包含用户名和密码SSH客户端推荐使用MobaXterm或Termius超算集群的登录节点地址连接步骤打开SSH客户端输入连接命令ssh usernamecluster.domain.edu.cn输入密码完成认证提示首次连接时可能会提示确认主机密钥输入yes即可。连接成功后你会看到类似这样的提示符[usernamelogin01 ~]$这表示你已经成功登录到超算集群的登录节点。2. 配置Python环境超算集群通常不预装Python环境我们需要自己安装Anaconda来管理Python环境。2.1 安装Anaconda首先检查你的存储配额/bin/myDiskQuota然后下载并安装Anacondawget https://repo.anaconda.com/archive/Anaconda3-2023.09-0-Linux-x86_64.sh bash Anaconda3-2023.09-0-Linux-x86_64.sh安装过程中需要注意选择安装路径时建议使用/home/username/anaconda3替换username为你的账号当询问是否初始化conda时选择yes安装完成后激活conda环境source ~/.bashrc验证安装conda --version2.2 创建PyTorch环境创建一个专门用于PyTorch的独立环境conda create -n pytorch_env python3.9 conda activate pytorch_env安装PyTorch以CUDA 11.8为例conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia验证PyTorch安装python -c import torch; print(torch.__version__); print(torch.cuda.is_available())3. 准备PyTorch示例代码我们将使用经典的MNIST分类作为第一个示例。创建一个名为mnist.py的文件import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms # 定义网络结构 class Net(nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 nn.Linear(784, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x x.view(-1, 784) x torch.relu(self.fc1(x)) x self.fc2(x) return x # 训练函数 def train(): transform transforms.Compose([transforms.ToTensor()]) trainset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) trainloader torch.utils.data.DataLoader(trainset, batch_size64, shuffleTrue) model Net() criterion nn.CrossEntropyLoss() optimizer optim.SGD(model.parameters(), lr0.01) for epoch in range(5): for data, target in trainloader: optimizer.zero_grad() output model(data) loss criterion(output, target) loss.backward() optimizer.step() print(fEpoch {epoch1}, Loss: {loss.item():.4f}) if __name__ __main__: train()4. 编写Slurm作业脚本创建一个名为job.sh的Slurm作业脚本#!/bin/bash #SBATCH --job-namemnist_pytorch #SBATCH --partitiongpu #SBATCH --nodes1 #SBATCH --ntasks-per-node1 #SBATCH --cpus-per-task4 #SBATCH --gresgpu:1 #SBATCH --time00:30:00 #SBATCH --outputmnist_%j.out # 加载必要的模块 module load cuda/11.8 # 激活conda环境 source /home/username/anaconda3/etc/profile.d/conda.sh conda activate pytorch_env # 运行Python脚本 python mnist.py关键参数说明参数说明推荐值--partition使用的计算分区gpu--nodes计算节点数量1--ntasks-per-node每个节点的任务数1--cpus-per-task每个任务的CPU核心数4--gresGPU数量1--time最大运行时间00:30:005. 提交和管理作业5.1 提交作业将mnist.py和job.sh上传到超算集群的同一目录下然后提交作业sbatch job.sh成功提交后会显示作业IDSubmitted batch job 1234565.2 监控作业状态查看作业状态squeue -u $USER输出示例JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON) 123456 gpu mnist_pt username R 0:05 1 gpu0015.3 查看作业输出作业完成后会生成输出文件mnist_123456.out数字为作业ID。查看输出cat mnist_123456.out预期输出类似Epoch 1, Loss: 0.1234 Epoch 2, Loss: 0.0987 Epoch 3, Loss: 0.0765 Epoch 4, Loss: 0.0654 Epoch 5, Loss: 0.05436. 常见问题与解决方案6.1 环境问题问题conda: command not found解决source ~/.bashrc问题No module named torch解决确保已激活正确的conda环境conda activate pytorch_env6.2 Slurm作业问题问题作业长时间处于PD(Pending)状态解决检查资源请求是否合理或尝试减少资源请求问题作业失败报错Out Of Memory解决增加内存请求或减少batch size6.3 数据传输问题上传文件到超算scp local_file usernamecluster.domain.edu.cn:remote_path从超算下载文件scp usernamecluster.domain.edu.cn:remote_file local_path7. 进阶技巧7.1 使用Jupyter Notebook在超算上运行Jupyter Notebookjupyter notebook --no-browser --port8889然后在本机建立SSH隧道ssh -N -L localhost:8888:localhost:8889 usernamecluster.domain.edu.cn7.2 使用TensorBoard在训练代码中添加TensorBoard记录from torch.utils.tensorboard import SummaryWriter writer SummaryWriter() for epoch in range(5): # ...训练代码... writer.add_scalar(Loss/train, loss.item(), epoch)在超算上启动TensorBoardtensorboard --logdirruns --port60067.3 批量提交作业创建多个作业脚本使用循环提交for i in {1..5}; do sbatch job_${i}.sh done7.4 资源使用优化监控GPU使用情况nvidia-smi优化Slurm参数组合任务类型--nodes--ntasks-per-node--cpus-per-task--gres单GPU训练114gpu:1多GPU数据并行118gpu:2多节点训练218gpu:1