Skip to content

04. Model and Context Partitioning | 模型与上下文切分 ​

学习问题 ​

状态分片主要改变训练状态的驻留方式;PP、TP 和 CP 则分别沿层、张量和序列切分模型内部工作。三者可以组合,但切分对象、通信位置和判断指标不能混用。请求到副本的动态分发进入下一节,与静态 Device Mesh 分开判断。

层、张量与序列三种模型内部切分维度

核心机制 ​

三种模型内部切分维度 ​

选择切分方案时,先写清每种方案切开什么、在哪里交换,以及应读取哪组证据,再判断多种切分能否放入同一个 Device Mesh。

机制切分对象主要通信首先观察的代价适用信号
Pipeline Parallel(PP)连续层 / stagestage 间 activation 与 backward gradientbubble、stage 不均衡模型很深,单卡放不下层序列
Tensor Parallel(TP)单层矩阵、特征维、Attention HeadAll-Reduce、All-Gather、Reduce-Scatter高频层内同步单层矩阵或 Attention 本身过大
Context Parallel(CP)sequence / contextK/V 或中间状态的环式交换通信随上下文增长长序列状态无法由单卡承载

CP 如何在不复制完整上下文时计算 Attention ​

CP 先让每个 rank 保留一段 Query 和对应的本地 K/V,再通过环式点对点交换或 All-Gather 逐步取得远端 K/V。每轮计算只处理当前可见的 K/V 分片,并把局部 Attention 统计合并为与未切分计算一致的结果。它减少单卡常驻的序列状态,但不会消除读取历史 K/V 的工作量。

阶段本地状态跨 rank 动作必须验证
初始布局本地 Query 与一段 K/V按序列位置建立分片归属分片连续、无遗漏且位置编码一致
分轮计算本地 Query 与当前 K/V 分片环式发送、接收下一段 K/V每轮 payload、调用次数与依赖顺序
统计合并局部最大值、指数和与加权输出合并分块 Softmax 统计输出与未切分 Attention 数值一致
结束与复用本地输出与下一层所需状态按后续布局保留或重分片峰值状态、交换字节和端到端延迟

组合并行先写 Device Mesh ​

组合策略应先写出 world size 如何分解,再为每个 mesh 维度指定切分对象和进程组。DP × TP × PP × CP / EP 不是越多越好;每增加一个维度,就增加一种通信模式和配置约束。

组合主要收益新的组合约束
PP + TP同时切深度和单层宽度stage 内 collective 与 stage 间 P2P 可能同时暴露
TP + CP同时扩展模型宽度和上下文Head 映射、序列布局和通信组必须兼容
DP/FSDP + TP扩吞吐并切单层状态数据并行组与 TP 组不能混淆

选择条件与验证证据 ​

切分问题机制页真实证据出口
stage 与 micro-batchPart 02 · 28 Pipeline87 PAR-MODEL-INTERNAL 的 PP 分支
Column / Row 与 Head 布局Part 02 · 29 Tensor Parallel87 PAR-MODEL-INTERNAL 的 TP 分支
长上下文序列切分PAR-CONTEXT87 PAR-MODEL-INTERNAL 的 CP 分支

相关阅读 ​

Released under the MIT License.