Distributed Training and Reasoning Optimization of LLM

Sufficient memory: Can we go faster? ¹

1. 数据并行 (Data Parallelism - DP)

Memory is a constraint. What to do? 2. 模型并行 (Model Parallelism - MP)

• 张量并行 (Tensor Parallelism - TP)
• 流水线并行 (Pipeline Parallelism - PP)

3. 序列与上下文并行 (Sequence & Context Parallelism)

• 序列并行 (Sequence Parallelism - SP)
• 上下文并行 (Context Parallelism - CP)

4. 专家并行 (Expert Parallelism - EP)

1. TPx (2D 张量并行 - 维1)：将模型权重和输入数据在模型维度（hidden dimension）上进行切分 。
2. TPy (2D 张量并行 - 维2)：将输入数据在序列长度（sequence length）维度上进行切分，与 TPx 共同构成 2D 张量并行，减少冗余并提高内存效率 。
3. SP (序列并行 - Sequence Parallelism)：通过 AlltoAll 操作将切分维度从序列轴转置到注意力头（attention head）轴，使每个设备独立处理部分注意力头 。
4. CP (上下文并行 - Context Parallelism)：在处理长序列时，通过环形（ring-based）通信在多个设备间切分并同步序列上下文 。
5. DP (数据并行 - Data Parallelism)：将训练数据（mini-batch）切分到不同设备上，每个设备维护模型的一个副本（或通过 ZeRO 进行分片）并同步梯度 。
6. PP (流水线并行 - Pipeline Parallelism)：将模型的不同层（layers）分布在不同的设备阶段，数据以微批次（micro-batches）形式在层间流水化执行 。
7. EP (专家并行 - Expert Parallelism)：针对 MoE（混合专家）模型，将不同的专家分配到不同设备，每个 Token 根据路由算法发送至对应的专家进行计算 。

SP 和 CP 的主要区别在于 SP 需要切分权重

Sequence Parallelism 完成 sequence 维度的并行，覆盖操作包括 LayerNorm、Dropout、FC，但不能切分 self-attention 模块

Attention 中 QKV 的计算需要用到一个完整 sequence 信息，计算上的耦合使得该模块不能先运算后进行简单拼接

Context Parallelism 解决 SP 中未完成的 self-attention 序列并行问题

ring attention

²

stripped ring attention

为解决 ring attention 在 casual mask 下的负载不均衡问题，将任务尽量均匀地分布到不同卡上

³

目前只关心 Q 的分布，不同卡计算 K/V 的顺序是不同的，这个是否会有什么影响？

⁴

Reference

• 序列并行视角下的各类研究