【DeepSpeed】多节点分布式训练

已在 3 个节点上使用 DeepSpeed 进行分布式训练为例，需要配置多节点环境、设置 DeepSpeed 的分布式策略，并确保通信正常。以下是详细步骤。

---

1. 前提条件

• 硬件：3 个节点，每节点有 GPU（例如每节点 4 张 NVIDIA A100）。
• 网络：节点间通过高带宽网络连接（如 InfiniBand 或高速以太网），支持 NCCL 通信。
• 软件：
  • 安装 PyTorch（推荐 torch>=2.0）和 DeepSpeed（pip install deepspeed）。
  • 安装 MPI（如 OpenMPI）用于多节点通信。
  • NCCL（NVIDIA Collective Communications Library）确保 GPU 间高效通信。
• 环境一致：所有节点需安装相同版本的 Python、PyTorch、DeepSpeed 和 CUDA。

验证环境：
运行以下命令检查 DeepSpeed 配置：

ds_report

确保 CUDA、NCCL 和 MPI 正常。

---

2. DeepSpeed 分布式训练原理

DeepSpeed 支持多节点分布式训练，主要依赖：

• 数据并行：每个节点处理不同数据批次。
• 模型并行（可选）：模型参数分布在节点间。
• 流水线并行（可选）：模型层分布在节点间。
• ZeRO：通过分区优化器状态、梯度和参数，降低显存占用，适合多节点。

对于 3 个节点，通常结合 数据并行 和 ZeRO-2/3，以平衡通信开销和显存效率。

---

3. 配置多节点环境

多节点训练需要一个主机文件和启动脚本。

主机文件（hostfile）

创建一个 hostfile，列出节点 IP 或主机名及 GPU 数量。例如：

node1 slots=4
node2 slots=4
node3 slots=4

• node1、node2、node3 是节点的主机名或 IP。
• slots 表示每节点 GPU 数量（这里假设每节点 4 张 GPU，总共 12 张）。

将 hostfile 放在可访问路径，例如 /path/to/hostfile。

SSH 配置

确保节点间支持无密码 SSH 登录：

1. 在主节点（例如 node1）生成 SSH 密钥：

   ssh-keygen
   

2. 复制公钥到其他节点：

   ssh-copy-id node2
   ssh-copy-id node3
   

3. 验证无密码登录：

   ssh node2
   

---

4. 编写训练代码

以下是一个简单的 PyTorch 模型，结合 DeepSpeed 进行多节点训练。

示例代码（train.py）

import torch
import torch.nn as nn
import deepspeed
import os

# 初始化分布式环境
deepspeed.init_distributed(dist_backend="nccl")

# 定义简单模型
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(128, 256)
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# DeepSpeed 配置
ds_config = {
    "train_batch_size": 192,  # 全局批大小（3 节点，每节点 64）
    "gradient_accumulation_steps": 1,
    "fp16": {
        "enabled": True,  # 混合精度（参考你对 FP16 的兴趣）
        "loss_scale": 0,
        "initial_scale_power": 16
    },
    "zero_optimization": {
        "stage": 2,  # ZeRO-2，适合多节点
        "allgather_partitions": True,
        "reduce_scatter": True
    },
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001
        }
    }
}

# 初始化模型
model = MLP()

# 初始化 DeepSpeed
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=ds_config,
    model_parameters=model.parameters()
)

# 模拟数据
input_data = torch.randn(64, 128).cuda()  # 每节点批大小 64
target = torch.randint(0, 10, (64,)).cuda()

# 训练
model_engine.train()
output = model_engine(input_data)
loss = nn.CrossEntropyLoss()(output, target)
model_engine.backward(loss)
model_engine.step()

if torch.distributed.get_rank() == 0:
    print(f"Rank {torch.distributed.get_rank()}: Loss = {loss.item()}")

说明：

• deepspeed.init_distributed 初始化多节点通信。
• train_batch_size=192：3 节点，每节点批大小 192/3=64。
• fp16：启用混合精度，参考你之前对 FP16 和 bfloat16 的兴趣。
• zero_optimization：使用 ZeRO-2，分区优化器状态和梯度。

---

5. DeepSpeed 配置文件

将配置保存为 ds_config.json，便于复用：

{
  "train_batch_size": 192,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "reduce_scatter": true
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001
    }
  }
}

然后修改代码，使用外部配置文件：

with open("ds_config.json") as f:
    ds_config = json.load(f)

---

6. 启动多节点训练

在主节点（例如 node1）运行以下命令启动训练：

deepspeed --num_nodes=3 --num_gpus=4 --hostfile=/path/to/hostfile train.py

参数说明：

• --num_nodes=3：使用 3 个节点。
• --num_gpus=4：每节点 4 张 GPU。
• --hostfile：指定主机文件路径。

环境变量（可选）：
如果遇到通信问题，设置以下变量：

export NCCL_IB_DISABLE=0
export NCCL_IB_GID_INDEX=3

---

7. 监控与调试

• 日志：DeepSpeed 会生成日志文件（默认在 ./output），检查每节点的输出。
• 通信检查：确保 NCCL 正常工作，运行：

  python -m torch.distributed.run --nproc_per_node=4 --nnodes=3 --node_rank=0 --master_addr=node1 --master_port=29500 -m torch.distributed.test
  

• 性能监控：使用 nvidia-smi 查看 GPU 利用率，或用 DeepSpeed 的 ds_report 检查配置。

---

8. 优化建议

• ZeRO 阶段：如果显存不足，尝试 ZeRO-3（stage: 3），但通信开销会增加。
• 流水线并行：对于大模型，将模型层分配到 3 个节点，设置 "pipeline_parallelism": {"enabled": true}。
• 批大小调整：根据节点数和 GPU 显存，调整 train_batch_size 和 gradient_accumulation_steps。
• 通信优化：确保节点间网络延迟低，使用 InfiniBand 或 RDMA。

---

9. 常见问题

• NCCL 错误：检查 CUDA 版本、NCCL 安装，或设置 NCCL_DEBUG=INFO 调试。
• 节点不同步：确保 hostfile 正确，所有节点代码和环境一致。
• 显存溢出：启用 ZeRO-3 或优化器卸载（offload_optimizer: {"device": "cpu"}）。