Repository navigation
fix(trainer-rank): price CP2 recompute using each rank’s layouts - #978
Draft
bradhilton wants to merge 37 commits into
Draft
bradhilton wants to merge 37 commits into
bradhilton wants to merge 37 commits into
Conversation
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
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
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>
bradhilton
deployed
to
trainer-rank-gpu-validation
September 26, 2026 00:31 — with
GitHub Actions
Active
This was referenced Sep 26, 2026
bradhilton
marked this pull request as ready for review
September 26, 2026 15:46
bradhilton
had a problem deploying
to
trainer-rank-gpu-validation
September 26, 2026 15:46 — with
GitHub Actions
Failure
bradhilton
marked this pull request as draft
September 26, 2026 16:08
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
bradhilton
had a problem deploying
to
trainer-rank-gpu-validation
September 26, 2026 23:14 — with
GitHub Actions
Failure
…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>
bradhilton
had a problem deploying
to
trainer-rank-gpu-validation
September 26, 2026 23:54 — with
GitHub Actions
Failure
This was referenced Sep 27, 2026
bradhilton
had a problem deploying
to
trainer-rank-gpu-validation
September 27, 2026 00:18 — with
GitHub Actions
Error
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton
had a problem deploying
to
trainer-rank-gpu-validation
September 27, 2026 00:32 — with
GitHub Actions
Failure
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
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
had a problem deploying
to
trainer-rank-gpu-validation
September 29, 2026 07:11 — with
GitHub Actions
Error
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
This was referenced Oct 2, 2026
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.