From 7d1d17e8922f556820d5e647a63e657a07468b6e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 16:48:30 +0000 Subject: [PATCH 1/5] Apply trainer-rank memory floors per context-parallel rank Outside TP1/CP1/EP1 the checkpoint, head and MoE floors were disabled, so a cold CP forward fell back to the generic activation factor. Measured on H200 (Qwen3.6-35B-A3B, 8 layers, EP1), cold peaks are 472-550 KB per local token at CP1, CP2 and CP4 alike, while main predicted 5.27 GB against 23-52 GB per CP rank. Price the floors from rows on the most loaded CP rank (an even share in the lower bound) and let CP, and EP for checkpoint retention, use them. EP dispatch working sets remain unmodeled. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 33 +++++++++++++++---- .../test_trainer_rank_admission_inputs.py | 5 +-- .../test_trainer_rank_checkpoint_memory.py | 19 ++++++----- tests/unit/test_trainer_rank_moe_memory.py | 11 ++++--- .../test_trainer_rank_recompute_memory.py | 16 ++++----- 5 files changed, 52 insertions(+), 32 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 4ae42575f..dea84d758 100644 --- a/src/art/trainer_rank/_impl.py +++ b/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 @@ -7724,7 +7738,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,)) ), ) diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index 92b5dbde4..a79d198e1 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/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) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 7a5e7d952..b0d7f6be7 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/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,16 @@ 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) + + def test_dp_empty_and_local_count(): r = rank() r._topology_key = lambda: (3, 1, 1, 1) @@ -427,10 +437,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) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index ccb1ed6a2..6b160e1c1 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/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 ) diff --git a/tests/unit/test_trainer_rank_recompute_memory.py b/tests/unit/test_trainer_rank_recompute_memory.py index ea145e4d4..39e87c7ea 100644 --- a/tests/unit/test_trainer_rank_recompute_memory.py +++ b/tests/unit/test_trainer_rank_recompute_memory.py @@ -428,20 +428,16 @@ def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind): ) 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") + # 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 ) From 606030f6cdbff543becea355de4f25dcbdd8c9e9 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 18:14:25 +0000 Subject: [PATCH 2/5] Defer context-parallel width probes to per-rank plan pricing Review follow-up. Width probing built group rows from global packed counts, so with the floors now enabled under CP it priced them about cp times too high, rejecting widths the exact planner would admit and recording a check that disagreed with the admitted plan's cost. Under CP every probe now uses the existing exact-plan fallback. A TP1/CP2 regression with uneven rank ownership checks that the probe defers, the per-rank floor is active, and the split lower bound stays below exact required and retained costs. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 12 ++++------- .../test_trainer_rank_checkpoint_memory.py | 20 +++++++++++++++++++ .../test_trainer_rank_recompute_memory.py | 7 ++++++- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index dea84d758..111ed71c6 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -5752,14 +5752,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 diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index b0d7f6be7..51bcd974e 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -200,6 +200,26 @@ def test_cp_floor_prices_rank_rows(cp, rows): assert r._checkpoint_memory_floor((rows,)) == single != (0, 0) +def test_cp_probe_defers_and_lower_bound_stays_below_exact(monkeypatch): + r = rank() + monkeypatch.setattr(r, "_topology_key", lambda: (1, 1, 2, 1)) + monkeypatch.setattr(r, "_topology", lambda: SimpleNamespace(cp=2, tp=1)) + # Uneven ownership: the busiest CP rank holds 3/4 of each group. + monkeypatch.setattr( + r, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() * 3 // 4 + ) + 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_dp_empty_and_local_count(): r = rank() r._topology_key = lambda: (3, 1, 1, 1) diff --git a/tests/unit/test_trainer_rank_recompute_memory.py b/tests/unit/test_trainer_rank_recompute_memory.py index 39e87c7ea..1d6eca3cb 100644 --- a/tests/unit/test_trainer_rank_recompute_memory.py +++ b/tests/unit/test_trainer_rank_recompute_memory.py @@ -421,7 +421,7 @@ 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 {}), @@ -433,6 +433,11 @@ def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind): 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._plan_group_rows(plan) == tuple( From 98570014411d3076e82cba3334cc7ff5e5b3cde8 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 18:42:21 +0000 Subject: [PATCH 3/5] Retry full sharing when a CP width falls outside profile trust Second review round. With every CP probe on the exact-plan fallback, width search retried memory-minimal layouts only on a memory failure; a cost-optimal layout that fit but exceeded the profile's trust window was rejected without trying full sharing, which main's cheap probe did. The fallback now retries in both cases (all conditions are DP-synchronized). A regression with unshared cost-optimal layouts reaches width 15 (width 8 without the retry), the CP bound test also covers an even per-rank share, and the probe docstring notes the CP deferral. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 38 +++++++++++------- .../test_trainer_rank_checkpoint_memory.py | 39 +++++++++++++++++-- 2 files changed, 59 insertions(+), 18 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 111ed71c6..11e1acfe4 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -5120,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: @@ -5729,6 +5736,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: diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 51bcd974e..7bcbf7fba 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -200,13 +200,14 @@ def test_cp_floor_prices_rank_rows(cp, rows): assert r._checkpoint_memory_floor((rows,)) == single != (0, 0) -def test_cp_probe_defers_and_lower_bound_stays_below_exact(monkeypatch): +@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)) - # Uneven ownership: the busiest CP rank holds 3/4 of each group. + # Even (tight) and uneven ownership of each group by the busiest CP rank. monkeypatch.setattr( - r, "_max_rank_model_tokens", lambda batch, **_: batch.tokens.numel() * 3 // 4 + 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. @@ -220,6 +221,38 @@ def test_cp_probe_defers_and_lower_bound_stays_below_exact(monkeypatch): 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 + + def test_dp_empty_and_local_count(): r = rank() r._topology_key = lambda: (3, 1, 1, 1) From 99640576de02ab38126b964ce1677e9b30c94104 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 19:58:23 +0000 Subject: [PATCH 4/5] Retry full sharing for an untrusted minimum wave under CP When the estimator is unavailable, the minimum wave used only the cost-optimal layout. A profiled layout outside the trust window was admitted cold even when full sharing was trusted. Apply the same DP-reduced retry the width probes use. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 13 +++++++++++++ tests/unit/test_trainer_rank_checkpoint_memory.py | 7 +++++++ tests/unit/test_trainer_rank_weird_shapes.py | 3 ++- 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 11e1acfe4..736bc951e 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -5218,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, diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 7bcbf7fba..c6e33ee9d 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -251,6 +251,13 @@ def test_cp_width_search_retries_full_sharing_inside_the_trust_window(monkeypatc 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(): diff --git a/tests/unit/test_trainer_rank_weird_shapes.py b/tests/unit/test_trainer_rank_weird_shapes.py index 174e56499..51b0c313a 100644 --- a/tests/unit/test_trainer_rank_weird_shapes.py +++ b/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 From 5585cd63818d2926be2f9c7f9ad49ee5560726b4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 20:56:10 +0000 Subject: [PATCH 5/5] Give the planning-status stub signature a CP1 topology Group rows now read the plan signature's CP size, so the None stub failed before the injected pricing error. Production signatures always carry a topology. Co-Authored-By: Claude Opus 5.5 (1M context) --- tests/unit/test_trainer_rank_planning_status.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_trainer_rank_planning_status.py b/tests/unit/test_trainer_rank_planning_status.py index 8b290a4a2..b684aadce 100644 --- a/tests/unit/test_trainer_rank_planning_status.py +++ b/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 )