Skip to content

fix(trainer-rank): credit staged adapter gradients and shared captures in graph admission - #1081

Merged
bradhilton merged 2 commits into
mainfrom
dalinar/925-admission-accounting
Oct 3, 2026
Merged

bradhilton merged 2 commits into
mainfrom
dalinar/925-admission-accounting

Conversation

@bradhilton

@bradhilton bradhilton commented Oct 3, 2026 •

Copy link
Copy Markdown
Collaborator

Agent: Dalinar

#925's graph-placement admission over-prices the first gradient wave of a slot by charging the same memory twice:

  1. Adapter gradients: it reserves 3× the slot's parameter bytes as gradient staging (current, staged and replacement), and also prices fix(trainer-rank): account for pending adapter gradients during recompute #1002's pending adapter gradients inside the unit's workspace. Those pending gradients are the staged gradients, the same bytes.
  2. LoRA version snapshot: the snapshot is charged once per unit or split child, but every unit with the same slot and staleness shares one capture.

This change credits the staged adapter gradients against the unit's workspace, so the full 3× staging reserve stays and rollback, the multi-forward prepare and the in-flight margin remain covered. It also charges each shared snapshot once. Nothing else in admission changes: CP no-release credit, custom heads, checkpoint sharing and learned split floors are untouched.

On a Qwen-like CPU fixture (40 layers, 920 MiB slot):

Case Admitted on main Admitted after
first wave, 1,024 rows 5,865 MiB 5,101 MiB
first wave, 4,096 rows 6,431 MiB 6,135 MiB
later waves unchanged unchanged
split, two same-slot children 6,244 MiB 4,716 MiB

Testing: new admission tests fail on main and pass here. 2,257 trainer-rank CPU tests pass. GPU CI runs on this PR.

Limits:

  • The arithmetic is CPU-derived from the code's allocation lifetimes, not a GPU measurement, and the fixture numbers are not a GPU qualification. The credit equals fix(trainer-rank): account for pending adapter gradients during recompute #1002's net pending adapter-gradient contribution, not the full staged tensor bytes.
  • Admission still over-prices in places. The replacement-gradient share never coincides with a backward, but removing it needs a measured in-flight gradient margin, and warm profiles may already include the snapshot. Both are left for a GPU-logged run.

🤖 Generated with Claude Code

…s in graph admission

Graph admission added the 3x gradient staging reserve beside each unit's
workspace, but the recompute walk's pending adapter gradients (#1002) are
the same _GradientBatch clones the reserve's staged share holds, so they
were charged twice. Credit them against the unit's workspace; the full
staging reserve (current, staged, replacement) stays for multi-forward
transactions.

Split children and groups of one slot and staleness reuse the first
capture's live LoRA weight snapshot; charge it once per admission.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation October 3, 2026 03:50 — with GitHub Actions Error
…tinct targets

The backward walk repeats a slot's pending inventory for each gradient
group priced together (routed and unrouted requests of one slot and
policy), so AG can reach 2P - B while one transaction stages each
distinct target once. Cap the credit at the pending bytes of the unit's
distinct gradient targets: the walk's pending gradients are a subset of
that staged aggregate.

A no-grad correction forward allocates no adapter gradients and can
coexist with the current gradient, staged head gradients and a
current-weight capture, so its workspace keeps no credit.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation October 3, 2026 04:15 — with GitHub Actions Active
@bradhilton
bradhilton marked this pull request as ready for review October 3, 2026 04:42
@bradhilton
bradhilton merged commit 6632b8d into main Oct 3, 2026
10 checks passed
bradhilton added a commit that referenced this pull request Oct 3, 2026
Brings in #1068 (eager head workspaces in admission and replay), #1067
(dense CP1 no-grad floor), #1081 (staged adapter-gradient credit), #1077,
#1080 and CI fixes. Conflicts and reconciliation:

- _group_head_workspace_bytes: TP1 keeps #1068 verbatim (capacity reserves
  the eager 7-buffer statistics, 3 buffers for requested logits). TP > 1
  goes through _tp_head_workspace_bytes, which uses #1068's model per chunk:
  the same 7-buffer eager charge where no kernel is attempted, the kernel
  path plus the bounded statistics' increment where one is (after a
  failure the bounded statistics, not the 7-buffer path, run), and #1068's
  logits copies gathered to the full vocabulary (1 + 2 TP buffers).
- Replay: one _frozen_head_bytes over #1068's version-3 facts
  (head_statistics, head_target_backward). TP1 is #1068's formula; TP > 1
  prices the kernel path, and capture declines TP > 1 waves with an eager
  chunk or with requested logits beside statistics, which the facts do not
  record. The optional rank facts (mixer_top_gaps, lora_modules_per_layer)
  sit beside #1067's defaulted dense_fc1_adapted.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
bradhilton added a commit that referenced this pull request Oct 10, 2026
…admission (#1081)

This reverts commit 6632b8d.

Dropped: this only prices, credits or places #925 graphs and outputs, which
this stack removes.

This branch was successfully deployed

1 active deployment
trainer-rank-gpu-validation — 3468be97 Deployed Oct 3, 2026 by bradhilton via Run on 2x H200 #1128
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