From a56535053d1016fd362942cb68d3c21d04c36be5 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 16:58:59 +0000
Subject: [PATCH 1/5] Price warm waves by the fit later plans measured, not the
 first plan's

A signature's memory profile max-merged every observed peak rate, so the
first executed plan's one-time costs (compilation, first-use workspaces)
priced every later wave. A small, cold first wave spreads those costs over
few tokens. On real Qwen3.6-35B-A3B (40 layers, CP2/EP2), the run's first
wave (15.6k tokens) peaked at 275 KB per packed token and later waves ran
206-247 KB; that profile bound 16 of 17 waves per run.

The profile now also keeps a warm fit over every plan after the first: the
max rate, the smallest size and the max sharing ratio. The first plan is
held by weak reference, so its caller-phase update stays provisional and a
later plan at its freed address is still warm. Waves at least as large as
the smallest later plan are priced at the lower of the warm fit and today's
fit; a rate learned under lighter sharing scales up for deeper-shared waves
as today's does. Smaller waves, and profiles without warm fields (replayed
reports), price as before.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 .github/workflows/prek.yml                    |   2 +
 src/art/trainer_rank/_impl.py                 |  60 ++++++-
 tests/unit/test_trainer_rank_active_memory.py |  15 +-
 tests/unit/test_trainer_rank_moe_memory.py    |   3 +-
 tests/unit/test_trainer_rank_profile_warm.py  | 157 ++++++++++++++++++
 5 files changed, 226 insertions(+), 11 deletions(-)
 create mode 100644 tests/unit/test_trainer_rank_profile_warm.py

diff --git a/.github/workflows/prek.yml b/.github/workflows/prek.yml
index d7991b079..18a49a302 100644
--- a/.github/workflows/prek.yml
+++ b/.github/workflows/prek.yml
@@ -230,6 +230,7 @@ jobs:
             tests/unit/test_trainer_rank_weird_shapes.py \
             tests/unit/test_trainer_rank_admission_inputs.py \
             tests/unit/test_trainer_rank_checkpoint_memory.py \
+            tests/unit/test_trainer_rank_profile_warm.py \
             tests/unit/test_trainer_rank_tp_floor.py \
             tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \
             tests/unit/test_trainer_rank_slot_memory.py \
@@ -275,6 +276,7 @@ jobs:
             --ignore=tests/unit/test_trainer_rank_weird_shapes.py \
             --ignore=tests/unit/test_trainer_rank_admission_inputs.py \
             --ignore=tests/unit/test_trainer_rank_checkpoint_memory.py \
+            --ignore=tests/unit/test_trainer_rank_profile_warm.py \
             --ignore=tests/unit/test_trainer_rank_tp_floor.py \
             --ignore=tests/unit/test_trainer_rank_checkpoint_gradient_memory.py \
             --ignore=tests/unit/test_trainer_rank_slot_memory.py \
diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index a7d67c072..016218b7f 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -614,6 +614,13 @@ class _MemoryProfile:
     # Separate the retained compute rate from the peak, which also learns
     # caller-owned backward workspace. Requested outputs are charged explicitly.
     retained_compute_bytes_per_token: float | None = None
+    # A signature's first executed plan also pays one-time costs (compilation,
+    # first-use workspaces), which a small first wave spreads over few tokens.
+    # The same fit over later plans only prices waves at least as large as the
+    # smallest of them; smaller waves keep the fit over every plan.
+    warm_bytes_per_token: float | None = None
+    warm_packed_tokens: int | None = None
+    warm_logical_per_packed: float | None = None
 
 
 @dataclass(frozen=True)
@@ -2031,6 +2038,10 @@ def memory_field(name: str, default: Any = None) -> Any:
         self._hybridep_rows_high_water = 0
         self._cache_recovery_state = _CacheRecoveryState()
         self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {}
+        # Each signature's first executed plan, while it is alive (in process).
+        self._profile_seed_plans: dict[
+            _MemorySignature, weakref.ReferenceType[_FlatForwardPlan]
+        ] = {}
         self._split_memory_floors: dict[bytes, int] = {}
         self._split_memory_floor_status = "not_observed"
         self._last_global_micro_batch_size: int | None = None
@@ -8051,27 +8062,46 @@ def _estimate_required_memory_bytes_from_values(
         # 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, self._one_layer_recompute()
         )
-        if profiled is not None and logical_tokens is not None:
-            profiled_tokens = max(
+
+        def profiled_tokens(logical_per_packed: float) -> int | float:
+            if logical_tokens is None:
+                return packed_tokens
+            return max(
                 packed_tokens,
                 logical_tokens
-                / profiled.logical_per_packed
+                / 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
         # was refused just inside it, even below a previously observed peak.
         if profiled is None:
             compute = static_compute
         else:
+            profiled_bytes = profiled.bytes_per_token * profiled_tokens(
+                profiled.logical_per_packed
+            )
+            if (
+                profiled.warm_bytes_per_token is not None
+                and profiled.warm_packed_tokens is not None
+                and profiled.warm_logical_per_packed is not None
+                and packed_tokens >= profiled.warm_packed_tokens
+            ):
+                # Later plans' own sharing: a rate learned under lighter
+                # sharing scales up for deeper-shared plans, as above.
+                profiled_bytes = min(
+                    profiled_bytes,
+                    profiled.warm_bytes_per_token
+                    * profiled_tokens(profiled.warm_logical_per_packed),
+                )
             compute = max(
                 static_compute,
                 int(
-                    profiled.bytes_per_token * profiled_tokens
+                    profiled_bytes
                     + (
                         _PACKED_PRICED_LOGICAL_ROW_BYTES * logical_tokens
                         # Branch states grow with segments, not packed rows;
@@ -8250,6 +8280,21 @@ def _update_memory_profile(
         compute_delta = max(0, peak_delta_bytes - plan.output_bytes)
         bytes_per_token = compute_delta / max(1, plan.packed_tokens)
         previous = self._memory_profiles.get(plan.signature)
+        logical_per_packed = plan.active_logical_tokens / max(1, plan.packed_tokens)
+        seeds = self.__dict__.setdefault("_profile_seed_plans", {})
+        if previous is None:
+            seeds[plan.signature] = weakref.ref(plan)
+        # The first plan's caller-phase update is still that plan. A weak
+        # reference, unlike ``id``, cannot match a later plan at a freed address.
+        seed = seeds.get(plan.signature)
+        warm = previous is not None and (seed is None or seed() is not plan)
+        warm_rate = None if previous is None else previous.warm_bytes_per_token
+        warm_tokens = None if previous is None else previous.warm_packed_tokens
+        warm_sharing = None if previous is None else previous.warm_logical_per_packed
+        if warm:
+            warm_rate = max(bytes_per_token, warm_rate or 0.0)
+            warm_tokens = min(plan.packed_tokens, warm_tokens or plan.packed_tokens)
+            warm_sharing = max(logical_per_packed, warm_sharing or 1.0)
         retained_fraction = None if previous is None else previous.retained_fraction
         retained_compute = (
             None if previous is None else previous.retained_compute_bytes_per_token
@@ -8278,11 +8323,14 @@ def _update_memory_profile(
                 0 if previous is None else previous.packed_tokens,
             ),
             logical_per_packed=max(
-                plan.active_logical_tokens / max(1, plan.packed_tokens),
+                logical_per_packed,
                 1.0 if previous is None else previous.logical_per_packed,
             ),
             retained_fraction=retained_fraction,
             retained_compute_bytes_per_token=retained_compute,
+            warm_bytes_per_token=warm_rate,
+            warm_packed_tokens=warm_tokens,
+            warm_logical_per_packed=warm_sharing,
         )
 
     def _forward_item(self, request: AnyForwardInput) -> _ForwardItem:
diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py
index a2033710d..f705147b7 100644
--- a/tests/unit/test_trainer_rank_active_memory.py
+++ b/tests/unit/test_trainer_rank_active_memory.py
@@ -107,7 +107,13 @@ def test_inactive_length_preserves_warm_cost_and_profile(monkeypatch, output, no
         )
         assert lower == rank._plan_cost(plan)
     rank._update_memory_profile(second, 10_000, retained_bytes=1000)
-    assert rank._memory_profiles[first.signature] == profile
+    # The second plan is warm; inactive rows change neither its rate nor sharing.
+    assert rank._memory_profiles[first.signature] == replace(
+        profile,
+        warm_bytes_per_token=profile.bytes_per_token,
+        warm_packed_tokens=second.packed_tokens,
+        warm_logical_per_packed=profile.logical_per_packed,
+    )
 
 
 def test_public_pair_avoids_inactive_only_split_and_keeps_total_telemetry(monkeypatch):
@@ -376,11 +382,12 @@ def cost(packed: int, logical: int = 8_000):
     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.
+    # Learning more sharing max-merges it; plans at or below the older ratio
+    # never get cheaper. (A later plan's lower rate may price larger plans;
+    # see test_trainer_rank_profile_warm.)
     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)
+    rank._update_memory_profile(wider, wider.output_bytes + 50_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))]
diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py
index 3fbdcee89..a31cd091d 100644
--- a/tests/unit/test_trainer_rank_moe_memory.py
+++ b/tests/unit/test_trainer_rank_moe_memory.py
@@ -103,8 +103,9 @@ def test_cp_memory_charges_local_and_gathered_outputs():
     compute = 16 * 2048 * 2 * 14
     assert estimate == int((compute + 2 * output_bytes) * 1.1)
     # A warm profile includes gather workspace already; do not add it twice.
+    # (A subclass: profiles hold a signature's first plan by weak reference.)
     rank._update_memory_profile(
-        SimpleNamespace(
+        type("Plan", (SimpleNamespace,), {})(
             signature=signature,
             packed_tokens=16,
             output_bytes=output_bytes,
diff --git a/tests/unit/test_trainer_rank_profile_warm.py b/tests/unit/test_trainer_rank_profile_warm.py
new file mode 100644
index 000000000..9bbb91e51
--- /dev/null
+++ b/tests/unit/test_trainer_rank_profile_warm.py
@@ -0,0 +1,157 @@
+"""A signature's first plan seeds its memory profile; later plans set a warm rate.
+
+CPU admission math, not a bound. On real Qwen3.6-35B-A3B (40 layers, CP2/EP2)
+the run's first, small, cold wave peaked at 275 KB per packed token while every
+later wave ran 206-239 KB; the max-merged rate priced every later wave with the
+first wave's one-time costs.
+"""
+
+from dataclasses import replace
+
+from test_trainer_rank_checkpoint_memory import rank, requests
+
+from art.trainer_rank._impl import _MemoryProfile
+
+FIRST, WARM = 3_000_000, 2_000_000  # Per packed token; far above the static floor.
+
+
+def _plans(r):
+    first, larger = (
+        r._plan_flat_forward(requests(1024, 64)),
+        r._plan_flat_forward(requests(8192, 64)),
+    )
+    assert first.signature == larger.signature
+    assert larger.packed_tokens > first.packed_tokens
+    return first, larger
+
+
+def _observe(r, plan, rate):
+    r._update_memory_profile(
+        plan, plan.output_bytes + rate * plan.packed_tokens, retained_bytes=None
+    )
+
+
+def _required(r, plan, logical_tokens=None):
+    return r._estimate_required_memory_bytes_from_values(
+        packed_tokens=plan.packed_tokens,
+        output_bytes=0,
+        signature=plan.signature,
+        logical_tokens=logical_tokens or plan.packed_tokens,
+    )
+
+
+def test_the_first_plans_one_time_costs_price_only_smaller_waves():
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    # The first plan's caller-phase update is still that plan: provisional.
+    _observe(r, first, FIRST)
+    profile = r._memory_profiles[first.signature]
+    assert profile.bytes_per_token == FIRST
+    assert profile.warm_bytes_per_token is None
+    assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
+    _observe(r, larger, WARM)
+    profile = r._memory_profiles[first.signature]
+    # Smaller waves keep today's max-merged rate; larger ones the warm rate.
+    assert profile.bytes_per_token == FIRST
+    assert profile.warm_bytes_per_token == WARM
+    assert profile.warm_packed_tokens == larger.packed_tokens
+    assert profile.warm_logical_per_packed == 1
+    assert _required(r, larger) == int(WARM * larger.packed_tokens * 1.1)
+    assert _required(r, first) == int(FIRST * first.packed_tokens * 1.1)
+
+
+def test_the_warm_rate_ratchets_and_its_extent_grows_down():
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    _observe(r, larger, WARM)
+    # A later, smaller plan: its higher rate is kept, and it extends the warm
+    # extent down to its own size.
+    middle = r._plan_flat_forward(requests(4096, 64))
+    _observe(r, middle, WARM + 100_000)
+    profile = r._memory_profiles[first.signature]
+    assert profile.warm_bytes_per_token == WARM + 100_000
+    assert profile.warm_packed_tokens == middle.packed_tokens
+    assert profile.bytes_per_token == FIRST
+    # A lower warm rate never lowers it.
+    _observe(r, larger, WARM - 100_000)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + 100_000
+
+
+def test_profiles_without_a_warm_rate_price_as_before():
+    """Replayed reports and callers that set profiles directly carry no warm rate."""
+    r = rank()
+    first, larger = _plans(r)
+    r._memory_profiles[first.signature] = _MemoryProfile(
+        bytes_per_token=FIRST, packed_tokens=first.packed_tokens
+    )
+    assert _required(r, first) == int(FIRST * first.packed_tokens * 1.1)
+    # A later plan's observation is warm: replay has no seed to exclude.
+    _observe(r, larger, WARM)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM
+
+
+def test_each_signature_and_a_cleared_profile_seeds_again():
+    """Profiles are keyed by memory signature; a new signature (a different
+    group count, slot shapes or topology) or a cleared profile starts cold."""
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    _observe(r, larger, WARM)
+    other = r._plan_flat_forward(requests(8192, 0)[:1])
+    assert other.signature != first.signature
+    _observe(r, other, FIRST)
+    assert r._memory_profiles[other.signature].warm_bytes_per_token is None
+    # Clearing a profile makes the next plan its first again.
+    del r._memory_profiles[first.signature]
+    _observe(r, larger, FIRST)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token is None
+    assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
+
+
+def test_a_later_plan_at_the_seeds_freed_address_is_warm():
+    """CPython reuses a freed plan's address for the next plan, so the seed is
+    held by weak reference rather than ``id``."""
+    r = rank()
+    _, larger = _plans(r)
+    seed = replace(larger)
+    _observe(r, seed, FIRST)
+    address = id(seed)
+    del seed
+    later = replace(larger)
+    assert id(later) == address
+    _observe(r, later, WARM)
+    assert r._memory_profiles[larger.signature].warm_bytes_per_token == WARM
+
+
+def test_a_warm_rate_learned_under_lighter_sharing_scales_for_deeper_sharing():
+    """A deeply shared first plan and unshared later plans: the later plans'
+    lower rate prices unshared waves, but deeper-shared waves scale it up by
+    the sharing gap, so they never price below today's fit over every plan.
+    A later plan that shares as deeply extends the warm rate to them."""
+    r = rank()
+    _, larger = _plans(r)
+    shared = replace(larger, logical_tokens=larger.packed_tokens * 64)
+    deep = shared.logical_tokens
+
+    def today():
+        profile = r._memory_profiles[larger.signature]
+        r._memory_profiles[larger.signature] = replace(
+            profile, warm_bytes_per_token=None, warm_packed_tokens=None
+        )
+        prices = _required(r, larger), _required(r, larger, deep)
+        r._memory_profiles[larger.signature] = profile
+        return prices
+
+    _observe(r, shared, FIRST)
+    _observe(r, larger, WARM)
+    profile = r._memory_profiles[larger.signature]
+    assert profile.logical_per_packed == 64
+    assert profile.warm_logical_per_packed == 1
+    unshared, deeper = today()
+    assert _required(r, larger) == int(WARM * larger.packed_tokens * 1.1) < unshared
+    assert _required(r, larger, deep) == deeper
+    _observe(r, replace(shared), WARM)
+    assert r._memory_profiles[larger.signature].warm_logical_per_packed == 64
+    assert _required(r, larger, deep) < today()[1]

From 107f9f53401e4f4eb4265d962dcc480551f33ccd Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 17:48:43 +0000
Subject: [PATCH 2/5] Fit the warm profile on caller-phase peaks and keep cost
 monotone

Review found two problems with the warm fit.

The size gate made cost drop at the smallest warm plan's size. The width
search accepts a width on the cheap no-sharing count and rejects on the
full-sharing count, so a shared layout below that size could execute priced
above the estimate that admitted it, and the lower bounds could refuse
feasible widths. A smaller wave is now priced by the warm fit as if it were
the smallest warm plan, and admission takes the lower of that and today's
fit. Cost is monotone in tokens again. This is sound where memory is a fixed
cost plus a per-token rate, since that plan's rate covers its share of the
fixed cost.

Split children and dp_rank_forward observe forward only, so they could set
a warm rate without the backward peak. Only a flat wave's caller-phase peak
now feeds the warm fit, and the profile's first such plan is the seed,
counted in the profile rather than held by weak reference.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/_impl.py                 |  52 ++--
 tests/unit/test_trainer_rank_active_memory.py |  15 +-
 tests/unit/test_trainer_rank_moe_memory.py    |   3 +-
 .../test_trainer_rank_physical_reserve.py     |   2 +-
 tests/unit/test_trainer_rank_profile_warm.py  | 224 +++++++++++++++---
 tests/unit/test_trainer_rank_validation.py    |  11 +-
 6 files changed, 239 insertions(+), 68 deletions(-)

diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index 016218b7f..89406e952 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -616,11 +616,13 @@ class _MemoryProfile:
     retained_compute_bytes_per_token: float | None = None
     # A signature's first executed plan also pays one-time costs (compilation,
     # first-use workspaces), which a small first wave spreads over few tokens.
-    # The same fit over later plans only prices waves at least as large as the
-    # smallest of them; smaller waves keep the fit over every plan.
+    # Admission uses the lower of the fit over every observation and the same
+    # fit over later plans' caller-phase peaks (which include backward), which
+    # prices a wave smaller than the smallest of them as if it were that large.
     warm_bytes_per_token: float | None = None
     warm_packed_tokens: int | None = None
     warm_logical_per_packed: float | None = None
+    caller_plans: int = 0
 
 
 @dataclass(frozen=True)
@@ -2038,10 +2040,6 @@ def memory_field(name: str, default: Any = None) -> Any:
         self._hybridep_rows_high_water = 0
         self._cache_recovery_state = _CacheRecoveryState()
         self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {}
-        # Each signature's first executed plan, while it is alive (in process).
-        self._profile_seed_plans: dict[
-            _MemorySignature, weakref.ReferenceType[_FlatForwardPlan]
-        ] = {}
         self._split_memory_floors: dict[bytes, int] = {}
         self._split_memory_floor_status = "not_observed"
         self._last_global_micro_batch_size: int | None = None
@@ -2968,7 +2966,9 @@ def _forward_micro_batches(
             # optimizer headroom. Peak only: the retained observation belongs
             # to the forward's return, already recorded for this same plan.
             if isinstance(candidate.plan, _FlatForwardPlan):
-                self._update_peak_memory_profile(candidate.plan, memory_baseline)
+                self._update_peak_memory_profile(
+                    candidate.plan, memory_baseline, caller_phase=True
+                )
             elif memory_baseline is not None:
                 self._record_split_memory_floor(
                     candidate.plan, memory_baseline, forward_peak
@@ -6288,6 +6288,8 @@ def _update_peak_memory_profile(
         plan: _FlatForwardPlan,
         baseline: int | None,
         retained_after: int | None = None,
+        *,
+        caller_phase: bool = False,
     ) -> None:
         if baseline is None:
             return
@@ -6301,6 +6303,7 @@ def _update_peak_memory_profile(
             retained_bytes=(
                 None if retained_after is None else max(0, retained_after - baseline)
             ),
+            caller_phase=caller_phase,
         )
 
     def _begin_planner_observation(
@@ -8089,14 +8092,20 @@ def profiled_tokens(logical_per_packed: float) -> int | float:
                 profiled.warm_bytes_per_token is not None
                 and profiled.warm_packed_tokens is not None
                 and profiled.warm_logical_per_packed is not None
-                and packed_tokens >= profiled.warm_packed_tokens
             ):
                 # Later plans' own sharing: a rate learned under lighter
-                # sharing scales up for deeper-shared plans, as above.
+                # sharing scales up for deeper-shared plans, as above. Pricing
+                # smaller waves as the smallest later plan keeps cost monotone
+                # in tokens, which the width search's bounds rely on; it is
+                # sound where a wave's memory is a fixed cost plus a per-token
+                # rate, as that plan's rate covers its share of the fixed cost.
                 profiled_bytes = min(
                     profiled_bytes,
                     profiled.warm_bytes_per_token
-                    * profiled_tokens(profiled.warm_logical_per_packed),
+                    * max(
+                        profiled.warm_packed_tokens,
+                        profiled_tokens(profiled.warm_logical_per_packed),
+                    ),
                 )
             compute = max(
                 static_compute,
@@ -8274,6 +8283,7 @@ def _update_memory_profile(
         peak_delta_bytes: int,
         *,
         retained_bytes: int | None,
+        caller_phase: bool = False,
     ) -> None:
         if plan.packed_tokens <= 0:
             return
@@ -8281,20 +8291,19 @@ def _update_memory_profile(
         bytes_per_token = compute_delta / max(1, plan.packed_tokens)
         previous = self._memory_profiles.get(plan.signature)
         logical_per_packed = plan.active_logical_tokens / max(1, plan.packed_tokens)
-        seeds = self.__dict__.setdefault("_profile_seed_plans", {})
-        if previous is None:
-            seeds[plan.signature] = weakref.ref(plan)
-        # The first plan's caller-phase update is still that plan. A weak
-        # reference, unlike ``id``, cannot match a later plan at a freed address.
-        seed = seeds.get(plan.signature)
-        warm = previous is not None and (seed is None or seed() is not plan)
         warm_rate = None if previous is None else previous.warm_bytes_per_token
         warm_tokens = None if previous is None else previous.warm_packed_tokens
         warm_sharing = None if previous is None else previous.warm_logical_per_packed
-        if warm:
-            warm_rate = max(bytes_per_token, warm_rate or 0.0)
-            warm_tokens = min(plan.packed_tokens, warm_tokens or plan.packed_tokens)
-            warm_sharing = max(logical_per_packed, warm_sharing or 1.0)
+        caller_plans = 0 if previous is None else previous.caller_plans
+        # Only a flat wave's caller-phase peak includes its backward; split
+        # children and dp_rank_forward observe forward only. The profile's
+        # first such plan also pays one-time costs, so it is not warm.
+        if caller_phase:
+            if caller_plans:
+                warm_rate = max(bytes_per_token, warm_rate or 0.0)
+                warm_tokens = min(plan.packed_tokens, warm_tokens or plan.packed_tokens)
+                warm_sharing = max(logical_per_packed, warm_sharing or 1.0)
+            caller_plans += 1
         retained_fraction = None if previous is None else previous.retained_fraction
         retained_compute = (
             None if previous is None else previous.retained_compute_bytes_per_token
@@ -8331,6 +8340,7 @@ def _update_memory_profile(
             warm_bytes_per_token=warm_rate,
             warm_packed_tokens=warm_tokens,
             warm_logical_per_packed=warm_sharing,
+            caller_plans=caller_plans,
         )
 
     def _forward_item(self, request: AnyForwardInput) -> _ForwardItem:
diff --git a/tests/unit/test_trainer_rank_active_memory.py b/tests/unit/test_trainer_rank_active_memory.py
index f705147b7..a2033710d 100644
--- a/tests/unit/test_trainer_rank_active_memory.py
+++ b/tests/unit/test_trainer_rank_active_memory.py
@@ -107,13 +107,7 @@ def test_inactive_length_preserves_warm_cost_and_profile(monkeypatch, output, no
         )
         assert lower == rank._plan_cost(plan)
     rank._update_memory_profile(second, 10_000, retained_bytes=1000)
-    # The second plan is warm; inactive rows change neither its rate nor sharing.
-    assert rank._memory_profiles[first.signature] == replace(
-        profile,
-        warm_bytes_per_token=profile.bytes_per_token,
-        warm_packed_tokens=second.packed_tokens,
-        warm_logical_per_packed=profile.logical_per_packed,
-    )
+    assert rank._memory_profiles[first.signature] == profile
 
 
 def test_public_pair_avoids_inactive_only_split_and_keeps_total_telemetry(monkeypatch):
@@ -382,12 +376,11 @@ def cost(packed: int, logical: int = 8_000):
     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 max-merges it; plans at or below the older ratio
-    # never get cheaper. (A later plan's lower rate may price larger plans;
-    # see test_trainer_rank_profile_warm.)
+    # 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, wider.output_bytes + 50_000, retained_bytes=None)
+    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))]
diff --git a/tests/unit/test_trainer_rank_moe_memory.py b/tests/unit/test_trainer_rank_moe_memory.py
index a31cd091d..3fbdcee89 100644
--- a/tests/unit/test_trainer_rank_moe_memory.py
+++ b/tests/unit/test_trainer_rank_moe_memory.py
@@ -103,9 +103,8 @@ def test_cp_memory_charges_local_and_gathered_outputs():
     compute = 16 * 2048 * 2 * 14
     assert estimate == int((compute + 2 * output_bytes) * 1.1)
     # A warm profile includes gather workspace already; do not add it twice.
-    # (A subclass: profiles hold a signature's first plan by weak reference.)
     rank._update_memory_profile(
-        type("Plan", (SimpleNamespace,), {})(
+        SimpleNamespace(
             signature=signature,
             packed_tokens=16,
             output_bytes=output_bytes,
diff --git a/tests/unit/test_trainer_rank_physical_reserve.py b/tests/unit/test_trainer_rank_physical_reserve.py
index bbf7ac3ef..b41a6e615 100644
--- a/tests/unit/test_trainer_rank_physical_reserve.py
+++ b/tests/unit/test_trainer_rank_physical_reserve.py
@@ -152,7 +152,7 @@ def forward(plan, **kwargs):
     monkeypatch.setattr(
         rank,
         "_update_peak_memory_profile",
-        lambda plan, baseline: profiles.append((plan, baseline)),
+        lambda plan, baseline, **_: profiles.append((plan, baseline)),
     )
     state = allocator(monkeypatch)
     original_phase = _impl._telemetry_phase
diff --git a/tests/unit/test_trainer_rank_profile_warm.py b/tests/unit/test_trainer_rank_profile_warm.py
index 9bbb91e51..95699ac3d 100644
--- a/tests/unit/test_trainer_rank_profile_warm.py
+++ b/tests/unit/test_trainer_rank_profile_warm.py
@@ -1,16 +1,20 @@
-"""A signature's first plan seeds its memory profile; later plans set a warm rate.
+"""A signature's first plan seeds its memory profile; later plans set a warm fit.
 
 CPU admission math, not a bound. On real Qwen3.6-35B-A3B (40 layers, CP2/EP2)
 the run's first, small, cold wave peaked at 275 KB per packed token while every
-later wave ran 206-239 KB; the max-merged rate priced every later wave with the
+later wave ran 206-247 KB; the max-merged rate priced every later wave with the
 first wave's one-time costs.
 """
 
 from dataclasses import replace
 
+from test_trainer_rank_active_memory import _rank as packed_rank
+from test_trainer_rank_active_memory import _requests as packed_requests
 from test_trainer_rank_checkpoint_memory import rank, requests
+import torch
 
-from art.trainer_rank._impl import _MemoryProfile
+from art.trainer_rank import _impl
+from art.trainer_rank._impl import _MemoryProfile, _packed_priced
 
 FIRST, WARM = 3_000_000, 2_000_000  # Per packed token; far above the static floor.
 
@@ -25,12 +29,21 @@ def _plans(r):
     return first, larger
 
 
-def _observe(r, plan, rate):
+def _update(r, plan, rate, *, caller_phase):
     r._update_memory_profile(
-        plan, plan.output_bytes + rate * plan.packed_tokens, retained_bytes=None
+        plan,
+        plan.output_bytes + rate * plan.packed_tokens,
+        retained_bytes=None,
+        caller_phase=caller_phase,
     )
 
 
+def _observe(r, plan, rate):
+    """One flat wave: its forward return, then its caller phase (backward)."""
+    _update(r, plan, rate, caller_phase=False)
+    _update(r, plan, rate, caller_phase=True)
+
+
 def _required(r, plan, logical_tokens=None):
     return r._estimate_required_memory_bytes_from_values(
         packed_tokens=plan.packed_tokens,
@@ -40,25 +53,42 @@ def _required(r, plan, logical_tokens=None):
     )
 
 
-def test_the_first_plans_one_time_costs_price_only_smaller_waves():
+def test_the_first_plans_one_time_costs_do_not_price_later_waves():
     r = rank()
     first, larger = _plans(r)
     _observe(r, first, FIRST)
-    # The first plan's caller-phase update is still that plan: provisional.
-    _observe(r, first, FIRST)
     profile = r._memory_profiles[first.signature]
     assert profile.bytes_per_token == FIRST
     assert profile.warm_bytes_per_token is None
+    assert profile.caller_plans == 1
     assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
     _observe(r, larger, WARM)
     profile = r._memory_profiles[first.signature]
-    # Smaller waves keep today's max-merged rate; larger ones the warm rate.
     assert profile.bytes_per_token == FIRST
     assert profile.warm_bytes_per_token == WARM
     assert profile.warm_packed_tokens == larger.packed_tokens
     assert profile.warm_logical_per_packed == 1
     assert _required(r, larger) == int(WARM * larger.packed_tokens * 1.1)
+    # Smaller waves: the lower of today's rate and the warm fit at the
+    # smallest later plan's size.
     assert _required(r, first) == int(FIRST * first.packed_tokens * 1.1)
+    near = r._plan_flat_forward(requests(6144, 64))
+    assert WARM * larger.packed_tokens < FIRST * near.packed_tokens
+    assert _required(r, near) == int(WARM * larger.packed_tokens * 1.1)
+
+
+def test_forward_only_observations_never_set_the_warm_fit():
+    """Split children and dp_rank_forward observe forward only; a flat wave's
+    caller phase also includes its backward."""
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    _update(r, larger, WARM, caller_phase=False)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token is None
+    assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
+    # A later caller phase is warm, and its peak includes the forward's.
+    _update(r, larger, WARM, caller_phase=True)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM
 
 
 def test_the_warm_rate_ratchets_and_its_extent_grows_down():
@@ -79,15 +109,17 @@ def test_the_warm_rate_ratchets_and_its_extent_grows_down():
     assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + 100_000
 
 
-def test_profiles_without_a_warm_rate_price_as_before():
-    """Replayed reports and callers that set profiles directly carry no warm rate."""
+def test_profiles_without_a_warm_fit_price_as_before():
+    """Replayed reports and callers that set profiles directly carry no warm
+    fit; the first caller phase after that seeds it."""
     r = rank()
     first, larger = _plans(r)
     r._memory_profiles[first.signature] = _MemoryProfile(
         bytes_per_token=FIRST, packed_tokens=first.packed_tokens
     )
     assert _required(r, first) == int(FIRST * first.packed_tokens * 1.1)
-    # A later plan's observation is warm: replay has no seed to exclude.
+    _observe(r, larger, WARM)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token is None
     _observe(r, larger, WARM)
     assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM
 
@@ -110,21 +142,6 @@ def test_each_signature_and_a_cleared_profile_seeds_again():
     assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
 
 
-def test_a_later_plan_at_the_seeds_freed_address_is_warm():
-    """CPython reuses a freed plan's address for the next plan, so the seed is
-    held by weak reference rather than ``id``."""
-    r = rank()
-    _, larger = _plans(r)
-    seed = replace(larger)
-    _observe(r, seed, FIRST)
-    address = id(seed)
-    del seed
-    later = replace(larger)
-    assert id(later) == address
-    _observe(r, later, WARM)
-    assert r._memory_profiles[larger.signature].warm_bytes_per_token == WARM
-
-
 def test_a_warm_rate_learned_under_lighter_sharing_scales_for_deeper_sharing():
     """A deeply shared first plan and unshared later plans: the later plans'
     lower rate prices unshared waves, but deeper-shared waves scale it up by
@@ -152,6 +169,157 @@ def today():
     unshared, deeper = today()
     assert _required(r, larger) == int(WARM * larger.packed_tokens * 1.1) < unshared
     assert _required(r, larger, deep) == deeper
-    _observe(r, replace(shared), WARM)
+    _observe(r, shared, WARM)
     assert r._memory_profiles[larger.signature].warm_logical_per_packed == 64
     assert _required(r, larger, deep) < today()[1]
+
+
+def test_cost_stays_monotone_across_the_warm_extent():
+    """The width search accepts on a larger layout's price and rejects on a
+    smaller one's, so a smaller wave must never price above a larger one."""
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    _observe(r, larger, WARM)
+    n = larger.packed_tokens
+
+    def required(packed, logical=None):
+        return r._estimate_required_memory_bytes_from_values(
+            packed_tokens=packed,
+            output_bytes=0,
+            signature=larger.signature,
+            logical_tokens=logical or packed,
+        )
+
+    sweep = [required(t) for t in (1, n // 2, (2 * n) // 3, n - 1, n, n + 1, 2 * n)]
+    assert sweep == sorted(sweep)
+    shared = [required(t, 4 * n) for t in (1, n // 4, n - 1, n, 2 * n, 4 * n)]
+    assert shared == sorted(shared)
+
+
+def test_packed_priced_cost_stays_monotone_across_the_warm_extent(monkeypatch):
+    # The fixture's requests are short; the short-request gate has its own tests.
+    monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1)
+    r = packed_rank()
+    signature = r._plan_flat_forward(packed_requests("target_tokens")).signature
+    assert _packed_priced(signature, r._one_layer_recompute())
+    today = _MemoryProfile(bytes_per_token=50_000, packed_tokens=8)
+    warm = replace(
+        today,
+        warm_bytes_per_token=20_000,
+        warm_packed_tokens=1000,
+        warm_logical_per_packed=4,
+    )
+
+    def cost(profile, packed, logical=8_000):
+        r._memory_profiles[signature] = profile
+        return r._subforward_cost(
+            packed_tokens=packed,
+            logical_tokens=logical,
+            output_bytes=0,
+            signature=signature,
+        ).required
+
+    sweep = [cost(warm, t) for t in (1, 500, 999, 1000, 1001, 2000, 8000)]
+    assert sweep == sorted(sweep)
+    # Just below the warm extent: priced as the smallest later plan, not today.
+    assert sweep[2] == sweep[3] < cost(today, 999)
+    rows = [cost(warm, 1000, logical) for logical in (1000, 8_000, 32_000, 64_000)]
+    assert rows == sorted(rows)
+
+
+def test_split_children_never_set_the_warm_fit(monkeypatch):
+    """The real split ladder: children observe forward only, and a split
+    wave's caller peak goes to split floors, never to the warm fit."""
+    from test_trainer_rank_split import _packed_budget, _request
+    from test_trainer_rank_split import _rank as split_rank
+
+    from art.trainer_rank import ForwardOutput
+
+    r = split_rank(monkeypatch)
+    monkeypatch.setattr(r, "_retained_memory_bytes", lambda *_args, **_kwargs: 0)
+    forward_rate, backward_rate = 100, 300
+    phases = []
+
+    def update(plan, baseline, retained_after=None, *, caller_phase=False):
+        # Stands in for the CUDA peak read: backward peaks above forward.
+        phases.append((plan.request_count, caller_phase))
+        rate = backward_rate if caller_phase else forward_rate
+        r._update_memory_profile(
+            plan,
+            plan.output_bytes + rate * plan.packed_tokens,
+            retained_bytes=None,
+            caller_phase=caller_phase,
+        )
+
+    def run(plan, **_kwargs):
+        # Production's forward-return update, then no CUDA baseline.
+        r._update_peak_memory_profile(plan, 0, 0)
+        return [ForwardOutput(None, None, None, None)] * plan.request_count, None
+
+    monkeypatch.setattr(r, "_update_peak_memory_profile", update)
+    monkeypatch.setattr(r, "_run_flat_plan_with_memory_tracking", run)
+    _packed_budget(monkeypatch, r, 20)
+    items = [[_request(0)], [_request(m) for m in range(1, 5)], [_request(5)]]
+    batches = r.forward_micro_batches(items)
+    next(batches)
+    assert next(batches).stats.subforward_count > 1
+    (signature,) = r._memory_profiles
+    # The seed's caller phase, then forward-only split children: not warm.
+    assert r._memory_profiles[signature].warm_bytes_per_token is None
+    assert [p for p in phases if p[1]] == [(1, True)]
+    list(batches)
+    # Only the later flat wave's caller phase fits the warm rate.
+    assert [p for p in phases if p[1]] == [(1, True), (1, True)]
+    assert r._memory_profiles[signature].warm_bytes_per_token == backward_rate
+
+
+def test_a_width_accepted_on_the_no_sharing_bound_never_executes_above_it(
+    monkeypatch,
+):
+    """The width search accepts a width on the cheap no-sharing count. With a
+    warm fit, that count can reach the warm size while the shared layout
+    executed falls below it; the executed plan must not price higher."""
+    from art.trainer_rank import ForwardInput, ForwardOutput
+
+    monkeypatch.setattr(_impl, "_PACKED_PRICED_MIN_REQUEST_TOKENS", 1)
+    r = packed_rank()
+    monkeypatch.setattr(r, "_dp_rank_and_size", lambda: (0, 1))
+    prefix = torch.arange(100)
+    items = [
+        [
+            ForwardInput(
+                input_tokens=torch.cat(
+                    [prefix, torch.arange(1000 + 10 * i, 1010 + 10 * i)]
+                ),
+                target_tokens=torch.cat(
+                    [prefix, torch.arange(1000 + 10 * i, 1010 + 10 * i)]
+                ),
+            )
+            for i in range(4)
+        ]
+    ]
+    plan = r._plan_flat_forward(items[0])
+    assert plan.packed_tokens < 200 <= plan.logical_tokens
+    r._memory_profiles[plan.signature] = _MemoryProfile(
+        bytes_per_token=100_000,
+        packed_tokens=plan.logical_tokens,
+        logical_per_packed=4,
+        warm_bytes_per_token=20_000,
+        warm_packed_tokens=200,
+        warm_logical_per_packed=4,
+    )
+    own = r._memory_check(plan).estimated_required_bytes
+    monkeypatch.setattr(r, "_available_memory_bytes", lambda: 10 * own)
+    monkeypatch.setattr(
+        r,
+        "_run_flat_plan_with_memory_tracking",
+        lambda plan, **_: (
+            [ForwardOutput(None, None, None, None)] * plan.request_count,
+            None,
+        ),
+    )
+    (batch,) = list(r.forward_micro_batches(items))
+    # Accepted on the no-sharing bound, which priced more tokens than ran.
+    assert batch.stats.packed_tokens == plan.packed_tokens
+    assert batch.stats.estimated_required_bytes > own
diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py
index 6212a0613..d58bf67ca 100644
--- a/tests/unit/test_trainer_rank_validation.py
+++ b/tests/unit/test_trainer_rank_validation.py
@@ -3639,12 +3639,12 @@ def test_forward_micro_batches_profiles_caller_peak_after_yield(
         "_run_flat_plan_with_memory_tracking",
         lambda *_args, **_kwargs: (_empty_outputs(plan), 123),
     )
-    profiles: list[tuple[int, int | None]] = []
+    profiles: list[tuple[int, int | None, bool]] = []
     monkeypatch.setattr(
         trainer,
         "_update_peak_memory_profile",
-        lambda candidate, baseline: profiles.append(
-            (candidate.packed_tokens, baseline)
+        lambda candidate, baseline, caller_phase=False: profiles.append(
+            (candidate.packed_tokens, baseline, caller_phase)
         ),
     )
 
@@ -3654,7 +3654,8 @@ def test_forward_micro_batches_profiles_caller_peak_after_yield(
     assert profiles == []
     with pytest.raises(StopIteration):
         next(batches)
-    assert profiles == [(plan.packed_tokens, 123)]
+    # The caller phase's peak includes backward: it may feed the warm fit.
+    assert profiles == [(plan.packed_tokens, 123, True)]
 
 
 @pytest.mark.parametrize("no_grad", [False, True])
@@ -3684,7 +3685,7 @@ def select_after_release(*args, **kwargs):
     monkeypatch.setattr(trainer, "_select_next_micro_batch", select_after_release)
     profiled: list[bool] = []
 
-    def profile(*_args):
+    def profile(*_args, **_kwargs):
         # Keep the completed wave through its caller peak observation.
         profiled.append(tensors[-1]() is not None)
 

From 5b44b77643b3e128b83b74888438ba5e7c67fcd8 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 18:39:07 +0000
Subject: [PATCH 3/5] Fit the warm profile only on whole caller phases

A nested tracked forward during a wave's yield resets the CUDA peak counter,
so the caller-phase reading then misses the wave's own forward and backward.
On main that reading was only max-merged; with the warm fit it could set the
warm rate. The micro-batch loop now captures the wave's interval (tracked
resets, forward peak) before yielding, and the caller phase fits the warm
profile only if no tracked forward reset the counter since and the counter
has not fallen below the wave's forward peak.

Other readings (forward-only, split children, interrupted caller phases)
cannot create or extend the warm fit, but a higher one now raises its rate,
as it raises the fit over every observation.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/_impl.py                |  39 +++++-
 tests/unit/test_trainer_rank_profile_warm.py | 136 ++++++++++++++++++-
 tests/unit/test_trainer_rank_validation.py   |  23 ++--
 3 files changed, 180 insertions(+), 18 deletions(-)

diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index 89406e952..15ef7795a 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -617,8 +617,9 @@ class _MemoryProfile:
     # A signature's first executed plan also pays one-time costs (compilation,
     # first-use workspaces), which a small first wave spreads over few tokens.
     # Admission uses the lower of the fit over every observation and the same
-    # fit over later plans' caller-phase peaks (which include backward), which
-    # prices a wave smaller than the smallest of them as if it were that large.
+    # fit over later flat waves' whole caller-phase peaks (forward plus the
+    # caller's loss and backward in the yield), which prices a wave smaller
+    # than the smallest of them as if it were that large.
     warm_bytes_per_token: float | None = None
     warm_packed_tokens: int | None = None
     warm_logical_per_packed: float | None = None
@@ -2040,6 +2041,9 @@ def memory_field(name: str, default: Any = None) -> Any:
         self._hybridep_rows_high_water = 0
         self._cache_recovery_state = _CacheRecoveryState()
         self._memory_profiles: dict[_MemorySignature, _MemoryProfile] = {}
+        # Tracked peak-counter resets, and the latest (resets, peak) reading.
+        self._peak_resets = 0
+        self._peak_reading: tuple[int, int] | None = None
         self._split_memory_floors: dict[bytes, int] = {}
         self._split_memory_floor_status = "not_observed"
         self._last_global_micro_batch_size: int | None = None
@@ -2891,6 +2895,7 @@ def _forward_micro_batches(
             outputs: list[Any] = []
             flat_outputs = iter(tracked_outputs)
             error: BaseException | None = None
+            interval: tuple[int, int] | None = None
             try:
                 if isinstance(candidate.plan, _FlatForwardPlan):
                     tracked_outputs, memory_baseline = (
@@ -2900,6 +2905,8 @@ def _forward_micro_batches(
                             context="forward_micro_batches",
                         )
                     )
+                    # This wave's peak interval, which its caller phase continues.
+                    interval = self.__dict__.get("_peak_reading")
                 else:
                     tracked_outputs, memory_baseline, forward_peak = (
                         self._execute_split_plan_with_memory_tracking(
@@ -2967,7 +2974,10 @@ def _forward_micro_batches(
             # to the forward's return, already recorded for this same plan.
             if isinstance(candidate.plan, _FlatForwardPlan):
                 self._update_peak_memory_profile(
-                    candidate.plan, memory_baseline, caller_phase=True
+                    candidate.plan,
+                    memory_baseline,
+                    caller_phase=True,
+                    interval=interval,
                 )
             elif memory_baseline is not None:
                 self._record_split_memory_floor(
@@ -6238,6 +6248,7 @@ def _run_flat_plan_with_memory_tracking(
             torch.cuda.synchronize(self.device)
             baseline = int(torch.cuda.memory_allocated(self.device))
             torch.cuda.reset_peak_memory_stats(self.device)
+            self._peak_resets = self.__dict__.get("_peak_resets", 0) + 1
         else:
             baseline = None
         observation = getattr(self, "_planner_observation", None)
@@ -6290,6 +6301,7 @@ def _update_peak_memory_profile(
         retained_after: int | None = None,
         *,
         caller_phase: bool = False,
+        interval: tuple[int, int] | None = None,
     ) -> None:
         if baseline is None:
             return
@@ -6297,13 +6309,21 @@ def _update_peak_memory_profile(
         observation = getattr(self, "_planner_observation", None)
         if observation is not None:
             observation["peak"] = max(observation["peak"], peak)
+        resets = self.__dict__.get("_peak_resets", 0)
+        self._peak_reading = (resets, peak)
         self._update_memory_profile(
             plan,
             max(0, peak - baseline),
             retained_bytes=(
                 None if retained_after is None else max(0, retained_after - baseline)
             ),
-            caller_phase=caller_phase,
+            # Only a whole wave: a nested forward resetting the counter during
+            # the yield, or an untracked reset (the counter fell below this
+            # wave's forward peak), leaves the peak since then short of it.
+            caller_phase=caller_phase
+            and interval is not None
+            and interval[0] == resets
+            and peak >= interval[1],
         )
 
     def _begin_planner_observation(
@@ -8295,15 +8315,20 @@ def _update_memory_profile(
         warm_tokens = None if previous is None else previous.warm_packed_tokens
         warm_sharing = None if previous is None else previous.warm_logical_per_packed
         caller_plans = 0 if previous is None else previous.caller_plans
-        # Only a flat wave's caller-phase peak includes its backward; split
-        # children and dp_rank_forward observe forward only. The profile's
-        # first such plan also pays one-time costs, so it is not warm.
+        # Only a flat wave's whole caller phase covers the caller's backward;
+        # split children and dp_rank_forward observe forward only. A caller that
+        # defers backward past the yield is not covered. The profile's first
+        # such plan also pays one-time costs, so it is not warm.
         if caller_phase:
             if caller_plans:
                 warm_rate = max(bytes_per_token, warm_rate or 0.0)
                 warm_tokens = min(plan.packed_tokens, warm_tokens or plan.packed_tokens)
                 warm_sharing = max(logical_per_packed, warm_sharing or 1.0)
             caller_plans += 1
+        elif warm_rate is not None:
+            # Other readings cannot fit the warm profile, but a higher one still
+            # raises its rate, as it raises the fit over every observation.
+            warm_rate = max(warm_rate, bytes_per_token)
         retained_fraction = None if previous is None else previous.retained_fraction
         retained_compute = (
             None if previous is None else previous.retained_compute_bytes_per_token
diff --git a/tests/unit/test_trainer_rank_profile_warm.py b/tests/unit/test_trainer_rank_profile_warm.py
index 95699ac3d..182639d41 100644
--- a/tests/unit/test_trainer_rank_profile_warm.py
+++ b/tests/unit/test_trainer_rank_profile_warm.py
@@ -241,7 +241,7 @@ def test_split_children_never_set_the_warm_fit(monkeypatch):
     forward_rate, backward_rate = 100, 300
     phases = []
 
-    def update(plan, baseline, retained_after=None, *, caller_phase=False):
+    def update(plan, baseline, retained_after=None, *, caller_phase=False, **_):
         # Stands in for the CUDA peak read: backward peaks above forward.
         phases.append((plan.request_count, caller_phase))
         rate = backward_rate if caller_phase else forward_rate
@@ -323,3 +323,137 @@ def test_a_width_accepted_on_the_no_sharing_bound_never_executes_above_it(
     # Accepted on the no-sharing bound, which priced more tokens than ran.
     assert batch.stats.packed_tokens == plan.packed_tokens
     assert batch.stats.estimated_required_bytes > own
+
+
+def _peak_reader(monkeypatch, r):
+    """Stand in for the CUDA peak counter, reset by each tracked forward."""
+    state = {"peak": 0}
+    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda *_: state["peak"])
+
+    def forward(plan, peak):
+        r._peak_resets = r.__dict__.get("_peak_resets", 0) + 1
+        state["peak"] = peak
+        r._update_peak_memory_profile(plan, 0, 0)
+        return r._peak_reading
+
+    return state, forward
+
+
+def _caller_phase(r, plan, interval):
+    r._update_peak_memory_profile(plan, 0, caller_phase=True, interval=interval)
+
+
+def test_only_a_whole_caller_phase_fits_the_warm_profile(monkeypatch):
+    r = rank()
+    first, larger = _plans(r)
+    state, forward = _peak_reader(monkeypatch, r)
+    n = larger.packed_tokens
+    _caller_phase(
+        r, first, forward(first, first.output_bytes + FIRST * first.packed_tokens)
+    )
+    signature = first.signature
+
+    def wave(interrupt):
+        interval = forward(larger, larger.output_bytes + WARM * n)
+        state["peak"] = larger.output_bytes + (WARM + 500_000) * n  # backward
+        interrupt()
+        _caller_phase(r, larger, interval)
+        return r._memory_profiles[signature].warm_bytes_per_token
+
+    # A nested forward during the yield resets the counter: not whole, even
+    # when its own peak is higher than this wave's forward.
+    def nested():
+        forward(first, larger.output_bytes + (WARM + 100_000) * n)
+
+    assert wave(nested) is None
+    # An untracked reset: the counter falls below this wave's forward peak.
+    assert wave(lambda: state.update(peak=0)) is None
+    # An uninterrupted wave: forward, then the caller's loss and backward.
+    assert wave(lambda: None) == WARM + 500_000
+
+
+def test_other_readings_raise_but_never_fit_the_warm_profile():
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    _update(r, larger, WARM, caller_phase=False)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token is None
+    _observe(r, larger, WARM)
+    profile = r._memory_profiles[first.signature]
+    assert profile.warm_bytes_per_token == WARM
+    # A higher forward-only or interrupted reading raises the warm rate, but
+    # neither extends the warm extent nor its sharing.
+    middle = r._plan_flat_forward(requests(4096, 64))
+    shared = replace(middle, logical_tokens=middle.packed_tokens * 8)
+    _update(r, shared, WARM + 300_000, caller_phase=False)
+    raised = r._memory_profiles[first.signature]
+    assert raised.warm_bytes_per_token == WARM + 300_000
+    assert raised.warm_packed_tokens == profile.warm_packed_tokens
+    assert raised.warm_logical_per_packed == profile.warm_logical_per_packed
+    _update(r, larger, WARM - 100_000, caller_phase=False)
+    assert r._memory_profiles[first.signature].warm_bytes_per_token == WARM + 300_000
+
+
+def test_a_nested_forward_during_the_yield_cannot_fit_the_warm_profile(monkeypatch):
+    """The real micro-batch loop: a caller that runs another tracked forward
+    after its backward leaves a peak counter reset since the wave's forward."""
+    from test_trainer_rank_split import _packed_budget, _request
+    from test_trainer_rank_split import _rank as split_rank
+
+    from art.trainer_rank import ForwardOutput
+
+    r = split_rank(monkeypatch)
+    monkeypatch.setattr(r, "_retained_memory_bytes", lambda *_args, **_kwargs: 0)
+    state, forward = _peak_reader(monkeypatch, r)
+
+    def run(plan, **_kwargs):
+        forward(plan, plan.output_bytes + 100 * plan.packed_tokens)
+        return [ForwardOutput(None, None, None, None)] * plan.request_count, 0
+
+    monkeypatch.setattr(r, "_run_flat_plan_with_memory_tracking", run)
+    _packed_budget(monkeypatch, r, 10)
+    items = [[_request(m)] for m in range(3)]
+    batches = r.forward_micro_batches(items)
+    for index, batch in enumerate(batches):
+        (signature,) = r._memory_profiles
+        state["peak"] += 200 * batch.stats.packed_tokens  # the caller's backward
+        if index == 1:
+            run(r._plan_flat_forward(items[0]))  # a nested tracked forward
+    profile = r._memory_profiles[signature]
+    assert profile.caller_plans == 2
+    # Only the uninterrupted third wave fit the warm rate, backward included.
+    assert profile.warm_bytes_per_token == 300
+
+
+def test_each_tracked_forward_starts_a_new_peak_interval(monkeypatch):
+    """The real tracked forward: its counter reset starts a new interval, so a
+    caller phase that continues an earlier one is not whole."""
+    r = rank()
+    first, larger = _plans(r)
+    state = {"peak": 0}
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "synchronize", lambda *_: None)
+    monkeypatch.setattr(torch.cuda, "memory_allocated", lambda *_: 0)
+    monkeypatch.setattr(
+        torch.cuda, "reset_peak_memory_stats", lambda *_: state.update(peak=0)
+    )
+    monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda *_: state["peak"])
+    monkeypatch.setattr(r, "device", torch.device("cuda", 0))
+
+    def execute(plan):
+        state["peak"] = plan.output_bytes + FIRST * plan.packed_tokens
+        return [None] * plan.request_count
+
+    monkeypatch.setattr(r, "_execute_flat_plan", execute)
+    check = _impl._MemoryCheck(0, 0, True)
+    _, baseline = r._run_flat_plan_with_memory_tracking(first, check=check, context="t")
+    seed = r._peak_reading
+    _caller_phase(r, first, seed)
+    _, baseline = r._run_flat_plan_with_memory_tracking(
+        larger, check=check, context="t"
+    )
+    interval = r._peak_reading
+    assert baseline == 0 and interval[0] == seed[0] + 1
+    r._run_flat_plan_with_memory_tracking(first, check=check, context="t")  # nested
+    _caller_phase(r, larger, interval)
+    assert r._memory_profiles[larger.signature].warm_bytes_per_token is None
diff --git a/tests/unit/test_trainer_rank_validation.py b/tests/unit/test_trainer_rank_validation.py
index d58bf67ca..0d85e276a 100644
--- a/tests/unit/test_trainer_rank_validation.py
+++ b/tests/unit/test_trainer_rank_validation.py
@@ -3634,17 +3634,19 @@ def test_forward_micro_batches_profiles_caller_peak_after_yield(
     trainer = TrainerRank(_runtime())
     _stub_forward(monkeypatch, trainer, profiled=True)
     plan = trainer._plan_flat_forward([_target_request(1)])
-    monkeypatch.setattr(
-        trainer,
-        "_run_flat_plan_with_memory_tracking",
-        lambda *_args, **_kwargs: (_empty_outputs(plan), 123),
-    )
-    profiles: list[tuple[int, int | None, bool]] = []
+
+    def run(*_args, **_kwargs):
+        # The forward's peak interval, which the caller phase continues.
+        trainer._peak_reading = (7, 99)
+        return _empty_outputs(plan), 123
+
+    monkeypatch.setattr(trainer, "_run_flat_plan_with_memory_tracking", run)
+    profiles: list[tuple[int, int | None, bool, object]] = []
     monkeypatch.setattr(
         trainer,
         "_update_peak_memory_profile",
-        lambda candidate, baseline, caller_phase=False: profiles.append(
-            (candidate.packed_tokens, baseline, caller_phase)
+        lambda candidate, baseline, caller_phase=False, interval=None: profiles.append(
+            (candidate.packed_tokens, baseline, caller_phase, interval)
         ),
     )
 
@@ -3654,8 +3656,9 @@ def test_forward_micro_batches_profiles_caller_peak_after_yield(
     assert profiles == []
     with pytest.raises(StopIteration):
         next(batches)
-    # The caller phase's peak includes backward: it may feed the warm fit.
-    assert profiles == [(plan.packed_tokens, 123, True)]
+    # The caller phase continues the wave's own interval: it may feed the
+    # warm fit if no other reset interrupted it.
+    assert profiles == [(plan.packed_tokens, 123, True, (7, 99))]
 
 
 @pytest.mark.parametrize("no_grad", [False, True])

From 8ae2585fc6e2d15f225117c113df2b8f96d3f017 Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 19:21:15 +0000
Subject: [PATCH 4/5] Keep higher readings after the seed for the first warm
 fit

A reading that cannot fit the warm profile (forward-only, a split child, an
interrupted caller phase) only raised an existing warm rate, so one taken
between the seed and the first whole warm plan was dropped when that plan
created the fit. After the seed, every reading now raises a pending warm
rate, which stays inert for pricing until a whole warm plan fits the size
extent and sharing. The forward_micro_batches docstrings state which waves
can lower later estimates.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/__init__.py             |  5 +++
 src/art/trainer_rank/_impl.py                | 13 ++++---
 tests/unit/test_trainer_rank_profile_warm.py | 36 ++++++++++++++------
 3 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py
index fd321a820..61dfa2d00 100644
--- a/src/art/trainer_rank/__init__.py
+++ b/src/art/trainer_rank/__init__.py
@@ -267,6 +267,11 @@ def forward_micro_batches(
         iterator before making collective calls after an early exit. Guards apply
         on the iterator's thread; raw torch.distributed calls are not guarded.
         Collective calls must still match across ranks.
+
+        Admission learns each wave's memory peak, including the caller's loss
+        and backward. Only waves whose backward runs inside the yield, with no
+        other TrainerRank forward there, can lower later estimates below the
+        first wave's.
         """
         forward = cast(
             Callable[..., Iterator[MicroBatch[ForwardInputs, ForwardOutputs]]],
diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index 15ef7795a..4c9623e54 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -2816,7 +2816,9 @@ def forward_micro_batches(
         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.
+        their estimate. Only waves whose loss and backward run inside the yield,
+        with no other TrainerRank forward there, can lower later estimates below
+        the signature's first wave.
         """
         if not isinstance(yield_empty, bool):
             raise TypeError("yield_empty must be a bool")
@@ -8325,10 +8327,11 @@ def _update_memory_profile(
                 warm_tokens = min(plan.packed_tokens, warm_tokens or plan.packed_tokens)
                 warm_sharing = max(logical_per_packed, warm_sharing or 1.0)
             caller_plans += 1
-        elif warm_rate is not None:
-            # Other readings cannot fit the warm profile, but a higher one still
-            # raises its rate, as it raises the fit over every observation.
-            warm_rate = max(warm_rate, bytes_per_token)
+        elif caller_plans:
+            # After the seed, other readings cannot fit the warm profile, but a
+            # higher one still raises its rate (inert until a whole warm plan
+            # fits the rest), as it raises the fit over every observation.
+            warm_rate = max(warm_rate or 0.0, bytes_per_token)
         retained_fraction = None if previous is None else previous.retained_fraction
         retained_compute = (
             None if previous is None else previous.retained_compute_bytes_per_token
diff --git a/tests/unit/test_trainer_rank_profile_warm.py b/tests/unit/test_trainer_rank_profile_warm.py
index 182639d41..f1fa1c38e 100644
--- a/tests/unit/test_trainer_rank_profile_warm.py
+++ b/tests/unit/test_trainer_rank_profile_warm.py
@@ -84,7 +84,7 @@ def test_forward_only_observations_never_set_the_warm_fit():
     first, larger = _plans(r)
     _observe(r, first, FIRST)
     _update(r, larger, WARM, caller_phase=False)
-    assert r._memory_profiles[first.signature].warm_bytes_per_token is None
+    assert r._memory_profiles[first.signature].warm_packed_tokens is None
     assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
     # A later caller phase is warm, and its peak includes the forward's.
     _update(r, larger, WARM, caller_phase=True)
@@ -265,8 +265,8 @@ def run(plan, **_kwargs):
     next(batches)
     assert next(batches).stats.subforward_count > 1
     (signature,) = r._memory_profiles
-    # The seed's caller phase, then forward-only split children: not warm.
-    assert r._memory_profiles[signature].warm_bytes_per_token is None
+    # The seed's caller phase, then forward-only split children: no warm fit.
+    assert r._memory_profiles[signature].warm_packed_tokens is None
     assert [p for p in phases if p[1]] == [(1, True)]
     list(batches)
     # Only the later flat wave's caller phase fits the warm rate.
@@ -358,18 +358,18 @@ def wave(interrupt):
         state["peak"] = larger.output_bytes + (WARM + 500_000) * n  # backward
         interrupt()
         _caller_phase(r, larger, interval)
-        return r._memory_profiles[signature].warm_bytes_per_token
+        return r._memory_profiles[signature]
 
     # A nested forward during the yield resets the counter: not whole, even
     # when its own peak is higher than this wave's forward.
     def nested():
-        forward(first, larger.output_bytes + (WARM + 100_000) * n)
+        forward(larger, larger.output_bytes + (WARM + 100_000) * n)
 
-    assert wave(nested) is None
+    assert wave(nested).warm_packed_tokens is None
     # An untracked reset: the counter falls below this wave's forward peak.
-    assert wave(lambda: state.update(peak=0)) is None
-    # An uninterrupted wave: forward, then the caller's loss and backward.
-    assert wave(lambda: None) == WARM + 500_000
+    assert wave(lambda: state.update(peak=0)).warm_packed_tokens is None
+    # Their readings raised the pending rate; an uninterrupted wave fits it.
+    assert wave(lambda: None).warm_bytes_per_token == WARM + 500_000
 
 
 def test_other_readings_raise_but_never_fit_the_warm_profile():
@@ -377,7 +377,7 @@ def test_other_readings_raise_but_never_fit_the_warm_profile():
     first, larger = _plans(r)
     _observe(r, first, FIRST)
     _update(r, larger, WARM, caller_phase=False)
-    assert r._memory_profiles[first.signature].warm_bytes_per_token is None
+    assert r._memory_profiles[first.signature].warm_packed_tokens is None
     _observe(r, larger, WARM)
     profile = r._memory_profiles[first.signature]
     assert profile.warm_bytes_per_token == WARM
@@ -456,4 +456,18 @@ def execute(plan):
     assert baseline == 0 and interval[0] == seed[0] + 1
     r._run_flat_plan_with_memory_tracking(first, check=check, context="t")  # nested
     _caller_phase(r, larger, interval)
-    assert r._memory_profiles[larger.signature].warm_bytes_per_token is None
+    assert r._memory_profiles[larger.signature].warm_packed_tokens is None
+
+
+def test_a_higher_reading_before_the_first_warm_plan_is_kept():
+    """A seed, then a higher forward-only reading, then a lower whole warm
+    plan: the warm rate keeps the higher post-seed reading."""
+    r = rank()
+    first, larger = _plans(r)
+    _observe(r, first, FIRST)
+    _update(r, larger, WARM + 300_000, caller_phase=False)
+    assert _required(r, larger) == int(FIRST * larger.packed_tokens * 1.1)
+    _observe(r, larger, WARM)
+    profile = r._memory_profiles[first.signature]
+    assert profile.warm_bytes_per_token == WARM + 300_000
+    assert _required(r, larger) == int((WARM + 300_000) * larger.packed_tokens * 1.1)

From d0c11d6a339d783eb2125d94978b81874b71576f Mon Sep 17 00:00:00 2001
From: Brad Hilton <brad.hilton.nw@gmail.com>
Date: Sat, 26 Sep 2026 19:57:56 +0000
Subject: [PATCH 5/5] State what caller-phase learning covers in the forward
 docstrings

Backward is learned only when it runs inside the yield; a deferred backward
is not detected. Only a nested TrainerRank forward is enforced. Also let the
cache-recovery test's profile stub accept the new keywords.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 src/art/trainer_rank/__init__.py               | 6 +++---
 src/art/trainer_rank/_impl.py                  | 7 ++++---
 tests/unit/test_trainer_rank_cache_recovery.py | 2 +-
 3 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/art/trainer_rank/__init__.py b/src/art/trainer_rank/__init__.py
index 61dfa2d00..04b78329e 100644
--- a/src/art/trainer_rank/__init__.py
+++ b/src/art/trainer_rank/__init__.py
@@ -269,9 +269,9 @@ def forward_micro_batches(
         Collective calls must still match across ranks.
 
         Admission learns each wave's memory peak, including the caller's loss
-        and backward. Only waves whose backward runs inside the yield, with no
-        other TrainerRank forward there, can lower later estimates below the
-        first wave's.
+        and backward when they run inside the yield; a backward deferred past the
+        yield is not learned. A wave with another TrainerRank forward inside its
+        yield cannot lower later estimates.
         """
         forward = cast(
             Callable[..., Iterator[MicroBatch[ForwardInputs, ForwardOutputs]]],
diff --git a/src/art/trainer_rank/_impl.py b/src/art/trainer_rank/_impl.py
index 4c9623e54..d4f7bde0f 100644
--- a/src/art/trainer_rank/_impl.py
+++ b/src/art/trainer_rank/_impl.py
@@ -2816,9 +2816,10 @@ def forward_micro_batches(
         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. Only waves whose loss and backward run inside the yield,
-        with no other TrainerRank forward there, can lower later estimates below
-        the signature's first wave.
+        their estimate. Backward is learned only when it runs inside the yield;
+        a backward deferred past it is not. A wave with another TrainerRank
+        forward inside its yield cannot lower later estimates below the
+        signature's first wave.
         """
         if not isinstance(yield_empty, bool):
             raise TypeError("yield_empty must be a bool")
diff --git a/tests/unit/test_trainer_rank_cache_recovery.py b/tests/unit/test_trainer_rank_cache_recovery.py
index e165346f2..ef34113c3 100644
--- a/tests/unit/test_trainer_rank_cache_recovery.py
+++ b/tests/unit/test_trainer_rank_cache_recovery.py
@@ -172,7 +172,7 @@ def make(self):
         self.addCleanup(observer_clock.stop)
         q = object.__new__(_impl.TrainerRank)
         q.device = types.SimpleNamespace(type="cuda")
-        q._update_peak_memory_profile = lambda *a: None
+        q._update_peak_memory_profile = lambda *a, **k: None
         q._execute_flat_plan = lambda p: [object() for _ in range(p.request_count)]
         q._telemetry_signature = lambda p: {}
         q._telemetry_plan_signature = lambda p: {}