GDN context parallel × tensor parallel fails: fla_cp broadcasts with a group-local src rank ("Global rank 1 is not part of group")

[bradhilton]

Summary

Combining context parallel with tensor parallel on a GatedDeltaNet model fails on the first forward: the native FLA context-parallel scan broadcasts its chain state with a group-local rank index passed as src=, but torch.distributed.broadcast interprets src as a global rank. It only works when the CP group happens to be global ranks 0..CP-1, i.e. TP 1 and DP 1. Any topology where CP groups are strided or offset (TP > 1, and by the same logic DP > 1 with CP > 1) raises:

ValueError: Global rank 1 is not part of group <torch.distributed.distributed_c10d.ProcessGroup object at 0x...>

Reproduction

• Qwen/Qwen3.8-27B (48 GDN + 16 attention layers), rank-1 LoRA, 4×H200, Trainer(cluster="H200:4", num_processes_per_node=4)
• env: ART_MEGATRON_CONTEXT_PARALLEL_SIZE=2 ART_MEGATRON_TENSOR_MODEL_PARALLEL_SIZE=2 (full recompute, otherwise defaults)
• ART 7496cc0 (main after Return full source-order context-parallel trainer outputs #912) via caladan main 32e00b6; caladan experiments/058-rank1-probe-bases.py probe … --gpus 4 --cp 2 --tp 2
• First forward_micro_batches call with hidden_states requested fails. The same probe code runs on CP 2 / CP 4 with TP 1 (64 s and 35 s per step) and on TP 2 / TP 4 with CP 1.

Traceback (rank 0, trimmed to ART frames):

art/trainer_rank/_impl.py:5425 _decoder_hidden
…megatron/core/transformer/transformer_block.py:784 forward → _checkpointed_forward
art/megatron/gdn/operator.py:220 _gdn_island_layer_forward
art/megatron/gdn/operator.py:409 gdn_prefix_tree_forward
art/megatron/gdn/operator.py:958 _run_cp_planned_prefixes_and_completions
art/megatron/gdn/operator.py:2314 run_gdn_bucket
art/megatron/gdn/fla_cp.py:106 chunk_gated_delta_rule_native_cp
art/megatron/gdn/fla_cp.py:216 forward
art/megatron/gdn/fla_cp.py:552 _broadcast_chain_final_state
torch/distributed/distributed_c10d.py:2956 broadcast → 1205 _canonicalize_group_rank → 1075 get_group_rank
ValueError: Global rank 1 is not part of group

Root cause

src/art/megatron/gdn/fla_cp.py:

def _broadcast_chain_final_state(final_state, group):
    ...
    owner = dist.get_world_size(group) - 1          # CP-local index
    dist.broadcast(final_state, src=owner, group=group)   # src must be a global rank

and the same pattern in _suffix_summary_exclusive_and_full:

    dist.broadcast(full, src=0, group=group)        # global rank 0 is not in every CP group

With TP 2 on four ranks, Megatron's CP groups are {0, 2} and {1, 3}; owner = 1 is not a member of {0, 2}, hence the error. With TP 1 / DP 2 / CP 2 the groups are {0, 1} and {2, 3} and src=0 fails for the second group the same way, so this also blocks DP × CP for GDN models, not just TP × CP.

The neighbouring _exchange_summary is fine: it uses all_to_all_single with per-rank split lists, which are indexed group-locally by design.

Fix

Either address the peer group-locally with the newer API:

dist.broadcast(final_state, group=group, group_src=owner)
dist.broadcast(full, group=group, group_src=0)

or translate explicitly:

dist.broadcast(final_state, src=dist.get_global_rank(group, owner), group=group)

The provider already validates GDN head counts against TP when CP > 1 (_validate_art_gdn_context_parallel_provider) and the oracle harnesses include Topology(tp=2, cp=2, sp=True), so TP × CP is clearly meant to be supported; this looks like the only remaining group-local src in the CP paths (I grepped src/art/megatron/gdn, src/art/megatron/context_parallel and src/art/trainer_rank for broadcast/send/recv; the checkpoint send/recv in trainer_rank/_checkpoint.py use global ranks).

Related: #911 / #912 (CP outputs now gathered to source order; TP 1 CP runs work end to end from caladan after that fix), #840 (CP 4 all-to-all hang, different code path).