[aoshen02]

Depends on #57733 and #62407 (merge those first); mypy/CI fail until they land. This PR's diff is only the 8 direct-reload files. They import helpers that #57733 adds (finalize_model_weights, discard_weight_parts, stash_weight_part) and rely on #62407's kept parameter objects.

Validated as the stack #57733 + #62407 + this PR + E1 (aoshen02#179) + derived-tensor in place (aoshen02#180), with same-weight direct reloads on GB200, NUPD=3. GSM8K is measured before and after, and every tensor is fingerprinted against cold start:

Model	Job	GSM8K pre → post	Tensors changed
DeepSeek-V4.1	40010	0.909 → 0.904	0/1709
MiniMax-M3.1	40013	0.951 → 0.951	0/3148
Muse Spark NVFP4	40042	0.839 → 0.838	0/3011
gpt-oss-120b TRTLLM MXFP4 (NUPD=2)	39950	0.938 → 0.936	0/796

Problem

RL weight sync reloads the checkpoint into a serving engine. The layerwise reload moves every layer to meta, materializes it again, loads the checkpoint and reruns post-load processing per layer. That costs a second copy of each layer at a time, and its correctness depends on per-layer completeness bookkeeping, so partial or reordered updates are hard to get right.

Fix

VLLM_RELOAD_MODE=direct (registered in vllm/envs.py; the default stays layerwise) writes checkpoint tensors straight into the live model and reruns cold-start post-load processing once, at the startup addresses:

• At cold start, before process_weights_after_loading (PWAL), record each tensor's checkpoint layout.
• begin_direct_reload, called from initialize_layerwise_reload, views every tensor whose layout PWAL changed back to its checkpoint layout, on the same storage. The loaders then write the checkpoint tensors in place.
• finish_direct_reload, called from finalize_layerwise_reload, clears the one-shot "already processed" guards and reruns PWAL for every module, then the model-level hook, as cold start does.
• replace_parameter keeps a parameter's address when PWAL changes its layout: under direct reload, the new tensor is written into a view of the existing storage when it fits.
• Refused at cold start: draft_lm_head_quantization. The quantized draft lm_head would be written into the target lm_head it shares.
• Refused at begin, before anything is modified:
  • no layout was recorded (also the case for pre-processed IPC weight-cache weights, already in runtime layout);
  • a storage is too small for its checkpoint layout.
• sharded_rdt refuses direct mode at init, because its replay plan is baked by a layerwise (meta) dry run.

Delivery contract. An update must deliver every checkpoint tensor of each module whose PWAL changes the tensor's layout or is not idempotent. Tensors without such processing, for example a frozen embedding table, may be left out and keep their bytes. Nothing in this PR checks the contract: until the completeness follow-up merges, a violation is not detected.

PWAL must keep every checkpoint tensor registered under its checkpoint name. This builds on #62407 (retire/alias instead of delete/rename).

Known limitations. Non-default MoE backends whose PWAL reallocates derived constants on a rerun are not yet refilled in place: TRTLLM FP8 per-tensor, FlashInfer CUTLASS NVFP4/MXFP4, CuteDSL, and compressed-tensors FP8 constants. They are fixed in the derived-tensor follow-up. Online quantization is supported by a later follow-up.

Tests

• tests/model_executor/model_loader/test_direct_reload.py uses a toy model with the hazards found on real models: a parent-derived tensor, a layout change, a loader bypass, a consumed scale and non-idempotent packing. It checks:
  • a reload equals a fresh cold start bit for bit, at unchanged addresses;
  • an unsent table without PWAL keeps its bytes;
  • refusals leave the model untouched;
  • an aborted update can be retried;
  • parked weight parts are handled.
• E2E (GB200, NCCL weight sync, 2 updates, reload diff on), this PR on [Model Loader] Run model-level post-load finalization from the loader hook #57733 + [Quantization] Keep deleted and renamed params registered after process_weights_after_loading #62407 alone: Qwen3.8 GSM8K 0.971 → 0.970, 0/1434 tensors changed (39863); Nemotron 0.853 → 0.854, 0/913 changed (39864). For the full stack, see the table at the top.

Dependencies / merge order

#57733 → #62407 → this PR. The follow-ups stack on this one: completeness check → moved-address/host-float check → side-tensor detector. The derived-tensor in-place PR, the NCCL empty_cache fix and the loader-bypass fix are independent.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

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

[hao-aaron]

Went through this against the RFC tracker. Leaving out anything F2/F4/F5/F7/F8 or PR 5 already covers. A few things I think belong here though:

1. Refuse online quant in this PR. F8 adds support, but until then nothing stops it, and finish_direct_reload pops _already_called_process_weights_after_loading, which is the only thing keeping online FP8 PWAL (online/fp8.py) from re-quantizing already-converted bytes. Same shape as the draft_lm_head_quantization refusal; a test next to that one would be nice.

2. Replaced param objects are silently treated as derived. _checkpoint_views skips anything without _ckpt_layout, so if cold-start PWAL registers a new Parameter under a checkpoint name, begin skips it and the loader writes checkpoint bytes into the runtime layout (a same-numel transpose passes default_weight_loader). F4 snapshots at first begin, after this already happened, so it won't catch it. Could we record the names in record_checkpoint_layout and raise in begin_direct_reload if one now points at an object with no recorded layout? That makes the #62407 contract enforced instead of assumed.

3. Failure state. test_retry_after_aborted_update_matches_cold_start says the worker refuses to run a partly updated model, but I don't see that anywhere. Either add a minimal in-progress/failed flag that begin/finish set, or fix the comment and add it to Known limitations.

Smaller:

• Dispatching direct mode from inside initialize_layerwise_reload / finalize_layerwise_reload reads oddly, and the env is checked separately at begin and finish. A small start_reload / finish_reload that stores the mode on the model would be cleaner (#63599 has one in reload/lifecycle.py if you want to borrow it).
• finish_direct_reload searches for a "cuda" param; the engines already have self.device, could just pass it through.
• Popping the guard on every module also hits MoERunnerInterface, which sets it in __init__ as a hack. Probably harmless, but limiting it to modules with a quant_method seems safer.
• _view_in_place bounds writes by storage.nbytes() rather than the tensor's own range, so it assumes nothing else lives in that storage. Fine if F4/F7 add an overlap check, just want to make sure it's tracked.
• Add VLLM_RELOAD_MODE to the ignored set in compile_factors so toggling it doesn't invalidate the compile cache.

This review/comment was written by an AI agent (Claude Code) on my behalf.

[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