[Refactor][Model Loader] Unify the weight loading lifecycle

[aoshen02]

Implements the lifecycle design proposed in #48920.

Related correctness taxonomy: #48312. Complementary storage-identity work: #48478 and #48902.

Problem

Initial model loading and checkpoint-format reload had the same three post-load requirements, but different callers assembled them independently:

1. quantization/repacking must run after checkpoint tensors are written;
2. Attention/MLA/MM encoder derived tensors must run after linear/quant processing;
3. reload must copy processed tensors back into the storage already referenced by CUDA Graphs.

That made a new loader or transfer path easy to get subtly wrong.

Design

flowchart LR
    I["Initial loader"] --> P["WeightLoadSession.prepare"]
    R["V2 checkpoint reload"] --> P
    T["TorchAO reload"] --> P
    W["Checkpoint IPC / NCCL"] --> P
    P --> L["Existing load / transfer implementation"]
    L --> F["WeightLoadSession.finish"]
    F --> LW["layerwise: materialize + process + copy back"]
    F --> D["utils: discover post-load module types"]
    D --> Q["Quant"]
    Q --> A["Attention / MLA / MM encoder"]
    A --> H["HPC-derived state"]

Loading

The boundary is deliberately small:

prepare -> existing load/transfer -> finish

The responsibilities are separated instead of putting all policy in the session:

Code	Single responsibility	Why it remains
WeightLoadSession.__init__	Store per-load exactly-once state and decide whether this load needs layerwise processing	Reload always needs storage-preserving copy-back; initial load needs it only for online quant methods using meta tensors
prepare()	Bind one active session and install the existing layerwise loading state	Prevents nested lifecycle ownership and ensures weight loaders see the same session
finish()	Coordinate layerwise finalization and post-load processing, then unbind	Every successful checkpoint-format load must cross one completion point
abort()	Restore temporary meta/wrapper structure, then let the original error propagate	Structural cleanup only; it is not rollback or retry
process_quant()	Invoke a formal QuantizeMethodBase hook at most once in this load	A layer can be reached eagerly and again by the final scan
finalize_attention_runtime()	Finalize already-classified Attention runtime state at most once in this load	This is separate from checkpoint/quant-method processing, and the session does not discover concrete attention classes
_process_quant_method()	Preserve the standalone online-quant path that has no session	Required by existing dummy/online quant loaders
utils._process_modules_after_loading()	Discover module types and preserve the established Quant -> Attention -> HPC order	Type policy stays next to the original model-wide traversal
layerwise._finish_layerwise_loading()	Finish delayed layers and restore original Parameter/Buffer storage	This is the CUDA Graph storage-identity guarantee

Attention, MLAAttention, and MMEncoderAttention now share one internal type tuple used by both the model-wide traversal and layerwise deferral. Adding another post-load attention wrapper requires changing one classification point, not the session.

This PR does not replace BaseModelLoader.load_weights, model.load_weights, IPC, or NCCL. They still own the source and transport of weights. It also adds no registry, generic object walk, runtime retry, or synchronization protocol.

Failure semantics

abort() restores temporary Python/module structure only. It cannot restore tensor values already written, does not retry the update, and does not make a partially updated model safe for inference. The original exception is re-raised.

The caller must already pause/drain inference during an update.

Simplification in the final revision

• load_session.py: 172 lines -> 150 lines
• final cleanup commit: 76 insertions / 82 deletions overall; production code is net -3 lines
• removed process_all_modules() and _process_hpc() from the session
• removed the unused standalone process_weights_after_loading() wrapper that constructed an unprepared session
• removed concrete Attention/HPC imports and isinstance policy from the session
• removed the duplicate finalize_layerwise_reload alias
• removed the unsupported promise that a failed partial update can be safely retried

Scope relative to #48312

This is lifecycle infrastructure, not the complete correctness checker:

• no category 1 pointer/lifetime/layout runtime check; that is complementary work in [RFC] Fail-Closed Graph Storage Contract for Weight Reload #48478
• no generic category 2 source-to-runtime value manifest
• no complete category 3 required-key, duplicate, ordering, or chunk-accounting checker
• no category 4 transactional rollback guarantee

Validation

Final commit: 20f3b063ff.

• all changed-file pre-commit hooks passed, including Ruff, format, mypy, SPDX, forbidden-import, and signoff checks
• lifecycle/reload/weight-transfer/V2-worker focused suite: 77 passed, 21 deselected
• final focused lifecycle suite after removing the obsolete wrapper: 9 passed
• V2 online FP8 + CPU offload on one H200: 1 passed
• git diff --check: passed

The final VIME E2E used Qwen2.5-0.5B-Instruct on h200-1: four trainer GPUs, four V2 vLLM engines, IPC weight transfer, FULL + PIECEWISE CUDA Graph capture, and two train/update/replay iterations.

Step	train_rollout_logprob_abs_diff	PPO KL	KL loss
0	0.013579	0	0
1	0.017642	0	0

The log confirms Using V2 Model Runner, IPCWeightTransferEngine, both CUDA Graph capture modes, successful start_weight_update -> update_weights -> finish_weight_update, and replay after each update. It ended with WEIGHT_LIFECYCLE_E2E_PASS variant=c.

AI assistance

AI assistance was used to research, implement, test, and draft this change. I reviewed the resulting diff and validation evidence and can explain and maintain the implementation.

[mergify]

Documentation preview: https://vllm--48908.org.readthedocs.build/en/48908/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[aoshen02]

The original PR cannot be reopened because its head repository was detached from the fork network after the visibility transition. The conflict-resolved continuation is now open as #49201, based on aoshen02/vllm-codex:codex/weight-lifecycle-plan-c.