Repository navigation
revert: restore pre-#925 trainer execution with API-name aliases - #1107
Merged
Merged
Conversation
bradhilton
had a problem deploying
to
trainer-rank-gpu-validation
October 9, 2026 02:35 — with
GitHub Actions
Failure
bradhilton
deployed
to
trainer-rank-gpu-validation
October 9, 2026 03:06 — with
GitHub Actions
Active
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
…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.
…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.
bradhilton
force-pushed
the
revert/trainer-925
branch
from
October 10, 2026 15:36
409634e to
fff6fcb
Compare
bradhilton
deployed
to
trainer-rank-gpu-validation
October 10, 2026 15:36 — with
GitHub Actions
Active
bradhilton
marked this pull request as ready for review
October 10, 2026 17:22
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: 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:
trainer_rank/andmegatron/match feat: add versioned trainer graphs and logical-rank APIs #925's parent652e0f085, apart from the independent fix(megatron): skip TE grouped GEMM when every group is empty #1089 and the fixes below. The five trainer commits betweendeb95f3fand that parent stay.46c911011). Without it, padded packed exports kept padding and dropped live rows.63b91ca90,b745db5a9).b745db5a9).forward_batches,forwardandreduceare the restored methods, andbackward(loss)matchesloss.backward().options=raisesTypeErroron every method and the constructor, even asoptions=None.ForwardOptions,ModuleHandle, CPU offload and replay, staleness corrections, and thecloudpickledependency.652e0f085, per Brad's decision. On current real paths it stays inactive, because Preserve literal reasoning tags in assistant history #1009's normalization already produces the template it gates on; three tests cover the legacy template directly. Otherwise tokenizer and chat-template behavior is unchanged from main.Validation:
quality-checkscommands):quality-checks, the only failure istest_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:
deb95f3fyet. The comparison long runs started from this revert are the first one.trainer.queuelists, nooptions=, and ports of #1501, #1663, #1741 and #1745.memory_placementfield.