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 と一致することを検証します。

プロセスグループの作成

parallel_state.initialize_model_parallel は、create_hierarchical_groups を介して HCP サイズが指定されると、階層型 CP サブグループを作成します。現在、Bridge は MPU ベースの ProcessGroupCollection を通じてこれらのグループを取得します。

TE との統合

TEDotProductAttention は、a2a+p2p を使用する場合に、階層型グループを Transformer Engine に渡します。Transformer Engine >= 1.12.0 が必要です。

注意点

1. 現在、Bridge の HCP は MPU 専用です：use_decentralized_pg=True の場合、Bridge はフラットな CP グループを初期化し、HCP は未設定のままです。
2. HCP を直接実行する、チェックイン済みの Bridge レシピはまだありません。
3. 単一 GPU のロード用ヘルパーは hierarchical_context_parallel_sizes をクリアします。
4. 古いスタックでは、エラーが表面化せず学習が壊れる可能性があります：hierarchical_context_parallel_sizes を設定せずに a2a+p2p を使用すると、現在の MCore ではアサーションが発生します。古いバージョンでは CP 通信が黙って無効になり、各ランクがローカルのチャンクだけにアテンションを適用していました。その結果、勾配が壊れた状態で、見かけ上のスループットが不自然に高くなっていました。
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 を参照）。代わりに、既存のユニットテストとログの確認を使用してください。

フラット CP の動作が維持されていることを確認するため、分散型 PG のユニットテストを実行します。

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 の作成が表示される
• エラーなく、少なくとも 1 ステップの学習が完了する
• CONTEXT_PARALLEL_GROUP しか表示されない場合、HCP は有効になっていません