Account for dense CP rank loads in retained-activation memory

dev/trainer_rank_recompute_memory.md

@@ -79,6 +79,12 @@ The existing static heuristic and routed FC2 bound remain floors. Profiles can
 only increase the estimate, and output storage and the 10% safety factor are
 applied afterward. Full-recompute and no-grad requests keep their existing path.
 
+The floor also applies with recompute disabled (`None`), including the `none`
+override and EP-overlap MoE trainers. Previously admitted workloads may now split
+or, for indivisible groups, be refused. Balanced MoE routing can be priced about
+3× above its measured peak because the allowance protects uneven dispatch.
+Caladan's full-recompute default is unaffected.
+
 ## Cross-checks and remaining conservatism
 
 The [CSV](trainer_rank_recompute_memory.csv) includes short-input failures that
@@ -175,7 +181,54 @@ runs in `scratch/recompute-memory-components/` and `scratch/recompute-memory-tig
 
 The [previous report](https://github.com/OpenPipe/ART/blob/57f9de2f9/dev/trainer_rank_recompute_memory.md)
 records the looser estimate and the earlier full-recompute underestimation.
-This work does not validate that legacy path, larger LoRA ranks, CP, pretrained
-routing distributions, arbitrary deep prefix trees, or other hardware/kernels.
+This initial campaign does not validate that legacy path, larger LoRA ranks, CP,
+pretrained routing distributions, arbitrary deep prefix trees, or other hardware/kernels.
 These measurements support the calibrated native paths, not a universal memory
 bound.
+
+## Context-parallel follow-up (#922)
+
+Dense selective and disabled-recompute floors now use the largest actual
+attention/GDN token load on any CP rank, summed across forward groups and padded
+as execution pads them. Balanced layouts approach a 1/CP token discount. Uneven
+layouts must use their actual load: the 27B 4k pair packs 6,964 global tokens but
+puts 4,096 on its busiest rank. Dividing by two would underprice its observed
+62.918 GiB peak. GDN segment states and cold workspace are not divided by CP;
+gathered outputs, the old static floor, and profile floors remain unchanged.
+
+Width selection uses the existing exact-plan fallback for these CP requests,
+reusing the native planner cache. This requires more CPU planning than a global
+token-count probe. The average CP load is used only as an optimistic split-search
+bound. MoE retains the previous estimate because expert dispatch can concentrate
+tokens from multiple CP ranks; it needs separate calibration before a discount.
+
+The follow-up measurements use two local H200s, TP1/CP2, eager native execution,
+bf16, rank-1 LoRA, random weights, and paired sequences with 30% shared prefixes.
+The [CP CSV](trainer_rank_recompute_memory_cp.csv) records maxima across ranks and
+both repetitions, including source hashes and forward/backward peaks. Values
+below are incremental allocated GiB and include the existing 10% estimate margin.
+
+| Model / tokens per sequence | Recompute | Previous estimate | CP estimate | Forward peak |
+| --- | --- | ---: | ---: | ---: |
+| 27B / 2,048 | selective | 69.123 | 34.954 | 31.518 |
+| 27B / 4,096 | selective | 117.374 | 69.524 | 62.918 |
+| 4B / 4,096 | selective | 41.478 | 20.922 | 18.727 |
+| 4B / 4,096 | none | 41.478 | 20.922 | 19.356 |
+| 4B / 4,096 | selective + mlp | 25.760 | 13.063 | 11.581 |
+| 1.7B / 4,096 | selective | 21.002 | 15.477 | 13.932 |
+
+All 13 cells (52 rank-samples) cover the measured forward peak, with at least 8.1%
+headroom; every backward completes with finite losses and adapter gradients.
+The 64-token 4B pair's cold backward peaks above its forward estimate (0.699 vs
+0.602 GiB); admission estimates forward memory, not arbitrary caller backward.
+Estimator source is `7d5fb47ce164fce3801cc8df1ce531c3a84832f5`; the initial 4B
+selective series uses `7467f35bb179ae11de6dfd960ff1b037ab4b8690`, before a TP-padding
+correction that does not change TP1. The CSV records both source and driver hashes.
+
+The 27B 4k pair now fits its 81.868 GiB incremental budget and completes backward;
+the previous estimate would refuse it. To reproduce, use the command above with
+TP=1, CP=2, `ART_DISABLE_MEGATRON_COMPILE=1`, and `--nproc-per-node=2`.
+Use `--model Qwen/Qwen3.5-4B` for the 4B checks and `--mode none` or
+`--modules core_attn mlp` for the other modes. Raw JSONL is retained under
+`scratch/recompute-memory-cp/`. This extends calibration to these dense CP2
+paths; larger CP sizes, combined TP/CP, and MoE CP remain uncalibrated.

dev/trainer_rank_recompute_memory.py

@@ -204,6 +204,7 @@ def gating(inputs, original=module.gating):
                     "logical_tokens": plan.logical_tokens,
                     "packed_tokens": plan.packed_tokens,
                     "grad_segment_count": plan.grad_segment_count,
+                    "retained_tokens": rank._plan_retained_tokens(plan),
                     "output_bytes": plan.output_bytes,
                     "selected_max_depth": plan.selected_max_depth,
                     "memory_minimal": memory_minimal,

dev/trainer_rank_recompute_memory_cp.csv

@@ -0,0 +1,14 @@
+case,source_sha,driver_sha256,model,mode,modules,compiled,tp,cp,lengths,shared_prefix,packed_tokens,retained_tokens,grad_segments,measured_samples,refused_samples,old_estimate_bytes,estimate_bytes,forward_peak_bytes,forward_backward_peak_bytes,finite,available_bytes_min
+27b-selective,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.8-27B,selective,core_attn,False,1,2,2048+2048,614,4096,2048,2,4,0,74220280217,37531864268,33842017792,33942171136,True,87907720909
+27b-selective,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.8-27B,selective,core_attn,False,1,2,4096+4096,1228,6964,4096,3,4,0,126029345587,74651231846,67557881856,67758206976,True,87905317069
+4b-mlp,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,selective,core_attn+mlp,False,1,2,256+256,76,512,256,2,4,0,2010434764,1158335692,933499392,970920448,True,133658705613
+4b-mlp,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,selective,core_attn+mlp,False,1,2,2048+2048,614,4096,2048,2,4,0,13980191948,7163399372,6243258880,6274177536,True,133658666701
+4b-mlp,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,selective,core_attn+mlp,False,1,2,4096+4096,1228,8192,4096,2,4,0,27659914444,14026329292,12435005952,12496841216,True,133658590925
+4b-none,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,none,,False,1,2,256+256,76,512,256,2,4,0,3065249792,1685743206,1478316032,1574740480,True,133658705613
+4b-none,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,none,,False,1,2,2048+2048,614,4096,2048,2,4,0,22418712166,11382659481,10417242624,10460772352,True,133658666701
+4b-none,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,none,,False,1,2,4096+4096,1228,8192,4096,2,4,0,44536954880,22464849510,20782973440,20870031360,True,133658590925
+4b-selective,7467f35bb179ae11de6dfd960ff1b037ab4b8690,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,selective,core_attn,False,1,2,64+64,19,128,64,2,4,0,991664537,646787891,457945600,750982144,True,133656611533
+4b-selective,7467f35bb179ae11de6dfd960ff1b037ab4b8690,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,selective,core_attn,False,1,2,2048+2048,614,4096,2048,2,4,0,22418712166,11382659481,10079601152,10135692800,True,133656572621
+4b-selective,7467f35bb179ae11de6dfd960ff1b037ab4b8690,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3.5-4B,selective,core_attn,False,1,2,4096+4096,1228,8192,4096,2,4,0,44536954880,22464849510,20107690496,20194748928,True,133656496845
+attention-selective,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3-1.7B,selective,core_attn,False,1,2,256+256,76,436,218,3,4,0,1481044787,779739136,713105408,752043520,True,138727714509
+attention-selective,7d5fb47ce164fce3801cc8df1ce531c3a84832f5,72f7ff4589b7af8b6a42d2eb8e3be152df2a67fa670770233c2dc50b03d61ad3,Qwen/Qwen3-1.7B,selective,core_attn,False,1,2,4096+4096,1228,6964,5120,3,4,0,22550816358,16618671308,14959463424,15113736704,True,138718277325

src/art/trainer_rank/_impl.py

@@ -2901,6 +2901,9 @@ def _split_chunk_lower_cost(
             output_bytes=output_bytes,
             signature=signature,
             logical_tokens=logical_tokens,
+            # The average CP load is an optimistic bound, not an admission cost.
+            retained_tokens=(packed_tokens + signature.topology[2] - 1)
+            // signature.topology[2],
         )
         profile = self._memory_profiles.get(signature)
         if (
@@ -2934,6 +2937,7 @@ def _plan_cost(self, plan: _FlatForwardPlan) -> _SubforwardCost:
             signature=plan.signature,
             logical_tokens=plan.active_logical_tokens,
             gdn_segments=plan.grad_segment_count,
+            retained_tokens=self._plan_retained_tokens(plan),
         )
 
     def _subforward_cost(
@@ -2944,13 +2948,15 @@ def _subforward_cost(
         signature: _MemorySignature,
         logical_tokens: int,
         gdn_segments: int = 0,
+        retained_tokens: int | None = None,
     ) -> _SubforwardCost:
         required = self._estimate_required_memory_bytes_from_values(
             packed_tokens=packed_tokens,
             output_bytes=output_bytes,
             signature=signature,
             logical_tokens=logical_tokens,
             gdn_segments=gdn_segments,
+            retained_tokens=retained_tokens,
         )
         retained = self._retained_memory_bytes(
             signature,
@@ -4630,6 +4636,15 @@ def _estimate_flat_forward(
                 checkpoint=checkpoint,
                 ensure_slots=not sync_planning_errors,
             )
+            if (
+                self._topology_key()[2] > 1
+                and self._recompute_granularity != "full"
+                and not self._geometry.moe_experts
+                and any(grad for (_, grad), _ in groups)
+            ):
+                # CP token ownership can be uneven. Use the existing exact-plan
+                # fallback; a global token count alone cannot price its peak.
+                return None
             packed_tokens = 0
             for (_slot, grad_enabled), group_indices in groups:
                 if exact:
@@ -5096,6 +5111,7 @@ def _memory_check(
                 signature=forward.signature,
                 logical_tokens=forward.active_logical_tokens,
                 gdn_segments=forward.grad_segment_count,
+                retained_tokens=self._plan_retained_tokens(forward),
             )
         return self._memory_check_required(required, sync_across_dp=sync_across_dp)
 
@@ -5150,6 +5166,30 @@ def _forward_memory_group() -> dist.ProcessGroup | None:
         except (AssertionError, ImportError, RuntimeError, ValueError):
             return None
 
+    def _plan_retained_tokens(self, plan: _FlatForwardPlan) -> int:
+        if (
+            plan.signature.topology[2] <= 1
+            or not plan.signature.grad_enabled
+            or self._recompute_granularity == "full"
+            or self._geometry.moe_experts
+        ):
+            return plan.packed_tokens
+        topology = self._topology()
+        # Bound each group's largest attention/GDN layout, including TP padding.
+        # Groups can place their peak on different ranks; summing is conservative.
+        return sum(
+            self._physical_tokens(
+                max(
+                    1,
+                    self._max_rank_model_tokens(
+                        _pad_packed_batch(group.packed, multiple=int(topology.tp)),
+                        topology=topology,
+                    ),
+                )
+            )
+            for group in plan.groups
+        )
+
     def _estimate_required_memory_bytes_from_values(
         self,
         *,
@@ -5158,6 +5198,7 @@ def _estimate_required_memory_bytes_from_values(
         signature: _MemorySignature,
         logical_tokens: int | None = None,
         gdn_segments: int = 0,
+        retained_tokens: int | None = None,
     ) -> int:
         if packed_tokens <= 0:
             return output_bytes
@@ -5256,7 +5297,13 @@ def _estimate_required_memory_bytes_from_values(
                 # native kernel initialization; a slope alone misses short inputs.
                 64 * 2**20
                 + gdn_state_bytes
-                + packed_tokens * self._param_dtype_size * retained_features,
+                + (
+                    packed_tokens
+                    if retained_tokens is None or geometry.moe_experts
+                    else retained_tokens
+                )
+                * self._param_dtype_size
+                * retained_features,
             )
         # Groups execute sequentially: summed packed rows conservatively bound
         # this FC2 component, not all workspace or retained graphs.
@@ -5954,7 +6001,9 @@ def _configure_hybridep(
             _pad_packed_batch(batch, multiple=int(topology.tp)) for batch in batches
         )
         sequence_length = max(int(batch.tokens.shape[1]) for batch in padded)
-        rows = tuple(self._hybridep_rows(batch, topology=topology) for batch in padded)
+        rows = tuple(
+            self._max_rank_model_tokens(batch, topology=topology) for batch in padded
+        )
         current = fused_a2a._hybrid_ep_buffer
         live = self._has_live_hybridep_graphs()
         # The buffer must hold the busiest rank's planned rows: cost-aware CP
@@ -6019,7 +6068,7 @@ def _set_hybridep_rows(rows: int) -> None:
 
         _set_hybridep_token_count(rows)
 
-    def _hybridep_rows(
+    def _max_rank_model_tokens(
         self,
         batch: PrefixTreePack,
         *,

tests/unit/test_trainer_rank_planning_status.py

@@ -97,6 +97,7 @@ def _worker(index: int, directory: Path) -> None:
             "estimate",
             "materialize",
             "price",
+            "cp_plan",
             "empty",
             "unequal",
             "unavailable",
@@ -150,6 +151,13 @@ def price(**kwargs):
                 return 0
 
             rank._estimate_required_memory_bytes_from_values = price
+
+            def retained_tokens(plan):
+                if index == 0 and mode == "cp_plan":
+                    raise primary
+                return plan.packed_tokens
+
+            rank._plan_retained_tokens = retained_tokens
             patches = pytest.MonkeyPatch()
             patches.setattr(_impl, "estimate_prefix_tree_packed_tokens", estimate)
             patches.setattr(_impl, "materialize_prefix_tree_layout", materialize)
@@ -181,14 +189,14 @@ def price(**kwargs):
                     if mode == "unavailable":
                         assert not rank._all_ranks_true(values is not None)
                     plan = rank._plan_flat_forward(requests, sync_planning_errors=True)
-                    if mode == "price":
+                    if mode in ("price", "cp_plan"):
                         rank._memory_check(
                             plan, sync_across_dp=True, sync_planning_errors=True
                         )
                     assert plan.request_count == len(requests)
                 except BaseException as caught:
                     error = caught
-                if mode in ("estimate", "materialize", "price"):
+                if mode in ("estimate", "materialize", "price", "cp_plan"):
                     if index == 0:
                         assert error is primary
                         assert error.__cause__ is cause and error.__context__ is context

tests/unit/test_trainer_rank_recompute_memory.py

@@ -368,3 +368,100 @@ def test_short_hybrid_pairs_pay_for_recurrent_states(monkeypatch):
     )
     assert inactive.grad_segment_count == 2
     assert rank._memory_check(inactive).estimated_required_bytes == estimate
+
+
+@pytest.mark.parametrize("granularity", (None, "selective"))
+def test_cp_prices_uneven_local_tokens_and_preserves_segment_states(
+    monkeypatch, granularity
+):
+    rank = _hybrid_rank(monkeypatch, 2)
+    rank._recompute_granularity = granularity
+    monkeypatch.setattr(rank, "_topology_key", lambda: (1, 2, 4, 1))
+    monkeypatch.setattr(rank, "_topology", lambda: SimpleNamespace(cp=4, tp=2))
+    # An uneven plan puts 3/4 of the rows on one rank, including TP padding.
+    monkeypatch.setattr(
+        rank, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() * 3 // 4
+    )
+    requests = [ForwardInput(input_tokens=torch.arange(4097), hidden_states=True)]
+    plan = rank._plan_flat_forward(requests)
+    assert rank._plan_retained_tokens(plan) == 3074
+    estimate = rank._memory_check(plan).estimated_required_bytes
+    assert rank._plan_cost(plan).required == estimate
+
+    def price(tokens=None, segments=plan.grad_segment_count):
+        return rank._estimate_required_memory_bytes_from_values(
+            packed_tokens=plan.packed_tokens,
+            output_bytes=plan.output_bytes,
+            signature=plan.signature,
+            gdn_segments=segments,
+            retained_tokens=tokens,
+        )
+
+    assert price(1026) < estimate < price()
+    state_costs = [
+        price(n) - price(n, segments=0) for n in (1026, 3074, plan.packed_tokens)
+    ]
+    assert max(state_costs) - min(state_costs) <= 1
+    assert min(state_costs) > 0
+    rank._memory_profiles[plan.signature] = _MemoryProfile(0, plan.packed_tokens)
+    assert rank._memory_check(plan).estimated_required_bytes == estimate
+    # Width selection must reach exact CP pricing even when the old global
+    # token bound would refuse. Both ordinary and memory-minimal probes defer.
+    for exact in (False, True):
+        assert rank._estimate_flat_forward(requests, exact=exact) is None
+    monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 1))
+    monkeypatch.setattr(rank, "_available_memory_bytes", lambda: estimate)
+    assert rank._search_next_micro_batch(requests, 0).check.fits
+    from art.trainer_rank._impl import Unset
+
+    lower = rank._split_chunk_lower_cost(
+        requests, [requests[0].input_tokens], checkpoint=Unset
+    )
+    assert lower.required <= estimate
+
+
+@pytest.mark.parametrize("kind", ("full", "no_grad", "moe"))
+def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind):
+    rank = _rank(
+        "full" if kind == "full" else "selective",
+        **({"num_moe_experts": 64} if kind == "moe" else {}),
+    )
+
+    monkeypatch.setattr(rank, "_topology_key", lambda: (1, 1, 2, 1))
+    monkeypatch.setattr(
+        rank,
+        "_max_rank_model_tokens",
+        lambda *a, **kw: pytest.fail("unexpected CP plan"),
+    )
+    plan = _plan(rank, no_grad=kind == "no_grad")
+    assert rank._plan_retained_tokens(plan) == plan.packed_tokens
+    assert rank._memory_check(plan).estimated_required_bytes == (
+        rank._estimate_required_memory_bytes_from_values(
+            packed_tokens=plan.packed_tokens,
+            output_bytes=plan.output_bytes,
+            signature=plan.signature,
+            gdn_segments=plan.grad_segment_count,
+        )
+    )
+
+
+def test_cp_floor_covers_recorded_uneven_27b_pair(monkeypatch):
+    # Native H200 CP2, 4k siblings: rank 0 peaks at 62.918 GiB with 4,096
+    # local tokens. Dividing the 6,964 global packed tokens by CP misses it.
+    rank = _hybrid_rank(monkeypatch, 1)
+    signature = replace(_plan(rank).signature, topology=(1, 1, 2, 1))
+    values: dict[str, Any] = dict(
+        packed_tokens=6964,
+        output_bytes=8192 * 5120 * 2,
+        signature=signature,
+        gdn_segments=3,
+    )
+    peak = 62.918 * 2**30
+    estimate = rank._estimate_required_memory_bytes_from_values(
+        **values, retained_tokens=4096
+    )
+    assert peak <= estimate <= 1.12 * peak
+    assert (
+        rank._estimate_required_memory_bytes_from_values(**values, retained_tokens=3482)
+        < peak
+    )