[yueming-yuan]

Co-authored with @yushengsu-thu.

Save and restore one adapter slot independently. Checkpoints contain adapter weight shards, per-rank Adam state and step counters, and FP32 master parameters; loading optimizer state requires the same parallel topology.

Checkpoint load/save execution failures, including filesystem errors, propagate to TrainerCell and terminate the current single-cell Tinker service. lora/actor.py and controller commands expose save_slot and load_slot(..., ckpt_path=..., load_optimizer=...). A TODO in checkpoint_io.py records deferred IO failure isolation.

Checkpoint names are relative symlinks to immutable version directories. The shared writer completes shards and metadata before atomically replacing the public link. A process crash leaves the old or new complete checkpoint reachable; older versions remain on disk for readers that already resolved them. Legacy directory checkpoints remain readable, but overwriting them is unsupported: save under a new name. Successful phases retain their collective barriers; execution errors still propagate to the trainer cell.

Stacked on #2841.

Validation:

• pre-commit run --all-files and incremental git diff --check passed on every updated component head. All previous heads remain ancestors; the final Unify HF snapshot and checkpoint writing #3198 tree exactly matches the reviewed implementation.
• Process-exit probes at both sides of publication lost the public path before the fix; afterward it resolves to the complete old or new version, and a subsequent save succeeds. A parameterized regression is included in the existing [multi-lora v2 7/7] Tinker gateway example and tests #2849 checkpoint test file.
• Sampler-to-training load probes with and without optimizer state now return a user failure before any backend load, while another tenant continues training. Repeated SAVE → LOAD verifies distinct retained versions and loading from a pinned version.
• Focused checkpoint continuation passed on 2 H200 GPUs (DP=2, TP/PP/EP=1), using real Megatron DDP, MultiLoRA layers, SlotOptimizer and distributed checkpoint IO. After three Adam steps, slot 0 was saved and restored into a fresh optimizer on slot 1; weights, FP32 masters, Adam moments, parameter-group metadata and step counters matched with zero tolerance before and after the next identical-gradient update. Rounding masters to BF16 and restarting from weights alone were both detected by negative controls.
• The enabled regression is tests/e2e/lora/test_slot_checkpoint_resume.py (Hopper, LoRA, 2-GPU stage). The tested final tree is Unify HF snapshot and checkpoint writing #3198 fd63c7ec8c41c2e0a61d273f0cd5a3504275ada7, with Megatron 73b54618f7e58e0f25f619bcfecbe2640765475a, Bridge 40b930897717941cfe2bd9806f417e28bf1bfa65, PyTorch 2.13.0+cu130 and Transformer Engine 2.17.0. Checkpoints were on GPFS. pre-commit run --all-files passed.
• This isolates checkpoint/Adam continuation with identical post-reduction gradients; it does not validate forward/backward numerics, topology changes, MoE/TP/PP layouts, server restart, or failure recovery. Fresh hosted CPU/NVIDIA CI is pending; broad GPU selection remains on Unify HF snapshot and checkpoint writing #3198. The separate Tinker cookbook acceptance test remains disabled pending GPU validation.

ci-megatron-pr: 73b54618f7e58e0f25f619bcfecbe2640765475a
ci-sglang-pr: 02b5e123c490c988d21526d0dd7837fefdd0e1da

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[nanjiangwill]

optimizer is the outer LayerWise/Chained optimizer, so this reload walks every slot’s child and copies BF16 model weights into FP32 masters. Loading slot S therefore rounds the masters of every live slot; only S is restored from masters below. The preceding lora_executor.load_slot() call does the same. Add a slot-local reload helper and use it in both paths, with a two-slot test that keeps slot B’s FP32 master different from its BF16 model value.

[nanjiangwill]

The read phase validates only optimizer list lengths. It does not validate master/moment tensor shapes or adapter weight keys and shapes. load_adapter() copies matching tensors first and reports missing or unused keys afterward, so a bad shard can mutate part of the live slot before raising. Validate the complete weight and optimizer tensor layout in read_shards() before any copy, or add a validation-only Bridge helper.

[nanjiangwill]

This code moved; the current comment is here: #2842 (comment)

[nanjiangwill]

These tests cover read-phase failures only. They never apply a valid checkpoint, so they do not verify FP32 masters, Adam moments and step counters, weights-only loading, or equivalence of the next Adam update. Add a CPU round-trip test for exact continuation. Include a second live slot whose FP32 master differs from its BF16 model value so the test also catches cross-slot reloads.

[nanjiangwill]

write_checkpoint_dir() now rolls back a caught failure from the second rename, but the published path is still absent between final_dir -> old_dir and tmp_dir -> final_dir. A process exit in that window leaves the last valid checkpoint only under _old_<name>, so readers lose the public path. Also, shutil.rmtree(old_dir) still runs inside the publish phase, so cleanup failure reports the save as failed after the replacement is committed. Use immutable version directories plus an atomically replaced pointer/exchange for crash atomicity, and make old-version cleanup best-effort.