把单卡换到多卡时踩过的坑

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup_ddp(rank, world_size):
    # 初始化进程组
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://127.0.0.1:12345',
        rank=rank,
        world_size=world_size
    )
    # 设置当前设备的 CUDA 设备
    torch.cuda.set_device(rank)

def cleanup_ddp():
    dist.destroy_process_group()

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

    # 创建模型并移动到当前设备
    model = MyModel()
    model = model.to(rank)
    model = DDP(model, device_ids=[rank])

    # 数据集和采样器
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)

    # 训练循环
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # 确保每个 epoch 数据打乱不同
        for batch in dataloader:
            # 训练逻辑
            pass

    cleanup_ddp()

if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size)

# 重新编译 NCCL
git clone https://github.com/NVIDIA/nccl.git
cd nccl
make -j8
sudo make install

# 重新安装 PyTorch
pip install torch --no-cache-dir

# 错误的梯度裁剪时机
for batch in dataloader:
    loss = model(batch)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 错误：在同步前裁剪
    optimizer.step()

# 正确的梯度裁剪时机
for batch in dataloader:
    loss = model(batch)
    loss.backward()
    # DDP 在这里自动同步梯度
    optimizer.step()
    # 梯度裁剪应该放在 optimizer.step() 之后
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

import math

# 线性缩放规则
base_lr = 1e-4
world_size = 2
adjusted_lr = base_lr * world_size

optimizer = torch.optim.AdamW(model.parameters(), lr=adjusted_lr)

from torch.distributed.pipeline.sync import Pipe

# 创建模型分段
model = nn.Sequential(
    nn.Linear(1024, 1024),
    nn.ReLU(),
    # ... 更多层
)

# 分割模型到多张卡
chunks = 4  # 分成 4 段
model = Pipe(model, chunks=chunks)

# 训练
for batch in dataloader:
    output = model(batch)  # 自动在多卡间流水线执行
    loss = criterion(output, target)
    loss.backward()

# 简化的张量并行示例
class ColumnParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.out_features_per_gpu = out_features // world_size
        self.weight = nn.Parameter(torch.randn(in_features, self.out_features_per_gpu))
        self.world_size = world_size

    def forward(self, x):
        # 本地计算
        local_output = torch.matmul(x, self.weight)
        # 聚合所有 GPU 的结果
        outputs = [torch.zeros_like(local_output) for _ in range(self.world_size)]
        dist.all_gather(outputs, local_output)
        return torch.cat(outputs, dim=-1)

import deepspeed

# DeepSpeed 配置
ds_config = {
    "train_batch_size": 32,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-4,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 1e-4,
            "warmup_num_steps": 1000
        }
    },
    "fp16": {
        "enabled": True,
        "loss_scale": 0,
        "initial_scale_power": 16,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": True
    },
    "gradient_clipping": 1.0
}

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

# 训练循环
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

# 混合精度训练 + DDP
scaler = torch.cuda.amp.GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        loss = model(batch)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

import torch
import numpy as np
import random

def set_seed(seed, rank):
    random.seed(seed + rank)
    np.random.seed(seed + rank)
    torch.manual_seed(seed + rank)
    torch.cuda.manual_seed(seed + rank)
    torch.cuda.manual_seed_all(seed + rank)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# 使用时
set_seed(42, rank)

# 只在主进程做验证
if rank == 0:
    model.eval()
    with torch.no_grad():
        for batch in val_dataloader:
            # 验证逻辑
            pass
    model.train()

# 定期清理缓存
if step % 100 == 0:
    torch.cuda.empty_cache()