[Refactor][Model Loader] Public start/finish/abort reload lifecycle

[aoshen02]

Why

Every caller that streams checkpoint-format weights into a live model names the mechanism directly — initialize_layerwise_reload / finalize_layerwise_reload — in five places: the NCCL, IPC and sharded-RDT weight-transfer engines, GPUModelRunner.reload_weights, and the torchao decorator. The RDT engine additionally reaches into layerwise internals for its dry-run bake and carried its own private "put the model back" routine, and its docstring asked for exactly this: "a public API should expose first-class ... an abort_layerwise_reload instead of _restore_after_dry_run."

RFC vllm-project#54477 (selective weight reload) needs to choose the reload mechanism in one place. With this PR that place exists.

What

• model_loader/reload/lifecycle.py: start_reload(model), finish_reload(model, model_config), abort_reload(model) — re-exported from model_loader.reload. Today they forward to the layerwise implementation; the old names stay exported.
• All five callers switched. Zero behavior change on the start/finish paths (pure forwarding; lazy imports kept lazy).
• abort_reload is new: undoes start_reload for every layer still waiting for weights — kernel tensors put back, online_process_loader wrappers removed, buffered weights dropped, torchao flag restored — so the model is loadable again. It is deliberately not a rollback: layers that completed keep their new weights, and tensors loaded in place (SKIP_TENSORS, e.g. bias) keep what was written. That matches the failure semantics in RFC [RFC]: Selective Weight Reload for RL Training vllm-project/vllm#54477 §11 (a failed update is terminal; recovery is a fresh reload or restart), and rollback would require the shadow copy the RFC is removing.
• Sharded RDT _bake now calls abort_reload and keeps only its own stamp unwrapping (_unwrap_recording_stamps); _restore_after_dry_run is deleted. Abort only touches online_process_loader wrappers, so RDT's stamps and loader identities are preserved exactly as before.
• Docs: docs/training/layerwise.md example uses the public API with finish_reload inside the protected block; docs/training/weight_transfer/base.md engine template updated.

Not a duplicate

vllm-project#49201 (closed, conflicts) proposed a full WeightLoadSession object. This is the minimal slice of that idea — three functions, no session state — sized to be mergeable and sufficient for the RFC.

Tests

.venv/bin/python -m pytest tests/model_executor/model_loader/test_reload.py -k "not test_reload_weights and not test_kv_scale_reload and not test_online_quantize_reload"
# 27 passed (26 existing + 1 new)
.venv/bin/python -m pytest tests/distributed/test_sharded_rdt_plan.py
# 124 passed; 3 GPU tests (TestBakeOnARealModel) could not start an engine here — GPU was fully occupied by another job (184/189 GiB). Needs a rerun on a free GPU.
pre-commit run --files <11 changed files>   # all pass

New test test_abort_reload_restores_pending_layers: a completed layer keeps new weights, a pending layer gets its original tensors/loaders back with weight untouched and in-place bias kept as written, and a subsequent start_reload→finish_reload on the same model succeeds with original tensor identity.

Audit

Reviewed by Codex (gpt-6-astra), two rounds. Round 1 caught: the new file missing from the commit, abort attaching a loader attribute to tensors that never had one (now only real online_process_loader wrappers are unwrapped), and finish_reload outside the try in the docs example. Round 2: PASS (independent CPU probes for RDT-cleanup equivalence, nested wrappers, first-load, state restoration).

AI assistance (Claude Code + Codex) was used; every line reviewed and tests run by the submitter.

🤖 Generated with Claude Code