AI分布式训练踩坑记录

import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)

    model = Model().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=1e-4)

    for data, label in dataloader:
        optimizer.zero_grad()
        output = ddp_model(data)
        loss = criterion(output, label)
        loss.backward()
        optimizer.step()

    cleanup()

class SplitModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.part0 = nn.Sequential(...).to('cuda:0')
        self.part1 = nn.Sequential(...).to('cuda:1')
        self.part2 = nn.Sequential(...).to('cuda:2')
        self.part3 = nn.Sequential(...).to('cuda:3')

    def forward(self, x):
        x = self.part0(x.to('cuda:0'))
        x = self.part1(x.to('cuda:1'))
        x = self.part2(x.to('cuda:2'))
        x = self.part3(x.to('cuda:3'))
        return x

# 原始矩阵 W shape: [hidden_size, hidden_size]

# 在两张卡上水平切分 W
W0 = W[:, :hidden_size//2]  # 在 GPU0
W1 = W[:, hidden_size//2:]  # 在 GPU1

# 输入 x 也需要对应切分
x0 = x[:, :hidden_size//2]
x1 = x[:, hidden_size//2:]

# 每张卡计算部分结果
y0 = x0 @ W0
y1 = x1 @ W1

# 最后 all-reduce 合并结果
y = y0 + y1

# 1. 在所有节点安装相同的 PyTorch 版本
pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121

# 2. 确保所有节点能无密码 SSH 互相访问
ssh-keygen -t rsa
ssh-copy-id user@node1
ssh-copy-id user@node2

# 3. 同步代码和数据
rsync -avz /data/model user@node1:/data/
rsync -avz /data/model user@node2:/data/

# 4. 配置环境变量
export MASTER_ADDR="node1"  # 主节点 IP
export MASTER_PORT="29500"
export WORLD_SIZE=8  # 总进程数（2 节点 x 4 卡）
export NCCL_DEBUG=INFO  # 方便调试

#!/bin/bash
# launch_cluster.sh

NNODES=2
NODE_RANK=0
GPUS_PER_NODE=4

# 在每个节点上运行
torchrun \
    --nproc_per_node=$GPUS_PER_NODE \
    --nnodes=$NNODES \
    --node_rank=$NODE_RANK \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    train.py \
    --config config.yaml

import deepspeed

ds_config = {
    "train_batch_size": 32,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 4,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-4,
            "betas": [0.9, 0.999],
            "eps": 1e-8
        }
    },
    "fp16": {
        "enabled": True,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 2,  # ZeRO-2
        "offload_optimizer": {
            "device": "cpu",  # 把优化器状态卸载到 CPU
            "pin_memory": True
        },
        "offload_param": {
            "device": "cpu"
        }
    },
    "gradient_clipping": 1.0,
    "prescale_gradients": False,
    "wall_clock_breakdown": False
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config
)