@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 146 lines 8.14 KB
Rendered Raw
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 ![Layerwise](../assets/training/layerwise.png)
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 Fp8PerTensorOnlineLinearMethod(LinearMethodBase):
32 """Online version of FP8 per-tensor quantization which loads a full
33 precision checkpoint 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/). Checkpoint-format weight transfer engines (e.g. the NCCL and IPC backends) run layerwise reloading automatically inside their `start_weight_update`/`finish_weight_update` lifecycle.
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 | ![Layerwise](../assets/training/layerwise_good_loading.png) | ![Layerwise](../assets/training/layerwise_bad_loading.png) |
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 ```