🚀 PyTorch 并行训练指南

1. 数据并行 (Data Parallelism)

• 每个设备持有完整的模型副本
• 将训练数据分割到不同设备上
• 各设备独立计算梯度，然后同步聚合（AllReduce）
• 适用场景 模型能放入单个GPU显存

2. 模型并行 / 张量并行 (Model/Tensor Parallelism)

• 将模型层内部切分到多个设备
• 例如：将一个大矩阵按行或列切分
• 需要设备间频繁通信
• 适用场景 单层参数太大，无法放入一个GPU

3. 流水线并行 (Pipeline Parallelism)

• 将模型按层切分到不同设备
• 各设备处理不同的层，形成流水线
• 通过 micro-batch 减少气泡（空闲时间）
• 适用场景 模型层数多，整体无法放入单个GPU

对比总结

💡 大规模训练（如GPT、LLaMA）通常三者混合使用，称为 3D并行。

训练与推理适用性

📌 推理场景总结：大模型推理框架（vLLM、TensorRT-LLM、HF TGI）主要使用张量并行和流水线并行，数据并行在推理中通常被简单的多实例部署替代。

1. 数据并行 - DistributedDataParallel (DDP)

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

dist.init_process_group(backend="nccl")
model = DDP(model, device_ids=[local_rank])

注意 旧版 DataParallel 已不推荐（单进程多线程，效率低）

2. 张量并行 - torch.distributed.tensor.parallel

from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel
)

parallelize_module(
    model,
    device_mesh,
    {
        "layer.weight": ColwiseParallel(),
        "output.weight": RowwiseParallel(),
    }
)

3. 流水线并行 - torch.distributed.pipelining

from torch.distributed.pipelining import (
    pipeline, SplitPoint, ScheduleGPipe
)

pipe = pipeline(
    model,
    mb_args=(microbatch,),
    split_spec={
        "layer4": SplitPoint.END,
        "layer8": SplitPoint.END,
    }
)
schedule = ScheduleGPipe(pipe, n_microbatches=4)
schedule.step(x)

3D 并行整合 - DeviceMesh

from torch.distributed.device_mesh import init_device_mesh

# 创建 2D mesh: (数据并行, 张量并行)
mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))

更高层封装工具

1. 环境变量（推荐）

# 只使用 GPU 0 和 2
CUDA_VISIBLE_DEVICES=0,2 python train.py

# 程序内看到的是 cuda:0 和 cuda:1（重新编号）

2. 代码中指定设备

# 指定默认设备
torch.cuda.set_device(0)

# 创建张量时指定
tensor = torch.randn(3, 3, device="cuda:1")

# 模型移动到指定GPU
model = model.to("cuda:2")
# 或
model = model.cuda(2)

3. 分布式训练 (DDP)

启动方式：

# torchrun 自动分配 (推荐)
torchrun --nproc_per_node=4 train.py

# 指定使用哪些 GPU
CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --nproc_per_node=4 train.py

代码中获取分配的 GPU：

import os
import torch.distributed as dist

dist.init_process_group(backend="nccl")

local_rank = int(os.environ["LOCAL_RANK"])  # torchrun 自动设置
torch.cuda.set_device(local_rank)

model = model.to(local_rank)
model = DDP(model, device_ids=[local_rank])

4. 多机分布式

# 机器1 (master)
torchrun --nnodes=2 --node_rank=0 --nproc_per_node=4 \
         --master_addr=192.168.1.1 --master_port=29500 train.py

# 机器2
torchrun --nnodes=2 --node_rank=1 --nproc_per_node=4 \
         --master_addr=192.168.1.1 --master_port=29500 train.py

5. DeviceMesh（3D 并行）

from torch.distributed.device_mesh import init_device_mesh

# 8卡: 2路数据并行 × 4路张量并行
mesh = init_device_mesh("cuda", (2, 4), mesh_dim_names=("dp", "tp"))

# 获取当前进程在各维度的位置
dp_rank = mesh.get_local_rank("dp")  # 0 或 1
tp_rank = mesh.get_local_rank("tp")  # 0, 1, 2, 或 3

常用查询命令

torch.cuda.device_count()        # 可用GPU数量
torch.cuda.current_device()      # 当前默认GPU
torch.cuda.get_device_name(0)    # GPU型号
torch.cuda.memory_summary(0)     # 显存使用情况

nvidia-smi                       # 查看所有GPU状态

总结

场景设定

• 模型：LLaMA-70B 级别（80层 Transformer，隐藏维度 8192）
• 硬件：128 节点 × 8 GPU/节点 = 1024 GPU（假设 H100 80GB）
• 目标：高效训练，最大化吞吐量

GPU 显存规格

LLaMA-70B 显存需求分析

3D 并行后每卡显存

💡 计算：1170 GB ÷ 64 (TP×PP) ≈ 18 GB/卡 + 激活值 ≈ 20-25 GB/卡

训练数据量参考

训练时间估算

理论计算公式：

实际估算示例

不同配置对比

⚠️ 注意：实际训练时间受集群网络、存储IO、故障恢复等因素影响，通常比理论值高 20-50%

⚡ 显存约束 vs 算力约束

大模型训练存在两个不同的瓶颈，需要区分对待：

✅ 总结：3D并行+ZeRO+Flash解决的是"能否训练"的问题，训练速度完全取决于集群总算力（GPU数量×单卡算力×利用率）

3D 并行配置方案

验证：TP × PP × DP = 8 × 8 × 16 = 1024 GPU ✅

GPU 拓扑结构图

启动脚本示例

#!/bin/bash
# run_1024gpu.sh - 在 SLURM 集群上启动训练

# SLURM 配置
#SBATCH --job-name=llama70b_train
#SBATCH --nodes=128
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=12

# 网络配置
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=29500
export WORLD_SIZE=1024

# 3D 并行配置
export TP_SIZE=8    # 张量并行度
export PP_SIZE=8    # 流水线并行度
export DP_SIZE=16   # 数据并行度（自动计算：1024/8/8=16）

# 训练超参数
export MICRO_BATCH_SIZE=1
export GLOBAL_BATCH_SIZE=2048  # DP_SIZE × gradient_acc_steps × micro_batch
export GRADIENT_ACC_STEPS=128  # 2048 / 16 / 1 = 128

# 启动训练
srun torchrun \
    --nnodes=128 \
    --nproc_per_node=8 \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train.py \
    --model llama-70b \
    --tp_size $TP_SIZE \
    --pp_size $PP_SIZE \
    --micro_batch_size $MICRO_BATCH_SIZE \
    --global_batch_size $GLOBAL_BATCH_SIZE

训练代码框架

# train.py - 3D并行训练主程序
import os
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor.parallel import (
    parallelize_module, ColwiseParallel, RowwiseParallel
)
from torch.distributed.pipelining import pipeline, SplitPoint, ScheduleGPipe
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

def main():
    # 1. 初始化分布式环境
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    world_rank = int(os.environ["RANK"])
    torch.cuda.set_device(local_rank)

    # 2. 创建 3D 设备网格 (DP, PP, TP)
    #    1024 GPU = 16 × 8 × 8
    mesh_3d = init_device_mesh(
        "cuda",
        (16, 8, 8),  # (DP, PP, TP)
        mesh_dim_names=("dp", "pp", "tp")
    )

    # 获取当前进程在各维度的位置
    dp_rank = mesh_3d.get_local_rank("dp")   # 0-15
    pp_rank = mesh_3d.get_local_rank("pp")   # 0-7 (哪个Stage)
    tp_rank = mesh_3d.get_local_rank("tp")   # 0-7

    # 3. 构建模型（仅当前PP Stage的层）
    layers_per_stage = 10  # 80层 / 8 stages
    start_layer = pp_rank * layers_per_stage
    end_layer = start_layer + layers_per_stage

    model = build_llama_layers(
        start_layer=start_layer,
        end_layer=end_layer,
        hidden_size=8192,
        num_heads=64
    )

    # 4. 应用张量并行（切分 Attention 和 FFN）
    tp_mesh = mesh_3d["tp"]
    parallelize_module(
        model,
        tp_mesh,
        {
            # Attention: Q/K/V 列切分，Output 行切分
            "attention.q_proj": ColwiseParallel(),
            "attention.k_proj": ColwiseParallel(),
            "attention.v_proj": ColwiseParallel(),
            "attention.o_proj": RowwiseParallel(),
            # FFN: gate/up 列切分，down 行切分
            "ffn.gate_proj": ColwiseParallel(),
            "ffn.up_proj": ColwiseParallel(),
            "ffn.down_proj": RowwiseParallel(),
        }
    )

    # 5. 设置流水线并行
    pp_group = mesh_3d.get_group("pp")

    # 6. 数据并行（使用FSDP进一步节省显存）
    dp_mesh = mesh_3d["dp"]
    model = FSDP(model, process_group=dp_mesh.get_group())

    # 7. 优化器
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

    # 8. 训练循环
    for step, batch in enumerate(dataloader):
        # 流水线执行（自动处理micro-batch）
        with pipeline_context(pp_group, num_microbatches=80):
            loss = model(batch)
            loss.backward()

        optimizer.step()
        optimizer.zero_grad()

        if world_rank == 0 and step % 10 == 0:
            print(f"Step {step}, Loss: {loss.item():.4f}")

if __name__ == "__main__":
    main()

通信与计算分析

GPU 互联带宽详解

1024 GPU 集群中，GPU 间带宽取决于物理连接层级：

带宽规格对比

1024 GPU (128节点) 实际带宽配置

带宽与并行策略匹配

💡 关键设计原则：高频通信（张量并行）放节点内用 NVLink；低频大数据量通信（数据并行）放节点间可重叠计算

ZeRO 与 Flash Attention 优化

这两项技术可以与 3D 并行同时使用，进一步提升训练效率：

ZeRO (Zero Redundancy Optimizer)

PyTorch 原生实现 - FSDP：

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    ShardingStrategy
)

# ZeRO-3 等效配置
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,  # ZeRO-3
    # ShardingStrategy.SHARD_GRAD_OP  # ZeRO-2
    # ShardingStrategy.NO_SHARD       # DDP (ZeRO-0)
)

# 或使用 DeepSpeed
import deepspeed
model, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config={
        "zero_optimization": {
            "stage": 3,  # ZeRO-3
            "offload_optimizer": {"device": "cpu"},  # 可选：卸载到CPU
        }
    }
)

Flash Attention

优化 Attention 计算，减少显存 5-20×，加速 2-4×

使用方式：

# 方式1: PyTorch 2.0+ 原生支持
with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=False
):
    output = F.scaled_dot_product_attention(q, k, v)

# 方式2: flash-attn 库 (推荐，功能更全)
from flash_attn import flash_attn_func
output = flash_attn_func(q, k, v, causal=True)

# 方式3: Transformers 库自动启用
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16
)

3D 并行 + ZeRO + Flash 组合效果

💡 最佳实践：3D并行 + ZeRO-2 + Flash Attention + 激活检查点 = 最优显存/性能平衡

关键优化建议

• 通信隐藏：数据并行的梯度 AllReduce 与下一轮计算重叠
• 激活检查点：节省显存，允许更大batch或更多层
• 混合精度：使用 BF16/FP16 减少通信量和显存
• Flash Attention：优化Attention计算，减少显存占用
• 梯度累积：80步累积实现大batch，减少通信频率

# 推荐的训练配置（含 ZeRO + Flash）
config = {
    # 精度设置
    "precision": "bf16",

    # 3D 并行配置
    "tensor_parallel_size": 8,
    "pipeline_parallel_size": 8,
    "data_parallel_size": 16,

    # ZeRO 配置 (与数据并行结合)
    "zero_stage": 2,              # 0=DDP, 1=优化器分片, 2=+梯度分片, 3=+参数分片
    "zero_offload": False,        # True=卸载到CPU (省显存但慢)

    # Flash Attention 配置
    "flash_attention": True,
    "flash_attention_version": 2, # 1 或 2 (推荐2)

    # 激活检查点
    "activation_checkpointing": True,
    "checkpoint_granularity": "selective",  # full / selective

    # Batch 配置
    "micro_batch_size": 1,
    "gradient_accumulation_steps": 128,
    "global_batch_size": 2048,    # 16 × 128 × 1
}

# DeepSpeed 完整配置示例
deepspeed_config = {
    "bf16": {"enabled": True},
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": True,           # 通信计算重叠
        "contiguous_gradients": True,
        "reduce_bucket_size": 5e8,
    },
    "activation_checkpointing": {
        "partition_activations": True,
        "contiguous_memory_optimization": True,
    },
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 128,
}

# PyTorch 原生配置 (FSDP + Flash)
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision

fsdp_config = {
    "sharding_strategy": ShardingStrategy.SHARD_GRAD_OP,  # ZeRO-2
    "mixed_precision": MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
    "use_orig_params": True,  # 兼容Flash Attention
}

📌 版本说明

• PyTorch 1.x：主要支持 DDP
• PyTorch 2.0+：原生支持张量并行、流水线并行
• 早期需依赖 FairScale、DeepSpeed、Megatron-LM 等第三方库
• PyTorch 2.x 已将这些能力逐步整合进核心库，API 仍在演进中