5D 并行配置#

RLinf 使用 NVIDIA Megatron-LM 作为大规模 Transformer 训练的后端，并提供五种相互独立的并行模式：

Tensor Parallelism (TP) 张量并行
Data Parallelism (DP) 数据并行
Pipeline Parallelism (PP) 流水线并行
Sequence Parallelism (SP) 序列并行
Context Parallelism (CP) 上下文并行

合适的组合可以让系统从单节点扩展到数百 GPU 的集群，同时在内存、通信和利用率之间取得平衡。

下面将详细介绍这五种并行模式，并给出启动时的配置方法。

1. 张量并行 (TP)#

定义

张量并行 将模型的 权重矩阵 切分到一个进程组中。每个 GPU 仅保存并计算自己负责的部分。

机制

线性层 – Linear(h, 4h) 按列切分 – Linear(4h, h) 按行切分
注意力投影 – \(Q,K,V\) 按 head 维度切分
每次切分计算后通过 all-reduce 将结果合并

YAML 示例

actor:
  model:
    tensor_model_parallel_size: 2    # tp_size = 2

优点

避免单 GPU 内存瓶颈
在 GPU 之间平衡计算

缺点

几乎每层后都需要 all-reduce → 高延迟
扩展性受限于隐藏层宽度

2. 数据并行 (DP)#

定义

数据并行 将 小批量数据 (minibatch) 切分，每个副本都存储完整的模型。

机制

全局 batch 在 DP rank 间切分
每个 rank 独立计算前向/反向
在优化器更新前通过 all-reduce 同步梯度

规模示例

cluster:
  num_nodes: 16

actor:
  model:
    tensor_model_parallel_size: 2     # TP
    pipeline_model_parallel_size: 2   # PP
    context_parallel_size: 2          # CP
  # dp_size = 128 / 2 / 2 / 2 = 16

优点

简单，模型无关
非常适合扩展数据集规模

缺点

每个 GPU 都保存完整模型 → 内存开销大
梯度 all-reduce 覆盖全部参数
通常需要与 TP/PP/CP 结合来支持更大模型

3. 流水线并行 (PP)#

定义

流水线并行 将 不同的层堆栈 分配到不同的 rank 上，组成计算流水线。

机制

将层均匀分配到 pp_size 个阶段
使用 1F1B （一前一后）或类似调度器来重叠计算

调度示意

GPU 0: [F1][F2][F3][F4][B4][B3][B2][B1]
GPU 1:     [F1][F2][F3][F4][B4][B3][B2][B1]
GPU 2:          [F1][F2][F3][F4][B4][B3][B2][B1]
GPU 3:               [F1][F2][F3][F4][B4][B3][B2][B1]

F = 前向 micro-batch，B = 反向 micro-batch，index = micro-batch 编号。

YAML 示例

actor:
  model:
    pipeline_model_parallel_size: 2

优点

降低超深模型的内存需求
只需相邻间通信（激活）

缺点

流水线气泡（空闲槽位）可能降低利用率

4. 序列并行 (SP)#

定义

Megatron 的 序列并行 在 TP 的基础上扩展，用于减少 长上下文 注意力和 MLP 模块的内存。

机制

必须与 TP 一起启用；两者使用相同进程组
注意力和 MLP 的输入/输出在序列维度上切分，而权重切分与 TP 保持一致

YAML 示例

actor:
  model:
    tensor_model_parallel_size: 2     # 启用 TP
    sequence_parallel: True           # 启用 SP

# 如果 TP = 1，则必须关闭 SP
actor:
  model:
    tensor_model_parallel_size: 1
    sequence_parallel: False

优点

对长序列显著减少内存占用

缺点

序列维度 shuffle 带来额外通信

5. 上下文并行 (CP)#

定义

上下文并行 面向 超长序列，通过将整个注意力计算在序列维度上切分，所有张量在该维度上分片。

机制

将 \(Q,K,V\) 和 logits 切分为 context 块
使用环形注意力进行通信，并逐步累积输出

YAML 示例

actor:
  model:
    context_parallel_size: 2

优点

打破 100k+ token 上下文的内存瓶颈
与动态 batch 大小搭配良好

缺点

带宽开销大；参数不被切分，模型内存依然完全复制

总结#

Megatron-LM 灵活组合 TP、DP、PP、SP 和 CP，使 RLinf 可以通过宽度（TP）、数据规模 （DP）、深度（PP）或 上下文长度 （SP / CP）来扩展模型。选择合适的规模取决于模型结构、目标序列长度、GPU 内存和互联拓扑，以获得最佳吞吐量。