Move TrainerRank memory accounting into trainer_rank/_memory

[bradhilton]

Summary

First of three stacked extractions from TrainerRank (src/art/trainer_rank/_impl.py, 9,890 lines, one class with 193 methods). This one moves the 26 memory-estimation, profiling and cache-recovery-accounting methods (1,229 lines) into src/art/trainer_rank/_memory.py.

Pattern — the same one _gdn_memory.py and _planner_evidence.py already use in this package: _memory.py holds module-level functions def _name(self: TrainerRank, ...) with verbatim bodies, and the class body binds them by assignment (_memory_check = _memory._memory_check; staticmethod(...) for the three statics). No collaborator object, no delegators, no __init__ change. Because bodies still call siblings via self., the ~100 per-instance monkeypatches in the test suite (rank._available_memory_bytes = …), the class-level calls (TrainerRank._split_required_memory(costs)) and the TrainerRank.__new__-built ranks all keep working, and co_name is preserved for _planner_evidence frame records. Names that tests patch on _impl (torch, dist, estimate_prefix_tree_packed_tokens, materialize_prefix_tree_layout, _HEAD_CHUNK_TOKENS, _language_model, …) are read as _impl.X at call time so those patches still intercept. _estimate_flat_forward has a def-time default = _impl.Unset, so _impl.py imports _memory immediately before class TrainerRank: (commented).

Moved: _split_required_memory, _split_memory_key, _record_split_memory_floor, _split_plan_memory_check, _head_workspace_bytes, _group_head_workspace_bytes, _plan_head_workspace_bytes, _plan_hybridep_growth_bytes, _checkpoint_moe_bytes_per_token, _moe_workspace_bytes, _checkpoint_memory_floor, _retained_memory_bytes, _estimate_flat_forward, _update_peak_memory_profile, _estimate_group_request_output_bytes, _memory_signature_from_requests, _slot_memory_shapes, _memory_check, _refresh_memory_check, _memory_check_required, _forward_memory_group, _estimate_required_memory_bytes_from_values, _gdn_segment_layer_bytes, _available_memory_bytes, _all_ranks_have_memory_profile, _update_memory_profile ( = static).

Deliberately left in _impl.py: the _backward_region-wrapped execution/cache-recovery cluster (_try_cache_recovery, _cache_recovery_episode, _release_cached_memory_for_backward, _memory_error_with_reduction_note, _execute_split_plan_with_memory_tracking, _run_flat_plan_with_memory_tracking), because tests/unit/test_trainer_rank_backward_work.py::actual_rank parses _impl.py source and re-executes those by name.

_memory.py is added to _planner_misses._SOURCE_NAMES so the planner-miss replay drift guard keeps fingerprinting the estimator it reruns (total fingerprinted bytes 552 KB < 1 MiB budget); the existing test asserts it.

Equivalence and rebase

• All 26 function bodies are AST-identical to main's current method bodies (modulo _impl. qualification, the self annotation, the dropped @staticmethod, and docstring re-indentation). The class body equals main's minus exactly these 26 methods plus the 26 bindings; module level equals main's plus the one pre-class import.
• Price warm waves by the fit later plans measured, not the first plan's #991 and Validate canonical checkpoint optimizer counters #992 landed on main after this was authored and modified three of the moved methods (_update_peak_memory_profile, _estimate_required_memory_bytes_from_values, _update_memory_profile); those changes are ported into _memory.py by re-running the same transformation on main's bodies. No new _impl. qualifications were needed.
• All three import orders work; _memory._impl is _impl.

Validation

• tests/unit/test_trainer_rank_*.py (minus gpu_ci) at -n 4: 1,256 passed on the branch vs 1,257 on a fresh main baseline; the symmetric difference is three known distributed/subprocess-timeout flakes that pass in isolation on both. main's new test_trainer_rank_profile_warm.py passes.
• ruff check, ruff format --check clean; ty check shows the same 6 pre-existing _checkpoint.py diagnostics as main.
• Independent review (thermo-nuclear rubric): every per-instance and class-level monkeypatch enumerated and verified effective; _impl.X table complete; full unit suite on branch matches base. Its one finding (the _SOURCE_NAMES gap) is fixed here.

Observable non-behavioral change: _planner_evidence.failure frame records for frames inside moved functions now report module: art.trainer_rank._memory (function/co_name unchanged; nothing branches on the module string).

Stack: refactor/trainer-rank-planner (34 planner methods → _micro_batch_planner.py) and refactor/trainer-rank-slots-optimizer (41 methods → _slots.py, _optimizer.py) follow on top of this. Combined, _impl.py goes from 9,890 to ~5,640 lines.

[bradhilton]

Schulman — no blocking findings in the reviewed scope.

Reviewed head f0a2d8863c91076a6b2529920e26a8b815aec928 against 3ad9a5f8d29ed4ea7f8e033e7e1f6cda3fe55607; both were rechecked immediately before posting. An independent delegate examined the stacked and cumulative extraction through #1000/#1001, and I reviewed its evidence. We did not author these extraction PRs; familiarity with earlier trainer fixes is separate.

All 26 moved memory methods match the original after narrowly accounting for module qualification, self annotations and docstring indentation. Remaining method bodies and class state are preserved. The warm-profile and optimizer-counter fixes survive, runtime globals remain late-bound through _impl, and instance monkeypatch dispatch and staticmethod bindings are retained. Packaging includes the new module, and the planner replay fingerprint includes its source.

Validation: five focused, framework-free CPU checks passed for method binding, exception identity, runtime-global rebinding, warm-profile values and source-fingerprint drift; independent AST comparisons and diff checks passed. No model, GPU, distributed execution, package build or full framework-import test was performed by this review.

Expected limitation: module/frame/line metadata changes. Frozen diagnostics that authenticate the old layout must remain pinned or be deliberately updated; these checks do not certify their automatic adoption. This is a source/CPU correctness review, not a new performance or GPU qualification.

[bradhilton]

Dalinar: memory/planner owner review. No blocking findings.

Reviewed head f0a2d8863c91076a6b2529920e26a8b815aec928 against base 3ad9a5f8d and current main c0b1296dc. Main's _impl.py is byte-identical to the base: #998 changed only _gdn_memory.py, and #994 only the Tau client. I didn't author this PR; its commits carry the shared git identity.

What I checked

• Equivalence. All 26 moved functions are AST-identical to main's methods, after normalizing _impl.X, the self annotation, docstrings and the dropped @staticmethod. The three statics are rebound with staticmethod(...). Every method left in the class is unchanged. The only other change to _impl.py is the pre-class from art.trainer_rank import _memory.
• Globals. _memory.py has no unresolved runtime names. torch and dist are TYPE_CHECKING-only and appear only in annotations. Of the names tests patch on _impl, _memory reads three (_HEAD_CHUNK_TOKENS, _language_model, dist), all through _impl., so those patches still intercept.
• Patching and replay. The bindings are class attributes. Per-instance overrides still win, and so do ranks built with TrainerRank.__new__ (the planner-miss replay does this and then sets _topology_key on the instance).
• Import order. Importing _memory, _impl, art.trainer_rank or _planner_misses first all work, and _memory._impl is _impl.
• Source fingerprints. _memory.py is in _SOURCE_NAMES. The fingerprinted total is 544 KiB of the 1 MiB limit (main: 545). Reports written before this change will need --allow-source-drift to replay, as after any planner source change.
• Price warm waves by the fit later plans measured, not the first plan's #991, Validate canonical checkpoint optimizer counters #992, Avoid repeated MoE workspace scans in GDN planning #998. The three methods Price warm waves by the fit later plans measured, not the first plan's #991/Validate canonical checkpoint optimizer counters #992 changed match main's current bodies. Avoid repeated MoE workspace scans in GDN planning #998's _gdn_memory.plan_floor calls rank._moe_workspace_bytes, which resolves through the new binding. The PR merges cleanly with current main.
• Packaging. The wheel packages src/art whole, and trainer-rank-gpu.yml triggers on all of src/art/trainer_rank/.
• Downstream (Caladan). Its source imports only TrainerRankSlotStateError and _CheckpointSlot from _impl; both stay. The experiment-local actor_memory.py observer fences on _impl.py's sha256, but that fence already pins an older ART, so this PR doesn't change its status.

Validation

• Full trainer-rank CPU suite on this head merged with main c0b1296dc: 1482 passed, 8 skipped, 0 failed. Main 0e0c31b32 gives the same count with the same command.
• ty and ruff are clean on the changed files.
• Separately, a private Fable review of both extraction PRs (source access, no shell) found no blocking issues. Its lows concern identity metadata and are listed on Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner #1000.

Low

1. _impl.py keeps two now-unused imports, nullcontext and MethodType. The repo's ruff config doesn't enable F401, and Move TrainerRank micro-batch planning into trainer_rank/_micro_batch_planner #1000 removes both.

Held planner-stack integration (my held drafts #963, #978, #981 and #986)

• The stack edits four methods this PR moves: _checkpoint_memory_floor, _estimate_required_memory_bytes_from_values, _memory_check and _moe_workspace_bytes. It also adds 18 methods.
• A dry-run merge of the stack into this head gives 5 conflict hunks in _impl.py. Resolving them toward this PR's side would silently drop the stack's edits, so the stack needs a hand port into _memory.py when it rebases.
• That work belongs to the stack, not this PR. After the port, the fingerprinted total would be about 590 KiB, still well under the limit.
• This review lifts no hold.

Limits: CPU tests only; no GPU smoke run. Behavioral equivalence comes from AST comparison and static reading.