Price trainer-rank memory profiles by packed tokens, not sharing ratio

[bradhilton]

Memory profiles were extrapolated by the prefix-sharing ratio: bytes_per_token * max(packed, logical / profile.logical_per_packed). So a profile learned without sharing priced a shared plan by its logical tokens. For grad-enabled single-target training requests of at least 64 tokens, this PR prices the profiled model memory by packed tokens, charges head and caller memory per logical token, and extrapolates sharing at most 8x beyond what the profile observed. Everything else keeps the existing extrapolation.

Evidence

Production report. Caladan 061 (Qwen3.6-35B-A3B, 1×H200, ART a655c051), planner report f3f203ed. The first multi-request batch (16,205 packed / 109,065 logical tokens) came right after a single-request calibration step (logical_per_packed = 1.0). The planner predicted a 56.3 GiB peak against 6.9 GiB observed. Its required estimate was 66.5 GiB with 67.7 GiB available, so it was admitted with 1.7% headroom. Across all shared steps in that run, observed peaks fit 66.6 MiB + 459,782 B/packed token.

Controlled measurement. One H200, 8 layers, random weights, 16,384 packed tokens with 512-token suffixes, sharing ratios 1/2/4/8, 3 shuffled repeats. The table shows peak bytes per packed token after warm-up:

model	family	ratio 1	ratio 8	old prediction on first ratio-8 call
Qwen3.6-35B-A3B	GDN-MoE	321.9K	323.5K	18.38 GiB (steady 8.46)
Qwen3-30B-A3B	attention-MoE	282.1K	282.0K	16.77 GiB (steady 7.72)
Qwen3-8B	dense attention	271.2K	271.2K	39.46 GiB on every repeat (steady 5.01)
Qwen3.5-4B	dense GDN	235.8K	237.4K	4.91 GiB

High-branch GDN measurement (same setup, 1 H200, 8 layers, raw log retained). This compares 101 twenty-token branches on a 5,100-token prompt (sharing ratio 72.6, 101 segments) with a single 7,119-token request. The warm peak per packed token rises from 323 KB to 378 KB (+17%), about 3.9 MB per extra segment. That matches one GDN layer's recurrent state (≈4.3 MB per segment), the amount live at a time under full recompute, and is now charged per segment.

The dense rows come from a retained raw log. The MoE rows were transcribed from a live log that was lost when the cluster auto-downed, so treat them as provisional. Ratios 2 and 4 are also flat. The only deviations are the first two calls of each run (cold allocation and two micro-batches), which aren't sharing effects.

One consequence of the old pricing: under a ratio-1 profile, memory-minimal (full-sharing) layouts cost as much as unshared ones, so the split ladder couldn't use sharing to fit a budget.

Changes

Supported invariant

Packed pricing assumes:

1. Model memory scales with packed rows at the profiled rate B.
2. Head memory plus the caller's loss memory (saved tensors and backward transients) peaks at no more than 6 KiB per active logical token.
3. Every packed-priced request has at least 64 tokens, so per-request constants up to 384 KiB, head blocks included, fit within that budget.

Callers heavier than item 2 are unsupported, and the forward_micro_batches docstring says so. main's own assumption, that memory scales with logical rows at the profile's density, is likewise unenforced. The invariant was agreed with both private reviewers before implementation.

Pricing

• Eligibility (_packed_priced(signature, one_layer_recompute)), all of:

  • ART's default one-layer full recompute, the only mode the evidence covers. That is exactly full / uniform / 1, read at pricing time from the live decoder config, with the decoder in training mode (Megatron checks the decoder's own mode and skips recompute in eval mode). A chunk without a GPT decoder falls back to the constructor settings and requires every submodule to be training.
  • Every grad mode is True.
  • Every request mix is target:single or inactive.

  There is no rate floor. An earlier floor made eligibility depend on the learned ratio, which could switch the logical charge off exactly where sharing is highest.

• Short requests. A batch containing a single-target request shorter than 64 tokens sets short_requests on its memory signature and keeps the logical extrapolation. A duplicate short request adds no packed row, and per-request caller memory can outgrow a few row charges. The flag is excluded from signature identity, so such a batch shares the profile learned from longer requests and is priced exactly as on main, rather than starting cold.

• Required = B·max(P, L/(8r)) + 6 KiB·L + grad_segments·GDN_state:

  • P is packed rows, L active logical rows and r the profile's sharing ratio.
  • The logical charge does not depend on layout, so the cost never falls as packed rows grow, and lower ≤ exact holds for split pruning; lower bounds pass zero segments.
  • Beyond 8x the observed sharing, packed rows are priced as if the sharing were 8r, so the model term grows at B/(8r) per logical token rather than main's B/r.
  • The GDN term is one layer's fp32 recurrent states plus conv history per grad segment; backward recomputes one layer at a time.

• Retained, inside the existing trust windows (ratio ≤ 8r, P ≤ 8x the profile's): R·P + 6 KiB·L. Outside them it falls back to required, as before. The existing split-search guard for retained trust crossings stays.

• The 6 KiB charge (12 allocator blocks). Each request's head buffers are separate allocations rounded up to 512 B blocks. On an H200, through the real head projection path (a stub bf16 output layer) and backward, a duplicate one-token request peaked at nine blocks plus 8 B (4,616 B) at the end of head forward, and 3,080 B stayed live into backward. The slope was the same at 20,000 and 100,000 duplicates. 061's CISPO caller saves one or two tensors per request (512 B–1 KiB, plus about 5 B per token).

• Cost. These figures come from the formula with the measured rates:

  • 061's first shared batch: about 8.3 GiB, against 6.9 GiB observed and the report's 56.3 GiB prediction; the logical charge adds 639 MiB, and the clamp is inactive at ratio 6.7.
  • Controlled ratio 8: about 6.3 GiB, against 4.9 GiB observed (the charge is 15% of B·P).
  • The ratio-72.6 GDN branch case: 25 GiB on the first call from a ratio-1 profile, then about 6.1 GiB once the ratio is learned. Observed: 2.5 GiB; main: 171 GiB.

• The recompute flag is recorded in planner reports and restored on replay. Reports written before this change fail replay as incomplete rank fields, as with earlier new fields.

• _request_mix_key also classifies labels by their normalized shape, matching _forward_item: an input of shape (1, N) with targets (N, K) is target:(K,), not target:single.

• Width probes still assume up to two GDN segments per request (an upper bound). Under packed pricing that adds a per-segment charge to the cheap feasibility bound, so a probe can reject a width the exact plan would admit. This is conservative: it can never admit more.

• Why the scope. Private review found memory that grows with logical rows or segments outside the head and caller budget: wide-label index copies, backward gradients of grad-enabled logits/hidden_states, and GDN branch states in no_grad forwards. The measured and production evidence covers grad-enabled single-target training, so only that path changes.

Tests

• CPU:
  • The packed-pricing scope test pins exact values for:
    • packed and extrapolated signatures (no_grad, mixed grad, dense output, wide labels) at profile ratio 2;
    • retained charges and the 8x clamp;
    • a packed-row sweep at rates of 1, 24,575 and 100,000 bytes per packed token, where cost never falls;
    • the per-segment GDN charge;
    • recompute detection from the live decoder config and mode;
    • the recorded replay flag.
  • The short-request gate: 63/64/65 tokens for (N,) and (1, N) inputs, with inactive unaffected. After calibrating on a 64-token request, a 63-token batch and a mixed batch reuse that profile (none pre-seeded) and price at main's extrapolation.
  • The clamp: a packed sweep across L/(8r) is monotone, and learning more sharing never cheapens a plan at or below the older ratio.
  • Replay rejects an unrecorded recompute mode.
  • Tests that exercise packed pricing with short fixture requests set the threshold to 1; tests of main's path keep main's expectations.
• CUDA (skipped without CUDA; they run in the 2xH200 CI job):
  • The per-request slope through the real head, a caller loss and backward, against the charge on the request's logical rows: head only at 1 token; 061's CISPO loss at 64 and 512 tokens; a caller saving a fixed 256 KiB per request at 64 tokens.
  • A negative control: that fixed caller on one-token requests exceeds the charge, which is why short requests are gated.
  • A whole-call check: a profile learned from 16 shared 4,096-token requests with the fixed caller must admit 20 shared 64-token requests at an estimate above their measured peak.

Validation

• Rebased onto main after Apply trainer-rank memory floors per context-parallel rank #951 merged.
• CPU-only (CUDA_VISIBLE_DEVICES=""), full tests/unit/test_trainer_rank_*.py plus test_prefix_tree_packing.py on the current head: 1,248 passed, 8 CUDA-only skips. The three failures also fail on current main in isolation: the live-graph collective test hits its wall-clock "hung" deadline, and test_native_checkpoint_gather_after_recovery fails a different pair of variants on each run.
• Private review: after twelve rounds, the supported-caller invariant was agreed with both reviewers before implementation. Round thirteen found nothing blocking at this head.
• CUDA on one H200: the six new tests pass (head slope; CISPO at 64 and 512 tokens; the fixed caller at 64 tokens; the one-token negative control; the calibration transition). CI's 2xH200 job also runs this file.
• Lint and types: ruff and ty are clean.
• Measurement mode: the controlled measurements ran forward plus loss.backward() (grad-enabled, target_tokens only), and the 061 reports come from grad-enabled training. This covers GDN models with up to 21 segments per batch.

🤖 Generated with Claude Code