nemo-mbridge-perf-hierarchical-context-parallel

分层上下文并行技能

本技能介绍分层上下文并行：由 cp_comm_type="a2a+p2p" 使用，并通过 hierarchical_context_parallel_sizes 配置的嵌套上下文并行进程组。

关于分层 CP 是什么、何时使用，以及决策树 （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_size
• seq_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 路径：

pg_collection = ProcessGroupCollection(
    ...
    cp=cp_pg,
    tp_cp=tp_cp_pg,
    hcp=None,
    ep=ep_pg,
    ...
)

实现概览

以上代码锚点展示了配置声明和参数验证。

验证（MCore）

TransformerConfig.__post_init__ 会强制要求 a2a+p2p 配置 HCP 尺寸，且其乘积必须与 CP 一致。

进程组创建

提供 HCP 尺寸时，parallel_state.initialize_model_parallel 会通过 create_hierarchical_groups 创建分层 CP 子组。目前，Bridge 通过基于 MPU 的 ProcessGroupCollection 获取这些组。

TE 集成

使用 a2a+p2p 时，TEDotProductAttention 会将分层组传递给 Transformer Engine。需要 Transformer Engine >= 1.12.0。

注意事项

1. 目前 Bridge 的 HCP 仅支持 MPU：如果设置了 use_decentralized_pg=True，Bridge 会初始化扁平 CP 组，并将 HCP 留空。
2. 目前没有纳入版本控制的 Bridge 配置会直接运行 HCP。
3. 单 GPU 加载辅助函数会清除 hierarchical_context_parallel_sizes。
4. 在旧版本堆栈上训练可能悄然出错：如果使用 a2a+p2p 却未设置 hierarchical_context_parallel_sizes，MCore 现在会触发断言。旧版本会悄然禁用 CP 通信，因此每个 rank 只会处理本地分块，并产生虚高的吞吐量，但梯度是错误的。
5. 乘积必须匹配：prod(hierarchical_context_parallel_sizes) 必须与 context_parallel_size 完全一致。不匹配会触发断言。
6. 检查日志：查找进程组初始化输出。应能看到创建了 HIERARCHICAL_CONTEXT_PARALLEL_GROUPS。如果只看到 CONTEXT_PARALLEL_GROUP，则 HCP 未启用。

验证

目前还没有专门针对 HCP 的 Bridge 端到端测试（参见 @skills/nemo-mbridge-perf-hierarchical-context-parallel/card.yaml follow_up_validation）。请改用现有单元测试并检查日志。

运行分散式 PG 单元测试，以确认扁平 CP 行为保持不变：

uv run python -m pytest tests/unit_tests/training/test_decentralized_pg.py -q

如需手动进行冒烟测试，可使用小型配置启动 4-GPU 运行，并设置 cp_comm_type=a2a+p2p 和 hierarchical_context_parallel_sizes=[2,2]：

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 未启用