Price warm waves by the fit later plans measured, not the first plan's by bradhilton · Pull Request #991 · OpenPipe/ART

bradhilton · 2026-09-26T16:59:14Z

A signature's first wave pays one-time costs (compilation, first-use workspaces). TrainerRank's memory profile keeps the highest per-token rate it has seen, so a small, cold first wave overprices every later wave. On Qwen3.6-35B-A3B that first wave ran 275 KB/token; later waves ran 206–247 KB/token.

Change: the profile also learns a warm rate from later waves, and admission uses the lower of the two.

Only a complete wave teaches the warm rate: its forward plus the caller's loss and backward inside the yield, with no other TrainerRank forward in between.
Other readings (split children, dp_rank_forward, interrupted waves) can raise the warm rate but never create it.
A wave smaller than the smallest warm wave is priced as if it were that large, so estimates never shrink as waves grow.

Before/after (Qwen3.6-35B-A3B, 40 layers, CP2, 8 steps; "today" prices the same waves without the warm rate):

	today	this PR	waves under their peak
EP1 (the profile sets the estimate)	+36% median over-estimate	+16% (min +9%)	0 of 25
EP2 (the static routing floor sets it)	+43%	+43%, unchanged	0 of 17

Testing: 15 new CPU tests, including integration tests through the real width search, split waves and a nested forward during a wave's caller phase, plus the GPU run above; the trainer-rank suite passes.

Limits:

Measured on one MoE model at CP2. Dense, TP>1, CP1 and PP are covered by CPU tests only.
Where no static floor applies, the 1.1× safety factor is the margin, as it already is for any signature whose first wave isn't small and cold.
A caller that runs backward after the yield, or resets CUDA peak stats itself, isn't detected; a TrainerRank forward inside every yield leaves this change inert.

🤖 Generated with Claude Code

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>

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>

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>

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>

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>

…d admission method bodies Pure code motion following the _memory / _gdn_memory precedent: 34 method bodies moved verbatim from _impl as module-level functions taking the rank as self. Names the bodies read from _impl's namespace (torch, dist, time, _telemetry_phase, plan/cost types, sibling helpers) are resolved as _impl.X at call time so tests that patch them on _impl keep intercepting; stdlib helpers and the unpatched prefix-tree / planner-cost functions are imported directly. @contextmanager stays on _planning_status and @_backward_region on _recover_admission. _forward_micro_batches carries the #991 caller-phase interval body.

bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 16:59 — with GitHub Actions Active

bradhilton had a problem deploying to trainer-rank-gpu-validation September 26, 2026 17:58 — with GitHub Actions Failure

bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 18:39 — with GitHub Actions Active

bradhilton and others added 4 commits September 26, 2026 19:21

bradhilton force-pushed the dalinar/profile-cold-seed branch from 4753dd4 to 8ae2585 Compare September 26, 2026 19:47

bradhilton had a problem deploying to trainer-rank-gpu-validation September 26, 2026 19:48 — with GitHub Actions Error

bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 19:58 — with GitHub Actions Active

bradhilton marked this pull request as ready for review September 26, 2026 20:18

bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 20:19 — with GitHub Actions Active

bradhilton merged commit 3ad9a5f into main Sep 26, 2026
10 checks passed

bradhilton mentioned this pull request Sep 26, 2026

Move TrainerRank memory accounting into trainer_rank/_memory #999

Draft

This was referenced Sep 26, 2026

Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner #1000

Draft

Move TrainerRank checkpoint-slot and dynamic-optimizer management into trainer_rank/_slots and _optimizer #1001

Draft

Navigation Menu

Price warm waves by the fit later plans measured, not the first plan's - #991

Price warm waves by the fit later plans measured, not the first plan's#991
bradhilton merged 5 commits into
mainfrom
dalinar/profile-cold-seed

bradhilton commented Sep 26, 2026 •

edited

Loading

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

bradhilton commented Sep 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

bradhilton commented Sep 26, 2026 •

edited

Loading