fix(trainer-rank): price CP2 recompute using each rank’s layouts

.github/workflows/prek.yml

@@ -246,6 +246,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 \
@@ -297,4 +299,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

src/art/megatron/context_parallel/executor.py

@@ -35,6 +35,7 @@
     DkvReducePlan,
     ExactMaskMetadata,
     FlexMaskSpec,
+    RankRuntimePlan,
     StageExecutionSpec,
     StagePlan,
     TokenRange,
@@ -1550,6 +1551,112 @@ 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 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
+    )
+    return (copies + q_heads * value_head_dim) * element_size + 2 * q_heads * 4
+
+
+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, 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
+    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:
+            # 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.
+        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:
+            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)
+        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:
+            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,

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,

src/art/trainer_rank/_impl.py

@@ -1069,6 +1069,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.
@@ -3041,6 +3052,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,
         head_backward_traced: bool = False,
@@ -3049,7 +3061,11 @@ def _subforward_cost(
         hybridep_growth_bytes: int = 0,
     ) -> _SubforwardCost:
         checkpoint_memory = self._checkpoint_memory_floor(
-            group_rows, slot_refs, gdn_segments, routed_rows=group_routed_rows
+            group_rows,
+            slot_refs,
+            gdn_segments,
+            routed_rows=group_routed_rows,
+            layouts=group_layouts,
         )
         required = self._estimate_required_memory_bytes_from_values(
             packed_tokens=packed_tokens,
@@ -3059,6 +3075,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,
@@ -3080,13 +3097,21 @@ def _subforward_cost(
         )
         # Input gradients live at the recomputed layer's peak; kept out of
         # forward retention, including the cold fallback above.
+        # Per-rank layouts price each rank's own boundaries; the input-gradient
+        # allowance keeps the busiest rank's (``_checkpoint_input_gradient_bytes``).
         gradient = self._checkpoint_input_gradient_bytes(
-            group_rows, slot_refs, retained=checkpoint_retained
+            group_rows,
+            slot_refs,
+            retained=checkpoint_retained if group_layouts is None else None,
         )
         gradient_slots = self._gradient_slots(group_rows, slot_refs)
         adapter_gradient = (
-            self._checkpoint_adapter_gradient_bytes(
-                self._checkpoint_gradient_groups(group_rows, slot_refs)
+            self._checkpoint_adapter_gradient_extra(
+                (checkpoint_retained, checkpoint_workspace),
+                group_rows,
+                slot_refs,
+                group_routed_rows,
+                group_layouts,
             )
             if gradient
             else 0
@@ -3099,7 +3124,7 @@ def _subforward_cost(
         forward_required = required
         if gradient:
             head_stage = self._checkpoint_head_stage_bytes(
-                head_workspace_bytes, gradient, group_rows, slot_refs
+                head_workspace_bytes, gradient, group_rows, slot_refs, group_layouts
             )
             peak = checkpoint_workspace + adapter_gradient
             if head_stage is not None:
@@ -4929,6 +4954,17 @@ def _cp_group_model_tokens(
         )
 
     _group_head_workspace_bytes = _memory._group_head_workspace_bytes
+    _layer_gdn_inputs = _memory._layer_gdn_inputs
+    _layout_checkpoint_floor = _memory._layout_checkpoint_floor
+    _layout_checkpoint_rank_floors = _memory._layout_checkpoint_rank_floors
+    _layout_pricing_supported = _memory._layout_pricing_supported
+    _minimum_layouts = _memory._minimum_layouts
+    _layout_layer_boundaries = _memory._layout_layer_boundaries
+    _layout_adapter_gradient_bytes = _memory._layout_adapter_gradient_bytes
+    _checkpoint_adapter_gradient_extra = _memory._checkpoint_adapter_gradient_extra
+    _recomputed_mixer_widths = _memory._recomputed_mixer_widths
+    _plan_group_layouts = _micro_batch_planner._plan_group_layouts
+    _compute_group_layouts = _micro_batch_planner._compute_group_layouts
     _triton_min_rows = _memory._triton_min_rows
     _head_backward_traced = _memory._head_backward_traced
     _te_workspace_growth_bytes = _memory._te_workspace_growth_bytes