From cdf4d2979e9febec0602c945ad6bdad9f253df44 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 22:28:04 +0000 Subject: [PATCH 1/9] Mirror the bytes a recomputed CP attention keeps for backward A size-only mirror of the executor's record_for_backward path, beside the code it mirrors: per stage, the Q/K/V flex consumes (padded copies, else contiguous copies of permuted views, or kept gathers and fetch buffers), flex's output and LSE at the execution length, logical copies when padded, and a merge-tape clone for every producing stage after the first. It reproduces the traced CP2 ranks: 0.971 GB for an aligned single stage and 3.07 GB for a padded local plus full-query remote stage. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/megatron/context_parallel/executor.py | 80 ++++++++++++ .../test_context_parallel_retained_bytes.py | 118 ++++++++++++++++++ 2 files changed, 198 insertions(+) create mode 100644 tests/unit/test_context_parallel_retained_bytes.py diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 4015011fe..78a65069f 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -35,6 +35,7 @@ DkvReducePlan, ExactMaskMetadata, FlexMaskSpec, + RankRuntimePlan, StageExecutionSpec, StagePlan, TokenRange, @@ -1550,6 +1551,85 @@ def _merge_stage_output_grads_from_tape( return stage_out_grads, stage_lse_grads +def retained_stage_record_bytes( + rank_plan: RankRuntimePlan, + *, + q_heads: int, + kv_heads: int, + head_dim: int, + value_head_dim: int, + element_size: int, + block_size: SparseBlockSize, +) -> int: + """Bytes a recomputed attention keeps for backward beyond its own-row tensors. + + A size-only mirror of ``_forward_stage_records`` with ``record_for_backward`` + and of ``_run_stage_attention``; keep the three in step. Every stage that + runs keeps the Q/K/V flex consumed: a padded copy when the execution length + differs, else a contiguous copy of the permuted ``q_flat``/``k_flat`` view; + partial-range gathers and remote fetch buffers are kept as the stage's + inputs. It keeps flex's output and LSE at the execution length, and + logical-length copies of them when padded. Every producing stage after the + first keeps a merge-tape clone of the accumulators. Accumulators themselves + are transient. + """ + own = int(rank_plan.local_valid_lengths[0]) if rank_plan.local_valid_lengths else 0 + accum_size = 4 if element_size < 4 else element_size + q_row = q_heads * head_dim * element_size + k_row = kv_heads * head_dim * element_size + v_row = kv_heads * value_head_dim * element_size + out_row = q_heads * value_head_dim * element_size + lse_row = q_heads * 4 + tape_row = q_heads * (value_head_dim + 1) * accum_size + total = 0 + local_produced = False + tapes: list[int] = [] + for stage in _ordered_stage_plans(rank_plan.stage_plans): + if not (stage.q_len > 0 and stage.k_len > 0 and stage.slices): + continue + q_len = _logical_stage_q_len(stage) + k_len = _logical_stage_k_len(stage) + q_pad, k_pad, _family = select_sparse_execution_family( + is_local_stage=bool(stage.is_local_stage), + q_len=int(stage.q_len), + k_len=int(stage.k_len), + block_size=block_size, + ) + q_full = _ranges_cover_full_length(stage.owner_local_q_ranges, length=own) + # Queries: the full range is a permuted view of q_flat; a partial range + # is gathered into a contiguous tensor kept as the stage's input. + if not q_full: + total += q_row * q_len + if q_pad != q_len: + total += q_row * q_pad + elif q_full and q_heads > 1: + total += q_row * q_len + # Keys and values: local ranges as for queries; remote ones land in + # contiguous head-major fetch buffers kept as the stage's inputs. + k_full = bool(stage.is_local_stage) and _ranges_cover_full_length( + stage.owner_local_k_ranges, length=own + ) + if not k_full: + total += (k_row + v_row) * k_len + if k_pad != k_len: + total += (k_row + v_row) * k_pad + elif k_full and kv_heads > 1: + total += (k_row + v_row) * k_len + total += (out_row + lse_row) * q_pad + if q_pad != q_len: + total += (out_row + lse_row) * q_len + tape = tape_row * (own if q_full else q_len) + if stage.is_local_stage: + local_produced = True + else: + tapes.append(tape) + # The first producing stage records no tape: the local stage when it ran, + # else whichever remote stage is ready first, so drop the smallest tape. + if tapes and not local_produced: + tapes.remove(min(tapes)) + return total + sum(tapes) + + def _forward_stage_records( *, q_flat: torch.Tensor, diff --git a/tests/unit/test_context_parallel_retained_bytes.py b/tests/unit/test_context_parallel_retained_bytes.py new file mode 100644 index 000000000..fd5e0e643 --- /dev/null +++ b/tests/unit/test_context_parallel_retained_bytes.py @@ -0,0 +1,118 @@ +"""Size-only retained-set mirror of the CP attention executor's recompute path.""" + +import pytest + +pytest.importorskip("triton") + +from art.megatron.context_parallel.executor import ( # noqa: E402 + retained_stage_record_bytes, +) +from art.megatron.context_parallel.types import ( # noqa: E402 + RankRuntimePlan, + StagePlan, + TokenRange, +) + +# Qwen3.6-35B-A3B attention: 16 query heads, 2 KV heads of 256, BF16, and the +# H200 flash block for a 256-wide head (128 query, 64 key rows). +GEOMETRY = dict( + q_heads=16, + kv_heads=2, + head_dim=256, + value_head_dim=256, + element_size=2, + block_size=(128, 64), +) +Q, KV, OUT, TAPE = 8192, 2048, 8192 + 64, 16 * 257 * 4 + + +def stage(index, *, local, q, k, q_len=None, k_len=None, own=None, source=0): + q_ranges = ( + (TokenRange(0, q),) if own is None or q == own else (TokenRange(1, 1 + q),) + ) + return StagePlan( + stage_index=index, + source_rank=source, + is_local_stage=local, + slices=("slice",), # ty: ignore[invalid-argument-type] + owner_local_q_ranges=q_ranges if q else (), + owner_local_k_ranges=(TokenRange(0, k),) if k else (), + q_len=q if q_len is None else q_len, + k_len=k if k_len is None else k_len, + ) + + +def plan(own, *stages): + return RankRuntimePlan( + rank=0, + original_seq_len=own, + token_layout_index=None, # ty: ignore[invalid-argument-type] + local_valid_lengths=(own,), + local_row_ranges=(TokenRange(0, own),), + stage_plans=stages, + remote_dkv_reduce_plan=None, # ty: ignore[invalid-argument-type] + ) + + +def test_aligned_single_stage_copies_views_without_output_copies(): + # Real-data rank 0: one aligned local stage. Contiguous copies of the + # permuted Q/K/V views, flex output and LSE; no padding, copies or tape. + rows = 52480 + retained = retained_stage_record_bytes( + plan(rows, stage(0, local=True, q=rows, k=rows)), **GEOMETRY + ) + assert retained == Q * rows + KV * rows + OUT * rows # 0.971 GB traced + + +def test_unaligned_single_stage_pads_and_copies_the_logical_output(): + rows = 52481 + q_pad, k_pad = 411 * 128, 821 * 64 + retained = retained_stage_record_bytes( + plan(rows, stage(0, local=True, q=rows, k=rows)), **GEOMETRY + ) + assert retained == Q * q_pad + KV * k_pad + OUT * q_pad + OUT * rows + + +def test_full_query_remote_stage_keeps_fetch_buffers_and_a_merge_tape(): + # Real-data rank 1: a local stage and a remote stage over all of its + # queries. Planner lengths round up, so both stages pad. + own, remote_k = 44314, 16504 + local = stage(0, local=True, q=own, k=own, q_len=44352, k_len=44352) + remote = stage(1, local=False, q=own, k=remote_k, q_len=44352, k_len=16512) + retained = retained_stage_record_bytes(plan(own, local, remote), **GEOMETRY) + q_pad = 347 * 128 # 44,416, as flex's traced output size shows + local_bytes = Q * q_pad + KV * 44352 + OUT * q_pad + OUT * own + remote_bytes = ( + Q * q_pad + KV * remote_k + KV * 16512 + OUT * q_pad + OUT * own + TAPE * own + ) + assert retained == local_bytes + remote_bytes + assert 3.05e9 < retained < 3.10e9 # 3.07 GB traced + + +def test_partial_query_remote_stage_keeps_its_gather_and_a_partial_tape(): + # Random-data rank 1: a large local stage and a 768-row remote stage. + own = 105153 + local = stage(0, local=True, q=own, k=own, q_len=105216, k_len=105216) + remote = stage(1, local=False, q=768, k=20608, own=own) + retained = retained_stage_record_bytes(plan(own, local, remote), **GEOMETRY) + local_bytes = Q * 105216 + KV * 105216 + OUT * 105216 + OUT * own + # Aligned: the query gather and fetch buffers feed flex without copies. + remote_bytes = Q * 768 + KV * 20608 + OUT * 768 + TAPE * 768 + assert retained == local_bytes + remote_bytes + + +def test_empty_remote_stage_and_missing_local_stage(): + rows = 1024 + empty = stage(1, local=False, q=0, k=0) + alone = retained_stage_record_bytes( + plan(rows, stage(0, local=True, q=rows, k=rows), empty), **GEOMETRY + ) + assert alone == Q * rows + KV * rows + OUT * rows + # Without a local stage, the first ready remote stage records no tape; the + # mirror drops the smallest so it never under-counts the order. + small = stage(1, local=False, q=256, k=256, own=rows) + full = stage(2, local=False, q=rows, k=512) + both = retained_stage_record_bytes(plan(rows, small, full), **GEOMETRY) + without_small_tape = retained_stage_record_bytes(plan(rows, full), **GEOMETRY) + assert both - without_small_tape == Q * 256 + KV * 256 + OUT * 256 + TAPE * rows + assert retained_stage_record_bytes(plan(rows), **GEOMETRY) == 0 From 5899a2dee8ce9b16013bf2afdf0827ec445f91b6 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 22:47:10 +0000 Subject: [PATCH 2/9] Price CP2 recompute per rank on each layer type's layout The checkpoint floor priced every rank on the busiest rank's rows with one mixer width, which over-counted boundaries and GDN and under-counted a rank whose attention runs two full-query stages. At CP2/TP1 with ART's CP core attention, price each rank on its own layouts instead: saved boundaries by each layer input's layout, a recomputed attention layer on attention rows plus what the executor keeps for backward, a GDN layer on GDN rows, and the MoE stage on that layer's rows with routed rows on the EP share. Admission takes the largest rank. Elsewhere the busiest-rank floor is unchanged. A rank can receive more routed rows than it holds, so routed rows are no longer clamped to local rows. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/megatron/context_parallel/runtime.py | 49 ++++ src/art/trainer_rank/_impl.py | 215 +++++++++++++++++- .../test_trainer_rank_admission_inputs.py | 1 + .../test_trainer_rank_checkpoint_memory.py | 9 +- tests/unit/test_trainer_rank_layout_memory.py | 161 +++++++++++++ tests/unit/test_trainer_rank_moe_memory.py | 2 +- 6 files changed, 420 insertions(+), 17 deletions(-) create mode 100644 tests/unit/test_trainer_rank_layout_memory.py diff --git a/src/art/megatron/context_parallel/runtime.py b/src/art/megatron/context_parallel/runtime.py index 531cb8829..2f7ac2dad 100644 --- a/src/art/megatron/context_parallel/runtime.py +++ b/src/art/megatron/context_parallel/runtime.py @@ -433,6 +433,55 @@ def context_parallel_rank_model_token_counts( ) +def context_parallel_rank_layouts( + *, + group_ids: torch.Tensor, + parent_ids: torch.Tensor, + topology: ParallelTopology, + config: ContextParallelConfig, + original_seq_len: int, + build_gdn_execution_spec: bool, + gdn_planner_config: Any | None = None, +) -> tuple[tuple[int, ...], tuple[int, ...] | None, tuple[RankRuntimePlan, ...]]: + """Each CP rank's attention rows, GDN rows and attention stage plan. + + Uses the cached planning bundle and per-rank runtime plans that execution + builds, so a memory estimate sees the layouts the ranks will run. + """ + planning_key, bundle, _group_ids_cpu, _parent_ids_cpu = ( + _get_or_build_planning_bundle( + group_ids=group_ids, + parent_ids=parent_ids, + topology=topology, + config=config, + original_seq_len=original_seq_len, + build_gdn_execution_spec=build_gdn_execution_spec, + ) + ) + attention = tuple(bundle.token_layout_index.token_counts_by_rank) + gdn = None + if build_gdn_execution_spec: + gdn = tuple( + _plan_gdn_global_execution( + planning_key=planning_key, + bundle=bundle, + topology=topology, + gdn_planner_config=gdn_planner_config, + ).gdn_token_counts_by_rank + ) + plans = tuple( + _get_or_build_bundle_rank_plan( + planning_key=planning_key, + bundle=bundle, + original_seq_len=original_seq_len, + target_rank=rank, + block_size=config.block_size, + ) + for rank in range(len(attention)) + ) + return attention, gdn, plans + + def context_parallel_model_token_total( *, group_ids: torch.Tensor, diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 8dfc41239..389a865c7 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1031,6 +1031,17 @@ def signature(self) -> _MemorySignature: _AnyForwardPlan = _FlatForwardPlan | _SplitForwardPlan +@dataclass(frozen=True) +class _GroupLayout: + """One packed group's CP layouts on every rank, for layout-aware pricing.""" + + attention_rows: tuple[int, ...] + gdn_rows: tuple[int, ...] | None + # What each rank's recomputed CP attention keeps for backward beyond its + # own-row activations (``retained_stage_record_bytes``). + attention_retained: tuple[int, ...] + + @dataclass(frozen=True) class _SubforwardCost: """Memory terms of one candidate subforward while all graphs stay live. @@ -4157,8 +4168,10 @@ def _moe_workspace_bytes( raise ValueError("Invalid constructor converted-weight stages") if type(shared) is not int or not 0 <= shared <= coefficient: raise ValueError("Invalid constructor shared-expert coefficient") - routed = rows if routed_rows is None else max(0, min(rows, routed_rows)) - return (rows - routed) * shared + ( + # A rank can receive more routed rows than it holds (an uneven CP + # split); the shared part then stays on the routed rows' count. + routed = rows if routed_rows is None else max(0, routed_rows) + return max(0, rows - routed) * shared + ( max( routed * coefficient, *(routed * per_row + fixed for per_row, fixed in stages), @@ -4172,11 +4185,14 @@ def _checkpoint_memory_floor( group_rows: tuple[tuple[int, bool], ...], slot_refs: tuple["LoRASlotRef | None", ...] | None = None, routed_rows: tuple[int, ...] | None = None, + layouts: tuple[_GroupLayout, ...] | None = None, ) -> tuple[int, int]: """Conservative saved-boundary charge and one recomputed layer's workspace. ``routed_rows`` are each group's balanced dispatched rows per rank - (``_plan_group_routed_rows``); by default, its local rows. + (``_plan_group_routed_rows``); by default, its local rows. With + ``layouts`` (``_plan_group_layouts``), price every rank on its own CP + layouts instead of the busiest rank's rows (``_layout_checkpoint_floor``). Count actual local full/uniform/1 boundaries, including aliases, rather than claiming measured distinct storage. Only this call's new groups enter the term; already-live graphs remain in the availability baseline. @@ -4241,9 +4257,12 @@ def _checkpoint_memory_floor( or getattr(decoder, "_forward_pre_hooks", None) ): return 0, 0 + refs = (None,) * len(group_rows) if slot_refs is None else slot_refs + routed = (None,) * len(group_rows) if routed_rows is None else routed_rows + if layouts is not None and all(grad for _, grad in group_rows): + return self._layout_checkpoint_floor(decoder.layers, refs, routed, layouts) retained = gradient_rows * layers * self._hidden_size * 2 moe = self._checkpoint_moe_bytes_per_token() if gradient_rows else 0 - refs = (None,) * len(group_rows) if slot_refs is None else slot_refs # Beside the mixer, the recomputed layer keeps its post-mixer residual # and pre-MLP norm output, and its MoE stage its routing state. mixer = ( @@ -4256,7 +4275,6 @@ def _checkpoint_memory_floor( if gradient_rows else 0 ) - routed = (None,) * len(group_rows) if routed_rows is None else routed_rows workspace = max( self._moe_workspace_bytes( rows, routed_rows=dispatched, checkpoint_grad=grad, slot_ref=ref @@ -4270,6 +4288,164 @@ def _checkpoint_memory_floor( workspace += self._te_workspace_growth_bytes() return retained, workspace + def _layout_checkpoint_floor( + self, + layers: Sequence[torch.nn.Module], + refs: tuple["LoRASlotRef | None", ...], + routed: tuple[int | None, ...], + layouts: tuple[_GroupLayout, ...], + ) -> tuple[int, int]: + """Each CP rank's boundaries and recomputed layer on its own layouts. + + A saved layer input arrives in the GDN layout when the layer follows a + GDN layer in its island (``_art_gdn_island_boundary``), and in the + attention layout otherwise. A recomputed attention layer keeps its + activations on its attention rows plus what the CP executor keeps for + backward; a GDN layer keeps the GDN width on its GDN rows. Either one's + residual, norm, routing state and MoE stage use that layer's rows, and + routed rows the EP share. Traced CP2 ranks: boundaries match exactly and + attention within 0.3%. Returns the largest rank's boundaries and the + rest of the largest rank total, so the two sum to that total. + """ + hidden = self._hidden_size * 2 + gdn_inputs = sum( + getattr( + getattr(layer, "_art_gdn_island_boundary", None), "input_layout", "" + ) + == "gdn" + for layer in layers + ) + attention_inputs = len(layers) - gdn_inputs + widths = self._recomputed_mixer_widths(stage_buffers=False) + moe = self._checkpoint_moe_bytes_per_token() + beside = 2 * hidden + self._moe_checkpoint_state_bytes_per_token() if moe else 0 + retained_by_rank: list[int] = [] + totals: list[int] = [] + for rank in range(len(layouts[0].attention_rows)): + retained = workspace = 0 + for ref, dispatched, layout in zip(refs, routed, layouts, strict=True): + attention = max(1, layout.attention_rows[rank]) + gdn = ( + attention + if layout.gdn_rows is None + else max(1, layout.gdn_rows[rank]) + ) + retained += hidden * (attention_inputs * attention + gdn_inputs * gdn) + for kind, rows, extra in ( + ("attention", attention, layout.attention_retained[rank]), + ("gdn", gdn, 0), + ): + if kind not in widths: + continue + stage = ( + rows * (widths[kind] + beside) + + extra + + self._moe_workspace_bytes( + rows, + routed_rows=dispatched, + checkpoint_grad=True, + slot_ref=ref, + ) + ) + workspace = max(workspace, stage) + retained_by_rank.append(retained) + totals.append(retained + workspace) + retained = max(retained_by_rank) + workspace = max(totals) - retained + if moe: + workspace += self._te_workspace_growth_bytes() + return retained, workspace + + def _plan_group_layouts( + self, plan: _FlatForwardPlan + ) -> tuple[_GroupLayout, ...] | None: + """Every rank's CP layouts per group, where layout pricing is modeled. + + Only CP2 at TP1/PP1 with gradient groups, ART's CP core attention with + no softmax offset, and GDN layers marked with island boundaries; the + executor's retained set is validated there. Elsewhere ``None`` keeps + the busiest-rank pricing. + """ + _dp, tp, cp, pp = plan.signature.topology + if (tp, cp, pp) != (1, 2, 1) or not plan.groups: + return None + if not all(group.grad_enabled for group in plan.groups): + return None + geometry = self._geometry + if not geometry.num_attention_heads or not geometry.kv_channels: + return None + try: + decoder = _language_model(self.runtime.model[0]).decoder + from art.megatron.context_parallel.core_attention import ( + ArtContextParallelCoreAttention, + ) + except (AttributeError, RuntimeError, ModuleNotFoundError): + return None + for layer in decoder.layers: + boundary = getattr(layer, "_art_gdn_island_boundary", None) + if boundary is not None and boundary.is_gdn: + continue + if self._gdn_layers and boundary is None: + return None + core = getattr( + getattr(layer, "self_attention", None), "core_attention", None + ) + if ( + type(core) is not ArtContextParallelCoreAttention + or getattr(core, "softmax_offset", None) is not None + ): + return None + from art.megatron.context_parallel.executor import retained_stage_record_bytes + from art.megatron.context_parallel.runtime import context_parallel_rank_layouts + from art.megatron.flex_attn.compiled import flash_sparse_block_size_for_head_dim + from art.megatron.training.microbatches import ( + _context_parallel_config_for_provider, + _gdn_planner_config_for_provider, + ) + + topology = self._topology() + handler = self.runtime.model_support_handler + config = _context_parallel_config_for_provider( + self.runtime.provider, self.device, handler + ) + head = int(geometry.kv_channels) + block = flash_sparse_block_size_for_head_dim( + head_dim=head, head_dim_v=head, device=self.device + ) + layouts = [] + for group in plan.groups: + batch = _pad_packed_batch(group.packed, multiple=int(topology.tp)) + attention, gdn, rank_plans = context_parallel_rank_layouts( + group_ids=batch.group_ids, + parent_ids=batch.parent_ids, + topology=topology, + config=config, + original_seq_len=int(batch.tokens.shape[1]), + build_gdn_execution_spec=handler.build_gdn_execution_spec, + gdn_planner_config=_gdn_planner_config_for_provider( + self.runtime.provider, handler + ), + ) + layouts.append( + _GroupLayout( + attention_rows=attention, + gdn_rows=gdn, + attention_retained=tuple( + retained_stage_record_bytes( + rank_plan, + q_heads=int(geometry.num_attention_heads), + kv_heads=int(geometry.num_query_groups), + head_dim=head, + value_head_dim=head, + element_size=self._param_dtype_size, + block_size=block, + ) + for rank_plan in rank_plans + ), + ) + ) + return tuple(layouts) + def _te_workspace_growth_bytes(self) -> int: """Transformer Engine's cuBLAS workspaces, until its GEMMs allocate them. @@ -4361,6 +4537,7 @@ def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost: gdn_segments=plan.grad_segment_count, group_rows=self._plan_group_rows(plan), group_routed_rows=self._plan_group_routed_rows(plan), + group_layouts=self._plan_group_layouts(plan), slot_refs=tuple(g.slot_ref for g in plan.groups), head_workspace_bytes=self._plan_head_workspace_bytes(plan), checkpoint_floor=_gdn_memory.plan_floor(self, plan), @@ -4378,6 +4555,7 @@ def _subforward_cost( gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), group_routed_rows: tuple[int, ...] | None = None, + group_layouts: tuple[_GroupLayout, ...] | None = None, slot_refs: tuple["LoRASlotRef | None", ...] | None = None, head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), @@ -4392,6 +4570,7 @@ def _subforward_cost( gdn_segments=gdn_segments, group_rows=group_rows, group_routed_rows=group_routed_rows, + group_layouts=group_layouts, slot_refs=slot_refs, head_workspace_bytes=head_workspace_bytes, checkpoint_floor=checkpoint_floor, @@ -4399,7 +4578,7 @@ def _subforward_cost( include_checkpoint_input_gradient=False, ) checkpoint_retained, checkpoint_workspace = self._checkpoint_memory_floor( - group_rows, slot_refs, group_routed_rows + group_rows, slot_refs, group_routed_rows, group_layouts ) retained = self._retained_memory_bytes( signature, @@ -7221,6 +7400,7 @@ def _memory_check( gdn_segments=forward.grad_segment_count, group_rows=self._plan_group_rows(forward), group_routed_rows=self._plan_group_routed_rows(forward), + group_layouts=self._plan_group_layouts(forward), slot_refs=tuple(g.slot_ref for g in forward.groups), head_workspace_bytes=self._plan_head_workspace_bytes(forward), checkpoint_floor=_gdn_memory.plan_floor(self, forward), @@ -8068,6 +8248,7 @@ def _estimate_required_memory_bytes_from_values( gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), group_routed_rows: tuple[int, ...] | None = None, + group_layouts: tuple[_GroupLayout, ...] | None = None, slot_refs: tuple["LoRASlotRef | None", ...] | None = None, head_workspace_bytes: int = 0, checkpoint_floor: tuple[int, int] = (0, 0), @@ -8160,7 +8341,7 @@ def _estimate_required_memory_bytes_from_values( ), ) retained, workspace = self._checkpoint_memory_floor( - group_rows, slot_refs, group_routed_rows + group_rows, slot_refs, group_routed_rows, group_layouts ) static_compute = max( static_compute, @@ -8299,18 +8480,26 @@ def _recomputed_mixer_bytes_per_token(self) -> int: more rows; its hidden-width input exchange and value-width output allowance price that (88 KB measured at CP2, 94 KB priced). """ + return max(self._recomputed_mixer_widths().values(), default=0) + + def _recomputed_mixer_widths(self, *, stage_buffers: bool = True) -> dict[str, int]: + """Recomputed attention and GDN mixer bytes per row, by layer type. + + ``stage_buffers=False`` leaves out the CP attention stage allowance, for + callers that price the executor's stage buffers from its stage plan. + """ geometry = self._geometry hidden = self._hidden_size tp = max(1, self._topology_key()[1]) cp = self._topology_key()[2] > 1 - widths = [] + widths: dict[str, float] = {} if self._gdn_layers < self._num_layers: attention, _gdn = self._mixer_activation_widths() - if cp: + if cp and stage_buffers: q = geometry.num_attention_heads * geometry.kv_channels or hidden kv = geometry.num_query_groups * geometry.kv_channels or hidden attention += (3 * q + 2 * kv) / tp - widths.append(attention) + widths["attention"] = attention if self._gdn_layers: key = geometry.gdn_key_heads * geometry.gdn_key_head_dim value = geometry.gdn_value_heads * geometry.gdn_value_head_dim @@ -8319,8 +8508,10 @@ def _recomputed_mixer_bytes_per_token(self) -> int: gdn = hidden + (2 * key + normalized + 6 * value + chunk) / tp if cp: gdn += hidden + value / tp - widths.append(gdn) - return int(max(widths, default=0) * self._param_dtype_size) + widths["gdn"] = gdn + return { + kind: int(width * self._param_dtype_size) for kind, width in widths.items() + } def _gdn_segment_layer_bytes(self) -> float: """Initial and final fp32 recurrent states plus convolution history.""" diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index 3481bfd6d..66b4bf812 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -33,6 +33,7 @@ def assert_plan_values(rank, plan, values): assert values["gdn_segments"] == plan.grad_segment_count assert values["group_rows"] == rank._plan_group_rows(plan) assert values["group_routed_rows"] == rank._plan_group_routed_rows(plan) + assert values["group_layouts"] == rank._plan_group_layouts(plan) assert values["head_workspace_bytes"] == rank._plan_head_workspace_bytes(plan) assert values["checkpoint_floor"] == _gdn_memory.plan_floor(rank, plan) assert values["retained_tokens"] == rank._plan_retained_tokens(plan) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 1d1087ab6..10ef298d1 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -534,10 +534,11 @@ def test_routed_rows_move_only_the_routed_moe_part(): fewer = r._checkpoint_memory_floor(((local, True),), None, (routed,)) assert fewer[0] == retained assert workspace - fewer[1] == (local - routed) * (188416 - 8192) - # Never more routed rows than local ones. - assert r._moe_workspace_bytes( - 10, routed_rows=20, checkpoint_grad=True - ) == r._moe_workspace_bytes(10, checkpoint_grad=True) + # A rank can receive more routed rows than it holds: all are priced, and + # the shared part stays on the routed count. + assert r._moe_workspace_bytes(10, routed_rows=20, checkpoint_grad=True) == ( + 20 * 188416 + ) r._moe_gradient_shared_bytes = 188417 with pytest.raises(ValueError, match="shared-expert"): r._moe_workspace_bytes(10, checkpoint_grad=True) diff --git a/tests/unit/test_trainer_rank_layout_memory.py b/tests/unit/test_trainer_rank_layout_memory.py new file mode 100644 index 000000000..a7ad7e6f1 --- /dev/null +++ b/tests/unit/test_trainer_rank_layout_memory.py @@ -0,0 +1,161 @@ +"""CP2 layout-aware recompute pricing: per-rank ledgers and gates, CPU only.""" + +from dataclasses import replace +from types import SimpleNamespace + +import pytest +from test_trainer_rank_checkpoint_memory import rank +import torch + +from art.trainer_rank import ForwardInput +from art.trainer_rank._impl import _TE_CUBLAS_WORKSPACE_BYTES, _GroupLayout + +H = 2048 * 2 + + +def qwen36(r): + """Qwen3.6-35B-A3B's 3:1 GDN/attention pattern and mixer geometry at CP2.""" + r._topology_key = lambda: (1, 1, 2, 1) + r._geometry = replace( + r._geometry, + num_attention_heads=16, + num_query_groups=2, + kv_channels=256, + gdn_key_heads=16, + gdn_key_head_dim=128, + gdn_value_heads=32, + gdn_value_head_dim=128, + ) + r._attention_output_gate = True + r._gdn_layers = 30 + for index, layer in enumerate(r.runtime.model[0].decoder.layers): + is_gdn = index % 4 != 3 + after_gdn = index > 0 and (index - 1) % 4 != 3 + layer._art_gdn_island_boundary = SimpleNamespace( + is_gdn=is_gdn, input_layout="gdn" if is_gdn and after_gdn else "attention" + ) + return r + + +def test_boundaries_follow_each_layer_inputs_layout(): + # Traced real-data CP2 ranks: 20 inputs in each layout, and boundaries of + # 8.248 and 7.612 GB. The busiest attention rank is not the busiest overall. + r = qwen36(rank()) + layers = r.runtime.model[0].decoder.layers + layout = _GroupLayout((52480, 44314), (48194, 48600), (0, 0)) + retained, _ = r._layout_checkpoint_floor(layers, (None,), (48397,), (layout,)) + assert retained == H * (20 * 52480 + 20 * 48194) + assert H * (20 * 44314 + 20 * 48600) < retained < 40 * 52480 * H + + +def test_largest_rank_total_not_a_sum_of_rank_maxima(): + r = qwen36(rank()) + layers = r.runtime.model[0].decoder.layers + widths = r._recomputed_mixer_widths(stage_buffers=False) + assert widths == {"attention": 2 * (2 * 2048 + 7 * 4096 + 3 * 512), "gdn": 94208} + # Rank 1's two full-query stages keep 3.08 GB against rank 0's 0.97 GB. + layout = _GroupLayout((52480, 44314), (48194, 48600), (970_720_000, 3_079_000_000)) + retained, workspace = r._layout_checkpoint_floor( + layers, (None,), (48397,), (layout,) + ) + + def total(rank_index): + attention = layout.attention_rows[rank_index] + gdn = layout.gdn_rows[rank_index] + ledger = H * (20 * attention + 20 * gdn) + stages = ( + attention * (widths["attention"] + 2 * H) + + layout.attention_retained[rank_index] + + r._moe_workspace_bytes( + attention, routed_rows=48397, checkpoint_grad=True + ), + gdn * (widths["gdn"] + 2 * H) + + r._moe_workspace_bytes(gdn, routed_rows=48397, checkpoint_grad=True), + ) + return ledger + max(stages) + + assert retained + workspace == max(total(0), total(1)) + _TE_CUBLAS_WORKSPACE_BYTES + # Rank 1 has fewer rows but is the largest total: its attention runs two + # full-query stages. Pricing it on its own rows is about neutral against + # the busiest-rank floor, which over-counts boundaries and GDN instead. + assert total(1) > total(0) + busiest = sum(r._checkpoint_memory_floor(((52480, True),), None, (48397,))) + assert abs(retained + workspace - busiest) < 0.01 * busiest + + +def test_routed_rows_above_a_ranks_own_rows_are_all_priced(): + r = qwen36(rank()) + layers = r.runtime.model[0].decoder.layers + few = _GroupLayout((100, 100), (100, 100), (0, 0)) + many = _GroupLayout((100, 100), (100, 100), (0, 0)) + low = sum(r._layout_checkpoint_floor(layers, (None,), (100,), (few,))) + high = sum(r._layout_checkpoint_floor(layers, (None,), (400,), (many,))) + assert high - low == 300 * 188416 + + +def test_no_grad_groups_keep_busiest_rank_pricing(): + r = qwen36(rank()) + layout = _GroupLayout((10, 8), (9, 9), (0, 0)) + groups = ((10, True), (12, False)) + assert r._checkpoint_memory_floor( + groups, None, (9, 12), (layout, layout) + ) == r._checkpoint_memory_floor(groups, None, (9, 12)) + + +def _plan(r, *, no_grad=False): + plan = r._plan_flat_forward( + [ + ForwardInput( + input_tokens=torch.arange(64), hidden_states=True, no_grad=no_grad + ) + ] + ) + return replace(plan, signature=replace(plan.signature, topology=(1, 1, 2, 1))) + + +@pytest.mark.parametrize( + "topology", [(1, 1, 1, 1), (1, 2, 2, 1), (1, 1, 4, 1), (1, 1, 2, 2)] +) +def test_layout_pricing_is_cp2_tp1_pp1_only(topology): + r = qwen36(rank()) + plan = _plan(r) + plan = replace(plan, signature=replace(plan.signature, topology=topology)) + assert r._plan_group_layouts(plan) is None + + +def test_layout_pricing_needs_gradient_groups_and_art_cp_attention(): + r = qwen36(rank()) + assert r._plan_group_layouts(_plan(r, no_grad=True)) is None + # The stub decoder's attention layers carry no ART CP core attention. + assert r._plan_group_layouts(_plan(r)) is None + # A GDN model without island boundaries is not modeled either. + for layer in r.runtime.model[0].decoder.layers: + del layer._art_gdn_island_boundary + assert r._plan_group_layouts(_plan(r)) is None + + +def test_plan_cost_and_admission_use_the_same_layouts(monkeypatch): + r = qwen36(rank()) + plan = _plan(r) + layout = _GroupLayout((40, 24), (32, 32), (1_000_000, 3_000_000)) + monkeypatch.setattr(r, "_plan_group_layouts", lambda plan: (layout,)) + monkeypatch.setattr(r, "_plan_group_rows", lambda plan: ((40, True),)) + monkeypatch.setattr(r, "_plan_group_routed_rows", lambda plan: (32,)) + monkeypatch.setattr(r, "_plan_hybridep_growth_bytes", lambda plan: 0) + monkeypatch.setattr(r, "_plan_retained_tokens", lambda plan: 32) + cost = r._plan_cost(plan) + # Admission adds CP output coexistence the plan cost leaves out (as it did + # before layouts); the layouts themselves enter both the same way. + gap = r._memory_check(plan).estimated_required_bytes - cost.required + monkeypatch.setattr(r, "_plan_group_layouts", lambda plan: None) + assert gap == r._memory_check(plan).estimated_required_bytes - ( + r._plan_cost(plan).required + ) + monkeypatch.setattr(r, "_plan_group_layouts", lambda plan: (layout,)) + retained, _ = r._layout_checkpoint_floor( + r.runtime.model[0].decoder.layers, + (None,), + r._plan_group_routed_rows(plan), + (layout,), + ) + assert cost.checkpoint_retained == plan.output_bytes + retained diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index d4e390ccf..c44363d4c 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -606,7 +606,7 @@ def held(rows): monkeypatch.setattr( rank, "_checkpoint_memory_floor", - lambda rows, refs=None, routed=None: (10**7, 10**6), + lambda rows, refs=None, routed=None, layouts=None: (10**7, 10**6), ) grown = rank._plan_cost(plan) monkeypatch.setattr(rank, "_plan_hybridep_growth_bytes", lambda plan: 0) From 0e7255c8b990d38ebf5af92689ba869b160f6ae0 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 23:30:55 +0000 Subject: [PATCH 3/9] Bound CP2 layout pricing from below and count flex's own LSE The split search's lower bound priced gated plans with the busiest-rank attention width, above what per-rank pricing can charge; price even-share layouts at the least attention state instead (one aligned local stage per rank). Count flex's saved LSE beside the normalized copy it returns. Share one gate between plan pricing and the bound, checking the topology before model state. Record group layouts in planner evidence. Tests cover the positive gate path against the executor's own plans, the runtime layouts, the softmax-offset gate, tiny and single-head stages, and the lower bound staying below plan cost. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/megatron/context_parallel/executor.py | 29 +++- src/art/trainer_rank/_impl.py | 95 +++++++++--- .../test_context_parallel_retained_bytes.py | 50 +++++-- tests/unit/test_trainer_rank_layout_memory.py | 137 +++++++++++++++++- 4 files changed, 276 insertions(+), 35 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 78a65069f..857b5ef2a 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -1551,6 +1551,26 @@ def _merge_stage_output_grads_from_tape( return stage_out_grads, stage_lse_grads +def minimum_retained_bytes_per_row( + *, + q_heads: int, + kv_heads: int, + head_dim: int, + value_head_dim: int, + element_size: int, +) -> int: + """The least ``retained_stage_record_bytes`` keeps per own row. + + Every rank with rows runs a local stage over all of them (each row attends + to itself); an aligned one keeps only the contiguous copies of multi-head + views, flex's output and its two LSEs. + """ + copies = (q_heads > 1) * q_heads * head_dim + (kv_heads > 1) * kv_heads * ( + head_dim + value_head_dim + ) + return (copies + q_heads * value_head_dim) * element_size + 2 * q_heads * 4 + + def retained_stage_record_bytes( rank_plan: RankRuntimePlan, *, @@ -1569,9 +1589,10 @@ def retained_stage_record_bytes( differs, else a contiguous copy of the permuted ``q_flat``/``k_flat`` view; partial-range gathers and remote fetch buffers are kept as the stage's inputs. It keeps flex's output and LSE at the execution length, and - logical-length copies of them when padded. Every producing stage after the - first keeps a merge-tape clone of the accumulators. Accumulators themselves - are transient. + logical-length copies of them when padded, plus flex's own LSE beside the + normalized one the FLASH backend returns (counted on every backend). Every + producing stage after the first keeps a merge-tape clone of the + accumulators. Accumulators themselves are transient. """ own = int(rank_plan.local_valid_lengths[0]) if rank_plan.local_valid_lengths else 0 accum_size = 4 if element_size < 4 else element_size @@ -1615,7 +1636,7 @@ def retained_stage_record_bytes( total += (k_row + v_row) * k_pad elif k_full and kv_heads > 1: total += (k_row + v_row) * k_len - total += (out_row + lse_row) * q_pad + total += (out_row + 2 * lse_row) * q_pad if q_pad != q_len: total += (out_row + lse_row) * q_len tape = tape_row * (own if q_full else q_len) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 389a865c7..5e0e9c88a 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -3680,6 +3680,7 @@ def _split_chunk_lower_cost( unshared_packed_tokens = 0 head_workspace_bytes = 0 group_rows: list[tuple[int, bool]] = [] + group_physical_rows: list[int] = [] for (_slot, grad_enabled), group_indices in groups: estimated = estimate_prefix_tree_packed_tokens( (rows[index] for index in group_indices), @@ -3687,6 +3688,7 @@ def _split_chunk_lower_cost( ) assert estimated is not None # rows are CPU copies physical_rows = self._physical_tokens(estimated) + group_physical_rows.append(physical_rows) packed_tokens += physical_rows # The most loaded CP rank holds at least an even share. cp = max(1, self._topology_key()[2]) @@ -3712,12 +3714,23 @@ def _split_chunk_lower_cost( slot_groups=tuple(key for key, _ in groups), ) logical_tokens = _active_logical_tokens(requests) + # Where exact plans price each rank's layouts, bound them from below + # the same way rather than with the busiest-rank widths. + layouts = ( + self._minimum_layouts(group_physical_rows, signature.topology[2]) + if self._layout_pricing_supported( + signature.topology, + gradient_groups=all(grad for _, grad in group_rows), + ) + else None + ) cost = self._subforward_cost( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, logical_tokens=logical_tokens, group_rows=tuple(group_rows), + group_layouts=layouts, slot_refs=tuple(ref for (ref, _), _ in groups), head_workspace_bytes=head_workspace_bytes, # The average CP load is an optimistic bound, not an admission cost. @@ -4356,37 +4369,29 @@ def _layout_checkpoint_floor( workspace += self._te_workspace_growth_bytes() return retained, workspace - def _plan_group_layouts( - self, plan: _FlatForwardPlan - ) -> tuple[_GroupLayout, ...] | None: - """Every rank's CP layouts per group, where layout pricing is modeled. - - Only CP2 at TP1/PP1 with gradient groups, ART's CP core attention with - no softmax offset, and GDN layers marked with island boundaries; the - executor's retained set is validated there. Elsewhere ``None`` keeps - the busiest-rank pricing. - """ - _dp, tp, cp, pp = plan.signature.topology - if (tp, cp, pp) != (1, 2, 1) or not plan.groups: - return None - if not all(group.grad_enabled for group in plan.groups): - return None + def _layout_pricing_supported( + self, topology: tuple[int, int, int, int], *, gradient_groups: bool + ) -> bool: + """Whether layout-aware pricing models this runtime and plan shape.""" + _dp, tp, cp, pp = topology + if (tp, cp, pp) != (1, 2, 1) or not gradient_groups: + return False geometry = self._geometry if not geometry.num_attention_heads or not geometry.kv_channels: - return None + return False try: decoder = _language_model(self.runtime.model[0]).decoder from art.megatron.context_parallel.core_attention import ( ArtContextParallelCoreAttention, ) except (AttributeError, RuntimeError, ModuleNotFoundError): - return None + return False for layer in decoder.layers: boundary = getattr(layer, "_art_gdn_island_boundary", None) if boundary is not None and boundary.is_gdn: continue if self._gdn_layers and boundary is None: - return None + return False core = getattr( getattr(layer, "self_attention", None), "core_attention", None ) @@ -4394,7 +4399,54 @@ def _plan_group_layouts( type(core) is not ArtContextParallelCoreAttention or getattr(core, "softmax_offset", None) is not None ): - return None + return False + return True + + def _minimum_layouts( + self, physical_rows: Sequence[int], cp: int + ) -> tuple[_GroupLayout, ...]: + """Even-share layouts keeping the least attention state: a lower bound. + + Some rank holds at least an even share of each layout's rows and runs + at least one aligned local stage over them. + """ + from art.megatron.context_parallel.executor import ( + minimum_retained_bytes_per_row, + ) + + geometry = self._geometry + per_row = minimum_retained_bytes_per_row( + q_heads=int(geometry.num_attention_heads), + kv_heads=int(geometry.num_query_groups), + head_dim=int(geometry.kv_channels), + value_head_dim=int(geometry.kv_channels), + element_size=self._param_dtype_size, + ) + return tuple( + _GroupLayout( + attention_rows=(rows // cp,) * cp, + gdn_rows=(rows // cp,) * cp if self._gdn_layers else None, + attention_retained=(rows // cp * per_row,) * cp, + ) + for rows in physical_rows + ) + + def _plan_group_layouts( + self, plan: _FlatForwardPlan + ) -> tuple[_GroupLayout, ...] | None: + """Every rank's CP layouts per group, where layout pricing is modeled. + + Only CP2 at TP1/PP1 with gradient groups, ART's CP core attention with + no softmax offset, and GDN layers marked with island boundaries; the + executor's retained set is validated there. Elsewhere ``None`` keeps + the busiest-rank pricing. + """ + if not plan.groups or not self._layout_pricing_supported( + plan.signature.topology, + gradient_groups=all(group.grad_enabled for group in plan.groups), + ): + return None + geometry = self._geometry from art.megatron.context_parallel.executor import retained_stage_record_bytes from art.megatron.context_parallel.runtime import context_parallel_rank_layouts from art.megatron.flex_attn.compiled import flash_sparse_block_size_for_head_dim @@ -6675,6 +6727,11 @@ def _fill_planner_snapshot( "retained_tokens": self._plan_retained_tokens(child), "group_rows": self._plan_group_rows(child), "group_routed_rows": self._plan_group_routed_rows(child), + "group_layouts": ( + None + if (layouts := self._plan_group_layouts(child)) is None + else [asdict(layout) for layout in layouts] + ), "hybridep_growth_bytes": ( self._plan_hybridep_growth_bytes(child) ), diff --git a/tests/unit/test_context_parallel_retained_bytes.py b/tests/unit/test_context_parallel_retained_bytes.py index fd5e0e643..a12b88a0c 100644 --- a/tests/unit/test_context_parallel_retained_bytes.py +++ b/tests/unit/test_context_parallel_retained_bytes.py @@ -5,6 +5,7 @@ pytest.importorskip("triton") from art.megatron.context_parallel.executor import ( # noqa: E402 + minimum_retained_bytes_per_row, retained_stage_record_bytes, ) from art.megatron.context_parallel.types import ( # noqa: E402 @@ -23,7 +24,9 @@ element_size=2, block_size=(128, 64), ) -Q, KV, OUT, TAPE = 8192, 2048, 8192 + 64, 16 * 257 * 4 +# Flex's output keeps its own LSE beside the normalized copy it returns; +# logical-length copies keep one. +Q, KV, FLEX, OUT, TAPE = 8192, 2048, 8192 + 2 * 64, 8192 + 64, 16 * 257 * 4 def stage(index, *, local, q, k, q_len=None, k_len=None, own=None, source=0): @@ -61,16 +64,41 @@ def test_aligned_single_stage_copies_views_without_output_copies(): retained = retained_stage_record_bytes( plan(rows, stage(0, local=True, q=rows, k=rows)), **GEOMETRY ) - assert retained == Q * rows + KV * rows + OUT * rows # 0.971 GB traced + assert retained == Q * rows + KV * rows + FLEX * rows # 0.971 GB traced + geometry = {k: v for k, v in GEOMETRY.items() if k != "block_size"} + assert retained == rows * minimum_retained_bytes_per_row(**geometry) def test_unaligned_single_stage_pads_and_copies_the_logical_output(): + # The tail chunk on the single-stage rank: the planner rounds both lengths + # up to its 128-row block, so both pad. rows = 52481 - q_pad, k_pad = 411 * 128, 821 * 64 + stage_len = 411 * 128 retained = retained_stage_record_bytes( - plan(rows, stage(0, local=True, q=rows, k=rows)), **GEOMETRY + plan( + rows, stage(0, local=True, q=rows, k=rows, q_len=stage_len, k_len=stage_len) + ), + **GEOMETRY, + ) + assert retained == Q * stage_len + KV * stage_len + FLEX * stage_len + OUT * rows + + +def test_tiny_stage_pads_to_two_blocks(): + retained = retained_stage_record_bytes( + plan(5, stage(0, local=True, q=5, k=5)), **GEOMETRY + ) + assert retained == Q * 256 + KV * 128 + FLEX * 256 + OUT * 5 + + +def test_single_head_views_are_already_contiguous(): + rows = 1024 + geometry = dict(GEOMETRY, q_heads=1, kv_heads=1) + retained = retained_stage_record_bytes( + plan(rows, stage(0, local=True, q=rows, k=rows)), **geometry ) - assert retained == Q * q_pad + KV * k_pad + OUT * q_pad + OUT * rows + assert retained == (256 * 2 + 2 * 4) * rows + del geometry["block_size"] + assert minimum_retained_bytes_per_row(**geometry) == 256 * 2 + 2 * 4 def test_full_query_remote_stage_keeps_fetch_buffers_and_a_merge_tape(): @@ -81,9 +109,9 @@ def test_full_query_remote_stage_keeps_fetch_buffers_and_a_merge_tape(): remote = stage(1, local=False, q=own, k=remote_k, q_len=44352, k_len=16512) retained = retained_stage_record_bytes(plan(own, local, remote), **GEOMETRY) q_pad = 347 * 128 # 44,416, as flex's traced output size shows - local_bytes = Q * q_pad + KV * 44352 + OUT * q_pad + OUT * own + local_bytes = Q * q_pad + KV * 44352 + FLEX * q_pad + OUT * own remote_bytes = ( - Q * q_pad + KV * remote_k + KV * 16512 + OUT * q_pad + OUT * own + TAPE * own + Q * q_pad + KV * remote_k + KV * 16512 + FLEX * q_pad + OUT * own + TAPE * own ) assert retained == local_bytes + remote_bytes assert 3.05e9 < retained < 3.10e9 # 3.07 GB traced @@ -95,9 +123,9 @@ def test_partial_query_remote_stage_keeps_its_gather_and_a_partial_tape(): local = stage(0, local=True, q=own, k=own, q_len=105216, k_len=105216) remote = stage(1, local=False, q=768, k=20608, own=own) retained = retained_stage_record_bytes(plan(own, local, remote), **GEOMETRY) - local_bytes = Q * 105216 + KV * 105216 + OUT * 105216 + OUT * own + local_bytes = Q * 105216 + KV * 105216 + FLEX * 105216 + OUT * own # Aligned: the query gather and fetch buffers feed flex without copies. - remote_bytes = Q * 768 + KV * 20608 + OUT * 768 + TAPE * 768 + remote_bytes = Q * 768 + KV * 20608 + FLEX * 768 + TAPE * 768 assert retained == local_bytes + remote_bytes @@ -107,12 +135,12 @@ def test_empty_remote_stage_and_missing_local_stage(): alone = retained_stage_record_bytes( plan(rows, stage(0, local=True, q=rows, k=rows), empty), **GEOMETRY ) - assert alone == Q * rows + KV * rows + OUT * rows + assert alone == Q * rows + KV * rows + FLEX * rows # Without a local stage, the first ready remote stage records no tape; the # mirror drops the smallest so it never under-counts the order. small = stage(1, local=False, q=256, k=256, own=rows) full = stage(2, local=False, q=rows, k=512) both = retained_stage_record_bytes(plan(rows, small, full), **GEOMETRY) without_small_tape = retained_stage_record_bytes(plan(rows, full), **GEOMETRY) - assert both - without_small_tape == Q * 256 + KV * 256 + OUT * 256 + TAPE * rows + assert both - without_small_tape == Q * 256 + KV * 256 + FLEX * 256 + TAPE * rows assert retained_stage_record_bytes(plan(rows), **GEOMETRY) == 0 diff --git a/tests/unit/test_trainer_rank_layout_memory.py b/tests/unit/test_trainer_rank_layout_memory.py index a7ad7e6f1..37030533d 100644 --- a/tests/unit/test_trainer_rank_layout_memory.py +++ b/tests/unit/test_trainer_rank_layout_memory.py @@ -8,7 +8,7 @@ import torch from art.trainer_rank import ForwardInput -from art.trainer_rank._impl import _TE_CUBLAS_WORKSPACE_BYTES, _GroupLayout +from art.trainer_rank._impl import _TE_CUBLAS_WORKSPACE_BYTES, Unset, _GroupLayout H = 2048 * 2 @@ -159,3 +159,138 @@ def test_plan_cost_and_admission_use_the_same_layouts(monkeypatch): (layout,), ) assert cost.checkpoint_retained == plan.output_bytes + retained + + +def art_cp(r, monkeypatch): + """Qwen3.6 at CP2 with ART's CP core attention and a CPU planning config.""" + from art.megatron.context_parallel.core_attention import ( + ArtContextParallelCoreAttention, + ) + from art.megatron.context_parallel.types import ParallelTopology + + r = qwen36(r) + for index, layer in enumerate(r.runtime.model[0].decoder.layers): + if index % 4 == 3: + core = ArtContextParallelCoreAttention.__new__( + ArtContextParallelCoreAttention + ) + torch.nn.Module.__init__(core) + core.softmax_offset = None + layer.self_attention = torch.nn.Module() + layer.self_attention.core_attention = core + provider = r.runtime.provider + provider.kv_channels = 256 + provider.ffn_hidden_size = 8192 + provider.linear_num_key_heads = 16 + provider.linear_num_value_heads = 32 + provider.linear_key_head_dim = 128 + provider.linear_value_head_dim = 128 + provider.params_dtype = torch.bfloat16 + r.runtime.model_support_handler = SimpleNamespace( + build_gdn_execution_spec=True, + context_parallel_workload_profile=lambda provider: None, + ) + monkeypatch.setattr(r, "_topology", lambda: ParallelTopology(tp=1, cp=2)) + return r + + +def _requests(lengths=(900, 700, 500)): + start, requests = 0, [] + for n in lengths: + requests.append( + ForwardInput( + input_tokens=torch.arange(start, start + n), hidden_states=True + ) + ) + start += n + return requests + + +def test_rank_layouts_are_the_executors_plans(monkeypatch): + from art.megatron.context_parallel.runtime import context_parallel_rank_layouts + from art.megatron.context_parallel.types import ( + ContextParallelConfig, + ParallelTopology, + ) + from art.megatron.prefix_tree_packing import prefix_tree_pack + + packed = prefix_tree_pack([r.input_tokens for r in _requests()], max_depth=1) + attention, gdn, plans = context_parallel_rank_layouts( + group_ids=packed.group_ids, + parent_ids=packed.parent_ids, + topology=ParallelTopology(tp=1, cp=2), + config=ContextParallelConfig(), + original_seq_len=int(packed.tokens.shape[1]), + build_gdn_execution_spec=True, + ) + total = int(packed.tokens.numel()) + assert sum(attention) == total and gdn is not None and sum(gdn) == total + # The ledger's attention rows and the mirror's own rows are the same count. + assert attention == tuple(plan.local_valid_lengths[0] for plan in plans) + + +def test_gated_plans_price_every_rank_from_its_stage_plan(monkeypatch): + from art.megatron.context_parallel.executor import retained_stage_record_bytes + from art.megatron.context_parallel.runtime import context_parallel_rank_layouts + from art.megatron.context_parallel.types import ParallelTopology + from art.megatron.training.microbatches import ( + _context_parallel_config_for_provider, + ) + + r = art_cp(rank(), monkeypatch) + plan = _plan_with(r, _requests()) + (layout,) = r._plan_group_layouts(plan) + assert sum(layout.attention_rows) == plan.packed_tokens + assert layout.gdn_rows is not None and sum(layout.gdn_rows) == plan.packed_tokens + # Each rank's retention is the executor mirror over that rank's own plan. + (group,) = plan.groups + _, _, rank_plans = context_parallel_rank_layouts( + group_ids=group.packed.group_ids, + parent_ids=group.packed.parent_ids, + topology=ParallelTopology(tp=1, cp=2), + config=_context_parallel_config_for_provider( + r.runtime.provider, r.device, r.runtime.model_support_handler + ), + original_seq_len=int(group.packed.tokens.shape[1]), + build_gdn_execution_spec=True, + ) + assert layout.attention_retained == tuple( + retained_stage_record_bytes( + rank_plan, + q_heads=16, + kv_heads=2, + head_dim=256, + value_head_dim=256, + element_size=2, + block_size=(128, 128), # the CPU device's flex block + ) + for rank_plan in rank_plans + ) + assert all(retained > 0 for retained in layout.attention_retained) + + +def test_softmax_offset_leaves_the_busiest_rank_floor(monkeypatch): + r = art_cp(rank(), monkeypatch) + plan = _plan_with(r, _requests()) + assert r._plan_group_layouts(plan) is not None + for layer in r.runtime.model[0].decoder.layers: + core = getattr(getattr(layer, "self_attention", None), "core_attention", None) + if core is not None: + core.softmax_offset = torch.zeros(16) + assert r._plan_group_layouts(plan) is None + + +def test_split_lower_bound_stays_below_the_layout_cost(monkeypatch): + r = art_cp(rank(), monkeypatch) + requests = _requests((2048, 1536, 1024, 512)) + plan = _plan_with(r, requests) + lower = r._split_chunk_lower_cost( + requests, tuple(item.input_tokens for item in requests), checkpoint=Unset + ) + assert lower.required <= r._plan_cost(plan).required + # It prices even-share layouts at the least attention state. + assert r._layout_pricing_supported((1, 1, 2, 1), gradient_groups=True) + + +def _plan_with(r, requests): + return r._plan_flat_forward(requests) From 9b4e9d709691b0dfdc33087761d34fee9dcf2307 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 23:43:04 +0000 Subject: [PATCH 4/9] Guard the layout gate and state why the lower bound rounds down The gate now declines models with several chunks or a decoder without layers before reading them. The bound holds because every rank's total grows with its own rows, so the largest is at least the total at the mean; rounding the even share up can exceed a split's exact cost. The lower-bound test covers even, skewed and odd splits and one long sequence. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 14 +++++++++++--- tests/unit/test_trainer_rank_layout_memory.py | 9 +++++++-- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 5e0e9c88a..bc47b03ed 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4379,6 +4379,8 @@ def _layout_pricing_supported( geometry = self._geometry if not geometry.num_attention_heads or not geometry.kv_channels: return False + if len(self.runtime.model) != 1: + return False try: decoder = _language_model(self.runtime.model[0]).decoder from art.megatron.context_parallel.core_attention import ( @@ -4386,7 +4388,10 @@ def _layout_pricing_supported( ) except (AttributeError, RuntimeError, ModuleNotFoundError): return False - for layer in decoder.layers: + layers = getattr(decoder, "layers", None) + if layers is None: + return False + for layer in layers: boundary = getattr(layer, "_art_gdn_island_boundary", None) if boundary is not None and boundary.is_gdn: continue @@ -4407,8 +4412,11 @@ def _minimum_layouts( ) -> tuple[_GroupLayout, ...]: """Even-share layouts keeping the least attention state: a lower bound. - Some rank holds at least an even share of each layout's rows and runs - at least one aligned local stage over them. + Every rank's total grows with its own rows, and every rank keeps at + least an aligned local stage's state per row (each row attends to + itself), so the largest rank's total is at least the total at the mean + rows. The mean is at least the floor of an even share, which is why + this rounds down; rounding up can exceed a split's exact cost. """ from art.megatron.context_parallel.executor import ( minimum_retained_bytes_per_row, diff --git a/tests/unit/test_trainer_rank_layout_memory.py b/tests/unit/test_trainer_rank_layout_memory.py index 37030533d..baf88257b 100644 --- a/tests/unit/test_trainer_rank_layout_memory.py +++ b/tests/unit/test_trainer_rank_layout_memory.py @@ -280,9 +280,14 @@ def test_softmax_offset_leaves_the_busiest_rank_floor(monkeypatch): assert r._plan_group_layouts(plan) is None -def test_split_lower_bound_stays_below_the_layout_cost(monkeypatch): +@pytest.mark.parametrize( + "lengths", + [(2048, 1536, 1024, 512), (4099, 3, 5, 7), (1, 2, 3, 4, 5, 6, 7), (8191,)], +) +def test_split_lower_bound_stays_below_the_layout_cost(monkeypatch, lengths): + # Even and skewed CP splits, odd row counts, and a single long sequence. r = art_cp(rank(), monkeypatch) - requests = _requests((2048, 1536, 1024, 512)) + requests = _requests(lengths) plan = _plan_with(r, requests) lower = r._split_chunk_lower_cost( requests, tuple(item.input_tokens for item in requests), checkpoint=Unset From 33185b49ed11ccbff0f5c433d5b8cee792e4e3d3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Fri, 25 Sep 2026 23:52:43 +0000 Subject: [PATCH 5/9] Run the CP layout memory tests in the Megatron CI lane Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/prek.yml | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml index 4ff185190..b66ac2c0e 100644 --- a/.github/workflows/prek.yml +++ b/.github/workflows/prek.yml @@ -239,6 +239,8 @@ jobs: tests/unit/test_trainer_rank_pending_memory.py \ tests/unit/test_trainer_rank_shared_memory.py \ tests/unit/test_trainer_rank_converted_memory.py \ + tests/unit/test_trainer_rank_layout_memory.py \ + tests/unit/test_context_parallel_retained_bytes.py \ tests/unit/test_trainer_rank_split.py \ tests/unit/test_megatron_compile_garbage.py \ tests/unit/test_trainer_rank_cache_recovery.py::test_dense_cp_exact_demand_fits_after_recovery \ @@ -283,4 +285,6 @@ jobs: --ignore=tests/unit/test_trainer_rank_pending_memory.py \ --ignore=tests/unit/test_trainer_rank_shared_memory.py \ --ignore=tests/unit/test_megatron_compile_garbage.py \ - --ignore=tests/unit/test_trainer_rank_converted_memory.py + --ignore=tests/unit/test_trainer_rank_converted_memory.py \ + --ignore=tests/unit/test_trainer_rank_layout_memory.py \ + --ignore=tests/unit/test_context_parallel_retained_bytes.py From c1500cafbb98cde712201f97f9768b63f9511445 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 00:24:13 +0000 Subject: [PATCH 6/9] Charge single-head view copies and partial-tape indices; time layout planning One head does not make a view of a fused QKV split contiguous, so the mirror now charges the copy flex makes of any full local view; the lower bound still counts only multi-head copies. A partial-query merge tape also keeps its int64 row index. The per-rank layout work now counts toward planning time (about 40-76 ms for a new layout's peer plan, under 1 ms once cached). Tests use typed geometry helpers. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/megatron/context_parallel/executor.py | 10 +- src/art/trainer_rank/_impl.py | 10 ++ .../test_context_parallel_retained_bytes.py | 92 ++++++++++--------- tests/unit/test_trainer_rank_layout_memory.py | 1 + 4 files changed, 68 insertions(+), 45 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 857b5ef2a..5d1295963 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -1623,7 +1623,10 @@ def retained_stage_record_bytes( total += q_row * q_len if q_pad != q_len: total += q_row * q_pad - elif q_full and q_heads > 1: + elif q_full: + # A view of the projection output: flex copies it unless it is + # already contiguous, which one head alone does not guarantee + # (a fused QKV split keeps the projection's token stride). total += q_row * q_len # Keys and values: local ranges as for queries; remote ones land in # contiguous head-major fetch buffers kept as the stage's inputs. @@ -1634,12 +1637,13 @@ def retained_stage_record_bytes( total += (k_row + v_row) * k_len if k_pad != k_len: total += (k_row + v_row) * k_pad - elif k_full and kv_heads > 1: + elif k_full: total += (k_row + v_row) * k_len total += (out_row + 2 * lse_row) * q_pad if q_pad != q_len: total += (out_row + lse_row) * q_len - tape = tape_row * (own if q_full else q_len) + # A partial-query tape also keeps its int64 row index. + tape = tape_row * own if q_full else (tape_row + 8) * q_len if stage.is_local_stage: local_produced = True else: diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index bc47b03ed..6d3d61b98 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4454,6 +4454,16 @@ def _plan_group_layouts( gradient_groups=all(group.grad_enabled for group in plan.groups), ): return None + started = time.perf_counter() + try: + return self._compute_group_layouts(plan) + finally: + # Planning work: every rank's CP plan, cached by planning key. + self._planning_seconds_accum += time.perf_counter() - started + + def _compute_group_layouts( + self, plan: _FlatForwardPlan + ) -> tuple[_GroupLayout, ...]: geometry = self._geometry from art.megatron.context_parallel.executor import retained_stage_record_bytes from art.megatron.context_parallel.runtime import context_parallel_rank_layouts diff --git a/tests/unit/test_context_parallel_retained_bytes.py b/tests/unit/test_context_parallel_retained_bytes.py index a12b88a0c..583de3a24 100644 --- a/tests/unit/test_context_parallel_retained_bytes.py +++ b/tests/unit/test_context_parallel_retained_bytes.py @@ -14,16 +14,31 @@ TokenRange, ) + # Qwen3.6-35B-A3B attention: 16 query heads, 2 KV heads of 256, BF16, and the # H200 flash block for a 256-wide head (128 query, 64 key rows). -GEOMETRY = dict( - q_heads=16, - kv_heads=2, - head_dim=256, - value_head_dim=256, - element_size=2, - block_size=(128, 64), -) +def retained(runtime_plan, *, q_heads=16, kv_heads=2): + return retained_stage_record_bytes( + runtime_plan, + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=256, + value_head_dim=256, + element_size=2, + block_size=(128, 64), + ) + + +def minimum(*, q_heads=16, kv_heads=2): + return minimum_retained_bytes_per_row( + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=256, + value_head_dim=256, + element_size=2, + ) + + # Flex's output keeps its own LSE beside the normalized copy it returns; # logical-length copies keep one. Q, KV, FLEX, OUT, TAPE = 8192, 2048, 8192 + 2 * 64, 8192 + 64, 16 * 257 * 4 @@ -61,12 +76,9 @@ def test_aligned_single_stage_copies_views_without_output_copies(): # Real-data rank 0: one aligned local stage. Contiguous copies of the # permuted Q/K/V views, flex output and LSE; no padding, copies or tape. rows = 52480 - retained = retained_stage_record_bytes( - plan(rows, stage(0, local=True, q=rows, k=rows)), **GEOMETRY - ) - assert retained == Q * rows + KV * rows + FLEX * rows # 0.971 GB traced - geometry = {k: v for k, v in GEOMETRY.items() if k != "block_size"} - assert retained == rows * minimum_retained_bytes_per_row(**geometry) + kept = retained(plan(rows, stage(0, local=True, q=rows, k=rows))) + assert kept == Q * rows + KV * rows + FLEX * rows # 0.971 GB traced + assert kept == rows * minimum() def test_unaligned_single_stage_pads_and_copies_the_logical_output(): @@ -74,31 +86,28 @@ def test_unaligned_single_stage_pads_and_copies_the_logical_output(): # up to its 128-row block, so both pad. rows = 52481 stage_len = 411 * 128 - retained = retained_stage_record_bytes( + kept = retained( plan( rows, stage(0, local=True, q=rows, k=rows, q_len=stage_len, k_len=stage_len) - ), - **GEOMETRY, + ) ) - assert retained == Q * stage_len + KV * stage_len + FLEX * stage_len + OUT * rows + assert kept == Q * stage_len + KV * stage_len + FLEX * stage_len + OUT * rows def test_tiny_stage_pads_to_two_blocks(): - retained = retained_stage_record_bytes( - plan(5, stage(0, local=True, q=5, k=5)), **GEOMETRY - ) - assert retained == Q * 256 + KV * 128 + FLEX * 256 + OUT * 5 + kept = retained(plan(5, stage(0, local=True, q=5, k=5))) + assert kept == Q * 256 + KV * 128 + FLEX * 256 + OUT * 5 -def test_single_head_views_are_already_contiguous(): +def test_single_head_views_may_still_be_copied(): + # One head does not make a view of a fused QKV split contiguous, so the + # mirror still charges the copies; only the lower bound leaves them out. rows = 1024 - geometry = dict(GEOMETRY, q_heads=1, kv_heads=1) - retained = retained_stage_record_bytes( - plan(rows, stage(0, local=True, q=rows, k=rows)), **geometry + kept = retained( + plan(rows, stage(0, local=True, q=rows, k=rows)), q_heads=1, kv_heads=1 ) - assert retained == (256 * 2 + 2 * 4) * rows - del geometry["block_size"] - assert minimum_retained_bytes_per_row(**geometry) == 256 * 2 + 2 * 4 + assert kept == (512 + 1024 + 512 + 2 * 4) * rows + assert minimum(q_heads=1, kv_heads=1) == 512 + 2 * 4 def test_full_query_remote_stage_keeps_fetch_buffers_and_a_merge_tape(): @@ -107,14 +116,14 @@ def test_full_query_remote_stage_keeps_fetch_buffers_and_a_merge_tape(): own, remote_k = 44314, 16504 local = stage(0, local=True, q=own, k=own, q_len=44352, k_len=44352) remote = stage(1, local=False, q=own, k=remote_k, q_len=44352, k_len=16512) - retained = retained_stage_record_bytes(plan(own, local, remote), **GEOMETRY) + kept = retained(plan(own, local, remote)) q_pad = 347 * 128 # 44,416, as flex's traced output size shows local_bytes = Q * q_pad + KV * 44352 + FLEX * q_pad + OUT * own remote_bytes = ( Q * q_pad + KV * remote_k + KV * 16512 + FLEX * q_pad + OUT * own + TAPE * own ) - assert retained == local_bytes + remote_bytes - assert 3.05e9 < retained < 3.10e9 # 3.07 GB traced + assert kept == local_bytes + remote_bytes + assert 3.05e9 < kept < 3.10e9 # 3.07 GB traced def test_partial_query_remote_stage_keeps_its_gather_and_a_partial_tape(): @@ -122,25 +131,24 @@ def test_partial_query_remote_stage_keeps_its_gather_and_a_partial_tape(): own = 105153 local = stage(0, local=True, q=own, k=own, q_len=105216, k_len=105216) remote = stage(1, local=False, q=768, k=20608, own=own) - retained = retained_stage_record_bytes(plan(own, local, remote), **GEOMETRY) + kept = retained(plan(own, local, remote)) local_bytes = Q * 105216 + KV * 105216 + FLEX * 105216 + OUT * own - # Aligned: the query gather and fetch buffers feed flex without copies. - remote_bytes = Q * 768 + KV * 20608 + FLEX * 768 + TAPE * 768 - assert retained == local_bytes + remote_bytes + # Aligned: the query gather and fetch buffers feed flex without copies; the + # partial tape also keeps its int64 row index. + remote_bytes = Q * 768 + KV * 20608 + FLEX * 768 + (TAPE + 8) * 768 + assert kept == local_bytes + remote_bytes def test_empty_remote_stage_and_missing_local_stage(): rows = 1024 empty = stage(1, local=False, q=0, k=0) - alone = retained_stage_record_bytes( - plan(rows, stage(0, local=True, q=rows, k=rows), empty), **GEOMETRY - ) + alone = retained(plan(rows, stage(0, local=True, q=rows, k=rows), empty)) assert alone == Q * rows + KV * rows + FLEX * rows # Without a local stage, the first ready remote stage records no tape; the # mirror drops the smallest so it never under-counts the order. small = stage(1, local=False, q=256, k=256, own=rows) full = stage(2, local=False, q=rows, k=512) - both = retained_stage_record_bytes(plan(rows, small, full), **GEOMETRY) - without_small_tape = retained_stage_record_bytes(plan(rows, full), **GEOMETRY) + both = retained(plan(rows, small, full)) + without_small_tape = retained(plan(rows, full)) assert both - without_small_tape == Q * 256 + KV * 256 + FLEX * 256 + TAPE * rows - assert retained_stage_record_bytes(plan(rows), **GEOMETRY) == 0 + assert retained(plan(rows)) == 0 diff --git a/tests/unit/test_trainer_rank_layout_memory.py b/tests/unit/test_trainer_rank_layout_memory.py index baf88257b..0cc147c6d 100644 --- a/tests/unit/test_trainer_rank_layout_memory.py +++ b/tests/unit/test_trainer_rank_layout_memory.py @@ -61,6 +61,7 @@ def test_largest_rank_total_not_a_sum_of_rank_maxima(): def total(rank_index): attention = layout.attention_rows[rank_index] + assert layout.gdn_rows is not None gdn = layout.gdn_rows[rank_index] ledger = H * (20 * attention + 20 * gdn) stages = ( From 4368998e490111574164a1ceb0b8d05b638a072f Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 00:30:13 +0000 Subject: [PATCH 7/9] Keep partial-tape indices when the first stage drops its tape The executor keeps a partial stage's int64 row index even when that stage produced first and recorded no accumulator tape, so charge indices apart from the tape the mirror drops. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/megatron/context_parallel/executor.py | 10 ++++++---- tests/unit/test_context_parallel_retained_bytes.py | 5 ++++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/art/megatron/context_parallel/executor.py b/src/art/megatron/context_parallel/executor.py index 5d1295963..c59751966 100644 --- a/src/art/megatron/context_parallel/executor.py +++ b/src/art/megatron/context_parallel/executor.py @@ -1562,8 +1562,8 @@ def minimum_retained_bytes_per_row( """The least ``retained_stage_record_bytes`` keeps per own row. Every rank with rows runs a local stage over all of them (each row attends - to itself); an aligned one keeps only the contiguous copies of multi-head - views, flex's output and its two LSEs. + to itself); an aligned one keeps at least the contiguous copies of + multi-head views, flex's output and its two LSEs. """ copies = (q_heads > 1) * q_heads * head_dim + (kv_heads > 1) * kv_heads * ( head_dim + value_head_dim @@ -1642,8 +1642,10 @@ def retained_stage_record_bytes( total += (out_row + 2 * lse_row) * q_pad if q_pad != q_len: total += (out_row + lse_row) * q_len - # A partial-query tape also keeps its int64 row index. - tape = tape_row * own if q_full else (tape_row + 8) * q_len + tape = tape_row * (own if q_full else q_len) + if not q_full: + # Its int64 row index is kept even by the first producing stage. + total += 8 * q_len if stage.is_local_stage: local_produced = True else: diff --git a/tests/unit/test_context_parallel_retained_bytes.py b/tests/unit/test_context_parallel_retained_bytes.py index 583de3a24..18f0f4b48 100644 --- a/tests/unit/test_context_parallel_retained_bytes.py +++ b/tests/unit/test_context_parallel_retained_bytes.py @@ -150,5 +150,8 @@ def test_empty_remote_stage_and_missing_local_stage(): full = stage(2, local=False, q=rows, k=512) both = retained(plan(rows, small, full)) without_small_tape = retained(plan(rows, full)) - assert both - without_small_tape == Q * 256 + KV * 256 + FLEX * 256 + TAPE * rows + # The dropped tape's int64 index stays: the executor keeps it regardless. + assert both - without_small_tape == ( + Q * 256 + KV * 256 + FLEX * 256 + 8 * 256 + TAPE * rows + ) assert retained(plan(rows)) == 0 From b4a619ec26b753a2a72a2f274607fd9a362d0c7c Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sat, 26 Sep 2026 23:35:12 +0000 Subject: [PATCH 8/9] Test the combine-extent floor through per-rank layouts Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_trainer_rank_moe_memory.py | 33 ++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 3644f0a90..a90b222c0 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -867,6 +867,39 @@ def test_hybridep_recompute_prices_fresh_dense_output_without_buffer_growth( assert not torch.cuda.is_initialized() +def test_hybridep_combine_extent_floors_the_layout_path( + hybrid_checkpoint_rank, monkeypatch +): + rank = hybrid_checkpoint_rank + groups = ((2, True),) + calls = [] + + def layout_floor(layers, refs, routed, layouts): + calls.append(layouts) + return 7, 11 + + def generic_floor(*args): + raise AssertionError("per-rank layouts must take the layout path") + + monkeypatch.setattr(rank, "_layout_checkpoint_floor", layout_floor) + monkeypatch.setattr(rank, "_generic_checkpoint_floor", generic_floor) + layouts = (object(),) + te = rank._te_workspace_growth_bytes() + # Two rows round up to four; the layout floor's small workspace loses. + assert rank._checkpoint_memory_floor(groups, layouts=layouts) == ( + 7, + 4 * 2048 * 2 + te, + ) + marker = torch.empty(0) + rank._pending_hybridep_graphs.append(weakref.ref(marker)) + rank._hybridep_rows_high_water = 218751 + assert rank._checkpoint_memory_floor(groups, layouts=layouts) == ( + 7, + 218752 * 2048 * 2 + te, + ) + assert calls == [layouts, layouts] + + @pytest.mark.parametrize("reference", ["absent", "expired", "smaller"]) def test_hybridep_high_water_needs_a_live_larger_graph( hybrid_checkpoint_rank, reference From 0f59d7c1c6d00b672f3df41a5c71a4c7ccfe47f4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Sun, 27 Sep 2026 00:22:07 +0000 Subject: [PATCH 9/9] Pin that the layout path's combine floor adds no second TE growth Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_trainer_rank_moe_memory.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index a90b222c0..1cc118bc3 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -897,7 +897,16 @@ def generic_floor(*args): 7, 218752 * 2048 * 2 + te, ) - assert calls == [layouts, layouts] + # A larger stage already carries its TE growth; the combine floor adds none. + stage = 218752 * 2048 * 2 + te + 1 + + def larger_layout_floor(layers, refs, routed, layouts): + calls.append(layouts) + return 7, stage + + monkeypatch.setattr(rank, "_layout_checkpoint_floor", larger_layout_floor) + assert rank._checkpoint_memory_floor(groups, layouts=layouts) == (7, stage) + assert calls == [layouts, layouts, layouts] @pytest.mark.parametrize("reference", ["absent", "expired", "smaller"])