A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 12:30 UTC
[Docs] [QeRL] Layerwise Reloading Documentation (#40317)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>
Kyle Sayers committed
Apr 29, 2026 at 00:06 UTC
a085b5257dd8cc8d6c255e9b92e4642ee12fc3aa
6 files changed
+214
-6
docs/assets/training/layerwise.png
Binary files /dev/null and b/docs/assets/training/layerwise.png differ
docs/assets/training/layerwise_bad_loading.png
Binary files /dev/null and b/docs/assets/training/layerwise_bad_loading.png differ
docs/assets/training/layerwise_good_loading.png
Binary files /dev/null and b/docs/assets/training/layerwise_good_loading.png differ
docs/training/layerwise.md
new
+146
@@ -0,0 +1,146 @@
1
+# What is Layerwise (Re)loading?
2
+
3
+Layerwise reloading is the system used to handle the loading of new weight data into existing weight data destinations without triggering recompilation of the cuda graph and other runtime artifacts. This system is used to enable [QeRL](https://arxiv.org/pdf/2510.11696)-style post training flows, where full-precision trainer weights are quantized and loaded into a target vLLM instance for fast, high-exploration rollouts. The core implementation can be found in [layerwise.py](../../vllm/model_executor/model_loader/reload/layerwise.py).
4
+
5
+
6
+
7
+## Layerwise Reloading for QeRL
8
+
9
+In order to load new weights into existing weight data destinations, a weight must undergo the following operations:
10
+
11
+- Transfer: weights must be transferred from trainer model to target node/device
12
+- Fuse: weight partitions must be fused, for example qkv/gate_up
13
+- Process: this typically means online quantization and kernel-specific padding or striding
14
+- Shard: weights must be sharded according to the selected parallelism strategy
15
+- Copy: weights must be copied into the existing weight data destinations
16
+
17
+Layerwise reloading achieves this using the following steps:
18
+
19
+1. Weights are **transferred** from the trainer to the target (see [weight_transfer](weight_transfer/README.md))
20
+2. Weights loaded via `model.load_weights`, during which they are **sharded** and **fused**
21
+3. Weights are **processed** in an online fashion as soon as all of a layer's weights are loaded
22
+4. Weights are **copied** into the existing weight data destinations
23
+
24
+For more information on implementation, see [Low Level `layerwise` API](#low-level-layerwise-api).
25
+
26
+## Layerwise Loading with Online Quantization
27
+
28
+Online quantization refers to when a user provides full precision weights and those weights are quantized on-the-fly as they are loaded into the model. The layerwise reloading system handles this by treating online quantization as a **processing** step, which is then handled in an online way both during first-time load and during reload. A typical online quantization method implementation should look like this:
29
+
30
+```python
31
+class Fp8OnlineLinearMethod(Fp8LinearMethod):
32
+ """Online version of Fp8LinearMethod which loads a full precision checkpoint
33
+ and quantizes weights during loading."""
34
+
35
+ uses_meta_device: bool = True
36
+
37
+ def create_weights(self, layer: torch.nn.Module, ...):
38
+ # weight is materialized and processed during loading
39
+ layer.weight = ModelWeightParameter(
40
+ data=torch.empty(..., device="meta"),
41
+ weight_loader=weight_loader,
42
+ )
43
+
44
+ # set up online processing
45
+ initialize_online_processing(layer)
46
+
47
+ def process_weights_after_loading(self, layer: Module) -> None:
48
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
49
+ return
50
+
51
+ layer.weight, layer.weight_scale = ops.scaled_fp8_quant(layer.weight)
52
+
53
+ # Prevent duplicate processing (e.g., during weight reload)
54
+ layer._already_called_process_weights_after_loading = True
55
+```
56
+
57
+## Example Usages
58
+
59
+### High Level Weight Transfer API
60
+
61
+The layerwise reloading system is integrated with the post-training weight transfer system. To use layerwise reloading in conjunction to the weight transfer system, follow the examples found [here](../../examples/rl/). Layerwise reloading is controlled by the `WeightTransferUpdateInfo.is_checkpoint_format` flag and is set to `True` by default.
62
+
63
+### Mid Level `reload_weights` API
64
+
65
+Layerwise reloading is also exposed via the `reload_weights` API. This interface can be called using the following code:
66
+
67
+```python
68
+from vllm import LLM
69
+
70
+llm = LLM("Qwen/Qwen3-0.6B")
71
+llm.collective_rpc("reload_weights")
72
+```
73
+
74
+This interface also allows specifying a `weights_path` which can be used to select a checkpoint path to load from:
75
+
76
+```python
77
+from vllm import LLM
78
+
79
+# fine tuned model checkpoints for testing
80
+mul_path = "inference-optimization/Qwen3-0.6B-debug-multiply"
81
+add_path = "inference-optimization/Qwen3-0.6B-debug-add"
82
+
83
+llm = LLM("Qwen/Qwen3-0.6B")
84
+llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_path})
85
+llm.generate("3 4 = ") # 12
86
+
87
+llm.collective_rpc("reload_weights", kwargs={"weights_path": add_path})
88
+llm.generate("3 4 = ") # 7
89
+```
90
+
91
+Finally, a `weights_iterator` can be provided directly. This iterator can be lazy or eagerly defined.
92
+
93
+```python
94
+from vllm import LLM
95
+
96
+weights_iterator = [("q_proj", ...), ("k_proj", ...), ...]
97
+
98
+llm = LLM("Qwen/Qwen3-0.6B")
99
+llm.collective_rpc("reload_weights", kwargs={"weights_iterator": weights_iterator})
100
+```
101
+
102
+### Low Level `layerwise` API
103
+
104
+[layerwise.py](../../vllm/model_executor/model_loader/reload/layerwise.py) Implements the following functions to execute its lifecycle:
105
+
106
+| Function | Purpose | Quantized Reload | Online Quantization |
107
+| - | - | - | - |
108
+| `record_metadata_for_reloading` | Record tensor metadata so that layers can be restored on the meta device | Called by `BaseModelLoader` | Called by `BaseModelLoader` |
109
+| `restore_layer_on_meta` | Restore layer to model format at start of reload | Called by `initialize_layerwise_reload` | Not called. Online quantized weights already start on meta device via `...OnlineLinearMethod.create_weights` |
110
+| `initialize_online_processing` | Wrap weight loaders with the `online_process_loader` wrapper, which buffers weights until all layer weights have been loaded | Called by `initialize_layerwise_reload` | Called by `...OnlineLinearMethod.create_weights` |
111
+| `_layerwise_process` | Process layer once all weights are loaded | Called by `online_process_loader` during loading | Called by `online_process_loader` during loading |
112
+| `_copy_and_restore_kernel_tensors` | Copy processed weights into original tensor locations to affect compiled cuda graphs, etc. | Called by `_layerwise_process` after `process_weights_after_loading` | Not called. There is no compiled cuda graph yet |
113
+| `finalize_layerwise_processing` | Catch any layers which did not load all weights (for example attention weights or weights with padding) | Called by `BaseModelLoader` | Called by `BaseModelLoader` |
114
+
115
+You can plug into this lifecycle directly by calling the `initialize_layerwise_reload`, loading weights, then calling `finalize_layerwise_processing`:
116
+
117
+```python
118
+from vllm import LLM
119
+from vllm.model_executor.model_loader.reload import initialize_layerwise_reload, finalize_layerwise_processing
120
+
121
+llm = LLM("Qwen/Qwen3-0.6B")
122
+
123
+# this model path requires `VLLM_ENABLE_V1_MULTIPROCESSING=0` and is not stable
124
+model = llm.llm_engine.engine_core.engine_core.model_executor.driver_worker.worker.get_model()
125
+
126
+# layerwise reload
127
+initialize_layerwise_reload(model)
128
+model.load_weights(...)
129
+finalize_layerwise_processing(model, llm.model_config)
130
+```
131
+
132
+## Troubleshooting Excessive Memory Usage
133
+
134
+Layerwise reloading allows users to incrementally load and process weights as they are loaded into the model. This system relies on buffering layer weights on device until all weights of a layer have been loaded. However, without offloading, this approach necessarily causes excessive buffering if weights are loaded out of order.
135
+
136
+For this reason, users must take care as to the order of weights when they are reloading into the model. Weight should be loaded "in order", meaning that each layer's weights are fully loaded before beginning to load the next layer's weights. "Out of order" loading can cause layer weights to stay buffered while other layer weights are loading, leading to excessive memory usage. In the example below, q_proj, k_proj, v_proj, and up_proj are all buffered at the same time, using more memory than if up_proj was loaded after q_proj, k_proj and v_proj.
137
+
138
+| Correct Loading | Incorrect Loading |
139
+| - | - |
140
+|  |  |
141
+
142
+Users will see a warning like the one below if weights are loaded out-of-order.
143
+
144
+```console
145
+WARNING [layerwise.py:198] Allocating 28.5 MB of device memory to buffers to load ["QKVParallelLinear", "MergedColumnParallelLinear"] layers. This extra memory usage can be avoided by ordering weights by their parent layer when reloading.
146
+```
vllm/model_executor/model_loader/reload/layerwise.py
+35
-5
@@ -3,7 +3,7 @@
3
import inspect
4
from collections.abc import Callable
5
from functools import wraps
6
-from weakref import WeakKeyDictionary
6
+from weakref import WeakKeyDictionary, WeakSet
7
8
import torch
9
@@ -21,7 +21,13 @@ from .meta import (
21
restore_layer_on_meta,
22
)
23
from .types import LayerReloadingInfo
24
-from .utils import get_layer_params_buffers, get_layer_size, get_layer_tensors
24
+from .utils import (
25
+ get_info_size,
26
+ get_layer_params_buffers,
27
+ get_layer_size,
28
+ get_layer_tensors,
29
+ has_device_tensors,
30
+)
31
32
logger = init_logger(__name__)
33
@@ -43,6 +49,9 @@ LAYERWISE_INFO: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
49
WeakKeyDictionary()
50
)
51
52
+# Global set used to track loading for logging purposes only
53
+LOADING_LAYERS: WeakSet[torch.nn.Module] = WeakSet()
54
+
55
56
def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:
57
"""
@@ -174,11 +183,30 @@ def make_online_process_loader(layer: torch.nn.Module, param_name: str) -> Calla
183
info.load_numel_total,
184
)
185
186
+ # Do not online process attention layers, must wait until finalize
187
+ if isinstance(layer, (Attention, MLAAttention)):
188
+ return ret
189
+
190
+ # Log warnings allocating excessive buffers on device
191
+ if has_device_tensors(bound_args):
192
+ LOADING_LAYERS.add(layer)
193
+ if len(LOADING_LAYERS) >= 2:
194
+ names = sorted([layer.__class__.__name__ for layer in LOADING_LAYERS])
195
+ mem_used = sum(
196
+ get_info_size(LAYERWISE_INFO[layer]) for layer in LOADING_LAYERS
197
+ )
198
+ logger.warning_once(
199
+ "Allocating %.1f MB of device memory to buffers to load %s layers. "
200
+ "This extra memory usage can be avoided by ordering weights "
201
+ "by their parent layer when reloading.",
202
+ mem_used / 1e6,
203
+ str(list(names)),
204
+ )
205
+
206
# Process and copy when all weights are loaded
178
- if info.load_numel >= info.load_numel_total and not isinstance( # type: ignore[operator]
179
- layer, (Attention, MLAAttention)
180
- ):
207
+ if info.load_numel >= info.load_numel_total: # type: ignore[operator]
208
_layerwise_process(layer, info)
209
+ LOADING_LAYERS.discard(layer)
210
211
return ret
212
@@ -240,6 +268,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
268
_finalize_attention_layer(layer, info, model_config)
269
info.reset()
270
271
+ LOADING_LAYERS.clear()
272
+
273
274
def finalize_layerwise_reload(*args, **kwargs):
275
finalize_layerwise_processing(*args, **kwargs)
vllm/model_executor/model_loader/reload/utils.py
+33
-1
@@ -1,14 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from inspect import BoundArguments
5
+
6
import torch
7
6
-from .types import LayerTensors
8
+from .types import LayerReloadingInfo, LayerTensors
9
10
__all__ = [
11
"get_layer_tensors",
12
"get_layer_params_buffers",
13
"get_layer_size",
14
+ "has_device_tensors",
15
+ "get_info_size",
16
]
17
18
@@ -39,3 +43,31 @@ def get_layer_size(layer: torch.nn.Module) -> int:
43
for name, tensor in get_layer_tensors(layer).items()
44
if name not in SKIP_TENSORS
45
)
46
+
47
+
48
+def has_device_tensors(bound_args: BoundArguments) -> bool:
49
+ """
50
+ Return True if the loaded weights exist on an accelerator device
51
+
52
+ :param bound_args: args to load weights
53
+ :return: True if weights are on accelerator device
54
+ """
55
+ return any(
56
+ isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu")
57
+ for value in bound_args.arguments.values()
58
+ )
59
+
60
+
61
+def get_info_size(info: LayerReloadingInfo) -> int:
62
+ """
63
+ Calculate the number of bytes used by loaded weights for a given layer
64
+
65
+ :param info: layerwise info to get size of
66
+ :return: number of bytes used by loaded weights
67
+ """
68
+ return sum(
69
+ value.nbytes
70
+ for _, args in info.loaded_weights
71
+ for value in args.arguments.values()
72
+ if isinstance(value, torch.Tensor) and value.device.type not in ("meta", "cpu")
73
+ )