5D 并行配置
================================

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

1. **Tensor Parallelism (TP)** 张量并行  
2. **Data Parallelism (DP)** 数据并行  
3. **Pipeline Parallelism (PP)** 流水线并行  
4. **Sequence Parallelism (SP)** 序列并行  
5. **Context Parallelism (CP)** 上下文并行  

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

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

1. 张量并行 (TP)
--------------------------

**定义**

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

**机制**

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

**YAML 示例**

.. code-block:: yaml

   actor:
     model:
       tensor_model_parallel_size: 2    # tp_size = 2

**优点**

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

**缺点**

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

2. 数据并行 (DP)
------------------------

**定义**

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

**机制**

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

**规模示例**

.. code-block:: yaml

   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 示例**

.. code-block:: yaml

   actor:
     model:
       pipeline_model_parallel_size: 2

**优点**

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

**缺点**

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

4. 序列并行 (SP)
----------------------------

**定义**

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

**机制**

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

**YAML 示例**

.. code-block:: 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)
---------------------------

**定义**

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

**机制**

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

**YAML 示例**

.. code-block:: yaml

   actor:
     model:
       context_parallel_size: 2

**优点**

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

**缺点**

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

总结
-------

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