在ParallelCluster 集群中通过 PyTorch DDP微调LLM模型的实践

当大语言模型的训练从单张消费级显卡迈向企业级集群，我们进入了一个全新的技术维度。这不仅仅是硬件规模的扩张，更是架构思维的根本转变------如何在多台机器之间协调计算资源，如何在网络延迟与计算效率之间寻找平衡，如何在异构环境中保持代码的可移植性。本文记录的是一次完整的技术探索：在AWS中国区（cn-north-1）使用AWS ParallelCluster搭建多节点GPU训练集群，并对Qwen2.5-0.5B模型进行微调的全过程。

这个项目的起点源于一个朴素的疑问：在中国区的云环境中，能否复现那些在国际社区中被广泛讨论的最佳实践？答案远比想象中复杂。从PyPI镜像的选择到SSH密钥格式的要求，从依赖版本的地狱到网络拓扑的限制，每一步都充满了意外的挑战。而当这些挑战被一一克服后，我们收获的不仅是一个可以运行的训练环境，更是对分布式机器学习工程本质的深刻理解。

架构设计

从手动配置到声明式基础设施

在云计算的早期时代，搭建一个GPU集群意味着繁琐的手动操作：登录每台机器，安装驱动，配置网络，同步文件。这种工作模式不仅效率低下，而且极易出错。AWS ParallelCluster的出现改变了这一切。作为AWS开源的集群管理工具，它采用了声明式配置的理念------你描述你想要的集群状态，工具负责实现它。

ParallelCluster的核心架构体现了一种优雅的分层思想。最底层是基础设施层，由AWS EC2实例、EBS存储、VPC网络构成；中间层是软件栈，包括操作系统镜像、调度器、共享文件系统；最上层是用户应用层，也就是我们的训练代码和数据。

在我们的实验中，集群由两类节点组成：

Head Node (c5.xlarge)：运行Slurm调度器的主控进程（slurmctld），负责作业的队列管理和资源分配。它不需要GPU，但需要足够的计算能力和存储空间。我们选择了c5.xlarge，这是一款基于Intel Xeon的通用计算实例，配备4个vCPU和8GB内存。更重要的是，我们将它的根卷扩展到了100GB------这个决定来自惨痛的经验教训，因为在最初的尝试中，默认的8GB磁盘在PyTorch安装过程中就宣告空间不足。
Compute Node (g4dn.xlarge)：真正的计算引擎，配备T4 GPU (16GB显存)。虽然与A100相比性能有限，但对于Qwen2.5-0.5B这样参数量仅5亿的模型而言，T4完全能够胜任微调任务。

分布式训练的网络拓扑

g4dn.xlarge不支持EFA（Elastic Fabric Adapter），这意味着我们只能使用普通TCP网络进行GPU之间的通信。PyTorch的分布式数据并行（DDP）底层使用NCCL作为通信后端。

NCCL（NVIDIA Collective Communications Library）是NVIDIA提供的多GPU/多节点通信库，专为深度学习优化。它实现了高效的集合通信原语：

Broadcast：一个节点广播到所有节点
AllReduce：所有节点reduce后广播结果（梯度同步的核心）
AllGather：收集所有节点的数据
ReduceScatter：reduce后分散到各节点

NCCL底层支持多种网络传输：

NVLink：GPU直连（单节点内）
InfiniBand：RDMA网络（如EFA）
TCP/IP：以太网（本项目使用）

注意：NCCL backend 不是"运行在 master" ，而是所有 rank 都使用 NCCL 进行通信。NCCL (NVIDIA Collective Communications Library) 是部署在每个节点上的通信库，所有参与训练的 GPU 都会初始化 NCCL 并参与集合通信。Master (Rank 0) 的特殊性在于"协调"（创建 TCP store、广播初始模型、保存 checkpoint），而非"独占 NCCL"。

组件	说明
NCCL Backend	所有 rank（0, 1, 2...）都初始化 NCCL，参与集合通信
Master (Rank 0)	创建 TCP store、广播初始模型、保存 checkpoint
Worker (Rank 1, 2...)	通过 NCCL 与 Master 和其他 worker 通信

DDP通信时序（3节点场景，展示Ring-AllReduce）：
NCCL Backend(Ring拓扑 - 所有节点) Node 2 (Rank 2) Node 1 (Rank 1) Node 0 (Rank 0) NCCL Backend(Ring拓扑 - 所有节点) Node 2 (Rank 2) Node 1 (Rank 1) Node 0 (Rank 0) #mermaid-svg-zkKcMQg05HLrLVni{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-zkKcMQg05HLrLVni .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-zkKcMQg05HLrLVni .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-zkKcMQg05HLrLVni .error-icon{fill:#552222;}#mermaid-svg-zkKcMQg05HLrLVni .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-zkKcMQg05HLrLVni .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-zkKcMQg05HLrLVni .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-zkKcMQg05HLrLVni .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-zkKcMQg05HLrLVni .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-zkKcMQg05HLrLVni .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-zkKcMQg05HLrLVni .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-zkKcMQg05HLrLVni .marker{fill:#333333;stroke:#333333;}#mermaid-svg-zkKcMQg05HLrLVni .marker.cross{stroke:#333333;}#mermaid-svg-zkKcMQg05HLrLVni svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-zkKcMQg05HLrLVni p{margin:0;}#mermaid-svg-zkKcMQg05HLrLVni .actor{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-zkKcMQg05HLrLVni text.actor>tspan{fill:black;stroke:none;}#mermaid-svg-zkKcMQg05HLrLVni .actor-line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-zkKcMQg05HLrLVni .innerArc{stroke-width:1.5;stroke-dasharray:none;}#mermaid-svg-zkKcMQg05HLrLVni .messageLine0{stroke-width:1.5;stroke-dasharray:none;stroke:#333;}#mermaid-svg-zkKcMQg05HLrLVni .messageLine1{stroke-width:1.5;stroke-dasharray:2,2;stroke:#333;}#mermaid-svg-zkKcMQg05HLrLVni #arrowhead path{fill:#333;stroke:#333;}#mermaid-svg-zkKcMQg05HLrLVni .sequenceNumber{fill:white;}#mermaid-svg-zkKcMQg05HLrLVni #sequencenumber{fill:#333;}#mermaid-svg-zkKcMQg05HLrLVni #crosshead path{fill:#333;stroke:#333;}#mermaid-svg-zkKcMQg05HLrLVni .messageText{fill:#333;stroke:none;}#mermaid-svg-zkKcMQg05HLrLVni .labelBox{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-zkKcMQg05HLrLVni .labelText,#mermaid-svg-zkKcMQg05HLrLVni .labelText>tspan{fill:black;stroke:none;}#mermaid-svg-zkKcMQg05HLrLVni .loopText,#mermaid-svg-zkKcMQg05HLrLVni .loopText>tspan{fill:black;stroke:none;}#mermaid-svg-zkKcMQg05HLrLVni .loopLine{stroke-width:2px;stroke-dasharray:2,2;stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-zkKcMQg05HLrLVni .note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-zkKcMQg05HLrLVni .noteText,#mermaid-svg-zkKcMQg05HLrLVni .noteText>tspan{fill:black;stroke:none;}#mermaid-svg-zkKcMQg05HLrLVni .activation0{fill:#f4f4f4;stroke:#666;}#mermaid-svg-zkKcMQg05HLrLVni .activation1{fill:#f4f4f4;stroke:#666;}#mermaid-svg-zkKcMQg05HLrLVni .activation2{fill:#f4f4f4;stroke:#666;}#mermaid-svg-zkKcMQg05HLrLVni .actorPopupMenu{position:absolute;}#mermaid-svg-zkKcMQg05HLrLVni .actorPopupMenuPanel{position:absolute;fill:#ECECFF;box-shadow:0px 8px 16px 0px rgba(0,0,0,0.2);filter:drop-shadow(3px 5px 2px rgb(0 0 0 / 0.4));}#mermaid-svg-zkKcMQg05HLrLVni .actor-man line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-zkKcMQg05HLrLVni .actor-man circle,#mermaid-svg-zkKcMQg05HLrLVni line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;stroke-width:2px;}#mermaid-svg-zkKcMQg05HLrLVni :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 1. 初始化阶段每个节点独立初始化NCCL 2. 建立Ring拓扑N0 ↔ N1 ↔ N2 ↔ N0所有节点平等参与 3. 模型广播Master广播到所有Worker 4. 前向传播各节点独立计算 5. 反向传播计算本地梯度 6. Ring-AllReduce所有节点参与梯度同步 6a. Reduce-Scatter沿Ring传递并累加 6b. All-Gather沿Ring传播完整梯度 7. 梯度平均各节点独立计算 avg = sum / 3 8. 参数更新各节点独立优化 loop $每个训练迭代$ init_process_group()MASTER_ADDR=172.31.2.1MASTER_PORT=29500init_process_group()init_process_group()Rank 0 (World=3)Rank 1 (World=3)Rank 2 (World=3)broadcast(model.state_dict())model paramsmodel paramsforward(data_0)forward(data_1)forward(data_2)backward()grad_0backward()grad_1backward()grad_2send(grad_0 $0$ )send(grad_1 $0$ +grad_0 $0$ )send(grad_2 $0$ +grad_1 $0$ +grad_0 $0$ )send(sum_grad)send(sum_grad)send(sum_grad)grad_0 = sum_grad / 3grad_1 = sum_grad / 3grad_2 = sum_grad / 3optimizer.step()optimizer.step()optimizer.step()

Master 的特殊性 vs NCCL 的平等性：

操作	Master (Rank 0)	Worker (Rank 1+)
创建 TCP store	只有Rank 0创建	其他rank连接
广播初始模型	发送方	接收方
前向/反向传播	计算	计算
梯度 AllReduce	参与（平等）	参与（平等）
Checkpoint 保存	仅Rank 0保存	无

Master 和 Head Node 的区别 ：Master (Rank 0) 不是 Head Node，它们是不同层级的概念。

概念	Master (Rank 0)	Head Node
层级	应用层（PyTorch DDP）	基础设施层（AWS ParallelCluster）
角色	DDP 训练中的协调者	Slurm 调度器的控制中心
运行位置	Compute Node (g4dn.xlarge)	Head Node (c5.xlarge)
识别方式	`SLURM_PROCID=0`	`hostname=ip-172-31-29-148`
GPU	有 T4 GPU	无 GPU
职责	创建 TCP store、广播模型、保存 checkpoint	运行 slurmctld、作业调度、用户登录入口

两个概念都用了"Master/Head"这种"主"的含义，但层级完全不同：

Head Node = 公司总部的调度中心（分配资源、管理任务）
Master (Rank 0) = 项目组的负责人（在实际干活的团队里协调）

实际部署：

Head Node (c5.xlarge) 没有 GPU，所以 Master (Rank 0) 必须在 Compute Node (g4dn.xlarge) 上运行
Master 跑在第一个 Compute Node 上，和 Head Node 是不同的机器

Ring-AllReduce原理

对于N个节点，每个节点有梯度G，目标：所有节点获得sum(G1,G2,...GN) / N。

假设有3个节点，每个梯度分成3块：

阶段1：Reduce-Scatter（沿Ring传递并累加）

节点0发送G0 $0$ 给节点1，节点1发送G1 $1$ 给节点2，节点2发送G2 $2$ 给节点0
节点1接收后计算G1[0]+G0[0]，节点2计算G2[1]+G1[1]，节点0计算G0[2]+G2[2]
经过3轮传递，每个节点持有完整的sum(G0,G1,G2)的一个块

阶段2：All-Gather（沿Ring传播完整梯度）

每个节点将持有的sum_grad块传给下一个节点
经过3轮传递，每个节点都拥有完整的sum(G0,G1,G2)

通信复杂度：

数据总量：每个梯度大小为S
传输数据量：2(N-1)S/N ≈ 2S（与节点数N无关！）
这就是Ring-AllReduce的魔力：扩展性好，节点数增加不会显著增加通信量

通信流程说明

初始化阶段 ：主节点（Rank 0）创建TCP store，工作节点通过MASTER_ADDR:MASTER_PORT加入通信组，NCCL建立跨节点连接
模型广播：训练开始前，Master将模型参数广播到所有Worker，确保各节点参数一致
独立计算：每个节点处理自己的数据批次，独立进行前向/反向传播，计算本地梯度
梯度同步 ：通过AllReduce操作（通常是Ring-AllReduce算法），各节点交换并平均梯度。对于N个节点，每个节点的梯度最终变为sum(grads) / N
参数更新：各节点使用同步后的梯度独立更新本地参数，由于初始参数相同且梯度相同，更新后的参数仍保持一致

关键环境变量（TCP网络场景）：

bash 复制代码

export NCCL_SOCKET_IFNAME=ens5    # 指定通信网卡
export NCCL_DEBUG=WARN              # 调试日志级别
export NCCL_IB_DISABLE=1           # 禁用InfiniBand（EFA不可用）

重要：如果NCCL选择错误的网络接口，会导致进程挂起或超时。使用以下命令查看可用网卡：
bash 复制代码
ip link show
在AWS g4dn实例上，实际发现应使用 ens5 而非 eth0：
bash 复制代码
export NCCL_SOCKET_IFNAME=ens5
NCCL网络接口选择错误的典型表现：

进程在 init_process_group() 处挂起，错误信息：NCCL WARN Connect to <ip> failed : Connection refused

性能对比：

方案	延迟	带宽	适用场景
EFA	<10μs	100Gbps	p3/p4/g5实例
TCP	~100μs	25Gbps	g4dn等入门级

对于Qwen2.5-0.5B这样的小模型，TCP网络完全足够。模型参数量仅5亿，梯度大小约2GB（FP32），AllReduce在25Gbps网络下仅需~0.6秒。

环境搭建

ParallelCluster配置

配置采用YAML格式，背后由AWS CloudFormation编排。以下是关键配置：

yaml 复制代码

Region: cn-north-1
Image:
  Os: ubuntu2204

HeadNode:
  InstanceType: c5.xlarge
  LocalStorage:
    RootVolume:
      Size: 100              # 避免8GB默认磁盘不足
      VolumeType: gp3
  Ssh:
    KeyName: cluster-key-ed25519  # Ubuntu 22.04要求ed25519
  CustomActions:
    OnNodeConfigured:
      Script: s3://llm-training-data-pcluster-v2/scripts/setup-head-node-v2.sh

Scheduling:
  Scheduler: slurm
  SlurmQueues:
    - Name: gpu
      CapacityType: ONDEMAND
      ComputeResources:
        - Name: g4dn-xlarge
          InstanceType: g4dn.xlarge
          MinCount: 0
          MaxCount: 2
          Efa:
            Enabled: false   # g4dn不支持EFA

SharedStorage:
  - Name: shared
    StorageType: Efs
    MountDir: /shared

第一次集群创建以失败告终，在安装PyTorch时pip报告磁盘空间不足。这是c5.xlarge实例默认配置的一个陷阱，根卷只有8GB而PyTorch的CUDA版本安装包就接近20GB。

错误信息 ：OSError: [Errno 28] No space left on device

在配置文件中明确设置RootVolume.Size: 100，并选择gp3作为卷类型。

Head Node初始化脚本

当Head Node启动完成后，ParallelCluster会执行CustomActions中指定的脚本：

bash 复制代码

#!/bin/bash
# setup-head-node-v2.sh

# 配置阿里云PyPI镜像
mkdir -p ~/.pip
cat > ~/.pip/pip.conf << 'EOF'
[global]
index-url = https://mirrors.aliyun.com/pypi/simple/
trusted-host = mirrors.aliyun.com
EOF

# 创建共享目录结构
mkdir -p /shared/{python,models,datasets,logs,output}

# 从S3同步预下载的模型和数据集
aws s3 sync s3://llm-training-data-pcluster-v2/models/ /shared/models/
aws s3 sync s3://llm-training-data-pcluster-v2/datasets/ /shared/datasets/

共享存储的设计权衡

在分布式训练环境中，存储的选择直接影响开发效率：

EBS：单实例挂载，这意味着每个Compute Node都需要维护数据副本
S3：高延迟（几十到几百毫秒），不适合训练过程中的实时数据读取
EFS：多实例共享，POSIX语义，自动扩展，是我们的最终选择

EFS挂载到/shared目录，承载着四个关键子目录：

路径	用途	大小
`/shared/models`	Qwen2.5-0.5B模型权重	954MB
`/shared/datasets`	古诗词数据集	113MB (388,599条)
`/shared/python`	共享Python环境	动态增长
`/shared/output`	训练输出和日志	动态增长

模型与数据的准备

我们选择了通义千问（Qwen）系列的Qwen2.5-0.5B模型，从ModelScope下载：

python 复制代码

from modelscope import snapshot_download

snapshot_download(
    model_id="Qwen/Qwen2.5-0.5B",
    cache_dir="/shared/models"
)

数据集采用chinese-poetry-collection，格式为Alpaca-style：

json 复制代码

{
  "instruction": "请续写或完成以下古诗词：",
  "input": "床前明月光",
  "output": "床前明月光，疑是地上霜。举头望明月，低头思故乡。"
}

训练实现

数据集类设计

PyTorch的Dataset类定义了如何从存储介质中读取样本。我们的实现如下：

python 复制代码

class ChinesePoetryDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=256):
        self.examples = []
        with open(data_path, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                text = f"{item.get('instruction', '')}\n{item.get('input', '')}\n{item.get('output', '')}"
                self.examples.append(text)
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, idx):
        text = self.examples[idx]
        encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding='max_length',
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': encoding['input_ids'].squeeze().clone()
        }

完整训练脚本

关键设计决策：

在实现训练脚本前，需要明确以下关键决策：

使用JSON Lines格式：适合流式处理，每行一个样本
一次性加载到内存：38万条记录数据量可控，避免频繁IO
truncation=True：截断超长序列，避免显存溢出
padding='max_length'：填充到统一长度，便于批处理

训练参数说明：

per_device_train_batch_size=1：T4显存限制下的保守选择
gradient_accumulation_steps=4：累积4个batch的梯度，等效batch size为4
fp16=True：混合精度训练，显存节省50%，利用Tensor Core加速
learning_rate=5e-5：Transformer微调的常用经验值

多节点 DDP 训练脚本

NCCL进程组初始化 ：dist.init_process_group("nccl")
DDP模型包装 ：model = DDP(model, device_ids=[local_rank])
DistributedSampler：数据自动分片到各节点
Slurm环境变量 ：自动检测SLURM_PROCID、SLURM_NTASKS等

python 复制代码

#!/usr/bin/env python3
import os
import json
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import Dataset

class ChinesePoetryDataset(Dataset):
    # 与单节点版本相同...
    def __init__(self, data_path, tokenizer, max_length=256):
        self.examples = []
        with open(data_path, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                text = f"{item.get('instruction', '')}\n{item.get('input', '')}\n{item.get('output', '')}"
                self.examples.append(text)
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, idx):
        text = self.examples[idx]
        encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            padding='max_length',
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),
            'attention_mask': encoding['attention_mask'].squeeze(),
            'labels': encoding['input_ids'].squeeze().clone()
        }

def setup_distributed():
    """初始化分布式训练环境"""
    if "SLURM_PROCID" in os.environ:
        # Slurm环境自动检测
        rank = int(os.environ["SLURM_PROCID"])
        world_size = int(os.environ["SLURM_NTASKS"])
        local_rank = int(os.environ.get("SLURM_LOCALID", 0))
        
        # 设置MASTER_ADDR和MASTER_PORT
        if rank == 0:
            os.environ.setdefault("MASTER_ADDR", "localhost")
            os.environ.setdefault("MASTER_PORT", "29500")
    else:
        rank = 0
        world_size = 1
        local_rank = 0
    
    # NCCL初始化
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
        
    return rank, world_size, local_rank

def main():
    # 分布式初始化
    rank, world_size, local_rank = setup_distributed()
    
    # 模型加载到对应GPU
    model_path = os.environ.get("MODEL_PATH", "/shared/models/Qwen/Qwen2.5-0.5B")
    model = AutoModelForCausalLM.from_pretrained(model_path)
    device = torch.device(f"cuda:{local_rank}")
    model = model.to(device)
    
    # DDP包装
    model = DDP(model, device_ids=[local_rank])
    
    # Tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token = tokenizer.eos_token
    
    # 数据加载
    data_path = os.environ.get("DATA_PATH", "/shared/datasets/chinese-poetry/train.jsonl")
    dataset = ChinesePoetryDataset(data_path, tokenizer)
    
    # DistributedSampler实现数据分片
    train_sampler = DistributedSampler(
        dataset, 
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )
    
    dataloader = DataLoader(
        dataset, 
        batch_size=1, 
        sampler=train_sampler
    )
    
    # 训练循环
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)
    num_epochs = 1
    
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)  # 确保每个epoch数据顺序不同
        
        for batch in dataloader:
            # 数据移到对应GPU
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)
            
            # 前向/反向传播
            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            loss.backward()
            
            optimizer.step()
            optimizer.zero_grad()
            
            # DDP自动同步梯度（AllReduce）
    
    # 仅Rank 0保存模型
    if rank == 0:
        output_dir = os.environ.get("OUTPUT_DIR", "/shared/output")
        model.save_pretrained(os.path.join(output_dir, "final"))
        tokenizer.save_pretrained(os.path.join(output_dir, "final"))

if __name__ == "__main__":
    main()

依赖版本冲突是Python生态中长期存在的问题。在我们的项目中，这个问题的复杂程度超出预期。以下是问题演进过程，验证通过的版本组合如下

复制代码

torch==2.0.1+cu117
transformers==4.35.2
accelerate==0.24.1
tokenizers==0.15.2
huggingface-hub==0.19.4
numpy==1.26.4
datasets==2.14.0

Slurm作业脚本与调度

作业调度器允许集群资源被多个用户公平共享，同时保证长时间任务稳定执行：

bash 复制代码

#!/bin/bash
#SBATCH --job-name=qwen-ddp
#SBATCH --partition=gpu
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:1
#SBATCH --output=/shared/logs/%j.out
#SBATCH --error=/shared/logs/%j.err
#SBATCH --time=02:00:00

# NCCL网络配置
export NCCL_SOCKET_IFNAME=ens5
export NCCL_DEBUG=INFO
export NCCL_IB_DISABLE=1

echo "Job started at $(date)"
echo "Node List: $SLURM_NODELIST"
echo "Total tasks: $SLURM_NTASKS"
echo "Master address: $(scontrol show hostname $SLURM_NODELIST | head -n 1)"
echo "Master port: 29500"

# 安装依赖
pip3 install --user torch==2.0.1 transformers==4.35.2 \
    accelerate==0.24.1 datasets==2.14.0 numpy==1.26.4 \
    --index-url https://mirrors.aliyun.com/pypi/simple/

# 使用srun启动多进程训练
srun --mpi=pmix --export=ALL python3 /shared/train-ddp.py

Slurm多节点调度机制：

Slurm作为集群调度器，在多节点作业启动时会自动注入环境变量，让训练进程能够互相发现并建立通信：

环境变量	Slurm注入的值	用途
`SLURM_NODELIST`	ip-172-31-2-1,ip-172-31-2-2	分配的所有节点列表
`SLURM_JOB_NODELIST`	ip-172-31-2- $1-2$	作业的节点列表（压缩格式）
`SLURM_NODEID`	0, 1, 2...	当前节点在作业中的索引
`SLURM_PROCID`	0, 1, 2...	当前进程的全局ID
`SLURM_LOCALID`	0, 1, 2...	当前节点上的本地进程ID
`SLURM_NTASKS`	2	任务总数
`SLURM_NNODES`	2	节点总数

节点发现流程（以2节点作业为例）：
DDP Process Node 2 Node 1 Slurmctld 用户 DDP Process Node 2 Node 1 Slurmctld 用户 #mermaid-svg-iwsdQgJ0lfkvjkxY{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-iwsdQgJ0lfkvjkxY .error-icon{fill:#552222;}#mermaid-svg-iwsdQgJ0lfkvjkxY .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-iwsdQgJ0lfkvjkxY .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-iwsdQgJ0lfkvjkxY .marker{fill:#333333;stroke:#333333;}#mermaid-svg-iwsdQgJ0lfkvjkxY .marker.cross{stroke:#333333;}#mermaid-svg-iwsdQgJ0lfkvjkxY svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-iwsdQgJ0lfkvjkxY p{margin:0;}#mermaid-svg-iwsdQgJ0lfkvjkxY .actor{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-iwsdQgJ0lfkvjkxY text.actor>tspan{fill:black;stroke:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .actor-line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-iwsdQgJ0lfkvjkxY .innerArc{stroke-width:1.5;stroke-dasharray:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .messageLine0{stroke-width:1.5;stroke-dasharray:none;stroke:#333;}#mermaid-svg-iwsdQgJ0lfkvjkxY .messageLine1{stroke-width:1.5;stroke-dasharray:2,2;stroke:#333;}#mermaid-svg-iwsdQgJ0lfkvjkxY #arrowhead path{fill:#333;stroke:#333;}#mermaid-svg-iwsdQgJ0lfkvjkxY .sequenceNumber{fill:white;}#mermaid-svg-iwsdQgJ0lfkvjkxY #sequencenumber{fill:#333;}#mermaid-svg-iwsdQgJ0lfkvjkxY #crosshead path{fill:#333;stroke:#333;}#mermaid-svg-iwsdQgJ0lfkvjkxY .messageText{fill:#333;stroke:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .labelBox{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-iwsdQgJ0lfkvjkxY .labelText,#mermaid-svg-iwsdQgJ0lfkvjkxY .labelText>tspan{fill:black;stroke:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .loopText,#mermaid-svg-iwsdQgJ0lfkvjkxY .loopText>tspan{fill:black;stroke:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .loopLine{stroke-width:2px;stroke-dasharray:2,2;stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);}#mermaid-svg-iwsdQgJ0lfkvjkxY .note{stroke:#aaaa33;fill:#fff5ad;}#mermaid-svg-iwsdQgJ0lfkvjkxY .noteText,#mermaid-svg-iwsdQgJ0lfkvjkxY .noteText>tspan{fill:black;stroke:none;}#mermaid-svg-iwsdQgJ0lfkvjkxY .activation0{fill:#f4f4f4;stroke:#666;}#mermaid-svg-iwsdQgJ0lfkvjkxY .activation1{fill:#f4f4f4;stroke:#666;}#mermaid-svg-iwsdQgJ0lfkvjkxY .activation2{fill:#f4f4f4;stroke:#666;}#mermaid-svg-iwsdQgJ0lfkvjkxY .actorPopupMenu{position:absolute;}#mermaid-svg-iwsdQgJ0lfkvjkxY .actorPopupMenuPanel{position:absolute;fill:#ECECFF;box-shadow:0px 8px 16px 0px rgba(0,0,0,0.2);filter:drop-shadow(3px 5px 2px rgb(0 0 0 / 0.4));}#mermaid-svg-iwsdQgJ0lfkvjkxY .actor-man line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;}#mermaid-svg-iwsdQgJ0lfkvjkxY .actor-man circle,#mermaid-svg-iwsdQgJ0lfkvjkxY line{stroke:hsl(259.6261682243, 59.7765363128%, 87.9019607843%);fill:#ECECFF;stroke-width:2px;}#mermaid-svg-iwsdQgJ0lfkvjkxY :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Slurm自动设置环境变量MASTER_ADDR, MASTER_PORTSLURM_NODELIST等开始DDP训练 sbatch --nodes=2 job.sh启动Task 0 (RANK=0)启动Task 1 (RANK=1)读取 $MASTER_ADDR (Node 1 IP)读取$ MASTER_PORT (29500)RANK=0作为Master启动TCP Store读取 $MASTER_ADDR读取$ MASTER_PORTTCP连接: RANK=1加入通信组连接成功, 进程组初始化完成

PyTorch自动检测Slurm：

这意味着用户不需要手动配置节点IP，Slurm自动处理了节点发现和服务发现的复杂性。

Slurm的slurmctld作为中央调度器，记录所有节点的IP地址和状态
作业启动时，slurmd代理在Compute Node上执行脚本
Slurm自动设置MASTER_ADDR为第一个节点的IP，MASTER_PORT为随机空闲端口
PyTorch的torchrun或torch.distributed.launch读取这些环境变量完成进程组初始化

python 复制代码

import os
import torch.distributed as dist

# PyTorch会自动读取Slurm注入的环境变量
# 如果检测到SLURM_PROCID，会自动使用Slurm进行初始化
if "SLURM_PROCID" in os.environ:
    rank = int(os.environ["SLURM_PROCID"])
    world_size = int(os.environ["SLURM_NTASKS"])
    dist.init_process_group(
        backend="nccl",
        rank=rank,
        world_size=world_size
    )

常用命令：

bash 复制代码

sbatch job-final.sh         # 提交作业
squeue                       # 查看作业队列
scancel <job_id>            # 取消作业
tail -f /shared/logs/*.out   # 实时查看日志

实际运行状态

多节点DDP训练（Job ID: 19）已成功启动并正在运行：

启动日志：

NCCL版本 2.14.3+cuda11.7 正确初始化
Rank 0 和 Rank 1 都成功初始化并识别到各自的GPU
数据自动分片：每个节点处理约 194,299 个样本
Checkpoint已保存到 /shared/output/checkpoint-500

Job started at Sun Jun 28 04:30:23 UTC 2026

Slurm Job Configuration:
Nodes: 2
Tasks per node: 1
GPUs per node: 1
Node List: gpu-dy-g4dn-xlarge-1 gpu-dy-g4dn-xlarge-2

Master address: gpu-dy-g4dn-xlarge-1
Master port: 29500
Launching DDP training with srun...
NCCL version 2.14.3+cuda11.7
[Rank 0/2] DDP initialized on cuda:0
[Rank 1/2] DDP initialized on cuda:0

Multi-Node DDP Training
World Size: 2 nodes
Model: Qwen2.5-0.5B

Loading model from: /shared/models/Qwen/Qwen2.5-0.5B
Loading dataset from: /shared/datasets/chinese-poetry/train.jsonl
Dataset size: 388599 samples
Samples per node: ~194299

Starting training...
Total steps: 194300

训练作业正在集群后台运行。可以通过以下方式监控：

bash 复制代码

# 查看作业状态
ssh -i ~/.ssh/pcluster-key-ed25519 ubuntu@<head-node-ip>
squeue

# 实时查看训练日志
tail -f /shared/logs/19.out

# 查看checkpoint
cd /shared/output && ls -lt

在ParallelCluster 集群中通过 PyTorch DDP微调LLM模型的实践

架构设计

从手动配置到声明式基础设施

分布式训练的网络拓扑

Ring-AllReduce原理

通信流程说明

环境搭建

ParallelCluster配置

Head Node初始化脚本

共享存储的设计权衡

模型与数据的准备

训练实现

数据集类设计

完整训练脚本

多节点 DDP 训练脚本

Slurm作业脚本与调度

实际运行状态

Job started at Sun Jun 28 04:30:23 UTC 2026

Slurm Job Configuration: Nodes: 2 Tasks per node: 1 GPUs per node: 1 Node List: gpu-dy-g4dn-xlarge-1 gpu-dy-g4dn-xlarge-2

Master address: gpu-dy-g4dn-xlarge-1 Master port: 29500 Launching DDP training with srun... NCCL version 2.14.3+cuda11.7 [Rank 0/2] DDP initialized on cuda:0 [Rank 1/2] DDP initialized on cuda:0

Multi-Node DDP Training World Size: 2 nodes Model: Qwen2.5-0.5B

Slurm Job Configuration:
Nodes: 2
Tasks per node: 1
GPUs per node: 1
Node List: gpu-dy-g4dn-xlarge-1 gpu-dy-g4dn-xlarge-2

Master address: gpu-dy-g4dn-xlarge-1
Master port: 29500
Launching DDP training with srun...
NCCL version 2.14.3+cuda11.7
[Rank 0/2] DDP initialized on cuda:0
[Rank 1/2] DDP initialized on cuda:0

Multi-Node DDP Training
World Size: 2 nodes
Model: Qwen2.5-0.5B