@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / da94c7c0eb
Fork 0

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 13, 00:45 UTC

Move online quantization to `model.load_weights` (#26327)

Signed-off-by: Jerry Zhang <jerryzh168@gmail.com>

Jerry Zhang committed Nov 18, 2025 at 16:52 UTC da94c7c0eb8dabea9c500dbd70fa042497497689
6 files changed +314 -113
examples/offline_inference/rlhf.py
+1 -1
@@ -62,7 +62,7 @@ ray.init()
62
63 # Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
64 # Learn more about Ray placement groups:
65 -# https://docs.ray.io/en/latest/placement-groups.html
65 +# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
66 pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
67 ray.get(pg_inference.ready())
68 scheduling_inference = PlacementGroupSchedulingStrategy(
examples/offline_inference/rlhf_online_quant.py new
+162
@@ -0,0 +1,162 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Demonstrates reinforcement learning from human feedback (RLHF) using vLLM and Ray.
5 +
6 +The script separates training and inference workloads onto distinct GPUs
7 +so that Ray can manage process placement and inter-process communication.
8 +A Hugging Face Transformer model occupies GPU 0 for training, whereas a
9 +tensor-parallel vLLM inference engine occupies GPU 1–2.
10 +
11 +The example performs the following steps:
12 +
13 +* Load the training model on GPU 0.
14 +* Split the inference model across GPUs 1–2 using vLLM's tensor parallelism
15 + and Ray placement groups.
16 +* Generate text from a list of prompts using the inference engine.
17 +* Update the weights of the training model and broadcast the updated weights
18 + to the inference engine by using a Ray collective RPC group. Note that
19 + for demonstration purposes we simply zero out the weights.
20 +
21 +For a production-ready implementation that supports multiple training and
22 +inference replicas, see the OpenRLHF framework:
23 +https://github.com/OpenRLHF/OpenRLHF
24 +
25 +This example assumes a single-node cluster with three GPUs, but Ray
26 +supports multi-node clusters. vLLM expects the GPUs are only used for vLLM
27 +workloads. Residual GPU activity interferes with vLLM memory profiling and
28 +causes unexpected behavior.
29 +"""
30 +
31 +import json
32 +import os
33 +
34 +import ray
35 +import torch
36 +from ray.util.placement_group import placement_group
37 +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
38 +from rlhf_utils import stateless_init_process_group
39 +from torchao.core.config import config_to_dict
40 +from torchao.quantization import (
41 + Float8DynamicActivationFloat8WeightConfig,
42 + PerRow,
43 +)
44 +from transformers import AutoModelForCausalLM
45 +
46 +from vllm import LLM, SamplingParams
47 +from vllm.utils.network_utils import get_ip, get_open_port
48 +
49 +
50 +class MyLLM(LLM):
51 + """Configure the vLLM worker for Ray placement group execution."""
52 +
53 + def __init__(self, *args, **kwargs):
54 + # Remove the top-level CUDA_VISIBLE_DEVICES variable set by Ray
55 + # so that vLLM can manage its own device placement within the worker.
56 + os.environ.pop("CUDA_VISIBLE_DEVICES", None)
57 + super().__init__(*args, **kwargs)
58 +
59 +
60 +# Load the OPT-125M model onto GPU 0 for the training workload.
61 +train_model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
62 +train_model.to("cuda:0")
63 +
64 +# Initialize Ray and set the visible devices. The vLLM engine will
65 +# be placed on GPUs 1 and 2.
66 +os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"
67 +ray.init()
68 +
69 +# Create a placement group that reserves GPU 1–2 for the vLLM inference engine.
70 +# Learn more about Ray placement groups:
71 +# https://docs.ray.io/en/latest/ray-core/scheduling/placement-group.html
72 +pg_inference = placement_group([{"GPU": 1, "CPU": 0}] * 2)
73 +ray.get(pg_inference.ready())
74 +scheduling_inference = PlacementGroupSchedulingStrategy(
75 + placement_group=pg_inference,
76 + placement_group_capture_child_tasks=True,
77 + placement_group_bundle_index=0,
78 +)
79 +
80 +# Launch the vLLM inference engine. The `enforce_eager` flag reduces
81 +# start-up latency.
82 +
83 +# generate torchao quantization config for RL rollout
84 +# see https://github.com/vllm-project/vllm/pull/23014 for instructions to
85 +# use serialized config files instead of passing around json string
86 +config = Float8DynamicActivationFloat8WeightConfig(granularity=PerRow())
87 +
88 +json_str = json.dumps(config_to_dict(config))
89 +
90 +llm = ray.remote(
91 + num_cpus=0,
92 + num_gpus=0,
93 + scheduling_strategy=scheduling_inference,
94 +)(MyLLM).remote(
95 + model="facebook/opt-125m",
96 + hf_overrides={"quantization_config_dict_json": json_str},
97 + enforce_eager=True,
98 + worker_extension_cls="rlhf_utils.WorkerExtension",
99 + tensor_parallel_size=2,
100 + distributed_executor_backend="ray",
101 +)
102 +
103 +# Generate text from the prompts.
104 +prompts = [
105 + "Hello, my name is",
106 + "The president of the United States is",
107 + "The capital of France is",
108 + "The future of AI is",
109 +]
110 +
111 +sampling_params = SamplingParams(temperature=0)
112 +
113 +outputs = ray.get(llm.generate.remote(prompts, sampling_params))
114 +
115 +print("-" * 50)
116 +for output in outputs:
117 + prompt = output.prompt
118 + generated_text = output.outputs[0].text
119 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
120 + print("-" * 50)
121 +
122 +# Set up the communication channel between the training process and the
123 +# inference engine.
124 +master_address = get_ip()
125 +master_port = get_open_port()
126 +
127 +handle = llm.collective_rpc.remote(
128 + "init_weight_update_group", args=(master_address, master_port, 1, 3)
129 +)
130 +
131 +model_update_group = stateless_init_process_group(
132 + master_address, master_port, 0, 3, torch.device("cuda:0")
133 +)
134 +ray.get(handle)
135 +
136 +# Simulate a training step by zeroing out all model weights.
137 +# In a real RLHF training loop the weights would be updated using the gradient
138 +# from an RL objective such as PPO on a reward model.
139 +for name, p in train_model.named_parameters():
140 + p.data.zero_()
141 +
142 +# Synchronize the updated weights to the inference engine.
143 +for name, p in train_model.named_parameters():
144 + dtype_name = str(p.dtype).split(".")[-1]
145 + handle = llm.collective_rpc.remote(
146 + "update_weight", args=(name, dtype_name, p.shape)
147 + )
148 + model_update_group.broadcast(p, src=0, stream=torch.cuda.current_stream())
149 + ray.get(handle)
150 +
151 +# Verify that the inference weights have been updated.
152 +assert all(ray.get(llm.collective_rpc.remote("check_weights_changed")))
153 +
154 +# Generate text with the updated model. The output is expected to be nonsense
155 +# because the weights are zero.
156 +outputs_updated = ray.get(llm.generate.remote(prompts, sampling_params))
157 +print("-" * 50)
158 +for output in outputs_updated:
159 + prompt = output.prompt
160 + generated_text = output.outputs[0].text
161 + print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
162 + print("-" * 50)
vllm/model_executor/model_loader/default_loader.py
+11 -35
@@ -22,6 +22,7 @@ from vllm.model_executor.model_loader.weight_utils import (
22 fastsafetensors_weights_iterator,
23 filter_duplicate_safetensors_files,
24 filter_files_not_needed_for_inference,
25 + get_quant_config,
26 maybe_download_from_modelscope,
27 multi_thread_pt_weights_iterator,
28 multi_thread_safetensors_weights_iterator,
@@ -273,42 +274,17 @@ class DefaultModelLoader(BaseModelLoader):
274 )
275
276 def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
276 - if model_config.quantization == "torchao" and torchao_version_at_least(
277 - "0.14.0"
278 - ):
279 - self.load_config.safetensors_load_strategy = "torchao"
280 - weights_to_load = {name for name, _ in model.named_parameters()}
281 -
282 - # if we don't have `model.weight_metadata_and_attr_saved` defined and
283 - # set to True, it means that this is either offline quantization case
284 - # or the first run of online quantization
285 - # see online_quantization.py for detailed notes
286 - offline_quantization_or_first_run_of_online_quantization = not getattr(
287 - model, "weight_metadata_and_attr_saved", False
288 - )
277 + if model_config.quantization == "torchao":
278 + quant_config = get_quant_config(model_config, self.load_config)
279 + if (
280 + hasattr(quant_config, "is_checkpoint_torchao_serialized")
281 + and quant_config.is_checkpoint_torchao_serialized
282 + and torchao_version_at_least("0.14.0")
283 + ):
284 + self.load_config.safetensors_load_strategy = "torchao"
285
290 - if model_config.quantization is None:
291 - # model is not quantized
292 - loaded_weights = model.load_weights(
293 - self.get_all_weights(model_config, model)
294 - )
295 - elif offline_quantization_or_first_run_of_online_quantization:
296 - # case 1: offline quantized checkpoint
297 - # case 2: Step I1 first run of weight loading with
298 - # online quantization
299 - # see online_quantization.py for detailed notes
300 - loaded_weights = model.load_weights(
301 - self.get_all_weights(model_config, model)
302 - )
303 - else:
304 - # to avoid circular dependency
305 - from vllm.model_executor.model_loader.online_quantization import (
306 - load_weights_and_online_quantize,
307 - )
308 -
309 - # subsequent runs of weight loading with online
310 - # quantization
311 - loaded_weights = load_weights_and_online_quantize(self, model, model_config)
286 + weights_to_load = {name for name, _ in model.named_parameters()}
287 + loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
288
289 self.counter_after_loading_weights = time.perf_counter()
290 logger.info_once(
vllm/model_executor/model_loader/online_quantization.py
+128 -77
@@ -2,13 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import types
5 +from collections.abc import Iterable
6
7 import torch
8 from torch import nn
9
10 from vllm.config import ModelConfig
11 from vllm.logger import init_logger
11 -from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
12 from vllm.model_executor.model_loader.utils import process_weights_after_loading
13
14 logger = init_logger(__name__)
@@ -56,6 +56,9 @@ logger = init_logger(__name__)
56 # R4. quantize weights (by calling process_weights_after_loading),
57 # also set `process_weights_after_loading_already_called` to
58 # True to stop it from running again
59 +# R5. (workaround for cudagraph), we restore the weight params to original quantized
60 +# weights params, and use original_weight_param.copy_(updated_weight_param) so that
61 +# the weight update work well with cudagraph
62 # process_weights_after_loading (if called):
63 # this will be skipped since it's already ran in
64 # load_weights
@@ -69,14 +72,6 @@ def maybe_save_metadata_and_attributes_for_weight_reloading(
72 if model_config.quantization != "torchao":
73 return
74
72 - if getattr(model, "process_weights_after_loading_already_called", False):
73 - # In case `process_weights_after_loading` is called multiple times
74 - # we'll skip it at later times
75 - logger.warning(
76 - "process_weights_after_loading already called for model %s", model
77 - )
78 - return
79 -
75 from vllm.model_executor.model_loader.weight_utils import get_quant_config
76
77 quant_config = get_quant_config(model_config, None)
@@ -137,6 +132,7 @@ def maybe_save_metadata_and_attributes_for_weight_reloading(
132 else:
133 model.recorded_weight_attr[name][key] = attr
134 # mark the metadata and attributes saved so we don't run it again
135 + model._model_config = model_config
136 model.weight_metadata_and_attr_saved = True
137
138
@@ -148,77 +144,132 @@ def _bond_method_to_cls(func, obj):
144 return types.MethodType(func, obj)
145
146
151 -def load_weights_and_online_quantize(
152 - model_loader: DefaultModelLoader, model: nn.Module, model_config: ModelConfig
153 -) -> set[str]:
147 +def support_quantized_model_reload_from_hp_weights(original_load_weights):
148 + """Decorator for `load_weights` method for AutoWeightsLoader.load_weights to support
149 + reloading high precision (bfloat16/float16/float32) weight for an already quantized
150 + model, this involves restoring the weights to a high precision weights and
151 + then online quantize the weights
152 + """
153 # online quantization, right now only enabled for
154 # torchao
156 - # R1, R2, R3, R4 in the Notes
157 -
158 - # TODO: Add fp8 support
159 - assert model_config.quantization == "torchao", (
160 - "online quantization is only enabled for torchao currently"
161 - )
162 - # TODO: use create_weights to restore the weights to original state
163 -
164 - # Step R1: First restore the quantized weights to original bfloat16
165 - # weights, with original metadata (shape, dtype, device)
166 - # and attributes, so that bfloat16 weights can be loaded properly
167 - existing_param_names = dict(model.named_parameters(remove_duplicate=False)).keys()
168 - named_modules = dict(model.named_modules(remove_duplicate=False))
169 - model_device = None
170 -
171 - # Step R2: recover the parameter to the state before first loading
172 - for name, d in model.original_weights_rebuild_keys.items():
173 - _shape = d["shape"]
174 - _dtype = d["dtype"]
175 - _device = d["device"]
155 + # R1, R2, R3, R4, R5 in the Notes
156 +
157 + def patched_model_load_weights(
158 + auto_weight_loader, weights: Iterable[tuple[str, torch.Tensor]], *, mapper=None
159 + ) -> set[str]:
160 + model = auto_weight_loader.module
161 + offline_quantization_or_first_run_of_online_quantization = not getattr(
162 + model, "weight_metadata_and_attr_saved", False
163 + )
164 +
165 + # if we don't have `model.weight_metadata_and_attr_saved` defined and
166 + # set to True, it means that this is either offline quantization case
167 + # or the first run of online quantization
168 + # see Notes in this file for more details
169 + if offline_quantization_or_first_run_of_online_quantization:
170 + # case 1: offline quantized checkpoint
171 + # case 2: Step I1 first run of weight loading with
172 + # online quantization
173 + return original_load_weights(auto_weight_loader, weights, mapper=mapper)
174 +
175 + model_config = model._model_config
176 +
177 + # TODO: Add fp8 support
178 + assert model_config.quantization == "torchao", (
179 + "online quantization is only enabled for torchao currently"
180 + )
181 + # TODO: use create_weights to restore the weights to original state
182 +
183 + # Step R1: First restore the quantized weights to original bfloat16
184 + # weights, with original metadata (shape, dtype, device)
185 + # and attributes, so that bfloat16 weights can be loaded properly
186 + # TODO: maybe set remove_duplicate to True?
187 + original_quantized_weight_dict = dict(
188 + model.named_parameters(remove_duplicate=False)
189 + )
190 + named_modules = dict(model.named_modules(remove_duplicate=False))
191 + model_device = None
192 +
193 + for name, d in model.original_weights_rebuild_keys.items():
194 + _shape = d["shape"]
195 + _dtype = d["dtype"]
196 + _device = d["device"]
197 + if model_device is not None:
198 + assert model_device == _device, (
199 + "Expecting all weights "
200 + "to be in the same device for now, got both: "
201 + f"{model_device} and {_device}"
202 + )
203 + else:
204 + model_device = _device
205 +
206 + if name in original_quantized_weight_dict:
207 + module_name, weight_name = name.rsplit(".", 1)
208 + module = named_modules[module_name]
209 + setattr(
210 + module,
211 + weight_name,
212 + torch.nn.Parameter(
213 + torch.empty(_shape, dtype=_dtype, device=_device),
214 + requires_grad=False,
215 + ),
216 + )
217 +
218 + # Step R2: recover the weight attributes to the state before first loading
219 + # recorded_weight_attr is
220 + # {"weight_name": {"weight_attr_key": attr}}
221 + # e.g.
222 + # {
223 + # {
224 + # "layer.0.weight": {
225 + # "weight_loader": weight_loader_function_object,
226 + # "input_dim": 0, ...
227 + # },
228 + # "layer.1.weight": ...,
229 + # }
230 + # }
231 + for full_weight_name, weight_attr_dict in model.recorded_weight_attr.items():
232 + for attr_name, attr in weight_attr_dict.items():
233 + module_name, weight_name = full_weight_name.rsplit(".", 1)
234 + module = named_modules[module_name]
235 + weight = getattr(module, weight_name)
236 + if not hasattr(weight, attr_name):
237 + setattr(weight, attr_name, _bond_method_to_cls(attr, weight))
238 +
239 + # Step R3: reload bfloat16 / high precision weights
240 + updated_params = original_load_weights(
241 + auto_weight_loader, weights, mapper=mapper
242 + )
243 +
244 + # Step R4: online quantize the weights
245 + # manually process weights after loading
246 + model.process_weights_after_loading_already_called = False
247 if model_device is not None:
177 - assert model_device == _device, (
178 - "Expecting all weights "
179 - "to be in the same device for now, got both: "
180 - f"{model_device} and {_device}"
181 - )
248 + process_weights_after_loading(model, model_config, model_device)
249 else:
183 - model_device = _device
184 -
185 - if name in existing_param_names:
186 - module_name, weight_name = name.rsplit(".", 1)
187 - module = named_modules[module_name]
188 - setattr(
189 - module,
190 - weight_name,
191 - torch.nn.Parameter(torch.empty(_shape, dtype=_dtype, device=_device)),
250 + logger.warning_once(
251 + "model_device is None, skip calling process_weights_after_loading"
252 )
253
194 - # recorded_weight_attr is
195 - # {"weight_name": {"weight_attr_key": attr}}
196 - # e.g.
197 - # {
198 - # {
199 - # "layer.0.weight": {
200 - # "weight_loader": weight_loader_function_object,
201 - # "input_dim": 0, ...
202 - # },
203 - # "layer.1.weight": ...,
204 - # }
205 - # }
206 - for full_weight_name, weight_attr_dict in model.recorded_weight_attr.items():
207 - for attr_name, attr in weight_attr_dict.items():
208 - module_name, weight_name = full_weight_name.rsplit(".", 1)
209 - module = named_modules[module_name]
210 - weight = getattr(module, weight_name)
211 - if not hasattr(weight, attr_name):
212 - setattr(weight, attr_name, _bond_method_to_cls(attr, weight))
213 -
214 - # Step I1: reload bfloat16 / high precision weights
215 - loaded_weights = model.load_weights(
216 - model_loader.get_all_weights(model_config, model)
217 - )
218 -
219 - # Step I2: online quantize the weights
220 - # manually process weights after loading
221 - model.process_weights_after_loading_already_called = False
222 - process_weights_after_loading(model, model_config, model_device)
223 - model.process_weights_after_loading_already_called = True
224 - return loaded_weights
254 + # Step R5 (workaround for cudagraph): restore the original quantized weights
255 + # and do a copy_ of the currents weights to the original weights
256 + updated_quantized_weights = dict(model.named_parameters(remove_duplicate=False))
257 + for name in model.original_weights_rebuild_keys:
258 + if name in original_quantized_weight_dict:
259 + original_quantized_weight = original_quantized_weight_dict[name]
260 + updated_quantized_weight = updated_quantized_weights[name]
261 +
262 + module_name, weight_name = name.rsplit(".", 1)
263 + module = named_modules[module_name]
264 + setattr(module, weight_name, original_quantized_weight)
265 + with torch.no_grad():
266 + original_quantized_weight.copy_(updated_quantized_weight)
267 +
268 + del original_quantized_weight_dict
269 + del named_modules
270 + del updated_quantized_weight
271 +
272 + model.process_weights_after_loading_already_called = True
273 + return updated_params
274 +
275 + return patched_model_load_weights
vllm/model_executor/model_loader/utils.py
+8
@@ -88,6 +88,14 @@ def initialize_model(
88 def process_weights_after_loading(
89 model: nn.Module, model_config: ModelConfig, target_device: torch.device
90 ) -> None:
91 + if getattr(model, "process_weights_after_loading_already_called", False):
92 + # In case `process_weights_after_loading` is called multiple times
93 + # we'll skip it at later times
94 + logger.debug_once(
95 + "process_weights_after_loading already called for model %s", model
96 + )
97 + return
98 +
99 # to avoid circular dependency
100 from vllm.model_executor.model_loader.online_quantization import (
101 maybe_save_metadata_and_attributes_for_weight_reloading,
vllm/model_executor/models/utils.py
+4
@@ -21,6 +21,9 @@ from vllm.logger import init_logger
21 from vllm.model_executor.layers.quantization.base_config import (
22 QuantizationConfig,
23 )
24 +from vllm.model_executor.model_loader.online_quantization import (
25 + support_quantized_model_reload_from_hp_weights,
26 +)
27 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28 from vllm.model_executor.models.interfaces import supports_any_eagle
29 from vllm.multimodal import NestedTensors
@@ -316,6 +319,7 @@ class AutoWeightsLoader:
319 )
320 raise ValueError(msg)
321
322 + @support_quantized_model_reload_from_hp_weights
323 def load_weights(
324 self,
325 weights: Iterable[tuple[str, torch.Tensor]],