模型并行与分布式训练

大模型训练需要多 GPU 甚至多机的计算资源，模型并行（Model Parallelism）和分布式训练是关键。

为什么需要并行？

单卡显存限制：

LLaMA-2-70B：需要 ~140GB 显存（FP16）
即使是 7B 模型，大 batch 也需要多卡

并行策略分类

1. 数据并行（Data Parallelism）

最常用：同一模型复制到多卡，数据分片。

GPU 1: [模型] ← batch[0:32]
GPU 2: [模型] ← batch[32:64]
GPU 3: [模型] ← batch[64:96]

前向+反向：每卡独立计算梯度 梯度同步：All-Reduce 聚合梯度（需通信） 参数更新：所有卡权重一致

实现（PyTorch DDP）

python

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

# 初始化进程组
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)

# 模型包装
model = MyModel().cuda()
model = DDP(model, device_ids=[local_rank])

# 数据分片（每个进程只取部分数据）
train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, sampler=train_sampler)

# 训练循环（每卡独立）
for batch in train_loader:
    loss = model(batch)
    loss.backward()
    optimizer.step()  # 梯度自动同步

优点：

扩展性好（线性加速）
实现简单

缺点：

梯度同步通信开销大
需要每卡存储完整模型副本

适用场景：模型能装下单卡，需要大批量训练。

2. 张量并行（Tensor Parallelism）

模型层内部分片：不同 GPU 计算不同的张量部分。

常见切分方式

行并行（Row-wise）：

Linear(W: [in, out]) → 切分输出维度 out
GPU 1: W[:, 0:out/2], 计算 out/2 个输出
GPU 2: W[:, out/2:out], 计算剩余输出
→ 拼接得到完整输出

列并行（Column-wise）：

GPU 1: W[0:in/2, :] → 部分输出
GPU 2: W[in/2:in, :] → 部分输出
→ 相加（或拼接）得到完整输出

Megatron-LM 示例

python

# 简化的张量并行 Linear
class TensorParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, tp_size, rank):
        super().__init__()
        self.tp_size = tp_size
        self.rank = rank
        # 每卡只存一部分权重
        self.weight = nn.Parameter(
            torch.randn(out_features // tp_size, in_features)
        )

    def forward(self, x):
        # All-Reduce 或其他通信收集结果
        out = F.linear(x, self.weight)
        # 跨卡通信（如 All-Gather 或 Reduce-Scatter）
        return out

框架支持：

Megatron-LM：NVIDIA 官方张量并行
DeepSpeed：ZeRO + 张量并行

混合并行示例

python

# DeepSpeed 配置（ZeRO-3 + 张量并行）
{
  "train_batch_size": 32,
  "fp16": {"enabled": true},
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {"device": "cpu"}
  },
  "tensor_parallel": {
    "tensor_model_parallel_size": 2  # 张量并行 2 卡
  }
}

优点：

单层大模型可训练（如 100B+）
通信主要发生在层内

缺点：

通信频繁（每个 layer 都需通信）
实现复杂

适用场景：单层参数巨大（如 GPT-3 175B 的 FFN 层）。

3. 流水线并行（Pipeline Parallelism）

模型按层切分到不同 GPU，形成流水线。

Stage 1 (GPU 0): Layer 1-10  + 微批0.1 + 微批0.2
Stage 2 (GPU 1): Layer 11-20 + 微批0.1 + 微批0.2
Stage 3 (GPU 2): Layer 21-30 + 微批0.1 + 微批0.2

前向：数据逐阶段流动反向：梯度反向传播

实现（GPipe）

python

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

# 划分模型
num_stages = 4
layers_per_stage = total_layers // num_stages
model = nn.Sequential(*layers)
model = model.to(local_rank)

# 将模型切分为 stages
model_stage = model[stage_start:stage_end]

# 数据分 micro-batch
micro_batches = torch.chunk(batch, num_micro_batches)

# 流水线执行
for micro_batch in micro_batches:
    output = model_stage(micro_batch)
    # 输出传递到下一 stage

通信：仅在 stage 边界通信（通信量小于张量并行）

挑战：

气泡（Bubble）：流水线空闲时间，降低效率
负载均衡：各 stage 计算量需均衡

优化：

微批（Micro-batching）：减小 bubble
重计算（Checkpointing）：节省内存，牺牲计算

4. 序列并行（Sequence Parallelism）

针对长序列训练，按序列维度切分。

Transformer 中，序列长度可能很长（如 32K tokens），按序列维度切分可降低单卡内存。

适用场景：长文本模型（如 GPT-3 175B 训练时序列长度达 2048）

主流框架

DeepSpeed（Microsoft）

最流行的分布式训练框架，支持：

ZeRO（Zero Redundancy Optimizer）：消除数据并行的参数冗余
ZeRO-3：参数、梯度、优化器状态全切分 → 显存降低 4× 以上
混合并行：组合数据、张量、流水线并行

python

# DeepSpeed 配置示例
ds_config = {
    "train_batch_size": 32,
    "gradient_accumulation_steps": 4,
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {"device": "cpu"},
        "offload_param": {"device": "cpu"},
        "overlap_comm": True,
        "contiguous_gradients": True
    }
}
model = deepspeed.initialize(model=model, config=ds_config)

ZeRO 阶段对比：

阶段	优化器状态	梯度	参数	显存节省
ZeRO-1	切分	不切分	不切分	4×
ZeRO-2	切分	切分	不切分	5×
ZeRO-3	切分	切分	切分	8×+

Megatron-LM（NVIDIA）

专为超大模型（>100B）设计的张量并行框架。

张量并行（Tensor Parallelism）
流水线并行（Pipeline Parallelism）
与 DeepSpeed 集成

bash

# 示例：Megatron-DeepSpeed 训练 GPT-3 175B
python pretrain_gpt.py \
  --tensor-model-parallel-size 2 \
  --pipeline-model-parallel-size 8 \
  --num-layers 96 \
  --hidden-size 12288 \
  --num-attention-heads 96

Fully Sharded Data Parallel (FSDP)

PyTorch 的 ZeRO-3 实现，语法简洁：

python

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision

model = MyModel().cuda()
model = FSDP(
    model,
    mixed_precision=MixedPrecision(param_dtype=torch.float16)
)

优势：与原生 PyTorch API 兼容性好。

组合策略（3D 并行）

最大模型通常组合三种并行：

示例：GPT-3 175B 训练

张量并行：2 卡/层 → 每层参数分 2 份
流水线并行：32 个 stage → 每 stage 3 层
数据并行：60×（960 卡总数 / 64 模型并行）

总卡数 = 张量并行 × 流水线并行 × 数据并行 = 2 × 32 × 60 = 3840 GPU

通信后端

后端	适用场景	协议
NCCL	NVIDIA GPU（推荐）	专为 GPU 优化
Gloo	CPU 多机	通用 CPU
MPI	传统 HPC	标准 MPI
MPI+NCCL	混合 CPU+GPU	组合

python

dist.init_process_group(backend="nccl")  # GPU 训练首选

实践建议

1. 环境准备

bash

# NCCL 调试（多卡通信）
export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0

# 多节点 SSH 免密登录
# 每个节点能互相访问

2. 启动命令

bash

# 单机多卡（DDP）
torchrun --nproc_per_node=4 train.py

# 多机多卡
torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --master_addr=192.168.1.1 \
  --master_port=29500 \
  train.py

3. 性能调优

增大 batch size：填满 GPU
使用混合精度（FP16/BF16）：减半内存，加速计算
梯度累积：模拟更大 batch
重叠通信：通信与计算重叠（DeepSpeed ZeRO-3）
检查点（Gradient Checkpointing）：用时间换空间，训练更大模型

常见问题

Q: 显存不足怎么办？

A: 检查并应用（按优先级）：

减小 batch size
启用梯度检查点
使用 ZeRO-2/3
增加模型并行或流水线并行

Q: 训练速度慢？

A: 查看卡利用率（nvidia-smi）：

利用率 < 30%：数据加载瓶颈（增加 dataloader workers）
利用率 30-70%：正常
利用率 ~100%：计算饱和

使用 profiler 定位瓶颈：

python

with torch.profiler.profile() as prof:
    train_step()
print(prof.export_chrome_trace("trace.json"))

Q: 多机训练无法通信？

A: 检查：

所有节点能互相 ping 通
端口（默认 29500）开放
使用相同版本的 PyTorch + NCCL

总结

数据并行：最常用，适合单卡能放下模型
张量并行：超大层参数（>10B）
流水线并行：非常大模型（>100B）
组合策略：三者结合训练最大模型

掌握这些技术，你可以训练任意规模的模型！

模型并行与分布式训练 ​

为什么需要并行？ ​

并行策略分类 ​

1. 数据并行（Data Parallelism） ​

实现（PyTorch DDP） ​

2. 张量并行（Tensor Parallelism） ​

常见切分方式 ​

Megatron-LM 示例 ​

混合并行示例 ​

3. 流水线并行（Pipeline Parallelism） ​

实现（GPipe） ​

4. 序列并行（Sequence Parallelism） ​

主流框架 ​

DeepSpeed（Microsoft） ​

Megatron-LM（NVIDIA） ​

Fully Sharded Data Parallel (FSDP) ​

组合策略（3D 并行） ​

通信后端 ​

实践建议 ​

1. 环境准备 ​

2. 启动命令 ​

3. 性能调优 ​

常见问题 ​

Q: 显存不足怎么办？ ​

Q: 训练速度慢？ ​

Q: 多机训练无法通信？ ​

总结 ​

模型并行与分布式训练

为什么需要并行？

并行策略分类

1. 数据并行（Data Parallelism）

实现（PyTorch DDP）

2. 张量并行（Tensor Parallelism）

常见切分方式

Megatron-LM 示例

混合并行示例

3. 流水线并行（Pipeline Parallelism）

实现（GPipe）

4. 序列并行（Sequence Parallelism）

主流框架

DeepSpeed（Microsoft）

Megatron-LM（NVIDIA）

Fully Sharded Data Parallel (FSDP)

组合策略（3D 并行）

通信后端

实践建议

1. 环境准备

2. 启动命令

3. 性能调优

常见问题

Q: 显存不足怎么办？

Q: 训练速度慢？

Q: 多机训练无法通信？

总结