分布式训练与并行策略：DDP、张量并行、流水线并行与 ZeRO

数据并行（Data Parallel）  ：每卡一份完整模型，切分数据 → 最常用
张量并行（Tensor Parallel） ：把单层矩阵切到多卡 → 层内切分
流水线并行（Pipeline Parallel）：把不同层放到不同卡 → 层间切分
专家并行（Expert Parallel） ：MoE 的不同专家分到不同卡 → 稀疏激活

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

def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def main():
    local_rank = setup()
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])

    dataset = MyDataset()
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = torch.utils.data.DataLoader(
        dataset, batch_size=32, sampler=sampler, num_workers=4
    )

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # 关键：每个 epoch 打乱不同
        for batch in loader:
            loss = model(**batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)

if __name__ == "__main__":
    main()

torchrun --nproc_per_node=8 --nnodes=2 --node_rank=0 \
  --master_addr=10.0.0.1 --master_port=29500 train.py

# 测量实际可用带宽（NCCL 自带测试）
./build/all_reduce_perf -b 8 -e 4G -f 2 -g 8

import torch
import torch.nn as nn
import torch.distributed as dist

class ColumnParallelLinear(nn.Module):
    """按输出维度切分：每卡持有一部分输出通道"""
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.out_per_rank = out_features // world_size
        self.weight = nn.Parameter(
            torch.empty(self.out_per_rank, in_features)
        )

    def forward(self, x):
        # x 在每卡相同，输出是完整输出的一段
        return torch.nn.functional.linear(x, self.weight)

class RowParallelLinear(nn.Module):
    """按输入维度切分：每卡持有一部分输入，输出 all-reduce 求和"""
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.in_per_rank = in_features // world_size
        self.weight = nn.Parameter(
            torch.empty(out_features, self.in_per_rank)
        )

    def forward(self, x):
        partial = torch.nn.functional.linear(x, self.weight)
        dist.all_reduce(partial)  # 各卡部分和相加
        return partial

时间 →
卡0: [F1][F2][F3][F4]            [B4][B3][B2][B1]
卡1:      [F1][F2][F3][F4]   [B4][B3][B2][B1]
卡2:           [F1][F2][F3][F4]  ...

# PyTorch 原生流水线（torch.distributed.pipeline.sync）
from torch.distributed.pipeline.sync import Pipe

model = nn.Sequential(layer1, layer2, layer3, layer4)
model = Pipe(model, chunks=8, checkpoint="except_last")
output = model(input)  # 自动切微批、调度 1F1B

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy

mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # 对应 ZeRO-3
    mixed_precision=mp_policy,
    device_id=torch.cuda.current_device(),
    use_orig_params=True,
)

{
  "train_batch_size": 256,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "initial_scale_power": 16
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu", "pin_memory": true },
    "offload_param": { "device": "cpu", "pin_memory": true },
    "overlap_comm": true,
    "contiguous_gradients": true,
    "stage3_gather_16bit_weights_on_model_save": true,
    "sub_group_size": 1e9,
    "reduce_bucket_size": "auto"
  },
  "gradient_clipping": 1.0,
  "steps_per_print": 100
}

deepspeed --num_gpus=8 --num_nodes=2 train.py --deepspeed ds_config_zero3.json \
  --per_device_train_batch_size 8 --gradient_accumulation_steps 4 --bf16

from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="./out",
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    bf16=True,
    deepspeed="ds_config_zero3.json",
    gradient_checkpointing=True,
)

deepspeed --num_gpus=64 train.py --deepspeed ds_config.json \
  --tensor_model_parallel_size 8 --pipeline_model_parallel_size 2 --zero_stage 1

accum_steps = 4
for i, batch in enumerate(loader):
    loss = model(**batch) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

with model.no_sync():        # 微批内不同步梯度
    loss.backward()
loss.backward()              # 最后一个微批才同步

export NCCL_DEBUG=INFO           # 打印通信细节
export NCCL_IB_DISABLE=1         # 无 InfiniBand 时禁用
export NCCL_SOCKET_IFNAME=eth0   # 指定网卡

1. 单卡能否跑通？→ 不能则先修模型代码
2. 单机 8 卡能否跑通？→ 不能则查 NCCL/拓扑
3. 多机能否跑通？→ 查网络与时钟同步
4. 加速比是否线性？→ 不线性则 profile 通信占比；收敛是否与单卡一致？对比前 100 步 loss

模型单卡装得下 + 想更快       → DDP 数据并行
模型单卡装不下 + 有 NVLink    → 张量并行（节点内）
模型很多层 + 想省通信         → 流水线并行（跨节点）
模型装不下 + 无高带宽         → ZeRO-1/2 + CPU offload
千亿参数 + 大规模集群         → 3D 并行（TP×PP×DP）+ ZeRO-1