Skip to content

fix(trainer-rank): price CP2 recompute using each rank’s layouts - #978

Draft
bradhilton wants to merge 37 commits into
dalinar/recompute-mixer-floorfrom
dalinar/cp-layout-recompute-floor
Draft

bradhilton wants to merge 37 commits into
dalinar/recompute-mixer-floorfrom
dalinar/cp-layout-recompute-floor

Conversation

@bradhilton

@bradhilton bradhilton commented Sep 25, 2026 •

Copy link
Copy Markdown
Collaborator

Agent: Dalinar

At CP2, the two ranks can hold different attention and GDN layouts. Price each rank’s saved boundaries, recomputed mixer, retained attention state and MoE stage from its own rows, then take the larger floor. Pair #1002’s pending adapter-gradient term with that same rank’s floor. This avoids combining one rank’s gradients with the other rank’s peak, which overcounted a short-sequence case by up to 632 MB.

Scope is CP2/TP1 gradient groups using ART’s CP attention without a softmax offset; other configurations retain the old floor. Building the peer layout costs about 40–76 ms initially and under 1 ms from cache, now counted in planning time. Reports retain the layouts for CPU replay.

Recorded trainer-rank, CP and replay tests passed. Local two-H200 Qwen3.6-35B-A3B traces validated the executor mirror/layout ledger within 0.3%. The layout correction left measured peaks and losses unchanged and accuracy roughly unchanged: real-data EP2 raw error was +18.2% cold/+20.2% warm. On the full stack with #1002, cold EP1 errors were +0.9–2.0% and EP2 +3.0–13.4%, with no measured underestimates.

The executor itself has not been tested directly against its mirror; that needs GPUs. CP4 and other executor configurations retain the old floor. GDN width and routing allowance remain unchanged, and routing still accounts for much of the EP2 excess. Stacked on unmerged #963 and includes #1002; review #1002 and #963 first.

bradhilton and others added 5 commits September 25, 2026 22:28
A size-only mirror of the executor's record_for_backward path, beside the
code it mirrors: per stage, the Q/K/V flex consumes (padded copies, else
contiguous copies of permuted views, or kept gathers and fetch buffers),
flex's output and LSE at the execution length, logical copies when padded,
and a merge-tape clone for every producing stage after the first. It
reproduces the traced CP2 ranks: 0.971 GB for an aligned single stage and
3.07 GB for a padded local plus full-query remote stage.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The checkpoint floor priced every rank on the busiest rank's rows with one
mixer width, which over-counted boundaries and GDN and under-counted a rank
whose attention runs two full-query stages. At CP2/TP1 with ART's CP core
attention, price each rank on its own layouts instead: saved boundaries by
each layer input's layout, a recomputed attention layer on attention rows
plus what the executor keeps for backward, a GDN layer on GDN rows, and the
MoE stage on that layer's rows with routed rows on the EP share. Admission
takes the largest rank. Elsewhere the busiest-rank floor is unchanged.

A rank can receive more routed rows than it holds, so routed rows are no
longer clamped to local rows.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The split search's lower bound priced gated plans with the busiest-rank
attention width, above what per-rank pricing can charge; price even-share
layouts at the least attention state instead (one aligned local stage per
rank). Count flex's saved LSE beside the normalized copy it returns.
Share one gate between plan pricing and the bound, checking the topology
before model state. Record group layouts in planner evidence.

Tests cover the positive gate path against the executor's own plans, the
runtime layouts, the softmax-offset gate, tiny and single-head stages, and
the lower bound staying below plan cost.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
The gate now declines models with several chunks or a decoder without
layers before reading them. The bound holds because every rank's total
grows with its own rows, so the largest is at least the total at the mean;
rounding the even share up can exceed a split's exact cost. The lower-bound
test covers even, skewed and odd splits and one long sequence.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 25, 2026 23:53 — with GitHub Actions Active
…planning

One head does not make a view of a fused QKV split contiguous, so the
mirror now charges the copy flex makes of any full local view; the lower
bound still counts only multi-head copies. A partial-query merge tape also
keeps its int64 row index. The per-rank layout work now counts toward
planning time (about 40-76 ms for a new layout's peer plan, under 1 ms once
cached). Tests use typed geometry helpers.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 26, 2026 00:24 — with GitHub Actions Error
The executor keeps a partial stage's int64 row index even when that stage
produced first and recorded no accumulator tape, so charge indices apart
from the tape the mirror drops.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Picks up main (0e0c31b) through the updated #963 branch. The layout floor
sits between #988's TP > 1 floor and the generic floor, and #971's combine-
extent floor applies after all three. Tests pass routed rows and layouts by
keyword.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 26, 2026
Picks up main (0e0c31b) through the updated #963 and #978 branches.
_checkpoint_floor_decoder stays TP1 by default, so the dense widths and
their discounts remain TP1/CP2 only; with sequence_parallel it applies main's
#988 checks for the traced TP x SP floor. _checkpoint_memory_floor prices
TP > 1 with #988's floor, then the layout and generic floors, and #971's
combine-extent floor after all three.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 26, 2026
Picks up main (0e0c31b) through the updated #963 and #978 branches. Clean
merge; the source equals the composition measured in M9 plus #998.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 26, 2026 23:14 — with GitHub Actions Failure
bradhilton and others added 2 commits September 26, 2026 23:27
…outs

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 27, 2026 00:32 — with GitHub Actions Failure
bradhilton and others added 3 commits September 28, 2026 06:48
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton and others added 15 commits September 28, 2026 13:24
Per-rank CP layouts: the head stage takes the largest rank's adapter term
over its own boundaries, sharing the per-rank walk with the decoder stage's
extra.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
…acted layout

Port #978 (through 6dc567b) onto the extracted planner modules: every
CP2 rank's checkpoint boundaries and recomputed layer are priced on its
own attention and GDN layouts plus the CP executor's retained stage
records, with adapter gradients paired per rank. Grouped planner-miss
replay freezes each group's layouts and the decoder's per-layer input
layouts (facts version 4).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Validate an empty GDN layout's length, require the TP1/CP2/PP1 topology
that layout pricing is limited to, and refuse a recorded group_layouts
argument in any estimate.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation September 29, 2026 07:11 — with GitHub Actions Error
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 29, 2026 07:43 — with GitHub Actions Active
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Sep 29, 2026
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton bradhilton changed the title Price CP2 recompute per rank on each layer type's layout fix(trainer-rank): price CP2 recompute using each rank’s layouts Oct 2, 2026

This branch was successfully deployed

1 active deployment
trainer-rank-gpu-validation — bf79e446 Deployed Sep 29, 2026 by bradhilton via Run on 2x H200 #976
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant