Apply trainer-rank memory floors per context-parallel rank

src/art/trainer_rank/_impl.py

@@ -1500,7 +1500,8 @@ def _moe_output_bytes_per_token(
     slot_ref: "LoRASlotRef | None" = None,
 ) -> int:
     """Known routed-expert working set, not a complete model/compiled bound."""
-    if shape != ParallelShape(tp=1, cp=1):
+    # CP shards rows, not the per-token working set; EP dispatch is not modeled.
+    if (shape.tp, shape.ep, shape.etp) != (1, 1, 1):
         return 0
     from megatron.core.extensions.transformer_engine import (
         TEColumnParallelGroupedLinear,
@@ -3461,7 +3462,9 @@ def _split_chunk_lower_cost(
             assert estimated is not None  # rows are CPU copies
             physical_rows = self._physical_tokens(estimated)
             packed_tokens += physical_rows
-            group_rows.append((physical_rows, grad_enabled))
+            # The most loaded CP rank holds at least an even share.
+            cp = max(1, self._topology_key()[2])
+            group_rows.append((-(-physical_rows // cp), grad_enabled))
             head_requests = tuple(requests[index] for index in group_indices)
             head_workspace_bytes = max(
                 head_workspace_bytes,
@@ -3536,7 +3539,7 @@ def _head_workspace_bytes(self, rows: int) -> int:
             rows <= 0
             or self._padded_vocab_size is None
             or len(self.runtime.model) != 1
-            or self._topology_key()[1:] != (1, 1, 1)
+            or self._topology_key()[1::2] != (1, 1)
         ):
             return 0
         try:
@@ -3776,9 +3779,21 @@ def _plan_head_workspace_bytes(self, plan: _FlatForwardPlan) -> int:
         return peak
 
     def _plan_group_rows(self, plan: _FlatForwardPlan) -> tuple[tuple[int, bool], ...]:
+        """Physical rows per group on the most loaded context-parallel rank."""
+        topology = self._topology() if plan.signature.topology[2] > 1 else None
         return tuple(
             (
-                self._physical_tokens(int(group.packed.tokens.numel())),
+                self._physical_tokens(
+                    int(group.packed.tokens.numel())
+                    if topology is None
+                    else max(
+                        1,
+                        self._max_rank_model_tokens(
+                            _pad_packed_batch(group.packed, multiple=int(topology.tp)),
+                            topology=topology,
+                        ),
+                    )
+                ),
                 group.grad_enabled,
             )
             for group in plan.groups
@@ -3902,8 +3917,7 @@ def _checkpoint_memory_floor(
             or config.params_dtype is not torch.bfloat16
             or self._param_dtype_size != 2
             or next(self.runtime.model[0].parameters()).dtype is not torch.bfloat16
-            or self._topology_key()[1:] != (1, 1, 1)
-            or _expert_parallel_shape(self.runtime.provider) != (1, 1)
+            or self._topology_key()[1::2] != (1, 1)
             or any(
                 type(getattr(config, name, None)) is not type(value)
                 or getattr(config, name) != value
@@ -5106,25 +5120,32 @@ def fits(width: int) -> tuple[bool, bool]:
             width = normalize(width)
             result = estimate(width)
             if result is None:
-                # Estimator unavailable (device inputs): admit on the
-                # materialized plan, trying the cost-optimal layouts first and
-                # the memory-minimal layouts if those do not fit.
+                # Estimator unavailable (device inputs, or CP per-rank floors):
+                # admit on the materialized plan, trying the cost-optimal
+                # layouts first and the memory-minimal layouts if those do not
+                # fit or fall outside the profile's trust window.
+                def price(plan: _FlatForwardPlan) -> tuple[_MemoryCheck, bool, bool]:
+                    check = self._memory_check(
+                        plan, sync_across_dp=True, sync_planning_errors=True
+                    )
+                    trusted = self._all_ranks_have_memory_profile(
+                        packed_tokens=plan.packed_tokens,
+                        signature=plan.signature,
+                    )
+                    profiled = self._all_ranks_true(
+                        plan.signature in self._memory_profiles
+                    )
+                    return check, trusted, profiled
+
                 plan = materialize(width)
-                check = self._memory_check(
-                    plan, sync_across_dp=True, sync_planning_errors=True
-                )
-                if not check.fits and not layout_modes.get(width, False):
+                check, trusted, profiled = price(plan)
+                if (
+                    not check.fits or (profiled and not trusted)
+                ) and not layout_modes.get(width, False):
                     layout_modes[width] = True
                     plans.pop(width, None)
                     plan = materialize(width)
-                    check = self._memory_check(
-                        plan, sync_across_dp=True, sync_planning_errors=True
-                    )
-                trusted = self._all_ranks_have_memory_profile(
-                    packed_tokens=plan.packed_tokens,
-                    signature=plan.signature,
-                )
-                profiled = self._all_ranks_true(plan.signature in self._memory_profiles)
+                    check, trusted, profiled = price(plan)
             else:
                 check, trusted, profiled = result
             if width in plans:
@@ -5197,6 +5218,19 @@ def candidate(width: int) -> _CandidateMicroBatch[ForwardInputsT]:
         first_estimate = estimate(min_width)
         if first_estimate is None or not (first_estimate[0].fits and first_estimate[1]):
             first = candidate(min_width)
+            if (
+                first_estimate is None
+                and first.check.fits
+                and first.cold_start
+                and not layout_modes.get(min_width, False)
+                and self._all_ranks_true(first.plan.signature in self._memory_profiles)
+            ):
+                # Materialized pricing (DP-uniform): a profiled cost-optimal
+                # layout outside trust; full sharing may be trusted, as the
+                # estimator path would find.
+                layout_modes[min_width] = True
+                plans.pop(min_width, None)
+                first = candidate(min_width)
             if not first.check.fits:
                 # The smallest wave cannot run unsplit: best effort is the
                 # bounded split ladder. Each DP rank runs it on its own share,
@@ -5715,6 +5749,7 @@ def _estimate_flat_forward(
         whose feasibility is monotone in width (valid for rejecting one).
         ``exact=True`` prices the planner's actual layouts (memoized by
         content) and is used only inside the band where those bounds disagree.
+        Under CP it returns None: per-rank floors need materialized layouts.
         """
 
         if sync_planning_errors:
@@ -5738,14 +5773,10 @@ def _estimate_flat_forward(
                 # Pending saves require the actual bucket/replayed-tail geometry.
                 # Existing unavailable handling materializes before admission.
                 return None
-            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.
+            if self._topology_key()[2] > 1:
+                # CP token ownership can be uneven and memory floors are priced
+                # per rank. Use the existing exact-plan fallback; a global token
+                # count alone cannot price its peak.
                 return None
             packed_tokens = 0
             head_workspace_bytes = 0
@@ -7724,7 +7755,12 @@ def _estimate_required_memory_bytes_from_values(
         static_compute = max(
             static_compute,
             *(
-                self._moe_workspace_bytes(packed_tokens, slot_ref=ref)
+                self._moe_workspace_bytes(
+                    sum(rows for rows, _ in group_rows)
+                    if signature.topology[2] > 1 and group_rows
+                    else packed_tokens,
+                    slot_ref=ref,
+                )
                 for ref in (slot_refs or (None,))
             ),
         )

tests/unit/test_trainer_rank_admission_inputs.py

@@ -93,7 +93,8 @@ def test_cp_gdn_segments_groups_and_retained_tokens_reach_exact_search(monkeypat
     assert _gdn_memory.model_shapes(rank) is None  # This exercises the CP fallback.
     plan = rank._plan_flat_forward(requests)
     assert plan.grad_segment_count == 2
-    assert rank._plan_group_rows(plan) == ((128, True), (80, False))
+    # Rows on the most loaded CP rank (mocked at 3/4 of each group).
+    assert rank._plan_group_rows(plan) == ((96, True), (60, False))
     assert rank._plan_retained_tokens(plan) == 156
     for exact in (False, True):
         for memory_minimal in (False, True):
@@ -108,7 +109,7 @@ def test_cp_gdn_segments_groups_and_retained_tokens_reach_exact_search(monkeypat
         requests, tuple(item.input_tokens for item in requests), checkpoint=Unset
     )
     assert len(calls) == 1
-    assert calls[0]["group_rows"] == ((128, True), (80, False))
+    assert calls[0]["group_rows"] == ((32, True), (20, False))  # Even CP share.
     assert calls[0]["retained_tokens"] == 52  # Optimistic CP average only here.
     calls.clear()
     cost = rank._plan_cost(plan)

tests/unit/test_trainer_rank_checkpoint_memory.py

@@ -181,7 +181,7 @@ def test_actual_config_revalidated(field, value):
     assert r._checkpoint_memory_floor(((10, True),)) == (0, 0)
 
 
-@pytest.mark.parametrize("axis", [1, 2, 3])
+@pytest.mark.parametrize("axis", [1, 3])
 def test_topology_revalidated(axis):
     r = rank()
     topology = [1, 1, 1, 1]
@@ -190,6 +190,76 @@ def test_topology_revalidated(axis):
     assert r._checkpoint_memory_floor(((10, True),)) == (0, 0)
 
 
+@pytest.mark.parametrize("rows", [(10, True), (11, False)])
+@pytest.mark.parametrize("cp", [2, 4])
+def test_cp_floor_prices_rank_rows(cp, rows):
+    # Callers pass rows on the most loaded CP rank; the per-row floor matches CP1.
+    r = rank()
+    single = r._checkpoint_memory_floor((rows,))
+    r._topology_key = lambda: (1, 1, cp, 1)
+    assert r._checkpoint_memory_floor((rows,)) == single != (0, 0)
+
+
+@pytest.mark.parametrize("share", [lambda n: -(-n // 2), lambda n: n * 3 // 4])
+def test_cp_probe_defers_and_lower_bound_stays_below_exact(monkeypatch, share):
+    r = rank()
+    monkeypatch.setattr(r, "_topology_key", lambda: (1, 1, 2, 1))
+    monkeypatch.setattr(r, "_topology", lambda: SimpleNamespace(cp=2, tp=1))
+    # Even (tight) and uneven ownership of each group by the busiest CP rank.
+    monkeypatch.setattr(
+        r, "_max_rank_model_tokens", lambda batch, **_: share(batch.tokens.numel())
+    )
+    reqs = requests()
+    # Global counts would price the per-rank floors about cp times too high.
+    assert r._estimate_flat_forward(reqs) is None
+    plan = r._plan_flat_forward(reqs)
+    assert r._checkpoint_memory_floor(r._plan_group_rows(plan))[0] > 0
+    exact = r._plan_cost(plan)
+    lower = r._split_chunk_lower_cost(
+        reqs, [q.input_tokens for q in reqs], checkpoint=Unset
+    )
+    assert lower.required <= exact.required and lower.retained <= exact.retained
+
+
+def test_cp_width_search_retries_full_sharing_inside_the_trust_window(monkeypatch):
+    r = rank()
+    r._dp_rank_and_size = lambda: (0, 1)
+    monkeypatch.setattr(r, "_topology_key", lambda: (1, 1, 2, 1))
+    monkeypatch.setattr(r, "_topology", lambda: SimpleNamespace(cp=2, tp=1))
+    monkeypatch.setattr(
+        r, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() * 3 // 4
+    )
+    r._available_memory_bytes = lambda: 1 << 50
+    # Cost-optimal layouts stay unshared; memory-minimal layouts share fully.
+    monkeypatch.setattr(
+        r,
+        "_layout_anchor",
+        lambda *, memory_minimal=False: (
+            "full_sharing" if memory_minimal else "no_sharing"
+        ),
+    )
+    items = [
+        ForwardInput(
+            input_tokens=torch.tensor([7, 100 + i]), hidden_states=True, no_grad=True
+        )
+        for i in range(16)
+    ]
+    signature = r._plan_flat_forward(items[:1]).signature
+    r._memory_profiles[signature] = _MemoryProfile(bytes_per_token=1, packed_tokens=2)
+    # Unshared layouts leave the 8x window above width 8; full sharing
+    # (width + 1 packed rows) stays trusted through width 15.
+    selected = r._search_next_micro_batch(items, 0)
+    assert not isinstance(selected, _ForwardRefusal)
+    assert selected.stats_global_count == 15
+    # The minimum wave retries too: one item's unshared layout (32 rows) is
+    # outside a 3-row profile's window; full sharing (17 rows) is inside it.
+    r._memory_profiles[signature] = _MemoryProfile(bytes_per_token=1, packed_tokens=3)
+    r._last_global_micro_batch_size = None
+    selected = r._search_next_micro_batch([items], 0)
+    assert not isinstance(selected, _ForwardRefusal)
+    assert selected.plan.packed_tokens == 17 and not selected.cold_start
+
+
 def test_dp_empty_and_local_count():
     r = rank()
     r._topology_key = lambda: (3, 1, 1, 1)
@@ -427,10 +497,3 @@ def test_no_grad_enclosure_config_guard(field, value):
     r = rank()
     setattr(r.runtime.model[0].decoder.config, field, value)
     assert r._checkpoint_memory_floor(((11, False),)) == (0, 0)
-
-
-@pytest.mark.parametrize("cp", [2, 4])
-def test_no_grad_enclosure_keeps_cp_fallback(cp):
-    r = rank()
-    r._topology_key = lambda: (1, 1, cp, 1)
-    assert r._checkpoint_memory_floor(((11, False),)) == (0, 0)

tests/unit/test_trainer_rank_moe_memory.py

@@ -174,11 +174,12 @@ def test_unknown_fc2_input_keeps_previous_pair(layer, mode):
 
 @pytest.mark.parametrize("field", ["tp", "cp", "ep", "etp"])
 def test_sharded_path_unchanged(layer, field):
-    assert (
-        _moe_output_bytes_per_token(
-            [layer], replace(ParallelShape(tp=1, cp=1), **{field: 2})
-        )
-        == 0
+    # CP shards rows, not the per-token working set; TP/EP/ETP are unmodeled.
+    single = _moe_output_bytes_per_token([layer], ParallelShape(tp=1, cp=1))
+    sharded = replace(ParallelShape(tp=1, cp=1), **{field: 2})
+    assert single > 0
+    assert _moe_output_bytes_per_token([layer], sharded) == (
+        single if field == "cp" else 0
     )
 
 

tests/unit/test_trainer_rank_planning_status.py

@@ -112,7 +112,10 @@ def _worker(index: int, directory: Path) -> None:
             rank._dp_rank_and_size = lambda: (index, 2)
             rank._physical_tokens = lambda tokens: tokens
             rank._estimate_group_request_output_bytes = lambda requests: 0
-            rank._memory_signature_from_requests = lambda *args, **kwargs: None
+            # Group rows read the plan's CP size; DP2/TP1/CP1/PP1 prices packed rows.
+            rank._memory_signature_from_requests = lambda *args, **kwargs: (
+                SimpleNamespace(topology=(2, 1, 1, 1))
+            )
             rank._forward_item = lambda request: SimpleNamespace(
                 input_ids=request.input_tokens, request=request
             )

tests/unit/test_trainer_rank_recompute_memory.py

@@ -421,27 +421,28 @@ def price(tokens=None, segments=plan.grad_segment_count):
 
 
 @pytest.mark.parametrize("kind", ("full", "no_grad", "moe"))
-def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind):
+def test_cp_prices_group_rows_at_the_most_loaded_rank(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, "_topology", lambda: SimpleNamespace(cp=2, tp=1))
     monkeypatch.setattr(
-        rank,
-        "_max_rank_model_tokens",
-        lambda *a, **kw: pytest.fail("unexpected CP plan"),
+        rank, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() - 1
     )
     plan = _plan(rank, no_grad=kind == "no_grad")
+    # Global counts cannot price per-rank floors: width probes defer to plans.
+    request = ForwardInput(
+        input_tokens=torch.tensor([1, 2]), hidden_states=True, no_grad=kind == "no_grad"
+    )
+    assert rank._estimate_flat_forward([request]) is None
+    # Retention stays global for these kinds; floors use the most loaded rank.
     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,
-        )
+    assert rank._plan_group_rows(plan) == tuple(
+        (int(group.packed.tokens.numel()) - 1, group.grad_enabled)
+        for group in plan.groups
     )
 
 

tests/unit/test_trainer_rank_weird_shapes.py

@@ -782,7 +782,8 @@ def test_adaptive_planner_globally_falls_back_when_one_rank_cannot_estimate(
     rank = TrainerRank(_runtime())
     monkeypatch.setattr(rank, "_dp_rank_and_size", lambda: (0, 2))
     # Planning succeeds; only estimator availability and profile trust are false.
-    outcomes = iter((True, False, True, True, True, False))
+    # The last outcome says no profile exists, so full sharing is not retried.
+    outcomes = iter((True, False, True, True, True, False, False))
     monkeypatch.setattr(rank, "_all_ranks_true", lambda _local: next(outcomes))
     plans = 0
     original = rank._plan_flat_forward