From 7467f35bb179ae11de6dfd960ff1b037ab4b8690 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 01:15:21 +0000 Subject: [PATCH 1/3] Price dense retained activations using actual CP rank loads --- dev/trainer_rank_recompute_memory.py | 1 + src/art/trainer_rank/_impl.py | 49 ++++++++++++- .../test_trainer_rank_recompute_memory.py | 73 +++++++++++++++++++ 3 files changed, 120 insertions(+), 3 deletions(-) diff --git a/dev/trainer_rank_recompute_memory.py b/dev/trainer_rank_recompute_memory.py index c888b3766..ba9fbc76e 100644 --- a/dev/trainer_rank_recompute_memory.py +++ b/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, diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 8b0603348..be34b5e93 100644 --- a/src/art/trainer_rank/_impl.py +++ b/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,6 +2948,7 @@ 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, @@ -2951,6 +2956,7 @@ def _subforward_cost( 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,24 @@ 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(group.packed, topology=topology)) + ) + for group in plan.groups + ) + def _estimate_required_memory_bytes_from_values( self, *, @@ -5158,6 +5192,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 +5291,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 +5995,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 +6062,7 @@ def _set_hybridep_rows(rows: int) -> None: _set_hybridep_token_count(rows) - def _hybridep_rows( + def _max_rank_model_tokens( self, batch: PrefixTreePack, *, diff --git a/tests/unit/test_trainer_rank_recompute_memory.py b/tests/unit/test_trainer_rank_recompute_memory.py index 56f03dd95..ec55e8520 100644 --- a/tests/unit/test_trainer_rank_recompute_memory.py +++ b/tests/unit/test_trainer_rank_recompute_memory.py @@ -368,3 +368,76 @@ 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)) + # 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) == 3072 + estimate = rank._memory_check(plan).estimated_required_bytes + assert rank._plan_cost(plan).required == estimate + values = dict( + packed_tokens=plan.packed_tokens, + output_bytes=plan.output_bytes, + signature=plan.signature, + gdn_segments=plan.grad_segment_count, + ) + price = rank._estimate_required_memory_bytes_from_values + assert price(**values, retained_tokens=1026) < estimate < price(**values) + state_costs = [ + price(**values, retained_tokens=n) + - price(**{**values, "gdn_segments": 0}, retained_tokens=n) + for n in (1026, 3072, 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, + ) + ) From 7d5fb47ce164fce3801cc8df1ce531c3a84832f5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 01:17:43 +0000 Subject: [PATCH 2/3] Match CP planning to execution padding --- src/art/trainer_rank/_impl.py | 8 +++++- .../test_trainer_rank_recompute_memory.py | 27 ++++++++++--------- 2 files changed, 21 insertions(+), 14 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index be34b5e93..6e2fc5a94 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -5179,7 +5179,13 @@ def _plan_retained_tokens(self, plan: _FlatForwardPlan) -> int: # Groups can place their peak on different ranks; summing is conservative. return sum( self._physical_tokens( - max(1, self._max_rank_model_tokens(group.packed, topology=topology)) + max( + 1, + self._max_rank_model_tokens( + _pad_packed_batch(group.packed, multiple=int(topology.tp)), + topology=topology, + ), + ) ) for group in plan.groups ) diff --git a/tests/unit/test_trainer_rank_recompute_memory.py b/tests/unit/test_trainer_rank_recompute_memory.py index ec55e8520..e3de11939 100644 --- a/tests/unit/test_trainer_rank_recompute_memory.py +++ b/tests/unit/test_trainer_rank_recompute_memory.py @@ -377,28 +377,29 @@ def test_cp_prices_uneven_local_tokens_and_preserves_segment_states( 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)) + 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) == 3072 + assert rank._plan_retained_tokens(plan) == 3074 estimate = rank._memory_check(plan).estimated_required_bytes assert rank._plan_cost(plan).required == estimate - values = dict( - packed_tokens=plan.packed_tokens, - output_bytes=plan.output_bytes, - signature=plan.signature, - gdn_segments=plan.grad_segment_count, - ) - price = rank._estimate_required_memory_bytes_from_values - assert price(**values, retained_tokens=1026) < estimate < price(**values) + + 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(**values, retained_tokens=n) - - price(**{**values, "gdn_segments": 0}, retained_tokens=n) - for n in (1026, 3072, plan.packed_tokens) + 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 From 44ea923bf07efe5b7bddf980e307209c72e2b400 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 17 Sep 2026 01:23:28 +0000 Subject: [PATCH 3/3] Record CP2 GPU evidence and cover collective planning failures --- dev/trainer_rank_recompute_memory.md | 57 ++++++++++++++++++- dev/trainer_rank_recompute_memory_cp.csv | 14 +++++ .../unit/test_trainer_rank_planning_status.py | 12 +++- .../test_trainer_rank_recompute_memory.py | 23 ++++++++ 4 files changed, 102 insertions(+), 4 deletions(-) create mode 100644 dev/trainer_rank_recompute_memory_cp.csv diff --git a/dev/trainer_rank_recompute_memory.md b/dev/trainer_rank_recompute_memory.md index 4f6ffe830..cebf59cb0 100644 --- a/dev/trainer_rank_recompute_memory.md +++ b/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. diff --git a/dev/trainer_rank_recompute_memory_cp.csv b/dev/trainer_rank_recompute_memory_cp.csv new file mode 100644 index 000000000..a4900af49 --- /dev/null +++ b/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 diff --git a/tests/unit/test_trainer_rank_planning_status.py b/tests/unit/test_trainer_rank_planning_status.py index 2054024ee..7cb996e50 100644 --- a/tests/unit/test_trainer_rank_planning_status.py +++ b/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 diff --git a/tests/unit/test_trainer_rank_recompute_memory.py b/tests/unit/test_trainer_rank_recompute_memory.py index e3de11939..ea145e4d4 100644 --- a/tests/unit/test_trainer_rank_recompute_memory.py +++ b/tests/unit/test_trainer_rank_recompute_memory.py @@ -426,6 +426,7 @@ def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind): "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, @@ -442,3 +443,25 @@ def test_cp_keeps_existing_full_no_grad_and_moe_costs(monkeypatch, kind): 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 + )