Skip to content

fix(trainer-rank): raise routed-row estimates from checkpoint observations - #981

Draft
bradhilton wants to merge 29 commits into
dalinar/cp-layout-recompute-floorfrom
dalinar/routed-share-observation
Draft

bradhilton wants to merge 29 commits into
dalinar/cp-layout-recompute-floorfrom
dalinar/routed-share-observation

Conversation

@bradhilton

@bradhilton bradhilton commented Sep 26, 2026 •

Copy link
Copy Markdown
Collaborator

Agent: Dalinar

A fixed expert-routing allowance can repeatedly underprice a skewed checkpoint. Observe each checkpoint’s worst routed share and raise later row estimates to that share plus 0.10 when it exceeds the allowance. Never lower a charge. This applies to HybridEP with matching EP/CP groups and TP/ETP1.

Only single-group micro-batch waves produce observations; split, multi-group and dp_rank_forward paths use prior observations. Ranks exchange shares with checkpoint epochs and record nothing on unsupported, failed, cancelled or mismatched observations. Loads, snapshots and discards check matching identities. Disk loads reset observations, snapshots inherit them, failed loads restore the previous epoch, and optimizer steps retain the high-water mark. Updated rows feed admission, splitting, costs and telemetry.

Recorded validation: 24 new CPU cases and the trainer-rank suite passed, including a two-process Gloo exchange and real checkpoint paths. A local two-H200 Qwen3.6-35B-A3B CP2/EP2 run made 12 real-data calls per arm. Zero-learning-rate shares stayed at 1.15–1.16; training reached 1.364 and raised later charges. No measured call was underestimated. These observations do not bound all routing or model configurations.

Includes #1002 via #963/#978; review those unmerged dependencies first. A warm discount is excluded. Planner replay still does not cover expert-parallel plans; EP1 records no share.

bradhilton and others added 4 commits September 26, 2026 04:20
At EP>1 TrainerRank prices the rows each rank's experts receive at a fixed
allowance over balanced routing (1.4x at EP2); skew beyond it went unseen.
Record each checkpoint's worst-layer routed share at the micro-batch
handoff, from HybridEP's retained tokens_per_expert, through the existing
handoff MAX reduce (EP>1 only: share, epoch, -epoch). A share is recorded
only when every rank observed the same checkpoint route epoch. Epochs are
assigned after a load or snapshot commits on every rank, never reused, and
forgotten on reload or discard.

When a checkpoint's share plus a 0.10 margin exceeds the allowance, its
balanced routed rows are scaled up to match. Charges never go down. The
share is reported in last_forward_telemetry() and planner-miss reports.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Review follow-ups:
- Loads and snapshot discards now check that every rank targets the same
  name and route epoch, so every rank forgets and assigns the same epochs.
- Cancellation during the observation still reaches the handoff exchange,
  as a failed forward.
- Epochs at or above 2**40 are not observed, so float64 transport stays
  exact.
- No-grad waves are observed: they dispatch the same way, and observations
  only raise charges.
- A snapshot starts from its source's share, since it has the same weights.

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 26, 2026 05:13 — with GitHub Actions Error
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 05:23 — with GitHub Actions Active
@bradhilton
bradhilton marked this pull request as ready for review September 26, 2026 15:53
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 26, 2026 15:54 — with GitHub Actions Active
@bradhilton
bradhilton marked this pull request as draft September 26, 2026 16:08
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 deployed to trainer-rank-gpu-validation September 26, 2026 23:14 — with GitHub Actions Active
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 Error
bradhilton and others added 8 commits September 28, 2026 06:53
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>
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 28, 2026 10:56 — with GitHub Actions Active
bradhilton and others added 12 commits September 28, 2026 13:25
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>
… the extracted layout

Port #981 (through e288179) onto the extracted planner modules: at the
EP>1 handoff every rank reads its HybridEP received rows, and a
checkpoint whose worst observed share plus a margin exceeds the cold
allowance is priced at that share; route epochs key the observations to
committed checkpoint content. Grouped planner-miss replay is unaffected
(capture refuses expert parallelism); the producers of recorded routed
rows join its stock-estimator check.

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>
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation September 29, 2026 09:15 — with GitHub Actions Active
@bradhilton bradhilton changed the title Raise EP routed-row pricing to each checkpoint's observed routing fix(trainer-rank): raise routed-row estimates from checkpoint observations Oct 2, 2026

This branch was successfully deployed

1 active deployment
trainer-rank-gpu-validation — 07bc450c Deployed Sep 29, 2026 by bradhilton via Run on 2x H200 #978
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