Move TrainerRank checkpoint-slot and dynamic-optimizer management into trainer_rank/_slots and _optimizer

[bradhilton]

Summary

Third and last of the stacked TrainerRank extractions (base: #1000, which is on #999). Moves 41 methods out of src/art/trainer_rank/_impl.py:

• 31 checkpoint-slot bookkeeping methods → src/art/trainer_rank/_slots.py (_resolve_slot_ref, _ensure_checkpoint_slots/_for, _register_checkpoint_prefetch, _load_checkpoint_slot, _load_registered_checkpoint, _guard_slot_can_load, _guard_checkpoint(s)_can_step, _validate_checkpoint_adapter_config, _validate_loaded_checkpoint_config, _validate_checkpoint_consistency, _selected_dynamic_checkpoints, _checkpoint_grad_flags, _track_slot_graph_outputs, _slot_graphs, _prune_slot_graphs, _has_live_slot_graph, _push_checkpoint(_sync), _discard_snapshot_checkpoint, _resolve_checkpoint_name, _checkpoint_group, _set_default_slot, _iter_slot_parameters, _resolve_custom_checkpoint, _checkpoint_prefetch_waiter, _prefetched_checkpoint, and the public prefetch_checkpoints, load_checkpoint, pop_checkpoint).
• 10 dynamic-optimizer methods → src/art/trainer_rank/_optimizer.py (optim_step, _dynamic_optim_step, _guard_optim_step_configuration, _dynamic_optimizer, _new_dynamic_optimizer, _extend_dynamic_optimizer, _restore_canonical_optimizer, _dynamic_param_step_flags, _zero_dynamic_optimizer_padding, _dynamic_optimizer_padding_masks).

Same pattern as #999/#1000 and the existing _gdn_memory.py/_planner_evidence.py: module-level functions def name(self: TrainerRank, ...) with verbatim bodies, bound in the class body by assignment. No collaborator object, no delegators, no __init__ change; sibling calls stay self._x(...) so the ~29 per-instance monkeypatches and the TrainerRank.__new__-built ranks in the test suite keep working, and the 42 call sites in _checkpoint.py (untouched) still resolve through the rank. None of the 41 originals was decorated. Names tests patch on _impl (torch, dist, _telemetry_phase) plus the _impl-defined classes/helpers the bodies use are read as _impl.X at call time. _resolve_slot_ref has a def-time = _impl.Unset default, so both modules are imported on the shared pre-class import line. The module was named _optimizer.py (not _dynamic_optimizer.py) to avoid colliding with the method of that name.

Stateless @staticmethods referenced widely from _checkpoint.py and tests (_slot_ref, _slot_state_error, _checkpoint_source, _checkpoint_source_key, _await_checkpoint_prefetch) and the thin public I/O wrappers (snapshot_checkpoint, save_checkpoint, export_lora, …) stay in _impl.py.

Equivalence and stacking

• All 41 function bodies are AST-identical to main's current method bodies (modulo _impl. qualification and the self annotation). None of the 41 was edited on main since the fork (Price warm waves by the fit later plans measured, not the first plan's #991/Validate canonical checkpoint optimizer counters #992/Preserve GPU image smoke and downstream dispatch #993 checked per method), so nothing needed porting.
• With Move TrainerRank memory accounting into trainer_rank/_memory #999 and Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner #1000 applied, the merged class has 82 remaining defs (83 counting the forward_micro_batches implementation separately from its @overload stubs), all AST-identical to main's, and 101 bindings (26 memory, 34 planner, 31 slots, 10 optimizer); defs ∩ bindings = ∅ and defs ∪ bindings = main's 183-method set. The rebase was done by regenerating both files from Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner #1000's _impl.py with the extraction generator rather than hand-resolving markers, so no stitched methods are possible. Newly dead imports (deepcopy, hashlib, LocalOptimizerState) removed; ruff --select F401 clean across _impl.py and all four extraction modules.
• All six import orders work with the TrainerRank.__new__ smoke passing under each.

Validation

• tests/unit/test_trainer_rank_*.py at -n 4: 1,311 passed / 19 failed on both the Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner #1000 baseline and this branch, identical failure sets (pre-existing megatron/distributed environment failures). test_trainer_rank_backward_work.py serially: 23 passed; the eight methods its actual_rank harness extracts from _impl.py by name are all still defs.
• ruff check, ruff format --check clean; ty check shows the same 6 pre-existing _checkpoint.py diagnostics as main.
• Independent review (thermo-nuclear rubric): 41/41 verbatim, no decorators, no sibling-module calls, _impl.X table cross-checked against every _impl name patched in tests, identical failure sets on base and branch across 1,613 tests. Verdict: safe to land; its one nit (the module/method name collision) is resolved by the _optimizer.py name.

_impl.py: 9,843 → 5,632 lines across the three PRs, with 101 methods relocated into four concern-named modules.

Note for xdist runs: test_trainer_rank_gpu_ci.py parametrizes from a set, so workers collect in different orders unless PYTHONHASHSEED is fixed — pre-existing on main.

[bradhilton]

Schulman — no blocking findings in the reviewed scope.

Reviewed head 5867bd43517da87e9a7fa123db0fafcfe8965816 against stacked base 35db95460bad370c1fd2222356e71a5327f96436, plus the cumulative stack from 3ad9a5f8d29ed4ea7f8e033e7e1f6cda3fe55607. Live head/base were rechecked immediately before posting. An independent delegate performed the review and I examined its evidence; we did not author these extraction PRs.

The method bodies match after narrowly accounting for _impl qualification, self annotations, docstring indentation and restored staticmethod bindings. Slot graph-lifetime guards, checkpoint/prefetch dispatch, optimizer creation and state reuse, hyperparameter updates, padding cleanup and failure cleanup survive the extraction. No unresolved runtime globals were found. Packaging covers the new modules; planner-estimator source fingerprints cover the moved estimator dependencies.

Four bounded framework-free CPU cases passed using the actual extracted source: sentinel identity, instance graph-guard dispatch, optimizer creation and existing optimizer/state preservation during configuration updates. Syntax and stacked/cumulative source comparisons passed. No full Torch/ART import-order test, broad pytest, real optimizer numerics, distributed/GPU execution or package build was performed by this review.

Nonblocking documentation nit: this head has 184 total/83 remaining methods under the review's count, rather than the description's 183/82; the extraction totals are correct. Changed function/frame metadata also means frozen source-authenticating diagnostics need explicit adoption rather than silent repinning.

[bradhilton]

Independent source review by Stark and his review delegate; posted by Jarvis at Brad's request. Scope and validation statements below describe that review.

Source review of ART #1001 at 5867bd4, stacked base 35db954: no introduced blocker found. All 41 extracted slot/optimizer method bodies match after explicit qualification/self-annotation normalization; ownership stays on TrainerRank and dispatch remains through self. Import order supplies definition-time defaults before module extraction bindings. This is source-only, not runtime qualification.

Trainer-v1 integration must port its consumed-graph markers into _slots, post-step DP0 head-buffer synchronization into optim_step, and accumulated-version validation/cleanup into _dynamic_optim_step. The full stack also predates main #998 and must preserve its GDN scan fix. Caladan observer source pins/digests need a reviewed update after any integration. These are future reconciliation obligations, not introduced #1001 defects.

Reviewed with a fresh task-local source reviewer; this lane did not author the extraction. No application imports, tests, GPU/provider work or GitHub posting occurred. Exact source evidence is recorded in the reviewer’s retained local report and source-proof file.