| 名称 | nemo-mbridge-perf-hierarchical-context-parallel |
| 描述 | 在 Megatron-Bridge 中启用层次化上下文并行的操作指南,包括配置旋钮、代码锚点、陷阱和验证。 |
| 开源协议 | Apache-2.0 when_to_use: 当需要将上下文并行扩展到 KV 头之外,或调查一个改变 CP 配置并导致 OOM 或回归的提交时使用;相关关键词包括 ‘hierarchical_context_parallel_sizes’、‘a2a+p2p’、‘hierarchical CP’、‘CP beyond KV heads’、‘multi-level CP’。 |
层次化上下文并行技能
本技能涵盖层次化上下文并行:由 cp_comm_type="a2a+p2p" 使用并由 hierarchical_context_parallel_sizes 配置的嵌套上下文并行进程组。
关于层次化上下文并行的定义、使用时机以及决策树(a2a+p2p 与纯 a2a 与 p2p 的对比),请参阅:
- @docs/training/hierarchical-context-parallel.md
- @skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml
启用方法
最小的 Bridge 覆盖配置:
cfg.model.context_parallel_size = 4
cfg.model.cp_comm_type = "a2a+p2p"
cfg.model.hierarchical_context_parallel_sizes = [2, 2]
cfg.dist.use_decentralized_pg = False
必需约束:
prod(hierarchical_context_parallel_sizes) == context_parallel_sizeseq_length % (2 * context_parallel_size) == 0- Transformer Engine
>= 1.12.0
代码锚点
上游配置与校验:
context_parallel_size: int = 1
"""Splits network input along sequence dimension across GPU ranks."""
hierarchical_context_parallel_sizes: Optional[list[int]] = None
"""Degrees of the hierarchical context parallelism. Users should provide a list to specify
the sizes for different levels. Taking the a2a+p2p cp comm type as example, it contains
groups of two levels, so the first value of the list indicates the group size of the a2a
communication type, and the second value indicates the group size of the p2p communication
type.
"""
if args.hierarchical_context_parallel_sizes:
from numpy import prod
assert args.context_parallel_size == prod(args.hierarchical_context_parallel_sizes)
if "a2a+p2p" in args.cp_comm_type:
assert args.hierarchical_context_parallel_sizes is not None, \
"--hierarchical-context-parallel-sizes must be set when a2a+p2p is used in cp comm"
Bridge MPU 路径:
parallel_state.initialize_model_parallel(
...
context_parallel_size=model_config.context_parallel_size,
hierarchical_context_parallel_sizes=model_config.hierarchical_context_parallel_sizes,
...
)
...
return ProcessGroupCollection.use_mpu_process_groups()
Bridge 去中心化进程组路径:
pg_collection = ProcessGroupCollection(
...
cp=cp_pg,
tp_cp=tp_cp_pg,
hcp=None,
ep=ep_pg,
...
)
实现映射
上述代码锚点展示了配置声明与参数校验。
校验(MCore)
TransformerConfig.__post_init__ 强制要求 a2a+p2p 需要 HCP 大小,且乘积与 CP 匹配。
进程组创建
当通过 create_hierarchical_groups 提供 HCP 大小时,parallel_state.initialize_model_parallel 会创建层次化 CP 子组。Bridge 目前通过基于 MPU 的 ProcessGroupCollection 获取这些组。
TE 集成
当使用 a2a+p2p 时,TEDotProductAttention 会将层次化组传递给 Transformer Engine。要求 Transformer Engine >= 1.12.0。
陷阱
- Bridge HCP 目前仅支持 MPU:如果
use_decentralized_pg=True,Bridge 会初始化扁平 CP 组,并将 HCP 留空。 - 目前没有签入的 Bridge 配方直接演练 HCP。
- 单 GPU 加载辅助器会清除
hierarchical_context_parallel_sizes。 - 旧技术栈上的静默训练损坏:如果在没有设置
hierarchical_context_parallel_sizes的情况下使用a2a+p2p,MCore 现在会断言。旧版本会静默禁用 CP 通信,因此每个 rank 只关注其本地块,产生人为的高吞吐量和损坏的梯度。 - 乘积必须匹配:
prod(hierarchical_context_parallel_sizes)必须恰好等于context_parallel_size。不匹配会触发断言。 - 在日志中验证:查找进程组初始化输出。您应该看到正在创建
HIERARCHICAL_CONTEXT_PARALLEL_GROUPS。如果只看到CONTEXT_PARALLEL_GROUP,则 HCP 未激活。
验证
目前还没有针对 HCP 的专门 Bridge 端到端测试(参见 @skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml 中的 follow_up_validation)。请改用现有的单元测试和日志检查。
运行 decentralized-PG 单元测试,确认扁平 CP 行为保持不变:
uv run python -m pytest tests/unit_tests/training/test_decentralized_pg.py -q
对于手动冒烟检查,使用一个小型配方和 cp_comm_type=a2a+p2p 以及 hierarchical_context_parallel_sizes=[2,2] 启动 4-GPU 运行:
CUDA_VISIBLE_DEVICES=0,1,2,3 uv run python -m torch.distributed.run --nproc_per_node=4 \
scripts/training/run_recipe.py \
--recipe llama32_1b_pretrain_config \
model.context_parallel_size=4 \
model.cp_comm_type=a2a+p2p \
"model.hierarchical_context_parallel_sizes=[2,2]" \
train.train_iters=2
成功标准:
- 日志显示正在创建
HIERARCHICAL_CONTEXT_PARALLEL_GROUPS - 训练至少完成一步且无错误
- 如果只看到
CONTEXT_PARALLEL_GROUP,则 HCP 未激活