NeMo-MBridge层次化上下文并行Skill nemo-mbridge-perf-hierarchical-context-parallel

本技能提供在Megatron-Bridge中启用、配置和验证层次化上下文并行(Hierarchical Context Parallelism)的操作指南,包括 cp_comm_type='a2a+p2p'、hierarchical_context_parallel_sizes 设置、必需约束、代码锚点、常见陷阱及日志验证方法,用于解决上下文并行规模扩展和CP配置导致的性能或OOM问题。关键词:层次化上下文并行、hierarchical context parallel、a2a+p2p、context parallel、Megatron-Bridge、CP配置、进程组、Transformer Engine、分布式训练调优

分布式训练调优 0 次安装 1 次浏览 更新于 9/7/2026
名称 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 与纯 a2ap2p 的对比),请参阅:

  • @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_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

陷阱

  1. Bridge HCP 目前仅支持 MPU:如果 use_decentralized_pg=True,Bridge 会初始化扁平 CP 组,并将 HCP 留空。
  2. 目前没有签入的 Bridge 配方直接演练 HCP。
  3. 单 GPU 加载辅助器会清除 hierarchical_context_parallel_sizes
  4. 旧技术栈上的静默训练损坏:如果在没有设置 hierarchical_context_parallel_sizes 的情况下使用 a2a+p2p,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)。请改用现有的单元测试和日志检查。

运行 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 未激活