From 8d70537ca5ee481c23c292660d59ff4c88b23b2e Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 14:13:22 +0000 Subject: [PATCH 01/13] Price trainer-rank memory profiles by packed tokens, not sharing ratio A memory profile learned at a lower logical/packed ratio was scaled up by the ratio gap, so a profile observed without prefix sharing priced a shared plan by its logical tokens. Measured forward+backward peaks are flat per packed token across sharing ratios 1-8 for GDN-MoE, attention-MoE, dense attention and dense GDN models, so apply profiled bytes per packed token in both the required and retained estimates. Sharing-ratio trust windows are unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 20 +----- tests/unit/test_trainer_rank_active_memory.py | 13 ++-- tests/unit/test_trainer_rank_moe_memory.py | 3 - tests/unit/test_trainer_rank_split.py | 63 ++++++++++--------- 4 files changed, 40 insertions(+), 59 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 736bc951e..d493dc22c 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -3976,7 +3976,6 @@ def _subforward_cost( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, - logical_tokens=logical_tokens, gdn_segments=gdn_segments, group_rows=group_rows, slot_refs=slot_refs, @@ -4054,8 +4053,7 @@ def _retained_memory_bytes( return required retained = output_bytes + max( checkpoint_retained_bytes, - profile.retained_compute_bytes_per_token - * max(packed_tokens, logical_tokens / profile.logical_per_packed), + profile.retained_compute_bytes_per_token * packed_tokens, ) return min(required, int(retained * _MEMORY_SAFETY_FACTOR)) @@ -5025,7 +5023,6 @@ def priced( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, - logical_tokens=logical_tokens, # A radix tree has fewer than twice as many segments as # active requests; the exact plan uses its actual count. gdn_segments=2 @@ -6790,7 +6787,6 @@ def _memory_check( packed_tokens=forward.packed_tokens, output_bytes=forward.output_bytes, signature=forward.signature, - logical_tokens=forward.active_logical_tokens, gdn_segments=forward.grad_segment_count, group_rows=self._plan_group_rows(forward), slot_refs=tuple(g.slot_ref for g in forward.groups), @@ -7636,7 +7632,6 @@ def _estimate_required_memory_bytes_from_values( packed_tokens: int, output_bytes: int, signature: _MemorySignature, - logical_tokens: int | None = None, gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), slot_refs: tuple["LoRASlotRef | None", ...] | None = None, @@ -7775,25 +7770,16 @@ def _estimate_required_memory_bytes_from_values( # Local head results coexist with full CP outputs during gathering. # Uneven rank plans can assign all of an item's rows to one rank. static_compute += output_bytes - # A profile learned under lighter sharing (lower logical/packed ratio) - # underestimates the per-packed-token footprint of a deeper-shared - # plan; scale the trusted estimate up by the ratio gap. - # Normalize before multiplying: cancelling packed tokens through two - # float operations can otherwise make a larger warm layout cheaper. - profiled_tokens: int | float = packed_tokens - if profiled is not None and logical_tokens is not None: - profiled_tokens = max( - packed_tokens, logical_tokens / profiled.logical_per_packed - ) # The trust window limits calibration growth, not the empirical floor. # Dropping that floor beyond the window can admit a larger request that # was refused just inside it, even below a previously observed peak. + # Measured peaks scale with packed tokens, not with the sharing ratio. if profiled is None: compute = static_compute else: compute = max( static_compute, - int(profiled.bytes_per_token * profiled_tokens), + int(profiled.bytes_per_token * packed_tokens), ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 5249e07f5..e8283b482 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -154,11 +154,10 @@ def test_inactive_observation_cannot_discount_later_shared_active_work(monkeypat monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 20_000) checks.append(rank._memory_check(candidate)) # Identical observed GPU work must produce identical future admission. - # Total-input ratios formerly discounted the second estimate to 11,985, - # admitting a plan that the equivalent short calibration refused. + # Sharing adds logical outputs, not packed activations. assert checks[0] == checks[1] - assert checks[0].estimated_required_bytes == 88_000 - assert not checks[0].fits + assert checks[0].estimated_required_bytes == 11_985 + assert checks[0].fits def test_warm_admission_rechecks_current_residency(monkeypatch): @@ -252,8 +251,7 @@ def unexpected_execution(*args, **kwargs): assert rank.last_forward_telemetry()["predicted_peak_bytes"] >= 10_000 -@pytest.mark.parametrize("logical_ratio", [1, 2, 10]) -def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): +def test_empirical_estimate_survives_packed_trust_boundary(): rank = _rank() observed = rank._plan_flat_forward(_requests("target_tokens")) rank._update_memory_profile(observed, 10_000, retained_bytes=1000) @@ -261,7 +259,6 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): values = [ estimate( packed_tokens=count, - logical_tokens=count * logical_ratio, output_bytes=count * 4, signature=observed.signature, ) @@ -269,7 +266,7 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): ] assert values == sorted(values) rate = rank._memory_profiles[observed.signature].bytes_per_token - assert values[-1] == int((800 * 4 + rate * 800 * logical_ratio) * 1.1) + assert values[-1] == int((800 * 4 + rate * 800) * 1.1) assert not rank._all_ranks_have_memory_profile( packed_tokens=800, signature=observed.signature ) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 6b160e1c1..594ee8abb 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -286,9 +286,6 @@ def test_profiles_outputs_and_empty_plan_preserve_empirical_floor(): assert not rank._all_ranks_have_memory_profile( packed_tokens=801, signature=signature ) - assert estimate( - packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature - ) == int((100 * 200000 + 123) * 1.1) def test_summed_group_envelope_and_retained_profile_unchanged(): diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 359700def..3284f6d8a 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -552,7 +552,7 @@ def test_retained_compute_keeps_growth_and_sharing_trust_limits( rank._update_memory_profile(plan, 100_000, retained_bytes=60_000) observed = rank._plan_cost(candidate) if trusted: - assert observed.retained == 220_000 + assert observed.retained == int((40_000 + 200 * packed_tokens) * 1.1) assert observed.retained < observed.required else: assert observed.retained == observed.required @@ -579,17 +579,21 @@ def _retained_ratio_requests() -> list[ForwardInput]: @pytest.mark.parametrize("admit", (False, True)) -@pytest.mark.parametrize(("profile_packed", "budget_gib"), ((8000, 40), (1000, 20))) +@pytest.mark.parametrize( + ("profile_packed", "budget_gib", "subforwards"), ((8000, 16, 2), (1000, 3, 4)) +) def test_retained_ratio_lower_bound_reaches_exact_split_admission( monkeypatch: pytest.MonkeyPatch, admit: bool, profile_packed: int, budget_gib: int, + subforwards: int, ) -> None: rank = _retained_ratio_rank(monkeypatch) # Actual packing: each child has 64k logical/no-sharing rows versus 4k # full-sharing rows. The latter crosses the retained-ratio limit; the # smaller profile also puts no-sharing beyond the packed-profile window. + # Sharing is priced by packed rows, so 17.19 GiB fits the whole forward. requests = _retained_ratio_requests() requests += [replace(request, no_grad=True) for request in requests] children = [rank._plan_flat_forward(requests[i : i + 16]) for i in (0, 16)] @@ -600,8 +604,15 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission( 4 * 131_072, profile_packed, retained_compute_bytes_per_token=128 ) costs = [rank._plan_cost(child) for child in children] - exact_bytes = rank._split_rung_check(costs).estimated_required_bytes - budget = budget_gib * 2**30 if admit else exact_bytes - 1 + rows = [request.input_tokens for request in requests] + lower_costs = [ + rank._split_chunk_lower_cost( + requests[i : i + 16], rows[i : i + 16], checkpoint=Unset + ) + for i in (0, 16) + ] + lower_bytes = rank._split_rung_check(lower_costs).estimated_required_bytes + budget = budget_gib * 2**30 if admit else lower_bytes - 1 monkeypatch.setattr(rank, "_available_memory_bytes", lambda: budget) exact = rank._split_rung_check(costs) if admit: @@ -609,9 +620,7 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission( requests, checkpoint=Unset, context="test" ) assert isinstance(plan, _SplitForwardPlan) - # The smaller profile cannot discount larger layouts after its trust - # window; the original 20 GiB budget now requires four subforwards. - assert plan.subforward_count == (2 if profile_packed == 8000 else 4) + assert plan.subforward_count == subforwards assert check == rank._split_rung_check( [rank._plan_cost(child) for child in plan.subforwards] ) @@ -637,30 +646,16 @@ def run(child, **kwargs): monkeypatch.setattr(rank, "_run_flat_plan_with_memory_tracking", run) outputs = rank._execute_admitted_plan(plan, check=check, context="test") assert [output.checkpoint for output in outputs] == list(map(str, range(32))) - lower = rank._split_rung_check( - [ - rank._split_chunk_lower_cost( - requests[i : i + 16], - [request.input_tokens for request in requests[i : i + 16]], - checkpoint=Unset, - ) - for i in (0, 16) - ] - ) + lower = rank._split_rung_check(lower_costs) assert [child.packed_tokens for child in children] == [64_000, 64_000] assert lower.estimated_required_bytes <= exact.estimated_required_bytes - assert lower.fits == (not admit or profile_packed == 8000) - assert exact.fits == (admit and profile_packed == 8000) - if not exact.fits: - # Reject this rung, at the lower bound or exact pricing; a later rung - # with smaller chunks may still fit this same budget. - split, rejected = rank._admit_split_rung( - [tuple(range(16)), tuple(range(16, 32))], - requests, - [request.input_tokens for request in requests], - checkpoint=Unset, - ) - assert split is None and not rejected.fits + assert lower.fits == admit and not exact.fits + # The lower bound passes whenever exact pricing admits this rung; smaller + # budgets reject it, though a later rung with smaller chunks may fit. + split, rejected = rank._admit_split_rung( + [tuple(range(16)), tuple(range(16, 32))], requests, rows, checkpoint=Unset + ) + assert (split is not None) == rejected.fits == (admit and subforwards == 2) @pytest.mark.parametrize( @@ -835,7 +830,7 @@ def test_retained_ratio_original_cost_witness_stays_conservative( ] # Exact retention keeps its conservative fallback. Only treating that # fallback as an optimistic split-search bound was incorrect. - assert small.required == large.required == 36_909_886_200 + assert (small.required, large.required) == (2_306_878_200, 4_613_745_400) assert small.retained == small.required and large.retained == 11_000 @@ -1090,7 +1085,13 @@ def test_warm_rounding_preserves_native_split_bound_and_exact_budget( rank._update_memory_profile( plan, peak_delta_bytes=7_864_390, retained_bytes=retained ) - costs = [rank._plan_cost(plan) for plan in children] + # Full sharing is the cheapest layout, so it sets the exact budget. + costs = [ + rank._plan_cost( + rank._plan_flat_forward(requests[a : a + 5], memory_minimal=True) + ) + for a in (0, 5) + ] lower = [ rank._split_chunk_lower_cost( requests[a : a + 5], rows[a : a + 5], checkpoint=Unset From 1cac203f0f59ad39bba04f6b0b33355acbc37185 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 15:02:38 +0000 Subject: [PATCH 02/13] Price logical-sized head copies and output gradients; drop stale callers Review follow-up. Profiles are per packed token, so memory that grows with logical rows is priced with outputs: the head's int64 copy of wide (2-D) labels, and the training gradient of dense logits/hidden-state outputs for requests not marked no_grad. Also remove the estimator logical_tokens argument from the landing harness and two tests that still passed it. Co-Authored-By: Claude Opus 5.5 (1M context) --- dev/trainer_rank_landing_acceptance.py | 2 -- src/art/trainer_rank/_impl.py | 15 ++++++++++++--- tests/unit/test_trainer_rank_active_memory.py | 2 +- tests/unit/test_trainer_rank_admission_inputs.py | 1 - tests/unit/test_trainer_rank_checkpoint_memory.py | 5 ++++- tests/unit/test_trainer_rank_head_memory.py | 1 - tests/unit/test_trainer_rank_split.py | 6 +++--- tests/unit/test_trainer_rank_validation.py | 6 +++--- 8 files changed, 23 insertions(+), 15 deletions(-) diff --git a/dev/trainer_rank_landing_acceptance.py b/dev/trainer_rank_landing_acceptance.py index dadd80c97..16a01bb7d 100644 --- a/dev/trainer_rank_landing_acceptance.py +++ b/dev/trainer_rank_landing_acceptance.py @@ -649,7 +649,6 @@ def unsplit_requirement() -> int: packed_tokens=plan.packed_tokens, output_bytes=plan.output_bytes, signature=plan.signature, - logical_tokens=plan.logical_tokens, ) def parity( @@ -1864,7 +1863,6 @@ def stream(arm: str, cap_bytes: int | None) -> dict[str, Any]: packed_tokens=plan.packed_tokens, output_bytes=plan.output_bytes, signature=plan.signature, - logical_tokens=plan.logical_tokens, ) torch.cuda.synchronize() cap = int(torch.cuda.memory_allocated()) + int(per_item * 2.5) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index d493dc22c..26b8c4322 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -6684,20 +6684,29 @@ def _estimate_group_request_output_bytes( total = 0 for request in requests: seq_len = int(request.input_tokens.numel()) + # Profiles are per packed token, so memory that grows with logical + # rows is priced here: wide label copies and training gradients of + # dense outputs. if request.target_tokens is not None: - total += int(request.target_tokens.numel()) * _dtype_size(torch.float32) + targets = int(request.target_tokens.numel()) + total += targets * _dtype_size(torch.float32) + if request.target_tokens.ndim > 1: + total += targets * _dtype_size(torch.long) if request.top_k is not None: total += ( seq_len * int(request.top_k) * (_dtype_size(torch.float32) + _dtype_size(torch.long)) ) + copies = 1 if request.no_grad else 2 if request.logits: if self._padded_vocab_size is None: raise RuntimeError("logits output memory requires a GPT model") - total += seq_len * self._padded_vocab_size * self._param_dtype_size + total += ( + copies * seq_len * self._padded_vocab_size * self._param_dtype_size + ) if request.hidden_states: - total += seq_len * self._hidden_size * self._param_dtype_size + total += copies * seq_len * self._hidden_size * self._param_dtype_size return total def _memory_signature_from_requests( diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index e8283b482..0b4bc2be8 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -156,7 +156,7 @@ def test_inactive_observation_cannot_discount_later_shared_active_work(monkeypat # Identical observed GPU work must produce identical future admission. # Sharing adds logical outputs, not packed activations. assert checks[0] == checks[1] - assert checks[0].estimated_required_bytes == 11_985 + assert checks[0].estimated_required_bytes == 12_971 assert checks[0].fits diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index a79d198e1..2fd5c0c49 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -28,7 +28,6 @@ def observed(**values): def assert_plan_values(rank, plan, values): assert values["packed_tokens"] == plan.packed_tokens assert values["output_bytes"] == plan.output_bytes - assert values["logical_tokens"] == plan.active_logical_tokens assert values["signature"] == plan.signature assert values["gdn_segments"] == plan.grad_segment_count assert values["group_rows"] == rank._plan_group_rows(plan) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index c6e33ee9d..5b74a450b 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -83,7 +83,10 @@ def test_same_old_signature_different_gradient_rows(): r = rank() a = r._estimate_flat_forward(requests()) b = r._estimate_flat_forward(requests(15360, 1024)) - assert a[:3] == b[:3] + # Same packed rows and signature; only gradient rows, and so the gradients + # of their hidden-state outputs, differ. + assert (a[0], a[2]) == (b[0], b[2]) + assert (a[1], b[1]) == (4096 * (2 * 1024 + 15360), 4096 * (2 * 15360 + 1024)) assert a[3] == ((1024, True), (15360, False)) assert b[3] == ((15360, True), (1024, False)) assert ( diff --git a/tests/unit/test_trainer_rank_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py index 50c16343f..744b64f09 100644 --- a/tests/unit/test_trainer_rank_head_memory.py +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -210,7 +210,6 @@ def test_exact_selector_estimate_matches_executed_layout(): packed_tokens=n, output_bytes=out, signature=sig, - logical_tokens=plan.active_logical_tokens, group_rows=groups, head_workspace_bytes=head, ) diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 3284f6d8a..7afd5e2f1 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -1109,7 +1109,7 @@ def test_warm_rounding_preserves_native_split_bound_and_exact_budget( @pytest.mark.parametrize("retained", (None, 0, 7864321 / 15)) -def test_normalized_warm_profile_is_monotone_through_ratio_floor( +def test_normalized_warm_profile_is_monotone_in_packed_tokens( monkeypatch: pytest.MonkeyPatch, retained: float | None ) -> None: rank = _retained_ratio_rank(monkeypatch) @@ -1117,8 +1117,8 @@ def test_normalized_warm_profile_is_monotone_through_ratio_floor( rank._memory_profiles[signature] = _MemoryProfile( 7864330 / 15, 1000, 15 / 13, retained_compute_bytes_per_token=retained ) - # All counts satisfy the retained guard. The logical term dominates until - # N=104, then packed-token growth dominates. Check every integer count. + # All counts satisfy the retained guard; logical rows no longer scale the + # profile, so cost must grow monotonically with packed rows. Check each count. costs = [ rank._subforward_cost( packed_tokens=packed, diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index d6ed71048..28798ebf7 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3817,10 +3817,10 @@ def _preprocess(self, *args: object, **kwargs: object) -> None: plan = trainer._plan_flat_forward([request]) estimate = trainer._estimate_flat_forward([request]) - target_bytes = 3 * 2 * 4 + target_bytes = 3 * 2 * (4 + 8) # logprobs plus the head's wide-label copy topk_bytes = 3 * 5 * (4 + 8) - logits_bytes = 3 * 10 * 4 - hidden_bytes = 3 * 4 * 4 + logits_bytes = 2 * 3 * 10 * 4 # dense outputs plus their training gradients + hidden_bytes = 2 * 3 * 4 * 4 assert estimate is not None and estimate[0] == plan.packed_tokens assert plan.output_bytes == target_bytes + topk_bytes + logits_bytes + hidden_bytes From bc5a1b607a0c2c3cb565b5a48c5419fa492d1488 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 15:23:52 +0000 Subject: [PATCH 03/13] Scope packed-token pricing to single-target training requests Second review round: moving logical-sized label copies and output gradients into output_bytes made profile learning subtract bytes that forward-only observations never held. Instead, keep main's output accounting and the estimator's logical_tokens argument, and skip the sharing-ratio extrapolation only when every request mix is single-target or output-free, the case the production reports and controlled measurements cover. Wide labels and dense or top-k outputs keep the conservative extrapolation. Co-Authored-By: Claude Opus 5.5 (1M context) --- dev/trainer_rank_landing_acceptance.py | 2 + src/art/trainer_rank/_impl.py | 48 +++++++++++++------ tests/unit/test_trainer_rank_active_memory.py | 11 +++-- .../test_trainer_rank_admission_inputs.py | 1 + .../test_trainer_rank_checkpoint_memory.py | 5 +- tests/unit/test_trainer_rank_head_memory.py | 1 + tests/unit/test_trainer_rank_moe_memory.py | 3 ++ tests/unit/test_trainer_rank_validation.py | 6 +-- 8 files changed, 51 insertions(+), 26 deletions(-) diff --git a/dev/trainer_rank_landing_acceptance.py b/dev/trainer_rank_landing_acceptance.py index 16a01bb7d..dadd80c97 100644 --- a/dev/trainer_rank_landing_acceptance.py +++ b/dev/trainer_rank_landing_acceptance.py @@ -649,6 +649,7 @@ def unsplit_requirement() -> int: packed_tokens=plan.packed_tokens, output_bytes=plan.output_bytes, signature=plan.signature, + logical_tokens=plan.logical_tokens, ) def parity( @@ -1863,6 +1864,7 @@ def stream(arm: str, cap_bytes: int | None) -> dict[str, Any]: packed_tokens=plan.packed_tokens, output_bytes=plan.output_bytes, signature=plan.signature, + logical_tokens=plan.logical_tokens, ) torch.cuda.synchronize() cap = int(torch.cuda.memory_allocated()) + int(per_item * 2.5) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 26b8c4322..8c6572362 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -3976,6 +3976,7 @@ def _subforward_cost( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, + logical_tokens=logical_tokens, gdn_segments=gdn_segments, group_rows=group_rows, slot_refs=slot_refs, @@ -4051,9 +4052,14 @@ def _retained_memory_bytes( ratio = logical_tokens / max(1, packed_tokens) if ratio > profile.logical_per_packed * _MEMORY_PROFILE_TRUST_GROWTH: return required + tokens = ( + packed_tokens + if _PACKED_PRICED_MIXES.issuperset(signature.request_mix) + else max(packed_tokens, logical_tokens / profile.logical_per_packed) + ) retained = output_bytes + max( checkpoint_retained_bytes, - profile.retained_compute_bytes_per_token * packed_tokens, + profile.retained_compute_bytes_per_token * tokens, ) return min(required, int(retained * _MEMORY_SAFETY_FACTOR)) @@ -5023,6 +5029,7 @@ def priced( packed_tokens=packed_tokens, output_bytes=output_bytes, signature=signature, + logical_tokens=logical_tokens, # A radix tree has fewer than twice as many segments as # active requests; the exact plan uses its actual count. gdn_segments=2 @@ -6684,29 +6691,20 @@ def _estimate_group_request_output_bytes( total = 0 for request in requests: seq_len = int(request.input_tokens.numel()) - # Profiles are per packed token, so memory that grows with logical - # rows is priced here: wide label copies and training gradients of - # dense outputs. if request.target_tokens is not None: - targets = int(request.target_tokens.numel()) - total += targets * _dtype_size(torch.float32) - if request.target_tokens.ndim > 1: - total += targets * _dtype_size(torch.long) + total += int(request.target_tokens.numel()) * _dtype_size(torch.float32) if request.top_k is not None: total += ( seq_len * int(request.top_k) * (_dtype_size(torch.float32) + _dtype_size(torch.long)) ) - copies = 1 if request.no_grad else 2 if request.logits: if self._padded_vocab_size is None: raise RuntimeError("logits output memory requires a GPT model") - total += ( - copies * seq_len * self._padded_vocab_size * self._param_dtype_size - ) + total += seq_len * self._padded_vocab_size * self._param_dtype_size if request.hidden_states: - total += copies * seq_len * self._hidden_size * self._param_dtype_size + total += seq_len * self._hidden_size * self._param_dtype_size return total def _memory_signature_from_requests( @@ -6796,6 +6794,7 @@ def _memory_check( packed_tokens=forward.packed_tokens, output_bytes=forward.output_bytes, signature=forward.signature, + logical_tokens=forward.active_logical_tokens, gdn_segments=forward.grad_segment_count, group_rows=self._plan_group_rows(forward), slot_refs=tuple(g.slot_ref for g in forward.groups), @@ -7641,6 +7640,7 @@ def _estimate_required_memory_bytes_from_values( packed_tokens: int, output_bytes: int, signature: _MemorySignature, + logical_tokens: int | None = None, gdn_segments: int = 0, group_rows: tuple[tuple[int, bool], ...] = (), slot_refs: tuple["LoRASlotRef | None", ...] | None = None, @@ -7779,16 +7779,29 @@ def _estimate_required_memory_bytes_from_values( # Local head results coexist with full CP outputs during gathering. # Uneven rank plans can assign all of an item's rows to one rank. static_compute += output_bytes + # Outputs that grow with logical rows make a profile learned under + # lighter sharing underestimate a deeper-shared plan; scale the trusted + # estimate up by the ratio gap for them. Normalize before multiplying: + # cancelling packed tokens through two float operations can otherwise + # make a larger warm layout cheaper. + profiled_tokens: int | float = packed_tokens + if ( + profiled is not None + and logical_tokens is not None + and not _PACKED_PRICED_MIXES.issuperset(signature.request_mix) + ): + profiled_tokens = max( + packed_tokens, logical_tokens / profiled.logical_per_packed + ) # The trust window limits calibration growth, not the empirical floor. # Dropping that floor beyond the window can admit a larger request that # was refused just inside it, even below a previously observed peak. - # Measured peaks scale with packed tokens, not with the sharing ratio. if profiled is None: compute = static_compute else: compute = max( static_compute, - int(profiled.bytes_per_token * packed_tokens), + int(profiled.bytes_per_token * profiled_tokens), ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) @@ -8698,6 +8711,11 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: ) +# Measured flat per packed row across sharing ratios; wide labels and dense or +# top-k outputs keep the conservative logical/packed ratio extrapolation. +_PACKED_PRICED_MIXES = frozenset({"target:single", "inactive"}) + + def _request_mix_key(request: AnyForwardInput) -> str: parts = [] if request.target_tokens is not None: diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 0b4bc2be8..9a95cd509 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -154,10 +154,11 @@ def test_inactive_observation_cannot_discount_later_shared_active_work(monkeypat monkeypatch.setattr(rank, "_available_memory_bytes", lambda: 20_000) checks.append(rank._memory_check(candidate)) # Identical observed GPU work must produce identical future admission. - # Sharing adds logical outputs, not packed activations. + # Total-input ratios formerly discounted the second estimate to 11,985, + # admitting a plan that the equivalent short calibration refused. assert checks[0] == checks[1] - assert checks[0].estimated_required_bytes == 12_971 - assert checks[0].fits + assert checks[0].estimated_required_bytes == 88_000 + assert not checks[0].fits def test_warm_admission_rechecks_current_residency(monkeypatch): @@ -251,7 +252,8 @@ def unexpected_execution(*args, **kwargs): assert rank.last_forward_telemetry()["predicted_peak_bytes"] >= 10_000 -def test_empirical_estimate_survives_packed_trust_boundary(): +@pytest.mark.parametrize("logical_ratio", [1, 2, 10]) +def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): rank = _rank() observed = rank._plan_flat_forward(_requests("target_tokens")) rank._update_memory_profile(observed, 10_000, retained_bytes=1000) @@ -259,6 +261,7 @@ def test_empirical_estimate_survives_packed_trust_boundary(): values = [ estimate( packed_tokens=count, + logical_tokens=count * logical_ratio, output_bytes=count * 4, signature=observed.signature, ) diff --git a/tests/unit/test_trainer_rank_admission_inputs.py b/tests/unit/test_trainer_rank_admission_inputs.py index 2fd5c0c49..a79d198e1 100644 --- a/tests/unit/test_trainer_rank_admission_inputs.py +++ b/tests/unit/test_trainer_rank_admission_inputs.py @@ -28,6 +28,7 @@ def observed(**values): def assert_plan_values(rank, plan, values): assert values["packed_tokens"] == plan.packed_tokens assert values["output_bytes"] == plan.output_bytes + assert values["logical_tokens"] == plan.active_logical_tokens assert values["signature"] == plan.signature assert values["gdn_segments"] == plan.grad_segment_count assert values["group_rows"] == rank._plan_group_rows(plan) diff --git a/tests/unit/test_trainer_rank_checkpoint_memory.py b/tests/unit/test_trainer_rank_checkpoint_memory.py index 5b74a450b..c6e33ee9d 100644 --- a/tests/unit/test_trainer_rank_checkpoint_memory.py +++ b/tests/unit/test_trainer_rank_checkpoint_memory.py @@ -83,10 +83,7 @@ def test_same_old_signature_different_gradient_rows(): r = rank() a = r._estimate_flat_forward(requests()) b = r._estimate_flat_forward(requests(15360, 1024)) - # Same packed rows and signature; only gradient rows, and so the gradients - # of their hidden-state outputs, differ. - assert (a[0], a[2]) == (b[0], b[2]) - assert (a[1], b[1]) == (4096 * (2 * 1024 + 15360), 4096 * (2 * 15360 + 1024)) + assert a[:3] == b[:3] assert a[3] == ((1024, True), (15360, False)) assert b[3] == ((15360, True), (1024, False)) assert ( diff --git a/tests/unit/test_trainer_rank_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py index 744b64f09..50c16343f 100644 --- a/tests/unit/test_trainer_rank_head_memory.py +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -210,6 +210,7 @@ def test_exact_selector_estimate_matches_executed_layout(): packed_tokens=n, output_bytes=out, signature=sig, + logical_tokens=plan.active_logical_tokens, group_rows=groups, head_workspace_bytes=head, ) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 594ee8abb..c51740349 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -286,6 +286,9 @@ def test_profiles_outputs_and_empty_plan_preserve_empirical_floor(): assert not rank._all_ranks_have_memory_profile( packed_tokens=801, signature=signature ) + assert estimate( + packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature + ) == int((100 * 100000 + 123) * 1.1) def test_summed_group_envelope_and_retained_profile_unchanged(): diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py index 28798ebf7..d6ed71048 100644 --- a/tests/unit/test_trainer_rank_validation.py +++ b/tests/unit/test_trainer_rank_validation.py @@ -3817,10 +3817,10 @@ def _preprocess(self, *args: object, **kwargs: object) -> None: plan = trainer._plan_flat_forward([request]) estimate = trainer._estimate_flat_forward([request]) - target_bytes = 3 * 2 * (4 + 8) # logprobs plus the head's wide-label copy + target_bytes = 3 * 2 * 4 topk_bytes = 3 * 5 * (4 + 8) - logits_bytes = 2 * 3 * 10 * 4 # dense outputs plus their training gradients - hidden_bytes = 2 * 3 * 4 * 4 + logits_bytes = 3 * 10 * 4 + hidden_bytes = 3 * 4 * 4 assert estimate is not None and estimate[0] == plan.packed_tokens assert plan.output_bytes == target_bytes + topk_bytes + logits_bytes + hidden_bytes From d243137590a5d5715812a971387843f943aa0ee4 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 16:11:58 +0000 Subject: [PATCH 04/13] Tighten packed-pricing scope and charge head rows beyond profile sharing Third review round: key wide labels by their normalized shape (a flattened token axis with trailing target dimensions is not single-target), restrict packed pricing to fully grad-enabled signatures (no_grad GDN branch states are not charged), and charge 128 B for each logical row beyond the profile's observed sharing to cover the head's per-row label, position and row-match buffers. Tests cover the excluded mixes and the wide-label key. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 44 ++++++++++---- tests/unit/test_trainer_rank_active_memory.py | 39 +++++++++++- tests/unit/test_trainer_rank_moe_memory.py | 3 +- tests/unit/test_trainer_rank_split.py | 60 +++++++++++-------- 4 files changed, 109 insertions(+), 37 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 8c6572362..b78959171 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4054,7 +4054,7 @@ def _retained_memory_bytes( return required tokens = ( packed_tokens - if _PACKED_PRICED_MIXES.issuperset(signature.request_mix) + if _packed_priced(signature) else max(packed_tokens, logical_tokens / profile.logical_per_packed) ) retained = output_bytes + max( @@ -7785,11 +7785,8 @@ def _estimate_required_memory_bytes_from_values( # cancelling packed tokens through two float operations can otherwise # make a larger warm layout cheaper. profiled_tokens: int | float = packed_tokens - if ( - profiled is not None - and logical_tokens is not None - and not _PACKED_PRICED_MIXES.issuperset(signature.request_mix) - ): + packed_priced = _packed_priced(signature) + if profiled is not None and logical_tokens is not None and not packed_priced: profiled_tokens = max( packed_tokens, logical_tokens / profiled.logical_per_packed ) @@ -7801,7 +7798,15 @@ def _estimate_required_memory_bytes_from_values( else: compute = max( static_compute, - int(profiled.bytes_per_token * profiled_tokens), + int(profiled.bytes_per_token * profiled_tokens) + + ( + _PACKED_PRICED_LOGICAL_ROW_BYTES + * max( + 0, logical_tokens - packed_tokens * profiled.logical_per_packed + ) + if packed_priced and logical_tokens is not None + else 0 + ), ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) @@ -8711,16 +8716,33 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: ) -# Measured flat per packed row across sharing ratios; wide labels and dense or -# top-k outputs keep the conservative logical/packed ratio extrapolation. +# Grad-enabled single-target training measured flat per packed row across +# sharing ratios. Wide labels, dense or top-k outputs and no_grad forwards (whose +# GDN branch states are uncharged) keep the logical/packed ratio extrapolation. _PACKED_PRICED_MIXES = frozenset({"target:single", "inactive"}) +# The head's label copies, positions, row-match vectors and saved masks (about +# 80-100 B) grow with logical rows. Under packed pricing, charge rows beyond the +# profile's observed sharing with margin. +_PACKED_PRICED_LOGICAL_ROW_BYTES = 128 + + +def _packed_priced(signature: "_MemorySignature") -> bool: + return signature.grad_modes == (True,) and _PACKED_PRICED_MIXES.issuperset( + signature.request_mix + ) def _request_mix_key(request: AnyForwardInput) -> str: parts = [] if request.target_tokens is not None: - target = request.target_tokens - tail_shape = tuple(target.shape[request.input_tokens.ndim :]) + target, inputs = request.target_tokens, request.input_tokens + # Match _forward_item: trailing target dims follow the input shape or a + # flattened token axis. + tail_shape = tuple( + target.shape[inputs.ndim :] + if target.shape[: inputs.ndim] == inputs.shape + else target.shape[1:] + ) parts.append(f"target:{tail_shape or 'single'}") if request.top_k is not None: parts.append(f"topk:{int(request.top_k)}") diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 9a95cd509..6b18d67a3 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -15,6 +15,7 @@ TrainerRankMemoryError, Unset, ) +from art.trainer_rank._impl import _PACKED_PRICED_LOGICAL_ROW_BYTES, _request_mix_key class _Model(torch.nn.Module): @@ -269,7 +270,8 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): ] assert values == sorted(values) rate = rank._memory_profiles[observed.signature].bytes_per_token - assert values[-1] == int((800 * 4 + rate * 800) * 1.1) + row = _PACKED_PRICED_LOGICAL_ROW_BYTES * 800 * (logical_ratio - 1) + assert values[-1] == int((800 * 4 + rate * 800 + row) * 1.1) assert not rank._all_ranks_have_memory_profile( packed_tokens=800, signature=observed.signature ) @@ -312,3 +314,38 @@ def importing(name, *args, **kwargs): with pytest.raises(type(error)) as caught: getattr(rank, method)(argument) assert caught.value is error + + +def test_packed_pricing_is_limited_to_grad_single_target_mixes(): + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + single = observed.signature + assert single.grad_modes == (True,) + assert rank._memory_profiles[single].logical_per_packed == 1 + excluded = ( + replace(single, grad_modes=(False,)), + replace(single, request_mix=("hidden",)), + replace(single, request_mix=("target:(2,)",)), + ) + for signature in excluded: + rank._memory_profiles[signature] = rank._memory_profiles[single] + + def estimate(signature, logical_tokens=64): + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=8, + logical_tokens=logical_tokens, + output_bytes=0, + signature=signature, + ) + + # No-grad, dense-output and wide-label signatures keep the logical/packed + # extrapolation; single-target sharing adds only per-row head buffers. + assert all(estimate(single) < estimate(signature) for signature in excluded) + assert estimate(single, 72) - estimate(single, 64) == pytest.approx( + 8 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 1.1, abs=1 + ) + # Flattened-axis wide labels are not single-target. + tokens = torch.arange(4).reshape(1, 4) + wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) + assert _request_mix_key(wide) == "target:(3,)" diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index c51740349..55ab7ec71 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -10,6 +10,7 @@ from art.trainer_rank import ForwardInput, TrainerRank from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, _MemoryProfile, _MemorySignature, _moe_output_bytes_per_token, @@ -288,7 +289,7 @@ def test_profiles_outputs_and_empty_plan_preserve_empirical_floor(): ) assert estimate( packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature - ) == int((100 * 100000 + 123) * 1.1) + ) == int((100 * 100000 + _PACKED_PRICED_LOGICAL_ROW_BYTES * 100 + 123) * 1.1) def test_summed_group_envelope_and_retained_profile_unchanged(): diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 7afd5e2f1..a3d46bc0e 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -51,6 +51,7 @@ TrainerRankSlotStateError, ) from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, Unset, _FlatForwardPlan, _MemoryCheck, @@ -579,21 +580,17 @@ def _retained_ratio_requests() -> list[ForwardInput]: @pytest.mark.parametrize("admit", (False, True)) -@pytest.mark.parametrize( - ("profile_packed", "budget_gib", "subforwards"), ((8000, 16, 2), (1000, 3, 4)) -) +@pytest.mark.parametrize(("profile_packed", "budget_gib"), ((8000, 40), (1000, 20))) def test_retained_ratio_lower_bound_reaches_exact_split_admission( monkeypatch: pytest.MonkeyPatch, admit: bool, profile_packed: int, budget_gib: int, - subforwards: int, ) -> None: rank = _retained_ratio_rank(monkeypatch) # Actual packing: each child has 64k logical/no-sharing rows versus 4k # full-sharing rows. The latter crosses the retained-ratio limit; the # smaller profile also puts no-sharing beyond the packed-profile window. - # Sharing is priced by packed rows, so 17.19 GiB fits the whole forward. requests = _retained_ratio_requests() requests += [replace(request, no_grad=True) for request in requests] children = [rank._plan_flat_forward(requests[i : i + 16]) for i in (0, 16)] @@ -604,15 +601,8 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission( 4 * 131_072, profile_packed, retained_compute_bytes_per_token=128 ) costs = [rank._plan_cost(child) for child in children] - rows = [request.input_tokens for request in requests] - lower_costs = [ - rank._split_chunk_lower_cost( - requests[i : i + 16], rows[i : i + 16], checkpoint=Unset - ) - for i in (0, 16) - ] - lower_bytes = rank._split_rung_check(lower_costs).estimated_required_bytes - budget = budget_gib * 2**30 if admit else lower_bytes - 1 + exact_bytes = rank._split_rung_check(costs).estimated_required_bytes + budget = budget_gib * 2**30 if admit else exact_bytes - 1 monkeypatch.setattr(rank, "_available_memory_bytes", lambda: budget) exact = rank._split_rung_check(costs) if admit: @@ -620,7 +610,9 @@ def test_retained_ratio_lower_bound_reaches_exact_split_admission( requests, checkpoint=Unset, context="test" ) assert isinstance(plan, _SplitForwardPlan) - assert plan.subforward_count == subforwards + # The smaller profile cannot discount larger layouts after its trust + # window; the original 20 GiB budget now requires four subforwards. + assert plan.subforward_count == (2 if profile_packed == 8000 else 4) assert check == rank._split_rung_check( [rank._plan_cost(child) for child in plan.subforwards] ) @@ -646,16 +638,32 @@ def run(child, **kwargs): monkeypatch.setattr(rank, "_run_flat_plan_with_memory_tracking", run) outputs = rank._execute_admitted_plan(plan, check=check, context="test") assert [output.checkpoint for output in outputs] == list(map(str, range(32))) - lower = rank._split_rung_check(lower_costs) + lower = rank._split_rung_check( + [ + rank._split_chunk_lower_cost( + requests[i : i + 16], + [request.input_tokens for request in requests[i : i + 16]], + checkpoint=Unset, + ) + for i in (0, 16) + ] + ) assert [child.packed_tokens for child in children] == [64_000, 64_000] assert lower.estimated_required_bytes <= exact.estimated_required_bytes - assert lower.fits == admit and not exact.fits - # The lower bound passes whenever exact pricing admits this rung; smaller - # budgets reject it, though a later rung with smaller chunks may fit. - split, rejected = rank._admit_split_rung( - [tuple(range(16)), tuple(range(16, 32))], requests, rows, checkpoint=Unset - ) - assert (split is not None) == rejected.fits == (admit and subforwards == 2) + assert lower.fits == (not admit or profile_packed == 8000) + assert exact.fits == (admit and profile_packed == 8000) + if not exact.fits: + # Otherwise reject this rung; a later rung with smaller chunks may still + # fit. Just under the untrusted 1000-row exact rung, the grad child's + # full-sharing layout (priced by packed rows) rescues the rung itself. + split, rejected = rank._admit_split_rung( + [tuple(range(16)), tuple(range(16, 32))], + requests, + [request.input_tokens for request in requests], + checkpoint=Unset, + ) + rescued = not admit and profile_packed == 1000 + assert (split is not None) == rejected.fits == rescued @pytest.mark.parametrize( @@ -830,7 +838,11 @@ def test_retained_ratio_original_cost_witness_stays_conservative( ] # Exact retention keeps its conservative fallback. Only treating that # fallback as an optimistic split-search bound was incorrect. - assert (small.required, large.required) == (2_306_878_200, 4_613_745_400) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES + assert (small.required, large.required) == tuple( + int((4 * 131_072 * packed + rows * (64_000 - packed) + 10_000) * 1.1) + for packed in (4000, 8000) + ) assert small.retained == small.required and large.retained == 11_000 From 73db1f9b9e79bfce928b54601526e2e8b3b74842 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 16:42:38 +0000 Subject: [PATCH 05/13] Charge retained head rows and keep packed pricing monotone Fourth review round: saved head indices and masks stay live until backward, so retained memory also charges logical rows beyond the profile's sharing. Packed pricing now requires a per-packed rate of at least the per-row charge times the profile ratio (and the retained row charge requires the same of the retained rate), keeping more-shared layouts no more expensive so lower-bound pruning stays sound. All-grad signatures with several slot groups qualify. The scope test pins exact values for packed and extrapolated signatures, ratios above one, mixed grad modes and both guards. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 51 +++++++++++------ tests/unit/test_trainer_rank_active_memory.py | 57 +++++++++++++------ tests/unit/test_trainer_rank_split.py | 3 +- 3 files changed, 77 insertions(+), 34 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index b78959171..c26266faf 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4052,15 +4052,20 @@ def _retained_memory_bytes( ratio = logical_tokens / max(1, packed_tokens) if ratio > profile.logical_per_packed * _MEMORY_PROFILE_TRUST_GROWTH: return required - tokens = ( - packed_tokens - if _packed_priced(signature) - else max(packed_tokens, logical_tokens / profile.logical_per_packed) - ) - retained = output_bytes + max( - checkpoint_retained_bytes, - profile.retained_compute_bytes_per_token * tokens, - ) + rate = profile.retained_compute_bytes_per_token + if _packed_priced(signature, profile): + # Saved head indices and masks stay live until backward; charge them + # while that keeps retention monotone in packed rows. + retained_compute = rate * packed_tokens + ( + _packed_row_bytes(profile, packed_tokens, logical_tokens) + if rate >= _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed + else 0 + ) + else: + retained_compute = rate * max( + packed_tokens, logical_tokens / profile.logical_per_packed + ) + retained = output_bytes + max(checkpoint_retained_bytes, retained_compute) return min(required, int(retained * _MEMORY_SAFETY_FACTOR)) def _split_request_order( @@ -7785,7 +7790,7 @@ def _estimate_required_memory_bytes_from_values( # cancelling packed tokens through two float operations can otherwise # make a larger warm layout cheaper. profiled_tokens: int | float = packed_tokens - packed_priced = _packed_priced(signature) + packed_priced = profiled is not None and _packed_priced(signature, profiled) if profiled is not None and logical_tokens is not None and not packed_priced: profiled_tokens = max( packed_tokens, logical_tokens / profiled.logical_per_packed @@ -7800,10 +7805,7 @@ def _estimate_required_memory_bytes_from_values( static_compute, int(profiled.bytes_per_token * profiled_tokens) + ( - _PACKED_PRICED_LOGICAL_ROW_BYTES - * max( - 0, logical_tokens - packed_tokens * profiled.logical_per_packed - ) + _packed_row_bytes(profiled, packed_tokens, logical_tokens) if packed_priced and logical_tokens is not None else 0 ), @@ -8726,9 +8728,24 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: _PACKED_PRICED_LOGICAL_ROW_BYTES = 128 -def _packed_priced(signature: "_MemorySignature") -> bool: - return signature.grad_modes == (True,) and _PACKED_PRICED_MIXES.issuperset( - signature.request_mix +def _packed_priced(signature: "_MemorySignature", profile: "_MemoryProfile") -> bool: + # Per-row charges must not outgrow the per-packed rate, or a more-shared + # layout could cost more and break lower-bound pruning. + return ( + bool(signature.grad_modes) + and all(signature.grad_modes) + and _PACKED_PRICED_MIXES.issuperset(signature.request_mix) + and profile.bytes_per_token + >= _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed + ) + + +def _packed_row_bytes( + profile: "_MemoryProfile", packed_tokens: int, logical_tokens: int +) -> float: + """Head buffers for logical rows beyond the profile's observed sharing.""" + return _PACKED_PRICED_LOGICAL_ROW_BYTES * max( + 0, logical_tokens - packed_tokens * profile.logical_per_packed ) diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 6b18d67a3..17e28eed6 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -322,29 +322,54 @@ def test_packed_pricing_is_limited_to_grad_single_target_mixes(): rank._update_memory_profile(observed, 10_000, retained_bytes=1000) single = observed.signature assert single.grad_modes == (True,) - assert rank._memory_profiles[single].logical_per_packed == 1 - excluded = ( - replace(single, grad_modes=(False,)), - replace(single, request_mix=("hidden",)), - replace(single, request_mix=("target:(2,)",)), + profile = replace( + rank._memory_profiles[single], + bytes_per_token=100_000, + retained_compute_bytes_per_token=50_000, + logical_per_packed=2, ) - for signature in excluded: - rank._memory_profiles[signature] = rank._memory_profiles[single] - - def estimate(signature, logical_tokens=64): + signatures = { + "single": single, + "multi_grad": replace(single, grad_modes=(True, True)), + "no_grad": replace(single, grad_modes=(False,)), + "mixed_grad": replace(single, grad_modes=(False, True)), + "hidden": replace(single, request_mix=("hidden",)), + "wide": replace(single, request_mix=("target:(2,)",)), + } + for signature in signatures.values(): + rank._memory_profiles[signature] = profile + + def estimate(name): return rank._estimate_required_memory_bytes_from_values( packed_tokens=8, - logical_tokens=logical_tokens, + logical_tokens=64, output_bytes=0, - signature=signature, + signature=signatures[name], ) - # No-grad, dense-output and wide-label signatures keep the logical/packed - # extrapolation; single-target sharing adds only per-row head buffers. - assert all(estimate(single) < estimate(signature) for signature in excluded) - assert estimate(single, 72) - estimate(single, 64) == pytest.approx( - 8 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 1.1, abs=1 + # Packed rows plus head buffers for logical rows beyond the profile's 2x. + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * (64 - 8 * 2) + assert ( + estimate("single") == estimate("multi_grad") == int((100_000 * 8 + rows) * 1.1) + ) + retained = rank._retained_memory_bytes( + single, packed_tokens=8, logical_tokens=64, output_bytes=0, required=1 << 40 + ) + assert retained == int((50_000 * 8 + rows) * 1.1) + # Others keep the logical/packed extrapolation: 64 / 2 profiled rows. + for name in ("no_grad", "mixed_grad", "hidden", "wide"): + assert estimate(name) == int(100_000 * 32 * 1.1) + # Per-row charges above the profiled rates would break lower-bound pruning: + # a small retained rate omits the rows, and a small per-packed rate + # extrapolates. + rank._memory_profiles[single] = replace( + profile, retained_compute_bytes_per_token=255 ) + assert rank._retained_memory_bytes( + single, packed_tokens=8, logical_tokens=64, output_bytes=0, required=1 << 40 + ) == int(255 * 8 * 1.1) + rank._memory_profiles[single] = replace(profile, bytes_per_token=255) + assert estimate("single") == int(255 * 32 * 1.1) # Flattened-axis wide labels are not single-target. tokens = torch.arange(4).reshape(1, 4) wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index a3d46bc0e..debce857e 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -553,7 +553,8 @@ def test_retained_compute_keeps_growth_and_sharing_trust_limits( rank._update_memory_profile(plan, 100_000, retained_bytes=60_000) observed = rank._plan_cost(candidate) if trusted: - assert observed.retained == int((40_000 + 200 * packed_tokens) * 1.1) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * (logical_tokens - packed_tokens) + assert observed.retained == int((40_000 + 200 * packed_tokens + rows) * 1.1) assert observed.retained < observed.required else: assert observed.retained == observed.required From 597061db7fdbf89b646f6be0e8e48322ae0b9e63 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 17:19:27 +0000 Subject: [PATCH 06/13] Charge GDN branch states and keep packed pricing monotone at the floor Fifth review round and a high-branch H200 measurement. Retained memory now falls back to the logical/packed extrapolation below the rate floor instead of dropping head rows. Costs are truncated once. Packed pricing requires per-packed rates of at least twice the per-row charge times the profile ratio, so fewer packed rows always cost less despite rounding. GDN branch states grow with segments, not packed rows: 101 twenty-token branches on a 5,100-token prompt (ratio 72.6, 8 layers) raised the warm peak per packed token from 323 KB to 378 KB, about 3.9 MB per extra segment. Packed pricing now charges one layer's recurrent state and convolution history per grad segment under full recompute (every GDN layer otherwise). The split lower bound passes no segments, so pruning stays sound. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 83 ++++++++++++------- tests/unit/test_trainer_rank_active_memory.py | 55 ++++++++++-- tests/unit/test_trainer_rank_split.py | 5 +- 3 files changed, 102 insertions(+), 41 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index c26266faf..17c50113b 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -4053,13 +4053,10 @@ def _retained_memory_bytes( if ratio > profile.logical_per_packed * _MEMORY_PROFILE_TRUST_GROWTH: return required rate = profile.retained_compute_bytes_per_token - if _packed_priced(signature, profile): - # Saved head indices and masks stay live until backward; charge them - # while that keeps retention monotone in packed rows. - retained_compute = rate * packed_tokens + ( - _packed_row_bytes(profile, packed_tokens, logical_tokens) - if rate >= _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed - else 0 + if _packed_priced(signature, profile) and rate >= _packed_rate_floor(profile): + # Saved head indices and masks stay live until backward. + retained_compute = rate * packed_tokens + _packed_row_bytes( + profile, packed_tokens, logical_tokens ) else: retained_compute = rate * max( @@ -7728,22 +7725,7 @@ def _estimate_required_memory_bytes_from_values( # state (fp32), plus convolution history. Unlike token activations, # these do not shrink with segment length. gdn_state_bytes = ( - 2 - * gdn_segments - * gdn_layers - / tp - * ( - 4 - * geometry.gdn_value_heads - * geometry.gdn_key_head_dim - * geometry.gdn_value_head_dim - + self._param_dtype_size - * ( - 2 * geometry.gdn_key_heads * geometry.gdn_key_head_dim - + geometry.gdn_value_heads * geometry.gdn_value_head_dim - ) - * max(0, geometry.gdn_conv_kernel - 1) - ) + gdn_segments * gdn_layers * self._gdn_segment_layer_bytes() ) static_compute = max( static_compute, @@ -7803,15 +7785,46 @@ def _estimate_required_memory_bytes_from_values( else: compute = max( static_compute, - int(profiled.bytes_per_token * profiled_tokens) - + ( - _packed_row_bytes(profiled, packed_tokens, logical_tokens) - if packed_priced and logical_tokens is not None - else 0 + int( + profiled.bytes_per_token * profiled_tokens + + ( + _packed_row_bytes(profiled, packed_tokens, logical_tokens) + # Branch states grow with segments, not packed rows: one + # recomputed layer at a time, or every GDN layer. + + gdn_segments + * self._gdn_segment_layer_bytes() + * ( + 1 + if self._recompute_granularity == "full" + else min(self._num_layers, self._gdn_layers) + ) + if packed_priced and logical_tokens is not None + else 0 + ) ), ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) + def _gdn_segment_layer_bytes(self) -> float: + """Initial and final fp32 recurrent states plus convolution history.""" + geometry = self._geometry + return ( + 2 + / max(1, self._topology_key()[1]) + * ( + 4 + * geometry.gdn_value_heads + * geometry.gdn_key_head_dim + * geometry.gdn_value_head_dim + + self._param_dtype_size + * ( + 2 * geometry.gdn_key_heads * geometry.gdn_key_head_dim + + geometry.gdn_value_heads * geometry.gdn_value_head_dim + ) + * max(0, geometry.gdn_conv_kernel - 1) + ) + ) + def _available_memory_bytes(self, sample: dict[str, Any] | None = None) -> int: if not (torch.cuda.is_available() and self.device.type == "cuda"): return 1 << 60 @@ -8728,15 +8741,21 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: _PACKED_PRICED_LOGICAL_ROW_BYTES = 128 +def _packed_rate_floor(profile: "_MemoryProfile") -> float: + """A per-packed rate at least twice the per-row charges it offsets. + + Fewer packed rows then always cost less, with margin for rounding, so + lower-bound pruning over more-shared layouts stays sound. + """ + return 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed + + def _packed_priced(signature: "_MemorySignature", profile: "_MemoryProfile") -> bool: - # Per-row charges must not outgrow the per-packed rate, or a more-shared - # layout could cost more and break lower-bound pruning. return ( bool(signature.grad_modes) and all(signature.grad_modes) and _PACKED_PRICED_MIXES.issuperset(signature.request_mix) - and profile.bytes_per_token - >= _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed + and profile.bytes_per_token >= _packed_rate_floor(profile) ) diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 17e28eed6..9c1a94d80 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -359,17 +359,58 @@ def estimate(name): # Others keep the logical/packed extrapolation: 64 / 2 profiled rows. for name in ("no_grad", "mixed_grad", "hidden", "wide"): assert estimate(name) == int(100_000 * 32 * 1.1) - # Per-row charges above the profiled rates would break lower-bound pruning: - # a small retained rate omits the rows, and a small per-packed rate - # extrapolates. + # Per-row charges must stay well below the profiled rates, or a more-shared + # layout could cost more: under 2 x 128 B x ratio, both costs extrapolate. + floor = 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 2 rank._memory_profiles[single] = replace( - profile, retained_compute_bytes_per_token=255 + profile, retained_compute_bytes_per_token=floor - 1 ) assert rank._retained_memory_bytes( single, packed_tokens=8, logical_tokens=64, output_bytes=0, required=1 << 40 - ) == int(255 * 8 * 1.1) - rank._memory_profiles[single] = replace(profile, bytes_per_token=255) - assert estimate("single") == int(255 * 32 * 1.1) + ) == int((floor - 1) * 32 * 1.1) + rank._memory_profiles[single] = replace(profile, bytes_per_token=floor - 1) + assert estimate("single") == int((floor - 1) * 32 * 1.1) + # At and above the floor, cost never falls as packed rows grow. + for rate in (floor, floor + 1, 100_000): + rank._memory_profiles[single] = replace( + profile, bytes_per_token=rate, retained_compute_bytes_per_token=rate + ) + costs = [ + rank._subforward_cost( + packed_tokens=packed, + logical_tokens=64, + output_bytes=0, + signature=single, + ) + for packed in range(1, 65) + ] + assert all(a.required <= b.required for a, b in zip(costs, costs[1:])) + assert all(a.retained <= b.retained for a, b in zip(costs, costs[1:])) + # GDN branch states grow with segments, not packed rows. + rank._memory_profiles[single] = profile + rank._geometry = replace( + rank._geometry, + gdn_key_heads=1, + gdn_key_head_dim=4, + gdn_value_heads=2, + gdn_value_head_dim=4, + gdn_conv_kernel=4, + ) + live = 1 if rank._recompute_granularity == "full" else rank._gdn_layers + assert live * rank._gdn_segment_layer_bytes() > 0 + + def segments(count): + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=8, + logical_tokens=64, + output_bytes=0, + signature=single, + gdn_segments=count, + ) + + assert segments(3) - segments(0) == pytest.approx( + 3 * live * rank._gdn_segment_layer_bytes() * 1.1, abs=1 + ) # Flattened-axis wide labels are not single-target. tokens = torch.arange(4).reshape(1, 4) wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index debce857e..83d3ba0fa 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -550,11 +550,12 @@ def test_retained_compute_keeps_growth_and_sharing_trust_limits( rank._update_memory_profile(plan, 100_000, retained_bytes=None) unknown = rank._plan_cost(candidate) assert unknown.retained == unknown.required - rank._update_memory_profile(plan, 100_000, retained_bytes=60_000) + # A retained rate of 260 B clears the 2 x 128 B floor for packed pricing. + rank._update_memory_profile(plan, 100_000, retained_bytes=66_000) observed = rank._plan_cost(candidate) if trusted: rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * (logical_tokens - packed_tokens) - assert observed.retained == int((40_000 + 200 * packed_tokens + rows) * 1.1) + assert observed.retained == int((40_000 + 260 * packed_tokens + rows) * 1.1) assert observed.retained < observed.required else: assert observed.retained == observed.required From c4a9427737588b76971b28ed986f4b702a3abd04 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 18:02:39 +0000 Subject: [PATCH 07/13] Limit packed pricing to one-layer full recompute Sixth review round. The per-segment GDN charge assumes backward recomputes one layer at a time, which holds only for ART's default full/uniform/1 recompute; block or multi-layer full recompute, selective recompute, and no recompute keep more GDN states and activations live, including across split subforwards. The measurements and production reports cover only that default, so packed pricing now requires it and every other mode keeps the logical/packed extrapolation. The flag is recorded in planner reports and restored on replay. Tests pin the per-layer segment bytes by hand and the fallback for other recompute modes. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 37 ++++++++++++------- src/art/trainer_rank/_planner_misses.py | 3 +- tests/unit/test_trainer_rank_active_memory.py | 9 ++++- .../unit/test_trainer_rank_planner_reports.py | 1 + 4 files changed, 34 insertions(+), 16 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 17c50113b..fbca63b4d 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1791,6 +1791,13 @@ def memory_field(name: str, default: Any = None) -> Any: ) self._recompute_granularity = memory_field("recompute_granularity", None) + # ART's default full/uniform/1: forward keeps only layer inputs and + # backward recomputes one layer at a time. Unset means that default. + self._one_layer_recompute = ( + self._recompute_granularity == "full" + and memory_field("recompute_method", None) in (None, "uniform") + and memory_field("recompute_num_layers", None) in (None, 1) + ) self._recompute_modules: frozenset[str] = frozenset( memory_field("recompute_modules", ()) or () ) @@ -4053,7 +4060,9 @@ def _retained_memory_bytes( if ratio > profile.logical_per_packed * _MEMORY_PROFILE_TRUST_GROWTH: return required rate = profile.retained_compute_bytes_per_token - if _packed_priced(signature, profile) and rate >= _packed_rate_floor(profile): + if _packed_priced( + signature, profile, self._one_layer_recompute + ) and rate >= _packed_rate_floor(profile): # Saved head indices and masks stay live until backward. retained_compute = rate * packed_tokens + _packed_row_bytes( profile, packed_tokens, logical_tokens @@ -6108,6 +6117,7 @@ def _fill_planner_snapshot( "hidden_size", "param_dtype_size", "recompute_granularity", + "one_layer_recompute", "sequence_parallel", "attention_output_gate", "mlp_activation_factor", @@ -7772,7 +7782,9 @@ def _estimate_required_memory_bytes_from_values( # cancelling packed tokens through two float operations can otherwise # make a larger warm layout cheaper. profiled_tokens: int | float = packed_tokens - packed_priced = profiled is not None and _packed_priced(signature, profiled) + packed_priced = profiled is not None and _packed_priced( + signature, profiled, self._one_layer_recompute + ) if profiled is not None and logical_tokens is not None and not packed_priced: profiled_tokens = max( packed_tokens, logical_tokens / profiled.logical_per_packed @@ -7789,15 +7801,9 @@ def _estimate_required_memory_bytes_from_values( profiled.bytes_per_token * profiled_tokens + ( _packed_row_bytes(profiled, packed_tokens, logical_tokens) - # Branch states grow with segments, not packed rows: one - # recomputed layer at a time, or every GDN layer. - + gdn_segments - * self._gdn_segment_layer_bytes() - * ( - 1 - if self._recompute_granularity == "full" - else min(self._num_layers, self._gdn_layers) - ) + # Branch states grow with segments, not packed rows; + # backward recomputes one layer at a time. + + gdn_segments * self._gdn_segment_layer_bytes() if packed_priced and logical_tokens is not None else 0 ) @@ -8750,9 +8756,14 @@ def _packed_rate_floor(profile: "_MemoryProfile") -> float: return 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed -def _packed_priced(signature: "_MemorySignature", profile: "_MemoryProfile") -> bool: +def _packed_priced( + signature: "_MemorySignature", profile: "_MemoryProfile", one_layer_recompute: bool +) -> bool: + # Measured only under one-layer full recompute: other recompute modes keep + # more GDN states and activations live per segment and logical row. return ( - bool(signature.grad_modes) + one_layer_recompute + and bool(signature.grad_modes) and all(signature.grad_modes) and _PACKED_PRICED_MIXES.issuperset(signature.request_mix) and profile.bytes_per_token >= _packed_rate_floor(profile) diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 2018c6186..7637fc253 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -521,7 +521,8 @@ def report( _RANK_FIELDS = frozenset( "num_layers hidden_size param_dtype_size recompute_granularity " - "sequence_parallel attention_output_gate mlp_activation_factor gdn_layers " + "one_layer_recompute sequence_parallel attention_output_gate " + "mlp_activation_factor gdn_layers " "checkpointed_moe_layers recompute_modules moe_output_bytes_per_token " "moe_forward_stages".split() ) diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 9c1a94d80..df4a66885 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -396,8 +396,9 @@ def estimate(name): gdn_value_head_dim=4, gdn_conv_kernel=4, ) - live = 1 if rank._recompute_granularity == "full" else rank._gdn_layers - assert live * rank._gdn_segment_layer_bytes() > 0 + # 2 states x 4 B x Hv=2 x dk=4 x dv=4, plus 2 B conv history of 16 rows x 3. + assert rank._gdn_segment_layer_bytes() == 2 * (4 * 2 * 4 * 4 + 2 * 16 * 3) + live = 1 def segments(count): return rank._estimate_required_memory_bytes_from_values( @@ -411,6 +412,10 @@ def segments(count): assert segments(3) - segments(0) == pytest.approx( 3 * live * rank._gdn_segment_layer_bytes() * 1.1, abs=1 ) + # Other recompute modes keep more live per segment and row: extrapolate. + rank._one_layer_recompute = False + assert estimate("single") == int(100_000 * 32 * 1.1) + rank._one_layer_recompute = True # Flattened-axis wide labels are not single-target. tokens = torch.arange(4).reshape(1, 4) wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) diff --git a/tests/unit/test_trainer_rank_planner_reports.py b/tests/unit/test_trainer_rank_planner_reports.py index a370f56fa..0617ad802 100644 --- a/tests/unit/test_trainer_rank_planner_reports.py +++ b/tests/unit/test_trainer_rank_planner_reports.py @@ -250,6 +250,7 @@ def test_replay_reruns_real_memory_estimator_and_prefix_layout(tmp_path): "hidden_size": 8, "param_dtype_size": 2, "recompute_granularity": "full", + "one_layer_recompute": True, "sequence_parallel": False, "attention_output_gate": False, "mlp_activation_factor": 3, From 81e0eb1919347475ccaa0f09644f95c07ea01756 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 18:32:19 +0000 Subject: [PATCH 08/13] Detect one-layer recompute from exact config and training mode Seventh review round. Megatron requires an explicit recompute method and layer count under full recompute and skips recompute in eval mode, so packed pricing now checks the exact full/uniform/1 settings and that every model chunk is training, each time it prices. Planner reports record the computed value and replay restores it. Test stubs declare ART's explicit settings, and a detector test covers selective, none, block, two-layer uniform, unset and eval mode against the extrapolated pricing. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 32 ++++++++++++------ src/art/trainer_rank/_planner_misses.py | 3 +- tests/unit/test_trainer_rank_active_memory.py | 33 ++++++++++++++++--- tests/unit/test_trainer_rank_moe_memory.py | 6 +++- tests/unit/test_trainer_rank_split.py | 6 +++- 5 files changed, 62 insertions(+), 18 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index fbca63b4d..6768dce05 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -1791,13 +1791,8 @@ def memory_field(name: str, default: Any = None) -> Any: ) self._recompute_granularity = memory_field("recompute_granularity", None) - # ART's default full/uniform/1: forward keeps only layer inputs and - # backward recomputes one layer at a time. Unset means that default. - self._one_layer_recompute = ( - self._recompute_granularity == "full" - and memory_field("recompute_method", None) in (None, "uniform") - and memory_field("recompute_num_layers", None) in (None, 1) - ) + self._recompute_method = memory_field("recompute_method", None) + self._recompute_num_layers = memory_field("recompute_num_layers", None) self._recompute_modules: frozenset[str] = frozenset( memory_field("recompute_modules", ()) or () ) @@ -4061,7 +4056,7 @@ def _retained_memory_bytes( return required rate = profile.retained_compute_bytes_per_token if _packed_priced( - signature, profile, self._one_layer_recompute + signature, profile, self._one_layer_recompute() ) and rate >= _packed_rate_floor(profile): # Saved head indices and masks stay live until backward. retained_compute = rate * packed_tokens + _packed_row_bytes( @@ -6117,7 +6112,6 @@ def _fill_planner_snapshot( "hidden_size", "param_dtype_size", "recompute_granularity", - "one_layer_recompute", "sequence_parallel", "attention_output_gate", "mlp_activation_factor", @@ -6127,6 +6121,7 @@ def _fill_planner_snapshot( ) } rank_fields["recompute_modules"] = sorted(self._recompute_modules) + rank_fields["one_layer_recompute"] = self._one_layer_recompute() rank_fields["moe_forward_stages"] = getattr(self, "_moe_forward_stages", ()) rank_fields["geometry"] = asdict(self._geometry) rank_fields["topology"] = list(plan.signature.topology) @@ -7783,7 +7778,7 @@ def _estimate_required_memory_bytes_from_values( # make a larger warm layout cheaper. profiled_tokens: int | float = packed_tokens packed_priced = profiled is not None and _packed_priced( - signature, profiled, self._one_layer_recompute + signature, profiled, self._one_layer_recompute() ) if profiled is not None and logical_tokens is not None and not packed_priced: profiled_tokens = max( @@ -7811,6 +7806,23 @@ def _estimate_required_memory_bytes_from_values( ) return int((output_bytes + compute) * _MEMORY_SAFETY_FACTOR) + def _one_layer_recompute(self) -> bool: + """ART's default full/uniform/1 recompute, which Megatron runs in training. + + Forward keeps only layer inputs and backward recomputes one layer at a + time. Eval mode skips recompute even with gradients enabled. + """ + recorded = self.__dict__.get("_recorded_one_layer_recompute") + if recorded is not None: + return recorded # Planner-report replay has no live model. + return ( + self._recompute_granularity, + self._recompute_method, + self._recompute_num_layers, + ) == ("full", "uniform", 1) and all( + chunk.training for chunk in self.runtime.model + ) + def _gdn_segment_layer_bytes(self) -> float: """Initial and final fp32 recurrent states plus convolution history.""" geometry = self._geometry diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 7637fc253..9c344b311 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -592,8 +592,9 @@ def replay( "incomplete replay: immutable rank fields differ (including MoE stages)" ) rank = _impl.TrainerRank.__new__(_impl.TrainerRank) - for name in _RANK_FIELDS: + for name in _RANK_FIELDS - {"one_layer_recompute"}: setattr(rank, "_" + name, values[name]) + rank._recorded_one_layer_recompute = values["one_layer_recompute"] rank._moe_forward_stages = tuple(tuple(row) for row in values["moe_forward_stages"]) rank._geometry = ModelGeometry(**values["geometry"]) dp, tp, cp, pp = values["topology"] diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index df4a66885..68fb5e62d 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -37,7 +37,11 @@ def _rank(): model=[_Model()], optimizer=None, provider=SimpleNamespace( - hidden_size=8, num_layers=4, recompute_granularity="full" + hidden_size=8, + num_layers=4, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, ), model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), ), @@ -412,10 +416,29 @@ def segments(count): assert segments(3) - segments(0) == pytest.approx( 3 * live * rank._gdn_segment_layer_bytes() * 1.1, abs=1 ) - # Other recompute modes keep more live per segment and row: extrapolate. - rank._one_layer_recompute = False - assert estimate("single") == int(100_000 * 32 * 1.1) - rank._one_layer_recompute = True + # Other recompute modes, and eval mode (which skips recompute), keep more + # live per segment and row: extrapolate. + for setting in ( + ("selective", "uniform", 1), + (None, None, None), + ("full", "block", 1), + ("full", "uniform", 2), + ("full", None, None), + ): + ( + rank._recompute_granularity, + rank._recompute_method, + rank._recompute_num_layers, + ) = setting + assert not rank._one_layer_recompute() + assert estimate("single") == estimate("hidden") + rank._recompute_granularity, rank._recompute_method = "full", "uniform" + rank._recompute_num_layers = 1 + assert rank._one_layer_recompute() + rank.runtime.model[0].eval() + assert not rank._one_layer_recompute() + assert estimate("single") == estimate("hidden") + rank.runtime.model[0].train() # Flattened-axis wide labels are not single-target. tokens = torch.arange(4).reshape(1, 4) wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 55ab7ec71..0513209e5 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -77,7 +77,11 @@ def _rank(layer=None): model=[model], optimizer=None, provider=SimpleNamespace( - hidden_size=2048, num_layers=40, recompute_granularity="full" + hidden_size=2048, + num_layers=40, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, ), model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), ), diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 83d3ba0fa..c2acbb9a4 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -85,7 +85,11 @@ def _runtime() -> "TrainingRuntime": model=[_FakeGPT()], optimizer=None, provider=SimpleNamespace( - hidden_size=8, num_layers=4, recompute_granularity="full" + hidden_size=8, + num_layers=4, + recompute_granularity="full", + recompute_method="uniform", + recompute_num_layers=1, ), model_support_handler=SimpleNamespace(build_gdn_execution_spec=False), ) # type: ignore From d2d4d8122ec069d74656f6b961a06f3c18add560 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 19:07:26 +0000 Subject: [PATCH 09/13] Read the decoder's own mode and recompute config when pricing Eighth review round. Megatron gates recompute on the decoder's training mode and config, which can diverge from the outer chunk, so the one-layer check now requires every submodule to be training and reads recompute settings from the live decoder config (stored fields only for stub models). Replay rejects a recorded flag that is not a bool. Tests cover a decoder in eval under a training chunk and show a recorded flag switches single-target pricing between packed and extrapolated. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 27 ++++++++++++++++--- src/art/trainer_rank/_planner_misses.py | 2 ++ tests/unit/test_trainer_rank_active_memory.py | 14 ++++++++++ 3 files changed, 39 insertions(+), 4 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 6768dce05..0adca3857 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -7810,17 +7810,36 @@ def _one_layer_recompute(self) -> bool: """ART's default full/uniform/1 recompute, which Megatron runs in training. Forward keeps only layer inputs and backward recomputes one layer at a - time. Eval mode skips recompute even with gradients enabled. + time. Megatron checks the decoder's own mode and config, and eval mode + skips recompute even with gradients enabled, so read both live. """ recorded = self.__dict__.get("_recorded_one_layer_recompute") if recorded is not None: return recorded # Planner-report replay has no live model. - return ( + stored = ( self._recompute_granularity, self._recompute_method, self._recompute_num_layers, - ) == ("full", "uniform", 1) and all( - chunk.training for chunk in self.runtime.model + ) + + def settings(chunk: torch.nn.Module) -> tuple[Any, ...]: + try: + config = _language_model(chunk).decoder.config + except (AttributeError, RuntimeError): + return stored + return tuple( + getattr(config, name, None) + for name in ( + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + ) + ) + + return all( + settings(chunk) == ("full", "uniform", 1) + and all(module.training for module in chunk.modules()) + for chunk in self.runtime.model ) def _gdn_segment_layer_bytes(self) -> float: diff --git a/src/art/trainer_rank/_planner_misses.py b/src/art/trainer_rank/_planner_misses.py index 9c344b311..61b23fc8c 100644 --- a/src/art/trainer_rank/_planner_misses.py +++ b/src/art/trainer_rank/_planner_misses.py @@ -594,6 +594,8 @@ def replay( rank = _impl.TrainerRank.__new__(_impl.TrainerRank) for name in _RANK_FIELDS - {"one_layer_recompute"}: setattr(rank, "_" + name, values[name]) + if type(values["one_layer_recompute"]) is not bool: + raise ValueError("incomplete replay: recompute mode is not recorded") rank._recorded_one_layer_recompute = values["one_layer_recompute"] rank._moe_forward_stages = tuple(tuple(row) for row in values["moe_forward_stages"]) rank._geometry = ModelGeometry(**values["geometry"]) diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 68fb5e62d..4bc72415a 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -439,6 +439,20 @@ def segments(count): assert not rank._one_layer_recompute() assert estimate("single") == estimate("hidden") rank.runtime.model[0].train() + # Megatron checks the decoder's own mode, which can diverge from the chunk. + rank.runtime.model[0].decoder = torch.nn.Module() + rank.runtime.model[0].decoder.eval() + assert rank.runtime.model[0].training and not rank._one_layer_recompute() + assert estimate("single") == estimate("hidden") + rank.runtime.model[0].decoder.train() + assert rank._one_layer_recompute() + # Replay trusts the recorded mode instead of a live model. + packed = estimate("single") + rank._recorded_one_layer_recompute = False + assert estimate("single") == estimate("hidden") != packed + rank._recorded_one_layer_recompute = True + assert estimate("single") == packed + del rank._recorded_one_layer_recompute # Flattened-axis wide labels are not single-target. tokens = torch.arange(4).reshape(1, 4) wide = ForwardInput(input_tokens=tokens, target_tokens=torch.zeros(4, 3).long()) From 960fae3dad6cbeecc3fe1a0af1d21a8e1bbe223d Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 20:00:44 +0000 Subject: [PATCH 10/13] Charge allocator blocks for shared rows; read the decoder's own mode Each request's head buffers are separate allocations rounded up to 512 B blocks, so duplicate one-token requests could exceed the 128 B row charge. Charge eight blocks per logical row beyond the profile's sharing; the rate floor scales with it, so lower bounds stay sound. Recompute detection now reads the decoder's training flag, which is what Megatron checks, instead of walking every submodule. Tests cover the live decoder config and replay of an unrecorded recompute mode. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 41 +++++++---------- tests/unit/test_trainer_rank_active_memory.py | 45 +++++++++++++++++-- .../unit/test_trainer_rank_planner_reports.py | 4 ++ tests/unit/test_trainer_rank_split.py | 9 ++-- 4 files changed, 69 insertions(+), 30 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 0adca3857..5cb960791 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -7816,31 +7816,22 @@ def _one_layer_recompute(self) -> bool: recorded = self.__dict__.get("_recorded_one_layer_recompute") if recorded is not None: return recorded # Planner-report replay has no live model. - stored = ( - self._recompute_granularity, - self._recompute_method, - self._recompute_num_layers, - ) + target = ("full", "uniform", 1) + names = ("recompute_granularity", "recompute_method", "recompute_num_layers") - def settings(chunk: torch.nn.Module) -> tuple[Any, ...]: + def active(chunk: torch.nn.Module) -> bool: try: - config = _language_model(chunk).decoder.config + decoder = _language_model(chunk).decoder + config = decoder.config except (AttributeError, RuntimeError): - return stored - return tuple( - getattr(config, name, None) - for name in ( - "recompute_granularity", - "recompute_method", - "recompute_num_layers", - ) - ) + # No decoder config (stub or non-GPT chunk): stored settings, + # and every module must be training. + stored = tuple(getattr(self, "_" + name) for name in names) + return stored == target and all(m.training for m in chunk.modules()) + settings = tuple(getattr(config, name, None) for name in names) + return settings == target and decoder.training is True - return all( - settings(chunk) == ("full", "uniform", 1) - and all(module.training for module in chunk.modules()) - for chunk in self.runtime.model - ) + return all(active(chunk) for chunk in self.runtime.model) def _gdn_segment_layer_bytes(self) -> float: """Initial and final fp32 recurrent states plus convolution history.""" @@ -8773,9 +8764,11 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: # GDN branch states are uncharged) keep the logical/packed ratio extrapolation. _PACKED_PRICED_MIXES = frozenset({"target:single", "inactive"}) # The head's label copies, positions, row-match vectors and saved masks (about -# 80-100 B) grow with logical rows. Under packed pricing, charge rows beyond the -# profile's observed sharing with margin. -_PACKED_PRICED_LOGICAL_ROW_BYTES = 128 +# 80-100 B) grow with logical rows. Each request's buffers are separate +# allocations rounded up to 512 B blocks, so a fully shared one-token request +# still holds several KiB. Under packed pricing, charge each row beyond the +# profile's observed sharing eight blocks. +_PACKED_PRICED_LOGICAL_ROW_BYTES = 8 * 512 def _packed_rate_floor(profile: "_MemoryProfile") -> float: diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 4bc72415a..a4cd1a174 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -15,7 +15,11 @@ TrainerRankMemoryError, Unset, ) -from art.trainer_rank._impl import _PACKED_PRICED_LOGICAL_ROW_BYTES, _request_mix_key +from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, + _packed_priced, + _request_mix_key, +) class _Model(torch.nn.Module): @@ -261,7 +265,7 @@ def unexpected_execution(*args, **kwargs): def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): rank = _rank() observed = rank._plan_flat_forward(_requests("target_tokens")) - rank._update_memory_profile(observed, 10_000, retained_bytes=1000) + rank._update_memory_profile(observed, 100_000, retained_bytes=1000) estimate = rank._estimate_required_memory_bytes_from_values values = [ estimate( @@ -281,6 +285,28 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): ) +def test_packed_pricing_covers_allocator_blocks_of_fully_shared_requests(): + # A duplicate one-token request adds no packed row but still allocates its + # head buffers, each rounded up to a 512 B block. Charge at least eight. + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 100_000, retained_bytes=1000) + profile = rank._memory_profiles[observed.signature] + assert _packed_priced(observed.signature, profile, rank._one_layer_recompute()) + duplicates = 100_000 + + def estimate(logical_tokens: int) -> int: + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=1, + logical_tokens=logical_tokens, + output_bytes=0, + signature=observed.signature, + ) + + base = profile.logical_per_packed + assert estimate(base + duplicates) - estimate(base) >= duplicates * 8 * 512 + + @pytest.mark.parametrize( "method,argument,fallback", [ @@ -364,7 +390,7 @@ def estimate(name): for name in ("no_grad", "mixed_grad", "hidden", "wide"): assert estimate(name) == int(100_000 * 32 * 1.1) # Per-row charges must stay well below the profiled rates, or a more-shared - # layout could cost more: under 2 x 128 B x ratio, both costs extrapolate. + # layout could cost more: under two row charges per ratio, both extrapolate. floor = 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 2 rank._memory_profiles[single] = replace( profile, retained_compute_bytes_per_token=floor - 1 @@ -446,6 +472,19 @@ def segments(count): assert estimate("single") == estimate("hidden") rank.runtime.model[0].decoder.train() assert rank._one_layer_recompute() + # With a decoder config, its live settings and the decoder's mode decide. + decoder = rank.runtime.model[0].decoder + decoder.config = SimpleNamespace( + recompute_granularity="full", recompute_method="uniform", recompute_num_layers=1 + ) + rank.runtime.model[0].eval() + decoder.train() + assert rank._one_layer_recompute() + rank.runtime.model[0].train() + decoder.config.recompute_num_layers = 2 + assert not rank._one_layer_recompute() + decoder.config.recompute_num_layers = 1 + assert rank._one_layer_recompute() # Replay trusts the recorded mode instead of a live model. packed = estimate("single") rank._recorded_one_layer_recompute = False diff --git a/tests/unit/test_trainer_rank_planner_reports.py b/tests/unit/test_trainer_rank_planner_reports.py index 0617ad802..842755e5a 100644 --- a/tests/unit/test_trainer_rank_planner_reports.py +++ b/tests/unit/test_trainer_rank_planner_reports.py @@ -346,6 +346,10 @@ def test_replay_reruns_real_memory_estimator_and_prefix_layout(tmp_path): changed = json.loads(path.read_bytes()) changed[field] += 1 assert reports.replay(changed)["aggregate"]["matches"] is False + unrecorded = reports.validate_report(path.read_bytes()) + unrecorded["replay"]["memory_replay"]["rank"]["one_layer_recompute"] = None + with pytest.raises(ValueError, match="recompute mode is not recorded"): + reports.replay(unrecorded) drifted = reports.validate_report(path.read_bytes()) drifted["replay"]["source_files"]["_impl.py"]["sha256"] = "0" * 64 with pytest.raises(ValueError, match="source differs"): diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index c2acbb9a4..88eb3c122 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -554,12 +554,15 @@ def test_retained_compute_keeps_growth_and_sharing_trust_limits( rank._update_memory_profile(plan, 100_000, retained_bytes=None) unknown = rank._plan_cost(candidate) assert unknown.retained == unknown.required - # A retained rate of 260 B clears the 2 x 128 B floor for packed pricing. - rank._update_memory_profile(plan, 100_000, retained_bytes=66_000) + # A retained rate just above the floor of two row charges keeps packed pricing. + rate = 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES + 4 + rank._update_memory_profile( + plan, 40_000 + 200 * rate, retained_bytes=40_000 + 100 * rate + ) observed = rank._plan_cost(candidate) if trusted: rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * (logical_tokens - packed_tokens) - assert observed.retained == int((40_000 + 260 * packed_tokens + rows) * 1.1) + assert observed.retained == int((40_000 + rate * packed_tokens + rows) * 1.1) assert observed.retained < observed.required else: assert observed.retained == observed.required From 740242c862419ae07de7aa5b7afbaac680acb9fc Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 20:35:28 +0000 Subject: [PATCH 11/13] Size the shared-row charge from a measured head peak On an H200, a duplicate one-token request peaked at nine 512 B blocks (4,616 B) at the end of head forward, above 4 KiB x 1.1. Charge twelve blocks per row, and add a CUDA test that runs the real head and backward for 20,001 duplicates and checks the per-request peak against the charge. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 8 ++-- tests/unit/test_trainer_rank_active_memory.py | 4 +- .../unit/test_trainer_rank_head_recompute.py | 45 +++++++++++++++++++ 3 files changed, 51 insertions(+), 6 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 5cb960791..9eaafc77b 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -8765,10 +8765,10 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: _PACKED_PRICED_MIXES = frozenset({"target:single", "inactive"}) # The head's label copies, positions, row-match vectors and saved masks (about # 80-100 B) grow with logical rows. Each request's buffers are separate -# allocations rounded up to 512 B blocks, so a fully shared one-token request -# still holds several KiB. Under packed pricing, charge each row beyond the -# profile's observed sharing eight blocks. -_PACKED_PRICED_LOGICAL_ROW_BYTES = 8 * 512 +# allocations rounded up to 512 B blocks: on an H200, a fully shared one-token +# request peaked at nine blocks (4,616 B) at the end of head forward. Under +# packed pricing, charge each row beyond the profile's observed sharing twelve. +_PACKED_PRICED_LOGICAL_ROW_BYTES = 12 * 512 def _packed_rate_floor(profile: "_MemoryProfile") -> float: diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index a4cd1a174..3e8717730 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -287,7 +287,7 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): def test_packed_pricing_covers_allocator_blocks_of_fully_shared_requests(): # A duplicate one-token request adds no packed row but still allocates its - # head buffers, each rounded up to a 512 B block. Charge at least eight. + # head buffers, each rounded up to a 512 B block: nine at the measured peak. rank = _rank() observed = rank._plan_flat_forward(_requests("target_tokens")) rank._update_memory_profile(observed, 100_000, retained_bytes=1000) @@ -304,7 +304,7 @@ def estimate(logical_tokens: int) -> int: ) base = profile.logical_per_packed - assert estimate(base + duplicates) - estimate(base) >= duplicates * 8 * 512 + assert estimate(base + duplicates) - estimate(base) >= duplicates * 9 * 512 @pytest.mark.parametrize( diff --git a/tests/unit/test_trainer_rank_head_recompute.py b/tests/unit/test_trainer_rank_head_recompute.py index 20e3d7c85..207b4f4c0 100644 --- a/tests/unit/test_trainer_rank_head_recompute.py +++ b/tests/unit/test_trainer_rank_head_recompute.py @@ -436,3 +436,48 @@ def exp(subtraction): assert model.output_layer.weight.grad is not None assert model.output_layer.weight.grad.isfinite().all() assert torch.cuda.is_initialized() == cuda_initialized + + +def test_shared_one_token_requests_fit_the_packed_row_charge(monkeypatch): + # A duplicate one-token request adds a logical row but no packed row. Its + # labels, indices, masks and output are separate CUDA allocations rounded + # to 512 B blocks; through backward they must fit the per-row charge. + if not torch.cuda.is_available(): + pytest.skip("allocator rounding needs CUDA") + _patch_local_head(monkeypatch) + device = torch.device("cuda") + weight = torch.randn(1024, 64, device=device, dtype=torch.bfloat16) / 8 + model = SimpleNamespace( + output_layer=_Head(weight), + vocab_size=1024, + share_embeddings_and_output_weights=False, + _scale_logits=lambda value: value, + ) + r = object.__new__(TrainerRank) + r.runtime = SimpleNamespace(model=[model]) + request = ForwardInput( + input_tokens=torch.tensor([7]), target_tokens=torch.tensor([5]) + ) + + def peak(count: int) -> int: + items = [r._forward_item(request) for _ in range(count)] + prepared = SimpleNamespace( + positions_by_item=tuple(torch.tensor([0]) for _ in range(count)), + source_positions_by_item=tuple(torch.arange(1) for _ in range(count)), + ) + hidden = torch.randn(1, 64, device=device, dtype=torch.bfloat16) + hidden.requires_grad_() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + loss = torch.zeros((), device=device) + for output in r._project_head(items, prepared, hidden): + loss = loss - output.target_logprobs.sum() + loss.backward() + torch.cuda.synchronize() + return torch.cuda.max_memory_allocated() - base + + peak(1) + extra = 20_000 + per_request = (peak(1 + extra) - peak(1)) / extra + assert per_request <= _impl._PACKED_PRICED_LOGICAL_ROW_BYTES From 058547e1c9111270b39a1d6b92e7d92b0861bdd3 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 22:22:21 +0000 Subject: [PATCH 12/13] Charge head and caller memory per logical row; clamp sharing to 8x Packed pricing now assumes an explicit, agreed invariant: model memory scales with packed rows, and head plus caller loss memory peaks at no more than 6 KiB per logical token. The estimate is B*max(P, L/(8r)) + 6 KiB*L plus GDN segment states, so the logical charge no longer vanishes at the profile's own density, and sharing beyond 8x the observed ratio is priced as if it were 8x. Single-target requests under 64 tokens keep the logical extrapolation. The rate floor is gone: it could switch the charge off where sharing is highest. CUDA tests run the real head with 061's CISPO loss and a fixed per-request caller, and check a 4,096 -> 64 token calibration transition against the estimator. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 89 ++++++------ tests/unit/test_trainer_rank_active_memory.py | 118 ++++++++++++--- tests/unit/test_trainer_rank_head_memory.py | 10 +- .../unit/test_trainer_rank_head_recompute.py | 135 ++++++++++++++---- tests/unit/test_trainer_rank_moe_memory.py | 2 +- tests/unit/test_trainer_rank_split.py | 21 +-- 6 files changed, 273 insertions(+), 102 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 9eaafc77b..53059f202 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -2709,6 +2709,16 @@ def forward_micro_batches( no_grad: bool | None = None, yield_empty: bool = False, ) -> Iterator[MicroBatch[ForwardInputs, ForwardOutputs]]: + """Yield admitted micro-batches; the caller runs its loss and backward. + + Admission learns each call's whole peak, including the caller's loss + and backward. For grad-enabled single-target requests of at least 64 + tokens under one-layer full recompute, it prices model memory by packed + rows and charges 6 KiB per logical token for the head plus the caller's + loss saves and backward transients. A caller whose per-token head and + loss memory peaks above that is unsupported: shared plans can exceed + their estimate. + """ if not isinstance(yield_empty, bool): raise TypeError("yield_empty must be a bool") enabled = torch.is_grad_enabled() if no_grad is None else not no_grad @@ -4055,12 +4065,11 @@ def _retained_memory_bytes( if ratio > profile.logical_per_packed * _MEMORY_PROFILE_TRUST_GROWTH: return required rate = profile.retained_compute_bytes_per_token - if _packed_priced( - signature, profile, self._one_layer_recompute() - ) and rate >= _packed_rate_floor(profile): - # Saved head indices and masks stay live until backward. - retained_compute = rate * packed_tokens + _packed_row_bytes( - profile, packed_tokens, logical_tokens + if _packed_priced(signature, self._one_layer_recompute()): + # Saved head indices, masks and caller saves stay live until + # backward. The ratio window above bounds the sharing extrapolation. + retained_compute = ( + rate * packed_tokens + _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens ) else: retained_compute = rate * max( @@ -7771,18 +7780,23 @@ def _estimate_required_memory_bytes_from_values( # Local head results coexist with full CP outputs during gathering. # Uneven rank plans can assign all of an item's rows to one rank. static_compute += output_bytes - # Outputs that grow with logical rows make a profile learned under + # Memory that grows with logical rows makes a profile learned under # lighter sharing underestimate a deeper-shared plan; scale the trusted - # estimate up by the ratio gap for them. Normalize before multiplying: - # cancelling packed tokens through two float operations can otherwise - # make a larger warm layout cheaper. + # estimate up by the ratio gap. Packed pricing instead charges head and + # caller memory per logical row and extrapolates the rest at most the + # usual trust growth. Normalize before multiplying: cancelling packed + # tokens through two float operations can otherwise make a larger warm + # layout cheaper. profiled_tokens: int | float = packed_tokens packed_priced = profiled is not None and _packed_priced( - signature, profiled, self._one_layer_recompute() + signature, self._one_layer_recompute() ) - if profiled is not None and logical_tokens is not None and not packed_priced: + if profiled is not None and logical_tokens is not None: profiled_tokens = max( - packed_tokens, logical_tokens / profiled.logical_per_packed + packed_tokens, + logical_tokens + / profiled.logical_per_packed + / (_MEMORY_PROFILE_TRUST_GROWTH if packed_priced else 1), ) # The trust window limits calibration growth, not the empirical floor. # Dropping that floor beyond the window can admit a larger request that @@ -7795,7 +7809,7 @@ def _estimate_required_memory_bytes_from_values( int( profiled.bytes_per_token * profiled_tokens + ( - _packed_row_bytes(profiled, packed_tokens, logical_tokens) + _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens # Branch states grow with segments, not packed rows; # backward recomputes one layer at a time. + gdn_segments * self._gdn_segment_layer_bytes() @@ -8763,26 +8777,19 @@ def _active_logical_tokens(requests: Sequence[AnyForwardInput]) -> int: # sharing ratios. Wide labels, dense or top-k outputs and no_grad forwards (whose # GDN branch states are uncharged) keep the logical/packed ratio extrapolation. _PACKED_PRICED_MIXES = frozenset({"target:single", "inactive"}) -# The head's label copies, positions, row-match vectors and saved masks (about -# 80-100 B) grow with logical rows. Each request's buffers are separate -# allocations rounded up to 512 B blocks: on an H200, a fully shared one-token -# request peaked at nine blocks (4,616 B) at the end of head forward. Under -# packed pricing, charge each row beyond the profile's observed sharing twelve. +# Packed pricing charges head and caller memory per active logical row: label +# copies, positions, row-match vectors, saved masks and the caller's loss saves +# and backward transients. Each request's buffers are separate allocations +# rounded up to 512 B blocks; on an H200, a fully shared one-token request's +# head peaked at nine blocks plus 8 B (4,616 B). Callers whose head and loss +# memory peaks above this charge per token are unsupported. _PACKED_PRICED_LOGICAL_ROW_BYTES = 12 * 512 +# Shorter single-target requests keep the logical extrapolation, so each +# packed-priced request brings at least 384 KiB for per-request constants. +_PACKED_PRICED_MIN_REQUEST_TOKENS = 64 -def _packed_rate_floor(profile: "_MemoryProfile") -> float: - """A per-packed rate at least twice the per-row charges it offsets. - - Fewer packed rows then always cost less, with margin for rounding, so - lower-bound pruning over more-shared layouts stays sound. - """ - return 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * profile.logical_per_packed - - -def _packed_priced( - signature: "_MemorySignature", profile: "_MemoryProfile", one_layer_recompute: bool -) -> bool: +def _packed_priced(signature: "_MemorySignature", one_layer_recompute: bool) -> bool: # Measured only under one-layer full recompute: other recompute modes keep # more GDN states and activations live per segment and logical row. return ( @@ -8790,16 +8797,6 @@ def _packed_priced( and bool(signature.grad_modes) and all(signature.grad_modes) and _PACKED_PRICED_MIXES.issuperset(signature.request_mix) - and profile.bytes_per_token >= _packed_rate_floor(profile) - ) - - -def _packed_row_bytes( - profile: "_MemoryProfile", packed_tokens: int, logical_tokens: int -) -> float: - """Head buffers for logical rows beyond the profile's observed sharing.""" - return _PACKED_PRICED_LOGICAL_ROW_BYTES * max( - 0, logical_tokens - packed_tokens * profile.logical_per_packed ) @@ -8821,7 +8818,15 @@ def _request_mix_key(request: AnyForwardInput) -> str: parts.append("logits") if request.hidden_states: parts.append("hidden") - return "+".join(parts) if parts else "inactive" + key = "+".join(parts) if parts else "inactive" + if ( + key == "target:single" + and int(request.input_tokens.numel()) < _PACKED_PRICED_MIN_REQUEST_TOKENS + ): + # Keep short requests out of packed pricing: a duplicate adds no packed + # row, and caller memory per request can outgrow its few row charges. + key += "+short" + return key def _pad_packed_batch( diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 3e8717730..1595b7921 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -14,6 +14,7 @@ TrainerRank, TrainerRankMemoryError, Unset, + _impl, ) from art.trainer_rank._impl import ( _PACKED_PRICED_LOGICAL_ROW_BYTES, @@ -22,6 +23,12 @@ ) +@pytest.fixture(autouse=True) +def _price_short_requests(monkeypatch): + # These fixtures use short requests; the short-request gate has its own tests. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) + + class _Model(torch.nn.Module): def __init__(self): super().__init__() @@ -277,9 +284,11 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): for count in (8, 63, 64, 65, 800) ] assert values == sorted(values) - rate = rank._memory_profiles[observed.signature].bytes_per_token - row = _PACKED_PRICED_LOGICAL_ROW_BYTES * 800 * (logical_ratio - 1) - assert values[-1] == int((800 * 4 + rate * 800 + row) * 1.1) + profile = rank._memory_profiles[observed.signature] + logical = 800 * logical_ratio + packed = max(800, logical / profile.logical_per_packed / 8) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * logical + assert values[-1] == int((800 * 4 + profile.bytes_per_token * packed + rows) * 1.1) assert not rank._all_ranks_have_memory_profile( packed_tokens=800, signature=observed.signature ) @@ -288,11 +297,12 @@ def test_empirical_estimate_survives_packed_trust_boundary(logical_ratio): def test_packed_pricing_covers_allocator_blocks_of_fully_shared_requests(): # A duplicate one-token request adds no packed row but still allocates its # head buffers, each rounded up to a 512 B block: nine at the measured peak. + # (The short-request gate keeps such requests out; the charge covers them.) rank = _rank() observed = rank._plan_flat_forward(_requests("target_tokens")) rank._update_memory_profile(observed, 100_000, retained_bytes=1000) profile = rank._memory_profiles[observed.signature] - assert _packed_priced(observed.signature, profile, rank._one_layer_recompute()) + assert _packed_priced(observed.signature, rank._one_layer_recompute()) duplicates = 100_000 def estimate(logical_tokens: int) -> int: @@ -307,6 +317,82 @@ def estimate(logical_tokens: int) -> int: assert estimate(base + duplicates) - estimate(base) >= duplicates * 9 * 512 +@pytest.mark.parametrize("shape", ["flat", "row"]) +def test_short_single_target_requests_keep_logical_pricing(monkeypatch, shape): + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 64) + + def request(length: int) -> ForwardInput: + tokens = torch.arange(length) + if shape == "row": + tokens = tokens[None] + return ForwardInput(input_tokens=tokens, target_tokens=tokens + 1) + + assert _request_mix_key(request(63)) == "target:single+short" + assert _request_mix_key(request(64)) == _request_mix_key(request(65)) + assert _request_mix_key(request(64)) == "target:single" + assert _request_mix_key(ForwardInput(input_tokens=torch.arange(3))) == "inactive" + rank = _rank() + long_plan = rank._plan_flat_forward([request(64), request(65)]) + mixed_plan = rank._plan_flat_forward([request(64), request(63)]) + recompute = rank._one_layer_recompute() + assert _packed_priced(long_plan.signature, recompute) + # One short request keeps the whole batch on the logical extrapolation. + assert not _packed_priced(mixed_plan.signature, recompute) + for plan in (long_plan, mixed_plan): + rank._memory_profiles[plan.signature] = _impl._MemoryProfile( + bytes_per_token=100_000, packed_tokens=1000 + ) + + def estimate(plan) -> int: + return rank._estimate_required_memory_bytes_from_values( + packed_tokens=10, + logical_tokens=1000, + output_bytes=0, + signature=plan.signature, + ) + + # Ratio-1 profiles: 1000 logical rows extrapolate to 1000 packed rows, or + # are clamped to 1000 / 8 = 125 rows plus the per-row charge. + assert estimate(mixed_plan) == int(100_000 * 1000 * 1.1) + assert estimate(long_plan) == int( + (100_000 * 125 + _PACKED_PRICED_LOGICAL_ROW_BYTES * 1000) * 1.1 + ) + + +def test_packed_sharing_clamp_is_monotone_and_learning_sharing_never_cheapens(): + rank = _rank() + observed = rank._plan_flat_forward(_requests("target_tokens")) + rank._update_memory_profile(observed, 100_000, retained_bytes=1000) + single = observed.signature + rank._memory_profiles[single] = replace( + rank._memory_profiles[single], bytes_per_token=50_000, logical_per_packed=1 + ) + + def cost(packed: int, logical: int = 8_000): + return rank._subforward_cost( + packed_tokens=packed, + logical_tokens=logical, + output_bytes=0, + signature=single, + ).required + + # Fixed logical rows; packed rows sweep across L / 8r = 1000. + sweep = [cost(packed) for packed in (1, 500, 999, 1000, 1001, 4000, 8000)] + assert sweep == sorted(sweep) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * 8_000 + assert sweep[0] == sweep[3] == int((50_000 * 1000 + rows) * 1.1) + assert sweep[4] == int((50_000 * 1001 + rows) * 1.1) + # Learning more sharing at a lower rate max-merges both; plans at + # or below the older ratio never get cheaper. + before = [cost(packed, logical) for packed, logical in ((100, 100), (50, 100))] + wider = replace(observed, packed_tokens=1, logical_tokens=8) + rank._update_memory_profile(wider, 1_000, retained_bytes=None) + profile = rank._memory_profiles[single] + assert profile.logical_per_packed > 1 and profile.bytes_per_token == 50_000 + after = [cost(packed, logical) for packed, logical in ((100, 100), (50, 100))] + assert all(b >= a for a, b in zip(before, after, strict=True)) + + @pytest.mark.parametrize( "method,argument,fallback", [ @@ -377,8 +463,8 @@ def estimate(name): signature=signatures[name], ) - # Packed rows plus head buffers for logical rows beyond the profile's 2x. - rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * (64 - 8 * 2) + # Packed rows plus head and caller memory for every logical row. + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * 64 assert ( estimate("single") == estimate("multi_grad") == int((100_000 * 8 + rows) * 1.1) ) @@ -389,19 +475,15 @@ def estimate(name): # Others keep the logical/packed extrapolation: 64 / 2 profiled rows. for name in ("no_grad", "mixed_grad", "hidden", "wide"): assert estimate(name) == int(100_000 * 32 * 1.1) - # Per-row charges must stay well below the profiled rates, or a more-shared - # layout could cost more: under two row charges per ratio, both extrapolate. - floor = 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 2 - rank._memory_profiles[single] = replace( - profile, retained_compute_bytes_per_token=floor - 1 + # Beyond 8x the profile's sharing, packed rows are priced as if sharing + # were 8x: 64 logical rows at ratio 2 x 8 = 16 is 4 rows, not 1. + clamped = rank._estimate_required_memory_bytes_from_values( + packed_tokens=1, logical_tokens=64, output_bytes=0, signature=single ) - assert rank._retained_memory_bytes( - single, packed_tokens=8, logical_tokens=64, output_bytes=0, required=1 << 40 - ) == int((floor - 1) * 32 * 1.1) - rank._memory_profiles[single] = replace(profile, bytes_per_token=floor - 1) - assert estimate("single") == int((floor - 1) * 32 * 1.1) - # At and above the floor, cost never falls as packed rows grow. - for rate in (floor, floor + 1, 100_000): + assert clamped == int((100_000 * 4 + rows) * 1.1) + # The logical charge does not depend on layout, so at any rate (none gates + # eligibility) cost never falls as packed rows grow. + for rate in (1, 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES * 2 - 1, 100_000): rank._memory_profiles[single] = replace( profile, bytes_per_token=rate, retained_compute_bytes_per_token=rate ) diff --git a/tests/unit/test_trainer_rank_head_memory.py b/tests/unit/test_trainer_rank_head_memory.py index 50c16343f..21b7ad071 100644 --- a/tests/unit/test_trainer_rank_head_memory.py +++ b/tests/unit/test_trainer_rank_head_memory.py @@ -8,7 +8,11 @@ import torch from art.trainer_rank import ForwardInput -from art.trainer_rank._impl import Unset, _MemoryProfile +from art.trainer_rank._impl import ( + _PACKED_PRICED_LOGICAL_ROW_BYTES, + Unset, + _MemoryProfile, +) def rank(): @@ -151,7 +155,9 @@ def test_outputs_retention_and_empirical_peak_are_counted_once(): retained_compute_bytes_per_token=1, ) cost = r._plan_cost(plan) - assert cost.required == int((plan.output_bytes + 512 * 2_000_000) * 1.1) + # Packed pricing adds head and caller memory for every logical row. + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * 512 + assert cost.required == int((plan.output_bytes + 512 * 2_000_000 + rows) * 1.1) assert cost.retained == int((plan.output_bytes + retained) * 1.1) diff --git a/tests/unit/test_trainer_rank_head_recompute.py b/tests/unit/test_trainer_rank_head_recompute.py index 207b4f4c0..3e40d2fa8 100644 --- a/tests/unit/test_trainer_rank_head_recompute.py +++ b/tests/unit/test_trainer_rank_head_recompute.py @@ -1,6 +1,7 @@ from __future__ import annotations from contextlib import nullcontext +from dataclasses import replace from datetime import timedelta import sys from types import SimpleNamespace @@ -438,13 +439,31 @@ def exp(subtraction): assert torch.cuda.is_initialized() == cuda_initialized -def test_shared_one_token_requests_fit_the_packed_row_charge(monkeypatch): - # A duplicate one-token request adds a logical row but no packed row. Its - # labels, indices, masks and output are separate CUDA allocations rounded - # to 512 B blocks; through backward they must fit the per-row charge. - if not torch.cuda.is_available(): - pytest.skip("allocator rounding needs CUDA") +_PRODUCTION_HEAD_CHUNK_TOKENS = _impl._HEAD_CHUNK_TOKENS + + +def _cispo(output, device): + # 061's caller: per-history sampled mask, detached clipped ratio. + sampled = torch.ones(output.target_logprobs.shape, dtype=torch.bool, device=device) + logp = output.target_logprobs[sampled] + ratio = (logp.detach() - 0.5).exp().clamp(max=5.0) + return -(ratio * logp).sum() + + +_CALLER_LOSSES = { + "sum": lambda output, device, refs: -output.target_logprobs.sum(), + "cispo": lambda output, device, refs: _cispo(output, device), + # Saves a fixed 256 KiB per request, whatever its length. + "fixed": lambda output, device, refs: ( + (output.target_logprobs.sum() - refs).square().mean() + ), +} + + +def _duplicate_request_peak(monkeypatch, count: int, length: int, loss: str) -> int: + """Peak bytes through the real head, a caller loss and backward.""" _patch_local_head(monkeypatch) + monkeypatch.setattr(_impl, "_HEAD_CHUNK_TOKENS", _PRODUCTION_HEAD_CHUNK_TOKENS) device = torch.device("cuda") weight = torch.randn(1024, 64, device=device, dtype=torch.bfloat16) / 8 model = SimpleNamespace( @@ -455,29 +474,85 @@ def test_shared_one_token_requests_fit_the_packed_row_charge(monkeypatch): ) r = object.__new__(TrainerRank) r.runtime = SimpleNamespace(model=[model]) - request = ForwardInput( - input_tokens=torch.tensor([7]), target_tokens=torch.tensor([5]) + tokens = torch.arange(length) + request = ForwardInput(input_tokens=tokens, target_tokens=tokens + 1) + items = [r._forward_item(request) for _ in range(count)] + prepared = SimpleNamespace( + positions_by_item=tuple(torch.arange(length) for _ in range(count)), + source_positions_by_item=tuple(torch.arange(length) for _ in range(count)), ) + hidden = torch.randn(length, 64, device=device, dtype=torch.bfloat16) + hidden.requires_grad_() + refs = torch.zeros(64 * 1024, device=device) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + total = torch.zeros((), device=device) + for output in r._project_head(items, prepared, hidden): + total = total + _CALLER_LOSSES[loss](output, device, refs) + total.backward() + torch.cuda.synchronize() + assert hidden.grad is not None + return torch.cuda.max_memory_allocated() - base - def peak(count: int) -> int: - items = [r._forward_item(request) for _ in range(count)] - prepared = SimpleNamespace( - positions_by_item=tuple(torch.tensor([0]) for _ in range(count)), - source_positions_by_item=tuple(torch.arange(1) for _ in range(count)), - ) - hidden = torch.randn(1, 64, device=device, dtype=torch.bfloat16) - hidden.requires_grad_() - torch.cuda.synchronize() - torch.cuda.reset_peak_memory_stats() - base = torch.cuda.memory_allocated() - loss = torch.zeros((), device=device) - for output in r._project_head(items, prepared, hidden): - loss = loss - output.target_logprobs.sum() - loss.backward() - torch.cuda.synchronize() - return torch.cuda.max_memory_allocated() - base - - peak(1) - extra = 20_000 - per_request = (peak(1 + extra) - peak(1)) / extra - assert per_request <= _impl._PACKED_PRICED_LOGICAL_ROW_BYTES + +@pytest.mark.parametrize( + ("length", "loss", "fits"), + [ + (1, "sum", True), + (64, "cispo", True), + (512, "cispo", True), + (64, "fixed", True), + # Why short requests keep logical pricing: per-request caller memory + # outgrows a one-token request's row charge. + (1, "fixed", False), + ], +) +def test_shared_requests_fit_the_per_row_charge(monkeypatch, length, loss, fits): + # A duplicate request adds logical rows but no packed rows. Its head buffers + # and the caller's saves are separate CUDA allocations rounded to 512 B + # blocks; through backward they must fit the charge on its logical rows. + if not torch.cuda.is_available(): + pytest.skip("allocator rounding needs CUDA") + _duplicate_request_peak(monkeypatch, 1, length, loss) + extra = 20_000 if length == 1 and loss == "sum" else 1_000 + single = _duplicate_request_peak(monkeypatch, 1, length, loss) + per_request = ( + _duplicate_request_peak(monkeypatch, 1 + extra, length, loss) - single + ) / extra + # Nine head blocks per one-token request at least; the rest is the caller. + assert per_request >= 9 * 512 + charge = _impl._PACKED_PRICED_LOGICAL_ROW_BYTES * length + assert (per_request <= charge) == fits + + +def test_calibrated_profile_admits_shorter_shared_requests_above_their_peak( + monkeypatch, +): + # A profile learned from 16 shared 4,096-token requests prices 20 shared + # 64-token requests (sharing 20 against 16) whose caller saves a fixed + # 256 KiB each: 64 B per calibration token but 4 KiB per new token. + if not torch.cuda.is_available(): + pytest.skip("allocator rounding needs CUDA") + from test_trainer_rank_active_memory import _rank + + rank = _rank() + plan = rank._plan_flat_forward( + [ForwardInput(input_tokens=torch.arange(64), target_tokens=torch.arange(64))] + ) + assert _impl._packed_priced(plan.signature, rank._one_layer_recompute()) + _duplicate_request_peak(monkeypatch, 1, 64, "fixed") + calibration = _duplicate_request_peak(monkeypatch, 16, 4096, "fixed") + observed = replace( + plan, packed_tokens=4096, logical_tokens=16 * 4096, output_bytes=16 * 4096 * 4 + ) + rank._update_memory_profile(observed, calibration, retained_bytes=None) + assert rank._memory_profiles[plan.signature].logical_per_packed == 16 + peak = _duplicate_request_peak(monkeypatch, 20, 64, "fixed") + estimate = rank._estimate_required_memory_bytes_from_values( + packed_tokens=64, + logical_tokens=20 * 64, + output_bytes=20 * 64 * 4, + signature=plan.signature, + ) + assert peak <= estimate diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py index 0513209e5..33b8ce9df 100644 --- a/tests/unit/test_trainer_rank_moe_memory.py +++ b/tests/unit/test_trainer_rank_moe_memory.py @@ -293,7 +293,7 @@ def test_profiles_outputs_and_empty_plan_preserve_empirical_floor(): ) assert estimate( packed_tokens=100, logical_tokens=200, output_bytes=123, signature=signature - ) == int((100 * 100000 + _PACKED_PRICED_LOGICAL_ROW_BYTES * 100 + 123) * 1.1) + ) == int((100 * 100000 + _PACKED_PRICED_LOGICAL_ROW_BYTES * 200 + 123) * 1.1) def test_summed_group_envelope_and_retained_profile_unchanged(): diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 88eb3c122..394cc3b48 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -49,6 +49,7 @@ TrainerRankMemoryError, TrainerRankPartialExecutionError, TrainerRankSlotStateError, + _impl, ) from art.trainer_rank._impl import ( _PACKED_PRICED_LOGICAL_ROW_BYTES, @@ -536,8 +537,13 @@ def test_changed_output_allocation_does_not_inflate_retained_compute() -> None: [(800, 800, True), (801, 801, False), (100, 800, True), (100, 801, False)], ) def test_retained_compute_keeps_growth_and_sharing_trust_limits( - packed_tokens: int, logical_tokens: int, trusted: bool + monkeypatch: pytest.MonkeyPatch, + packed_tokens: int, + logical_tokens: int, + trusted: bool, ) -> None: + # Price the short test request as a packed-priced one. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = TrainerRank(_runtime()) plan = replace( rank._plan_flat_forward([_request(0)]), @@ -554,14 +560,13 @@ def test_retained_compute_keeps_growth_and_sharing_trust_limits( rank._update_memory_profile(plan, 100_000, retained_bytes=None) unknown = rank._plan_cost(candidate) assert unknown.retained == unknown.required - # A retained rate just above the floor of two row charges keeps packed pricing. - rate = 2 * _PACKED_PRICED_LOGICAL_ROW_BYTES + 4 + rate = 260 rank._update_memory_profile( plan, 40_000 + 200 * rate, retained_bytes=40_000 + 100 * rate ) observed = rank._plan_cost(candidate) if trusted: - rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * (logical_tokens - packed_tokens) + rows = _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens assert observed.retained == int((40_000 + rate * packed_tokens + rows) * 1.1) assert observed.retained < observed.required else: @@ -847,11 +852,7 @@ def test_retained_ratio_original_cost_witness_stays_conservative( ] # Exact retention keeps its conservative fallback. Only treating that # fallback as an optimistic split-search bound was incorrect. - rows = _PACKED_PRICED_LOGICAL_ROW_BYTES - assert (small.required, large.required) == tuple( - int((4 * 131_072 * packed + rows * (64_000 - packed) + 10_000) * 1.1) - for packed in (4000, 8000) - ) + assert small.required == large.required == 36_909_886_200 assert small.retained == small.required and large.retained == 11_000 @@ -1018,6 +1019,8 @@ def loss(indices: tuple[int, ...]) -> torch.Tensor: def test_retained_ratio_bound_uses_original_guard_at_trusted_endpoint( monkeypatch: pytest.MonkeyPatch, direction: float | None ) -> None: + # Short and long requests share one signature here. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) def request(tokens: list[int]) -> ForwardInput: From ccca0dc7661fc1e8531ee841ec05fff976f5f29b Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 24 Sep 2026 22:43:59 +0000 Subject: [PATCH 13/13] Keep short-request batches on the calibrated profile The short-request gate was a request-mix key, which gave any batch with a short request its own profile signature: it started cold on the static estimate instead of main's calibrated extrapolation. Carry the gate as a signature flag excluded from identity, so those batches share the profile and price exactly as on main. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/art/trainer_rank/_impl.py | 23 +++++---- tests/unit/test_trainer_rank_active_memory.py | 48 ++++++++----------- tests/unit/test_trainer_rank_split.py | 10 ++-- 3 files changed, 43 insertions(+), 38 deletions(-) diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py index 53059f202..5f0006c4a 100644 --- a/src/art/trainer_rank/_impl.py +++ b/src/art/trainer_rank/_impl.py @@ -941,6 +941,9 @@ class _MemorySignature: grad_enabled: bool grad_modes: tuple[bool, ...] slot_shapes: tuple[tuple[bool, tuple[tuple[int, ...], ...]], ...] = () + # Short single-target requests keep the logical extrapolation, but share + # the profile learned from longer requests of the same signature. + short_requests: bool = dataclass_field(default=False, compare=False) @dataclass(frozen=True) @@ -3352,6 +3355,7 @@ def feed(value: Any) -> None: signature.grad_enabled, signature.grad_modes, signature.slot_shapes, + signature.short_requests, p.packed_tokens, p.logical_tokens, p.inactive_logical_tokens, @@ -6745,6 +6749,7 @@ def _memory_signature_from_requests( grad_enabled=any(modes), grad_modes=modes, slot_shapes=shapes if any(any(shape) for _, shape in shapes) else (), + short_requests=any(_short_request(request) for request in requests), ) def _slot_memory_shapes( @@ -8797,6 +8802,7 @@ def _packed_priced(signature: "_MemorySignature", one_layer_recompute: bool) -> and bool(signature.grad_modes) and all(signature.grad_modes) and _PACKED_PRICED_MIXES.issuperset(signature.request_mix) + and not signature.short_requests ) @@ -8818,15 +8824,16 @@ def _request_mix_key(request: AnyForwardInput) -> str: parts.append("logits") if request.hidden_states: parts.append("hidden") - key = "+".join(parts) if parts else "inactive" - if ( - key == "target:single" + return "+".join(parts) if parts else "inactive" + + +def _short_request(request: AnyForwardInput) -> bool: + """Too short for packed pricing: a duplicate adds no packed row, and caller + memory per request can outgrow its few logical-row charges.""" + return ( + _request_mix_key(request) == "target:single" and int(request.input_tokens.numel()) < _PACKED_PRICED_MIN_REQUEST_TOKENS - ): - # Keep short requests out of packed pricing: a duplicate adds no packed - # row, and caller memory per request can outgrow its few row charges. - key += "+short" - return key + ) def _pad_packed_batch( diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py index 1595b7921..a2033710d 100644 --- a/tests/unit/test_trainer_rank_active_memory.py +++ b/tests/unit/test_trainer_rank_active_memory.py @@ -327,36 +327,30 @@ def request(length: int) -> ForwardInput: tokens = tokens[None] return ForwardInput(input_tokens=tokens, target_tokens=tokens + 1) - assert _request_mix_key(request(63)) == "target:single+short" - assert _request_mix_key(request(64)) == _request_mix_key(request(65)) - assert _request_mix_key(request(64)) == "target:single" - assert _request_mix_key(ForwardInput(input_tokens=torch.arange(3))) == "inactive" + assert [_impl._short_request(request(n)) for n in (63, 64, 65)] == [ + True, + False, + False, + ] + assert not _impl._short_request(ForwardInput(input_tokens=torch.arange(3))) rank = _rank() - long_plan = rank._plan_flat_forward([request(64), request(65)]) - mixed_plan = rank._plan_flat_forward([request(64), request(63)]) recompute = rank._one_layer_recompute() - assert _packed_priced(long_plan.signature, recompute) - # One short request keeps the whole batch on the logical extrapolation. - assert not _packed_priced(mixed_plan.signature, recompute) - for plan in (long_plan, mixed_plan): - rank._memory_profiles[plan.signature] = _impl._MemoryProfile( - bytes_per_token=100_000, packed_tokens=1000 + # Calibrate on a 64-token request, then price a short one: the short batch + # shares the calibrated profile (no cold start) but keeps main's pricing. + long_plan = rank._plan_flat_forward([request(64)]) + rank._update_memory_profile(long_plan, 64 * 100_000, retained_bytes=None) + for requests in ([request(63)], [request(64), request(63)]): + short = rank._plan_flat_forward(requests) + assert short.signature.short_requests and short.signature == long_plan.signature + assert _packed_priced(long_plan.signature, recompute) + assert not _packed_priced(short.signature, recompute) + profile = rank._memory_profiles[short.signature] + cost = rank._plan_cost(short) + extrapolated = profile.bytes_per_token * max( + short.packed_tokens, + short.active_logical_tokens / profile.logical_per_packed, ) - - def estimate(plan) -> int: - return rank._estimate_required_memory_bytes_from_values( - packed_tokens=10, - logical_tokens=1000, - output_bytes=0, - signature=plan.signature, - ) - - # Ratio-1 profiles: 1000 logical rows extrapolate to 1000 packed rows, or - # are clamped to 1000 / 8 = 125 rows plus the per-row charge. - assert estimate(mixed_plan) == int(100_000 * 1000 * 1.1) - assert estimate(long_plan) == int( - (100_000 * 125 + _PACKED_PRICED_LOGICAL_ROW_BYTES * 1000) * 1.1 - ) + assert cost.required >= int((short.output_bytes + extrapolated) * 1.1) def test_packed_sharing_clamp_is_monotone_and_learning_sharing_never_cheapens(): diff --git a/tests/unit/test_trainer_rank_split.py b/tests/unit/test_trainer_rank_split.py index 394cc3b48..c9d5b3c9f 100644 --- a/tests/unit/test_trainer_rank_split.py +++ b/tests/unit/test_trainer_rank_split.py @@ -1019,7 +1019,7 @@ def loss(indices: tuple[int, ...]) -> torch.Tensor: def test_retained_ratio_bound_uses_original_guard_at_trusted_endpoint( monkeypatch: pytest.MonkeyPatch, direction: float | None ) -> None: - # Short and long requests share one signature here. + # Price the short test requests as packed-priced ones. monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) @@ -1095,6 +1095,8 @@ def select_leaf_sharing(input_ids, *, memory_minimal=False, grad_enabled=True): def test_warm_rounding_preserves_native_split_bound_and_exact_budget( monkeypatch: pytest.MonkeyPatch, retained: int | None, admit: bool ) -> None: + # Price the short test requests as packed-priced ones. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) tokens = torch.arange(3) part = [ForwardInput(input_tokens=tokens, target_tokens=tokens) for _ in range(5)] @@ -1136,13 +1138,15 @@ def test_warm_rounding_preserves_native_split_bound_and_exact_budget( def test_normalized_warm_profile_is_monotone_in_packed_tokens( monkeypatch: pytest.MonkeyPatch, retained: float | None ) -> None: + # Price the short test request as a packed-priced one. + monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1) rank = _retained_ratio_rank(monkeypatch) signature = rank._plan_flat_forward([_request(0)]).signature rank._memory_profiles[signature] = _MemoryProfile( 7864330 / 15, 1000, 15 / 13, retained_compute_bytes_per_token=retained ) - # All counts satisfy the retained guard; logical rows no longer scale the - # profile, so cost must grow monotonically with packed rows. Check each count. + # All counts satisfy the retained guard; the logical charge does not depend + # on layout, so cost must grow monotonically with packed rows. Check each. costs = [ rank._subforward_cost( packed_tokens=packed,