From 3074453671401b55b0494a30fc0f235704adf782 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 00:02:49 +0000 Subject: [PATCH 1/2] fla_cp: broadcast chain state from the CP group's global owner rank torch.distributed.broadcast takes a global rank as src. _broadcast_chain_final_state passed cp_size - 1 and _suffix_summary_exclusive_and_full passed 0, both CP-group-local indices, which only coincide with global ranks when the context-parallel group is ranks 0..cp_size-1 (TP 1, DP 1). With TP 2 the CP groups are {0, 2} and {1, 3} and the first forward raised 'Global rank 1 is not part of group'; DP x CP fails the same way. Translate through dist.get_global_rank. Fixes #919. Co-Authored-By: Claude Fable 5.1 --- src/art/megatron/gdn/fla_cp.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/art/megatron/gdn/fla_cp.py b/src/art/megatron/gdn/fla_cp.py index bb053018e..893452202 100644 --- a/src/art/megatron/gdn/fla_cp.py +++ b/src/art/megatron/gdn/fla_cp.py @@ -468,7 +468,11 @@ def _suffix_summary_exclusive_and_full( recv_from=rank + 1 if rank + 1 < world_size else None, ) full = inclusive if rank == 0 else torch.empty_like(summary) - dist.broadcast(full, src=0, group=group) # ty: ignore[possibly-missing-attribute] + dist.broadcast( # ty: ignore[possibly-missing-attribute] + full, + src=dist.get_global_rank(group, 0), # ty: ignore[possibly-missing-attribute] + group=group, + ) return exclusive, full @@ -547,7 +551,13 @@ def _scan_fwd_initial_state(summary: Tensor | None, h0: Tensor) -> Tensor: def _broadcast_chain_final_state(final_state: Tensor | None, group: Any) -> Tensor: if final_state is None: raise RuntimeError("native FLA CP did not produce a local final state") - owner = dist.get_world_size(group) - 1 # ty: ignore[possibly-missing-attribute] + # ``src`` is a global rank; the chain owner is the last rank *of the CP + # group*, whose global rank only equals ``cp_size - 1`` when the group is + # ranks 0..cp_size-1 (TP 1, DP 1). Translate explicitly. + owner = dist.get_global_rank( # ty: ignore[possibly-missing-attribute] + group, + dist.get_world_size(group) - 1, # ty: ignore[possibly-missing-attribute] + ) final_state = final_state.contiguous() dist.broadcast(final_state, src=owner, group=group) # ty: ignore[possibly-missing-attribute] return final_state From ba959aa8f8f3e5b95966e511b92ef29a3da9476e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 00:07:51 +0000 Subject: [PATCH 2/2] Test FLA CP collectives on strided and offset context-parallel groups Gloo regression test on four CPU ranks with CP groups {0,2}/{1,3} (TP 2) and {0,1}/{2,3} (DP 2): _broadcast_chain_final_state must deliver the last group member's state to every member, and _suffix_summary_exclusive_and_full must agree on the full-chain summary across the group. Fails on main with 'Global rank N is not part of group'. Co-Authored-By: Claude Fable 5.1 --- tests/unit/test_fla_cp_group_ranks.py | 78 +++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 tests/unit/test_fla_cp_group_ranks.py diff --git a/tests/unit/test_fla_cp_group_ranks.py b/tests/unit/test_fla_cp_group_ranks.py new file mode 100644 index 000000000..a472476a0 --- /dev/null +++ b/tests/unit/test_fla_cp_group_ranks.py @@ -0,0 +1,78 @@ +"""Native FLA CP collectives must address peers by global rank. + +Context-parallel groups are only ``range(cp_size)`` when TP = DP = 1. With +tensor parallelism the groups are strided (TP 2 on four ranks: {0, 2} and +{1, 3}); with data parallelism they are offset ({2, 3}). ``torch.distributed`` +``src`` arguments are global ranks, so a CP-local index raises +``ValueError: Global rank N is not part of group`` (OpenPipe/ART#919). +""" + +from __future__ import annotations + +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from art.megatron.gdn.fla_cp import ( + _broadcast_chain_final_state, + _suffix_summary_exclusive_and_full, +) + + +@pytest.mark.parametrize( + "layout", + ("tp2_cp2", "dp2_cp2"), +) +def test_fla_cp_collectives_use_global_ranks(layout: str, tmp_path) -> None: + mp.spawn( + _worker, + args=(layout, f"file://{tmp_path / 'init'}"), + nprocs=4, + join=True, + ) + + +def _cp_groups(layout: str) -> list[list[int]]: + if layout == "tp2_cp2": + # Megatron orders TP fastest: CP peers are two ranks apart. + return [[0, 2], [1, 3]] + if layout == "dp2_cp2": + # Two CP groups of contiguous ranks; the second never contains rank 0. + return [[0, 1], [2, 3]] + raise AssertionError(layout) + + +def _worker(rank: int, layout: str, init_method: str) -> None: + dist.init_process_group( + "gloo", + init_method=init_method, + rank=rank, + world_size=4, + timeout=timedelta(seconds=90), + ) + try: + groups = [dist.new_group(ranks) for ranks in _cp_groups(layout)] + members = next(r for r in _cp_groups(layout) if rank in r) + group = groups[_cp_groups(layout).index(members)] + owner = members[-1] + + # Chain final state: every member must receive the last member's state. + final_state = torch.full((2, 3, 3), float(rank)) + received = _broadcast_chain_final_state(final_state, group) + assert torch.equal(received, torch.full((2, 3, 3), float(owner))) + + # Suffix summaries: the full-chain summary is broadcast from the first + # member; all members must agree on it. + summary = torch.zeros((2, 3, 5)) + summary[..., :3] = torch.eye(3) * (1.0 + rank) + summary[..., 3:] = float(rank) + _exclusive, full = _suffix_summary_exclusive_and_full(summary, group) + gathered = [torch.empty_like(full) for _ in members] + dist.all_gather(gathered, full, group=group) + for other in gathered: + assert torch.equal(other, full) + finally: + dist.destroy_process_group()