Best-effort internal splitting for TrainerRank forwards

[bradhilton]

Summary

Follow-up to #826. When the unsplit plan (cost-optimal, then memory-minimal) cannot be admitted, TrainerRank.dp_rank_forward now tries a bounded, deterministic split ladder — 2, 4, … subforwards, cutting the requests in prefix-local depth-first order into token-balanced chunks — and executes the fewest subforwards that fit, keeping every returned graph live together and reconstructing outputs in caller order. The same ladder rescues forward_micro_batches when even the minimum wave cannot fit unsplit.

Contract (as agreed): try not to raise when splitting would make execution feasible; account for all returned graphs being live together; if finding out is too expensive or fragile, refuse — worded as "unable to find a feasible split", never as a claim that none exists.

Design

• Rung check = the research thread's cumulative memory invariant, reduced to one inequality. Every returned graph stays live, so subforward j needs its own transient peak plus the memory retained by the subforwards before it. Each of those sums is bounded by all retained memory plus the largest ephemeral share, which therefore decides a rung by itself, in any order (chunks still execute larger-ephemeral-first to minimize the running peak). The same quantity is the headroom the caller's backward can count on; that is a heuristic for backward workspace, not a bound, so the GPU gate measures it under real pressure instead of claiming it.
• Bounded planning cost. The cheap full-sharing lower bound (one O(tokens) CPU scan per chunk) rejects a rung without planning anything; a surviving rung is priced exactly with cost-optimal layouts and, failing that, memory-minimal ones (whose packed tokens equal the lower bound). The planner runs for at most one rung — the one that executes.
• Retained fraction learned online per memory signature: memory still allocated after a forward returns ÷ that forward's observed peak — a physical ratio that needs no trusted denominator (the cold call's static estimate is far below the real peak, so a ratio against the estimate would clamp to 1.0 and, max-merged, pin the signature to "everything retained" forever). None until observed, then max-merged (an observed 1.0 is distinct from "unobserved"); trusted only within the profile's packed-token trust range and near its observed sharing ratio, so a small profiled forward cannot authorize a much larger split. Unobserved means 1.0, so a cold oversized call refuses until a profile exists. The observation is taken at forward return and says nothing about backward (documented limitation).
• Collective safety. Ensuring checkpoint slots is a world collective and the ladder's length depends on this rank's DP-local inputs, so slots are ensured exactly once per call and all further planning skips the ensure. In the minimum-wave path every DP rank runs its own ladder, then all ranks agree on the outcome with one collective, so a refusal is raised everywhere or nowhere. (Found while restructuring: per-chunk planning calls in the first draft would have deadlocked DP>1.)
• No retry after the first forward. The complete ordered split is admitted before any model execution; any execution-time memory failure of an admitted split raises TrainerRankPartialExecutionError (a TrainerRankMemoryError) naming how many subforwards completed. Each subforward's outputs carry their own slot-graph sentinel, so slot load/step stays blocked until every subforward's graph is released.
• Splitting is disabled under expert parallelism in this release (HybridEP capacity must not be resized between subforwards while earlier graphs are live); the refusal says so.
• Observability: last_forward_telemetry() gains subforward_count, subforward_request_indices, predicted_peak_bytes, usable_limit_bytes; MicroBatchStats gains subforward_count. A test-only usable-memory cap (ART_TRAINER_RANK_TEST_MEMORY_LIMIT_BYTES, gated by ART_TRAINER_RANK_TEST_HOOKS) drives the deterministic GPU arm.

Design brief: dev/trainer_rank_planner_design.md ("Best-effort internal splitting").

Evidence

CPU gates (tests/unit/test_trainer_rank_split.py, 15 tests; the original 8 were written and fail-verified before implementation): split instead of raise; unsplit reports one subforward; nested caller-order reconstruction; bounded ladder + honest refusal wording with the planner never running for a failed rung; cumulative live-graph admission refuses when halves fit alone but not together; retained profile enables the split; profile trusted only near its observed scale; observations max-merged once observed; minimum-wave splitting; deterministic partitions; one slot-ensure collective per call; execution-time failure reported as partial execution (first and later subforward); independent slot-graph sentinels per subforward (load/step blocked until the last subforward's graph is released).

GPU gate (dev/trainer_rank_landing_acceptance.py --phase split-conversion), mirroring the research thread's sealed split-conversion cell (Qwen3.5-4B, 4 layers, CP1, 4 × 12,288-token inputs), fresh process per arm, H200:

Arm	Result
--pressure cap unlimited	unsplit (subforward_count 1); forward peak 26.0 GB, forward+backward 26.3 GB
--pressure cap conversion (cap just below the unsplit requirement)	2 subforwards; parity vs unsplit: mean 0.041% relative (max-abs 0.138, bf16 reduction order); combined backward and reverse-order per-subforward backward both succeed with every graph live
--pressure cap bounded-decline (cap below one request)	refused before any model execution, bounded-search wording
--pressure ballast training forward (no test hooks)	measured retained fraction f = 0.990 of the forward's peak stays allocated for backward, so the 2-way conversion window is only (1−f)·R/2 = 141 MB on an unsplit requirement R = 28.6 GB. 112.6 GB of live ballast puts the real usable budget inside that window (28.55 GB; 4.3 GB free on the device): unsplit refused before execution, 2 subforwards run under the reduced headroom (parity mean 0.041%), combined backward succeeds with the ballast still live; observed forward+backward peak 26.19 GB ≤ predicted 28.48 GB ≤ admitted budget 28.55 GB (forward-only peak 25.82 GB; no reserved-pool growth)
--pressure ballast bounded-decline (budget 3.3 GiB, below one request)	refused before any model execution
--pressure ballast no_grad forward	retained fraction 2.6e-5 (outputs only), so splitting pays: with the real budget at 60% of the unsplit requirement (8.37 GB → 5.02 GB) the call converts to 2 subforwards (parity exact); observed peak 3.80 GB ≤ predicted 4.18 GB ≤ budget 5.02 GB

The ballast arm's main lesson: splitting cannot shrink retained activations, only the transient share. no_grad forwards (reference/old-policy logprobs) are the demonstrated high-value case, with wide-margin physical-safety evidence. The training benefit is workload-dependent and small in this sealed landing cell (~0.5%); the research thread's full-height cell retained closer to 92%, so CP/GDN, output-heavy or workspace-heavy training shapes may have several gigabytes of splittable transient memory, and grad-enabled support is kept for them.

Regression: 537 CPU tests (acceptance, tree-equivalence, trainer-rank shard, split gates), prek --all-files.

Known limitations

• Cold calls that cannot fit unsplit refuse until a retained-bytes profile exists (conservative by design; no universal prior is defensible).
• The backward headroom term is forward-ephemeral memory: a heuristic that held on one real-pressure cell, not a bound.
• Splitting is disabled under EP>1; per-subforward execution overhead is not modeled.
• The full-sharing lower bound can conservatively reject a rung whose cost-optimal layouts would have fit if retained-profile trust changes with the sharing ratio: a false refusal, never an unsafe admission. A cold oversized no_grad call likewise refuses until a compatible profile exists (a later simplification could model no_grad retained memory directly from the known output bytes).
• No claim of exhaustive feasibility or broadly calibrated split optimality; the GPU seal is one cell.

🤖 Generated with Claude Code

[bradhilton]

Second commit series addresses the research-thread review. Point by point:

1. Retained-profile hardening — done. _MemoryProfile.retained_fraction is now float | None (None = never observed), and merging is: unobserved → take the observation; observed → max(previous, observed). So an observed 1.0 can no longer be replaced by a later lower observation. That fix immediately exposed a second problem on the GPU cell: the fraction was measured against the estimate, and the first (cold) call's estimate is the untrusted static formula (~1.4 GB vs a real 28.5 GB peak), so its ratio clamped to 1.0 and, max-merged, pinned the signature to "everything retained" — both GPU arms refused instead of converting. The fraction is now a physical ratio, retained bytes ÷ the same forward's observed peak, which needs no trusted denominator; admission applies it to a subforward's estimated peak (≥ the real peak whenever the estimate is trusted). Applicability is size-aware: _retained_fraction(signature, packed_tokens=…, logical_tokens=…) returns 1.0 unless the chunk is within _MEMORY_PROFILE_TRUST_GROWTH of the profile's packed-token scale (the same range that already gates bytes_per_token) and its logical/packed ratio is within that range of the observed ratio. Retained and bytes_per_token are co-observed on every CUDA forward, so the profile's packed_tokens is the retained scale; the micro-batch path's post-yield re-profile is peak-only and leaves the retained observation alone. Tests: test_retained_profile_is_trusted_only_near_its_observed_scale, test_retained_observations_are_max_merged_once_observed.

2. Backward term — kept as a heuristic, worded as one, and measured. Admission is now a single rung check, all retained memory plus the largest ephemeral share. That inequality implies every per-subforward check (transient peak plus retained-before-it) in any execution order, so it is the whole forward invariant; as a backward term it is the headroom the caller's backward can count on, which the code, design brief and PR now call a heuristic for backward workspace rather than an upper bound. The ballast arm gates on it: observed forward+backward peak must stay within both the admitted budget and the predicted peak (numbers below).

3. Ballast arm — done (--phase split-conversion --pressure ballast; fresh process, warm kernels via the unlimited pass, no test hooks). Live ballast brings the planner's real usable budget 2% under the unsplit requirement; the split runs under that headroom; the combined backward runs with the ballast still live; deeper ballast refuses before execution. Sealed evidence in the PR description table.

Cold policy — kept: refuse until profiled; no prior baked in. The forward-only nature of the observation is documented in _MemoryProfile and the design brief.

Checkpoint-slot lifetime — added test_split_subforwards_track_independent_slot_graphs: two subforwards on one slot, backward the first → load and step still blocked, backward the second → allowed. It runs the real _execute_flat_plan (only the model call is faked) so the sentinels come from production code.

EP>1 — unchanged (explicit refusal).

Planning-cost bound — done via the cheap precheck. The full-sharing lower bound per chunk (one O(tokens) CPU scan, from CPU copies of the rows made once per call) rejects a rung with no planning at all; a surviving rung is priced exactly with cost-optimal layouts and, failing that, memory-minimal ones, whose packed tokens equal the lower bound — so the planner runs for at most one rung, the one that executes. test_split_ladder_is_bounded_and_refuses_when_one_request_cannot_fit now asserts the planner ran exactly twice (the two unsplit attempts) across a three-rung refusal. No work-limit knob was needed.

Minor — any execution-time memory failure of an admitted split (first subforward included) is TrainerRankPartialExecutionError, naming how many subforwards completed; the class is in the public __module__ tuple; wording is now "prefix-local depth-first order" everywhere (a token-balanced cut can still divide a sibling subtree).

One thing the restructuring surfaced that the review did not: _group_active_request_indices(ensure_slots=True) runs _ensure_checkpoint_slots, an all_gather_object over a world-wide gloo group, on every call — even with nothing to ensure. The first draft's ladder called it per chunk per rung, a count that depends on each rank's DP-local inputs, so DP>1 would have deadlocked the moment two ranks stopped at different rungs (the CP=2 CI job and the 1-GPU gate could not see this). Slots are now ensured exactly once per call and all further planning passes ensure_slots=False (test_split_ladder_ensures_checkpoint_slots_once). Relatedly, the minimum-wave path now has every DP rank run its own ladder and then agree on the outcome with one collective, so a refusal is raised on all ranks or none.

[bradhilton]

Wording updated per the review: no_grad is the demonstrated high-value case; the training benefit is workload-dependent and small in the sealed landing cell (the full-height research cell retained ~92%, so other shapes may have gigabytes of splittable transient memory). Grad-enabled support stays. The two accepted limitations (conservative lower-bound refusals when profile trust changes with sharing ratio; cold no_grad refusal until a compatible profile exists) are recorded in the design brief and PR description. Marked ready for review.