@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d28d86e8a3

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 5, 21:00 UTC

[QeRL] Fix online quantized reloading (#38442)

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Kyle Sayers committed Mar 29, 2026 at 16:56 UTC d28d86e8a34bf2617be294c235d6e6ef3321917b

9 files changed +104 -62

.buildkite/test-amd.yaml

+3 -3

@@ -812,7 +812,7 @@ steps:

812 commands:

813 - apt-get update && apt-get install -y curl libsodium23

814 - export VLLM_WORKER_MULTIPROC_METHOD=spawn

815 - - pytest -v -s model_executor

815 + - pytest -v -s model_executor -m '(not slow_test)'

816 - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py

817

818

@@ -1242,7 +1242,7 @@ steps:

1242 - vllm/platforms/rocm.py

1243 commands:

1244 - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

1245 - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py

1245 + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'

1246 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'

1247 - pytest models/language -v -s -m 'distributed(num_gpus=2)'

1248 - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py

@@ -2501,7 +2501,7 @@ steps:

2501 - tests/models/

2502 commands:

2503 - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

2504 - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py

2504 + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'

2505 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'

2506 - pytest models/language -v -s -m 'distributed(num_gpus=2)'

2507 - pytest models/multimodal -v -s -m 'distributed(num_gpus=2)' --ignore models/multimodal/generation/test_whisper.py

.buildkite/test_areas/model_executor.yaml

+1 -1

@@ -13,5 +13,5 @@ steps:

13 commands:

14 - apt-get update && apt-get install -y curl libsodium23

15 - export VLLM_WORKER_MULTIPROC_METHOD=spawn

16 - - pytest -v -s model_executor

16 + - pytest -v -s model_executor -m '(not slow_test)'

17 - pytest -v -s entrypoints/openai/completion/test_tensorizer_entrypoint.py

.buildkite/test_areas/models_distributed.yaml

+1 -1

@@ -14,7 +14,7 @@ steps:

14 - tests/models/

15 commands:

16 - TARGET_TEST_SUITE=L4 pytest basic_correctness/ -v -s -m 'distributed(num_gpus=2)'

17 - - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py

17 + - CUDA_VISIBLE_DEVICES=0,1 pytest -v -s model_executor/model_loader/test_sharded_state_loader.py -m '(not slow_test)'

18 # Avoid importing model tests that cause CUDA reinitialization error

19 - pytest models/test_transformers.py -v -s -m 'distributed(num_gpus=2)'

20 - pytest models/language -v -s -m 'distributed(num_gpus=2)'

tests/model_executor/model_loader/test_reload.py

+45 -20

@@ -38,7 +38,10 @@ def test_move_metatensors():

39 def test_reload_lifecycle():

40 layer = torch.nn.Linear(2, 3)

41 - info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))

41 + info = LayerReloadingInfo(

42 + restore_metadata=capture_layer_to_meta(layer),

43 + restore_device=torch.device("cpu"),

44 + )

46 restore_layer_on_meta(layer, info)

47 for name, tensor in get_layer_tensors(layer).items():

@@ -48,7 +51,7 @@ def test_reload_lifecycle():

51 assert tensor.__class__ == meta_tensor.__class__

52 assert tensor.__dict__ == meta_tensor.__dict__

51 - materialize_layer(layer)

54 + materialize_layer(layer, info)

55 for name, tensor in get_layer_tensors(layer).items():

56 materialized_tensor = getattr(layer, name)

57 assert tensor.dtype == materialized_tensor.dtype

@@ -60,7 +63,10 @@ def test_reload_lifecycle():

63 def test_model_cleanup(dist_init, default_vllm_config):

64 layer = QKVParallelLinear(2, 3, 4)

65 assert layer.weight.weight_loader.__self__ is layer

63 - info = LayerReloadingInfo(restore_metadata=capture_layer_to_meta(layer))

66 + info = LayerReloadingInfo(

67 + restore_metadata=capture_layer_to_meta(layer),

68 + restore_device=torch.device("cpu"),

69 + )

71 mock_info_dict: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (

72 WeakKeyDictionary()

@@ -90,39 +96,46 @@ def test_get_numel_loaded():

96 assert ret == "value"

93 -@pytest.mark.parametrize("tp_size", [2])

99 +@pytest.mark.parametrize(

100 + "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]

101 +)

102 @pytest.mark.parametrize(

103 "base_model,mul_model,add_model",

104 [

97 - (

105 + pytest.param(

106 "Qwen/Qwen3-0.6B",

107 "inference-optimization/Qwen3-0.6B-debug-multiply",

108 "inference-optimization/Qwen3-0.6B-debug-add",

109 + marks=[pytest.mark.slow_test],

110 ),

102 - (

111 + pytest.param(

112 "inference-optimization/Qwen3-0.6B-FP8_BLOCK",

113 "inference-optimization/Qwen3-0.6B-debug-multiply-FP8_BLOCK",

114 "inference-optimization/Qwen3-0.6B-debug-add-FP8_BLOCK",

115 + marks=[pytest.mark.slow_test],

116 ),

107 - (

117 + pytest.param(

118 "inference-optimization/Qwen3-0.6B-W4A16-G128",

119 "inference-optimization/Qwen3-0.6B-debug-multiply-W4A16-G128",

120 "inference-optimization/Qwen3-0.6B-debug-add-W4A16-G128",

121 + marks=[pytest.mark.slow_test],

122 ),

112 - (

123 + pytest.param(

124 "inference-optimization/DeepSeek-V3-debug-empty",

125 "inference-optimization/DeepSeek-V3-debug-multiply",

126 "inference-optimization/DeepSeek-V3-debug-add",

127 + marks=[pytest.mark.slow_test],

128 ),

117 - (

129 + pytest.param(

130 "inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC",

131 "inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC",

132 "inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC",

133 ),

122 - (

134 + pytest.param(

135 "inference-optimization/DeepSeek-V3-debug-empty-NVFP4A16",

136 "inference-optimization/DeepSeek-V3-debug-multiply-NVFP4A16",

137 "inference-optimization/DeepSeek-V3-debug-add-NVFP4A16",

138 + marks=[pytest.mark.slow_test],

139 ),

140 ],

141 )

@@ -138,6 +151,8 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):

151 tensor_parallel_size=tp_size,

152 enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),

153 enable_prefix_caching=False,

154 + max_model_len=16,

155 + max_num_seqs=1,

156 ) as llm:

157 llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})

158 mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]

@@ -150,34 +165,42 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):

165 assert add_perp < mul_perp

166

167

153 -@pytest.mark.parametrize("tp_size", [2])

168 +@pytest.mark.parametrize(

169 + "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]

170 +)

171 @pytest.mark.parametrize(

172 "base_model,mul_model,add_model,quantization",

173 [

157 - (

174 + pytest.param(

175 "Qwen/Qwen3-0.6B",

176 "inference-optimization/Qwen3-0.6B-debug-multiply",

177 "inference-optimization/Qwen3-0.6B-debug-add",

178 "fp8",

179 ),

163 - (

180 + pytest.param(

181 "inference-optimization/DeepSeek-V3-debug-empty",

182 "inference-optimization/DeepSeek-V3-debug-multiply",

183 "inference-optimization/DeepSeek-V3-debug-add",

184 "fp8",

185 + marks=[pytest.mark.slow_test],

186 ),

169 - (

187 + pytest.param(

188 "Qwen/Qwen3-0.6B",

189 "inference-optimization/Qwen3-0.6B-debug-multiply",

190 "inference-optimization/Qwen3-0.6B-debug-add",

191 "mxfp8",

192 + marks=[pytest.mark.slow_test],

193 + ),

194 + pytest.param(

195 + "inference-optimization/DeepSeek-V3-debug-empty",

196 + "inference-optimization/DeepSeek-V3-debug-multiply",

197 + "inference-optimization/DeepSeek-V3-debug-add",

198 + "mxfp8",

199 + marks=[

200 + pytest.mark.slow_test,

201 + pytest.mark.xfail(reason="mxfp4 & mla is not supported yet"),

202 + ],

203 ),

175 - # ( TODO: support mxfp4 & mla

176 - # "inference-optimization/DeepSeek-V3-debug-empty",

177 - # "inference-optimization/DeepSeek-V3-debug-multiply",

178 - # "inference-optimization/DeepSeek-V3-debug-add",

179 - # "mxfp8",

180 - # ),

204 ],

205 )

206 def test_online_quantize_reload(

@@ -195,6 +218,8 @@ def test_online_quantize_reload(

218 tensor_parallel_size=tp_size,

219 enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),

220 enable_prefix_caching=False,

221 + max_model_len=16,

222 + max_num_seqs=1,

223 ) as llm:

224 llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})

225 mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]

vllm/model_executor/layers/quantization/fp8.py

+5 -2

@@ -1006,14 +1006,17 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1006 initialize_online_processing(layer)

1007

1008 def process_weights_after_loading(self, layer: Module) -> None:

1009 + # TODO(@ksayers): inplace fp8 quant kernel, initialize scales with ones

1010 if getattr(layer, "_already_called_process_weights_after_loading", False):

1011 return

1012

1013 fp8_dtype = current_platform.fp8_dtype()

1014 w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)

1015 w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)

1015 - w13_scale = torch.ones(layer.num_experts, dtype=torch.float32)

1016 - w2_scale = torch.ones(layer.num_experts, dtype=torch.float32)

1016 + w13_scale = torch.ones(

1017 + layer.num_experts, device=w13.device, dtype=torch.float32

1018 + )

1019 + w2_scale = torch.ones(layer.num_experts, device=w2.device, dtype=torch.float32)

1020 layer.w13_input_scale = None

1021 layer.w2_input_scale = None

1022

vllm/model_executor/model_loader/reload/layerwise.py

+15 -3

@@ -49,7 +49,10 @@ def get_layerwise_info(layer: torch.nn.Module) -> LayerReloadingInfo:

49 information existed, a new entry is constructed

50 """

51 if layer not in LAYERWISE_INFO:

52 - LAYERWISE_INFO[layer] = LayerReloadingInfo()

52 + LAYERWISE_INFO[layer] = LayerReloadingInfo(

53 + restore_metadata=({}, {}),

54 + restore_device=torch.get_default_device(),

55 + )

57 return LAYERWISE_INFO[layer]

@@ -64,6 +67,7 @@ def record_metadata_for_reloading(model: torch.nn.Module):

67 for layer in model.modules():

68 info = get_layerwise_info(layer)

69 info.restore_metadata = capture_layer_to_meta(layer)

70 + info.restore_device = torch.get_default_device()

73 @torch.no_grad()

@@ -99,10 +103,18 @@ def initialize_layerwise_reload(model: torch.nn.Module):

103 # Restore layer parameters/buffers onto meta device

104 restore_layer_on_meta(layer, info)

105

106 + # Wrap weight loaders to buffer loading

107 initialize_online_processing(layer)

108

109

110 def initialize_online_processing(layer: torch.nn.Module):

111 + """

112 + Wrap a layer's weight loaders with online processing loaders.

113 + Called by either `initialize_layerwise_reload` or an online quantization scheme,

114 + prevents double wrapping in the case of online quantization + reloading

115 +

116 + :param layer: layer whose parameter weight loaders will be wrapped

117 + """

118 info = get_layerwise_info(layer)

119

120 # Track loading progress to determine when to process/copy

@@ -211,7 +223,7 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon

223 elif info.load_numel <= 0:

224 # first load but received no weights. This happens on dummy load

225 if info.kernel_tensors is None:

214 - materialize_layer(layer)

226 + materialize_layer(layer, info)

227

228 # reloading: place kernel tensors back as a fallback

229 else:

@@ -244,7 +256,7 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):

256 4. Copies processed values back to original tensor storage

257 """

258 # Materialize layer tensors onto device

247 - materialize_layer(layer)

259 + materialize_layer(layer, info)

260

261 # Reset online quantization flag so process_weights_after_loading

262 # will run again during reload

vllm/model_executor/model_loader/reload/meta.py

+5 -4

@@ -94,14 +94,15 @@ def restore_layer_on_meta(layer: torch.nn.Module, info: LayerReloadingInfo):

94 layer.register_buffer(name, buffer)

97 -def materialize_layer(layer: torch.nn.Module) -> None:

97 +def materialize_layer(layer: torch.nn.Module, info: LayerReloadingInfo):

98 """Materialize all meta tensors in a layer to actual tensors."""

99 if layer.__class__.__name__ in SKIP_MODULES:

100 return

101

102 - for name, tensor in get_layer_tensors(layer).items():

103 - if name not in SKIP_TENSORS:

104 - setattr(layer, name, materialize_meta_tensor(tensor))

102 + with info.restore_device:

103 + for name, tensor in get_layer_tensors(layer).items():

104 + if name not in SKIP_TENSORS:

105 + setattr(layer, name, materialize_meta_tensor(tensor))

106

107

108 class CopyCounter(TorchDispatchMode):

vllm/model_executor/model_loader/reload/types.py

+12 -7

@@ -13,21 +13,26 @@ LayerTensors = tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]

14 @dataclass

15 class LayerReloadingInfo:

16 - # model format (meta), populated by `record_metadata_for_reloading`

17 - restore_metadata: LayerTensors = field(default_factory=lambda: ({}, {}))

16 + # model format metadata, recorded by `record_metadata_for_reloading`

17 + restore_metadata: LayerTensors

19 - # kernel format (device), used to copy into when reloading only

20 - kernel_tensors: LayerTensors | None = None

19 + # device to materialize layers with, recorded by `record_metadata_for_reloading`

20 + restore_device: torch.device

22 - # track how many restored elements are ready for loading

22 + # track how many elements are ready for loading, used by `online_process_loader`

23 load_numel: int = 0

24 load_numel_total: int | None = None

26 - # stores arguments and tensors ready for loading

26 + # used by `online_process_loader` to buffer args and tensors until ready to load

27 loaded_weights: list[tuple[str, BoundArguments]] = field(default_factory=list)

29 + # kernel formatted tensors, copied into by `_layerwise_process` when reloading

30 + kernel_tensors: LayerTensors | None = None

31 +

32 def reset(self):

30 - self.__init__(restore_metadata=self.restore_metadata) # type: ignore[misc]

33 + self.__init__( # type: ignore[misc]

34 + restore_metadata=self.restore_metadata, restore_device=self.restore_device

35 + )

37 def can_load(self) -> bool:

38 return self.load_numel_total is not None

vllm/v1/worker/gpu_model_runner.py

+17 -21

@@ -4943,28 +4943,24 @@ class GPUModelRunner(

4943

4944 # begin loading weights

4945 logger.info_once("Reloading weights inplace...", scope="local")

4946 - load_device = (

4947 - self.vllm_config.load_config.device or self.vllm_config.device_config.device

4948 - )

4949 - with torch.device(load_device):

4950 - if is_checkpoint_format:

4951 - # load weights from checkpoint/ original model format

4952 - initialize_layerwise_reload(model)

4953 - loaded_weights = model.load_weights(weights_iterator)

4954 - finalize_layerwise_reload(model, self.model_config)

4946 + if is_checkpoint_format:

4947 + # load weights from checkpoint/ original model format

4948 + initialize_layerwise_reload(model)

4949 + loaded_weights = model.load_weights(weights_iterator)

4950 + finalize_layerwise_reload(model, self.model_config)

4951

4956 - else:

4957 - # load weights from kernel format

4958 - logger.warning_once(

4959 - "Reloading with `is_checkpoint_format=True` requires that "

4960 - "weights be in kernel format and already sharded",

4961 - scope="local",

4962 - )

4963 - loaded_weights = set()

4964 - for name, loaded_weight in weights_iterator:

4965 - param = model.get_parameter(name) # TODO: buffers?

4966 - param.copy_(loaded_weight)

4967 - loaded_weights.add(name)

4952 + else:

4953 + # load weights from kernel format

4954 + logger.warning_once(

4955 + "Reloading with `is_checkpoint_format=True` requires that "

4956 + "weights be in kernel format and already sharded",

4957 + scope="local",

4958 + )

4959 + loaded_weights = set()

4960 + for name, loaded_weight in weights_iterator:

4961 + param = model.get_parameter(name) # TODO: buffers?

4962 + param.copy_(loaded_weight)

4963 + loaded_weights.add(name)

4964

4965 # logging and validation

4966 counter_after_reloading = time.perf_counter()