Skip to content

revert: restore pre-#925 trainer execution with API-name aliases - #1107

Merged
bradhilton merged 24 commits into
mainfrom
revert/trainer-925
Oct 10, 2026
Merged

bradhilton merged 24 commits into
mainfrom
revert/trainer-925

Conversation

@bradhilton

@bradhilton bradhilton commented Oct 9, 2026 •

Copy link
Copy Markdown
Collaborator

Agent: Tarkin

Problem: #925 sends every CUDA forward through versioned graphs, CPU snapshots and a placement ladder, with no way to turn it off. It cost throughput, caused the 062 MemoryError (#1096), and needed #1081, #1086 and #1099 to price itself. No current experiment uses the shared-service features it was built for.

Behavior after this change:

Validation:

  • Local checks (GPUs hidden, using the quality-checks commands):
    • prek (ruff with the locked 0.15.15, ruff-format, ty, uv-lock-check) passes.
    • Megatron lightweight: 1,604 passed, 14 skipped, with only the known memory test failing.
    • Tokenizer and trajectory suites: 1,966 passed.
    • Alias tests pass.
  • Caladan: all 16 GPT-OSS export cases from caladan#1995 pass, including the 4 padded cases that failed on the previous head.
  • GitHub CI: the required 2×H200 TrainerRank GPU validation passes. In quality-checks, the only failure is test_tensor_parallel_heads_fit_their_charge_without_a_spare_chunk, which fails the same way on main. It's a real head-checkpoint memory bug that fix(trainer-rank): keep head row normalization outside the statistics checkpoint #1120 is fixing, and CI skips the unit step after it.

Limitations:

@bradhilton
bradhilton had a problem deploying to trainer-rank-gpu-validation October 9, 2026 02:35 — with GitHub Actions Failure
@bradhilton
bradhilton deployed to trainer-rank-gpu-validation October 9, 2026 03:06 — with GitHub Actions Active
…1105)

This reverts commit e2a406a.

Temporarily reverted so #925 reverts cleanly; re-applied on the restored
pre-#925 path later in this stack.
This reverts commit f8ec63c.

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

#1099 fixes the #1096 MemoryError raised by #925's output placement check.
This reverts commit 564e1b6.

Temporarily reverted so #925 reverts cleanly; re-applied on the restored
pre-#925 path later in this stack.
This reverts commit 87c6464.

Temporarily reverted so #925 reverts cleanly; re-applied on the restored
pre-#925 path later in this stack.
This reverts commit faa173a.

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

Without the graph memory policy, candidate() already falls back to plans[width],
so this is a no-op on the restored path.
This reverts commit 1f173b4.

Temporarily reverted so #925 reverts cleanly; re-applied with adaptations
to the restored pre-#925 path later in this stack.
…admission (#1081)

This reverts commit 6632b8d.

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

Temporarily reverted so #925 reverts cleanly; re-applied with adaptations
to the restored pre-#925 path later in this stack.
This reverts commit 3886916.

Temporarily reverted so #925 reverts cleanly; re-applied on the restored
pre-#925 path later in this stack.
…1067)

This reverts commit 6ee08f6.

Temporarily reverted so #925 reverts cleanly; re-applied on the restored
pre-#925 path later in this stack.
This reverts commit 6ac64d4.

Temporarily reverted so #925 reverts cleanly; re-applied with adaptations
to the restored pre-#925 path later in this stack.
This reverts commit 1006597.

Restores pre-#925 SPMD trainer execution: trainer_rank/ and megatron/ now
match 652e0f0 apart from the independent #1089. The cloudpickle dependency
goes with it; the uv.lock conflict was the adjacent datrie line #1104 removed.

Rendering is intentionally unchanged. #925 also removed #967's
_preserve_literal_thinking_off_content and part of #1009's chat-template
parser handling. Those, and the tests that cover them, stay byte-identical to
main until a separate investigation decides:

  src/art/trajectories/_tokenize.py
  src/art_inference/chat_template.py
  tests/unit/test_literal_reasoning_content.py
  tests/unit/test_reasoning_parser_joined_branches.py
  tests/unit/trajectories/test_literal_thinking_off.py
  tests/unit/trajectories/test_recorded_boundaries.py
  tests/unit/trajectories/test_recorded_boundary_source_guard.py
)

Re-applies 6ee08f6 onto the pre-#925 planner. The src hunks apply as-is.
The only conflict was in test_trainer_rank_checkpoint_memory.py, where
pre-#925's test_no_grad_enclosure_config_guard sits beside the new tests;
both are kept.
…red path (#1068)

Re-applies 3886916 onto the pre-#925 planner. The src hunks apply as-is.
The only conflict was the prek.yml deselect list: #925's
test_trainer_rank_commands.py no longer exists, so only the eager-head
deselect is added.
…ng on restored path (#1075)

Re-applies 87c6464 onto the pre-#925 planner. The src hunks apply as-is,
plus the functools.lru_cache import #925 had added to _impl.py.

test_trainer_rank_tp_floor.py takes #1075's version, whose only #925
dependency was a helper refactor; tests/unit/trainer_rank_test_support.py is
restored for those helpers. One assertion that priced the split through
#925's placement_cost graph path is dropped; the restored admission compares
`required` directly, which the test already asserts.
…tored path (#1088)

Re-applies 564e1b6 onto the pre-#925 planner without changes.
…ored path (#1105)

Re-applies e2a406a onto the pre-#925 head. The only conflict was the
_impl.py import block: keep gather_target_logits, drop #925's _rng import.

Pricing was rechecked against the restored code rather than carried over.
The head implementation (_impl.py head/statistics functions, topk.py,
_targets.py) and its pricing (_memory.py head terms, _planner_replay.py) are
byte-identical to main. The restored forward calls the same _forward_packed
that #925 wrapped; #925's saved-tensor capture copies were priced in
_memory_policy, not in the head workspace terms, so the index-vector and
dense-buffer terms hold as written.
Re-applies 6ac64d4 adapted to the pre-#925 planner. Under the oversized
opt-in, an expert-parallel refusal is overridable again and keeps the
lower-estimate unsplit plan; internal EP splitting stays disabled.

Dropped with #925's CPU offload: the cpu_fits clause that preferred a
host-admissible plan, its test, and the graph host-budget case added to
#925's test_trainer_rank_memory_admission.py, which stays deleted. The
re-applied tests call the restored dp_rank_forward name.
… path (#1084)

Re-applies 1f173b4 adapted to the pre-#925 planner. The src change applies
as-is: the DP-local EP path cannot bypass admission; the WORLD-coordinated
forward_micro_batches path still can.

Dropped: the graph host-budget case in #925's
test_trainer_rank_memory_admission.py (still deleted) and the cpu_fits test
already removed with #1080. The re-applied tests call the restored
dp_rank_forward and forward_micro_batches names.
…tored path (#1077)

Re-applies da26eb1 adapted to the pre-#925 slots. Caladan's automatic
frozen reload (#1501) calls prepare_checkpoint(forward_only=True),
_register_checkpoint_source and load_checkpoint(_source_owner=...), which
are all kept, along with the snapshot cache and slot-write reservations.

Dropped: the _graph_cache lookup in _has_live_slot_graph; restored slot-graph
markers alone protect residency. Two tests that drove #925's graph cache and
checkpoint versions are rewritten against the same contracts: a live slot
graph blocks eviction until backward, and a native forward-only reload
restores frozen bytes without touching the live training slot. The
distributed recovery test keeps its pre-#925 fixture plus the new fields.
…ethods

forward_batches, forward and reduce are the restored forward_micro_batches,
dp_rank_forward and dp_reduce; backward(loss, gradient, retain_graph=) runs
torch.autograd.backward exactly as loss.backward() does. Old names keep
working. There is no options= keyword on these methods or on the
constructor: passing it, even options=None, raises TypeError instead of
silently dropping a staleness or placement request.

The pre-#925 guard that TrainerRank has no `forward` now checks that
`forward` is dp_rank_forward.
The restored #925 test helper module now carries only what the re-applied
#1075 and #1105 tests import: fake_rank, recompute_model, process_group,
AllReduceSum, all_reduce_max and spawn_and_join. The rest served #925 tests
this stack removes.
@bradhilton
bradhilton merged commit 739dc73 into main Oct 10, 2026
8 of 9 checks passed

This branch was successfully deployed

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