Price dense recompute by its traced stage and one input gradient

[bradhilton]

For dense models TrainerRank charges one input gradient per saved layer boundary, on top of the boundaries themselves. It also prices several no-grad groups as if they ran at once. On dense Qwen3.8-27B at CP2 this over-prices gradient waves by about 40% and multi-group no-grad waves by up to about 2.5x. In q062 the planner refused a wider wave in 242 of 269 decisions.

Allocator traces of dense Qwen3.8-27B at CP2 (random init, reproducing q062's shapes exactly) show:

• the recomputed layer's peak holds one input gradient;
• its MLP FC1 stage holds the base output, the LoRA output and their sum, plus one F-wide tensor: 7F per row;
• a no-grad group holds 6F + 6H per row, and groups run one after another.

What changes (dense models at CP2 only, gated)

• Structural gate. _dense_mlp_recompute_bytes_per_token returns the per-row gradient stage and no-grad transient only when every decoder layer matches the traced execution:

  • the TE layer-norm FC1 and row-parallel FC2 with their LoRA wrappers;
  • the fused SwiGLU that Megatron Bridge's Qwen3.5 providers configure, with no clamp or offset; BF16; no bias or FP8/FP4; TP1/PP1; no sequence parallelism or CUDA graphs;
  • LoRA ranks up to 256, rechecked per active slot;
  • the checkpoint floor's own decoder conditions (full/uniform/1 recompute in training), so a discount always carries its TE workspace growth;
  • Megatron TransformerLayer with exactly the traced mixer types: SelfAttention, Bridge's Qwen3VLSelfAttention or GatedDeltaNet;
  • nothing the layer runs, the mixer's children included, has hooks or forward overrides besides ART's own GDN layer, mixer and empty-safe norm wrappers, which must still call the class's forward (through Dynamo's compile wrapper, as in training);
  • every adapter in the layer, the mixer's too, is an exact LoRA that selects its tensors the way execution does, with every adapter's rank priced.

  The traced model is a GDN/attention hybrid. A pure-attention dense model with the same MLP passes the gate on extrapolated widths.

  Anything else returns 0 and keeps today's pricing.

• Eligibility. Decided once over all of a wave's slots, so one unsupported slot keeps both allowances.

• Gradient groups. Priced with the recomputed mixer, residual and norm, and the MLP stage. The stage also carries one more FC1 triplet (6F) for the early-run recompile residue seen in q062. The gradient is one row of H instead of one per boundary. This applies to both the busiest-rank floor and Price CP2 recompute per rank on each layer type's layout #978's per-layout floor, plus TE workspace growth (no-grad-only waves included).

• No-grad groups. Charged their 6F + 6H transient, including beside retained gradient boundaries in mixed waves.

• Multi-group no-grad waves. Priced at the largest group's own rows at that width, and at least its share of today's per-token floor. A single no-grad group also gains its per-row transient from the checkpoint floor, which is conservative. That share falls as another group's rows grow, so the split planner's lower bound keeps only the largest group's rows.

• Not changed. CP1, CP>2, TP>1, MoE, the 1.10 factor and learned profiles.

Before/after (q062, 269 real waves held out from the trace fit)

• Gradient waves: the floor alone is never under an observed peak (min +8.1%). The median raw error moves only from +40.9% to +35.8%, because the always-on recompile allowance keeps cold waves safe without a profile.
• No-grad waves: median from +137% to +32%, never under.

Testing: 65 new unit tests cover the gate (43 structural negatives, including a compiled delegate, the Qwen attention class, mixer children and adapters, and LoRA selector overrides), named slots through the real slot lookup, per-slot ranks and mixed-slot fallback, the checkpoint-floor conditions, the split lower bound, one gradient, both floors, the no-grad transient in mixed waves, unequal group ratios and the CP2-only scope. Replay of older reports stays compatible. The trainer-rank suite passes.

On dense Qwen3.8-27B compiled as in training, the gate is active on the real model, including its walk over each layer's children. The GPU results below come from the previous revision, whose widths were 0.06% lower. On 2×H200 (CP2, random init), held-out single waves:

• gradient waves of 35k and 70k tokens: +28% raw;
• a two-group no-grad wave: +13% (main's formula: +119%).

None is under. On q062-shaped batches, #986 packs 10 waves where main needs 11, but runs no faster. The binding limit is the ~36 GB of free memory, not the estimate.

Stacked on #978 (itself on #963); none of the stack is merged, so review in that order.

🤖 Generated with Claude Code