Price TP4 sequence-parallel boundary shards in the checkpoint floor

[bradhilton]

At TP>1 the full-recompute checkpoint floor returns nothing, so a cold wave on a TP4 sequence-parallel model is priced only by the static per-token floor. That floor covers the saved layer boundaries but not the recomputed layer's workspace, until a memory profile is learned.

In two 062 SFT runs on dense Qwen3.8-27B at TP4, the first 25,728-row wave admitted 4.64 GB (4.22 GB raw) and peaked at 7.06 GB on all four ranks: +67.5% raw, +52.3% admitted. That was one wave per run, replicated across the TP ranks. It didn't OOM because about 77 GB was free. Later, profiled waves were covered.

Trace. A random-init trace of that exact wave on four H200s reproduces it: 4.637 GB admitted, 7.13 GB peak on every rank (+69.2% raw, +53.8% admitted). At the peak:

• the saved boundary shards are 4.22 GB (per rank, 64 × H × 2 bytes per sharded row);
• the rest, 2.92 GB, is the recomputed layer's transient workspace over the SP-gathered rows:
  • the MLP FC1 stage (6F/TP);
  • the mixer;
  • the gathered norm input (2H);
  • norms and other workspace;
  • one input-gradient shard.

Change.

• The floor now applies at TP4 with sequence parallelism, charging each rank's boundary shard (ceil(rows/TP) rows per boundary). As at TP1, the repeated boundaries serve as the input-gradient term, and here they cover that workspace.
• GDN recurrent states grow with segments rather than rows, so for gradient waves the recomputed layer's initial and final states are priced per segment.
  • The count covers the gradient groups' segments plus up to TP − 1 padding roots per group, since padding gives each pad token its own root.
  • Width probing now counts only gradient groups' segments:
    • twice their requests as the no-sharing upper bound, for accepting a width;
    • one per group as the full-sharing lower bound, for rejecting one;
    • the selected layouts' actual counts when pricing exactly, matching the materialized plan.
  • Kernel-internal chunk states aren't bounded by this term.
• For the traced wave that is 8.43 GB raw (9.28 GB admitted) against the 7.13 GB peak.

Scope. Only where that cover is justified by the traced decomposition:

• TP exactly 4, CP1, PP1, sequence parallel, full/uniform/1 recompute;
• dense models only, with readable mixer geometry (including the GDN conv kernel) and at least TP KV groups (no replicated QKV);
• deep enough that the repeated shards cover the traced workspace at the model's widths: layers × H ≥ 4 × (2H + H + H) + max(6, activation factor) × F + the wider mixer's projection width + H. That is about 49 layers for this model.

TP2, TP8, TP with CP, MoE, and shallower or wider models keep today's pricing; this trace doesn't establish their coverage. No-grad waves at TP4 are priced as before.

Before/after (the traced wave):

• cold admission: 4.64 GB → 9.28 GB (peak 7.13 GB);
• warm admission: 8.02 GB (profile) → 9.28 GB.
  The floor is also a minimum for profiled waves and for each child of a split plan, as main already does at TP1; cold split plans sum their children's floors. So qualifying TP4 gradient waves roughly double their cold static floor, and steady-state packing gets about 16% more conservative here.

Testing: 22 new CPU tests:

• the traced wave's floor and admission;
• ceiling sharding and gradient-only groups;
• segment states priced with 4,096 segments, TP-padding roots, and no charge for no-grad waves;
• width probes counting only gradient segments (acceptance and rejection bounds, and the exact count);
• an ungated attention-only model;
• fallbacks for TP2, TP8, CP2, PP2, no SP, SP at TP1, selective recompute, MoE (layer and geometry), replicated QKV, missing attention geometry, a missing conv kernel, a shallow model and a wide FFN;
• the depth boundary (49 vs 48 layers).

Each gate condition was mutation-tested. The trainer-rank suite passes.

🤖 Generated with Claude Code