online fp8 quant with streaming weight post-processing by vkuzo · Pull Request #29196 · vllm-project/vllm

vkuzo · 2025-11-21T18:58:53Z

Summary:

updates the fp8.py vllm online quant entry point to post-process weights as soon as they are ready, to prevent from holding the entire high precision model in memory. Specifically:

during weight loading, keep track of how many elements we have loaded
when we have loaded all the elements, call post-processing

I'm putting up this PR to demonstrate the approach. Since for now the # of use cases is 2 (linear and moe in fp8.py), keeping the logic simple. If this is extended to other online quant providers such as compressed-tensors/torchao/etc in the future, might be good to make the code more reusable in future PRs.

Test Plan:

tested locally with facebook/opt-125m for linear and Qwen/Qwen1.5-MoE-A2.7B for MoE and fp8 online quantization. Verified (via prints) that layer.process_weights_after_loading is called immediately after the last weight chunk is loaded in both cases.

Reviewers:

Subscribers:

Tasks:

Tags:

Purpose

Test Plan

Test Result

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

vllm/model_executor/layers/quantization/fp8.py

gemini-code-assist

Code Review

This pull request introduces a proof-of-concept for online FP8 quantization with streaming weight post-processing. The approach is clever, patching the weight loader to trigger post-processing as soon as a weight tensor is fully loaded. This can help reduce peak memory usage during model loading.

My main feedback is to improve the robustness of the state management. The flag _already_called_process_weights_after_loading is currently stored on the Fp8LinearMethod instance. While this works with the current code structure, it's fragile. Attaching this state to the layer object instead would make the implementation more robust against future changes, such as instance reuse for optimization. I've added specific comments with code suggestions to address this.

gemini-code-assist · 2025-11-21T19:01:30Z

vllm/model_executor/layers/quantization/fp8.py

+                    del param._loaded_numel
+                    # Prevent the usual `process_weights_after_loading` call from doing
+                    # anything
+                    self._already_called_process_weights_after_loading = True


Storing _already_called_process_weights_after_loading on self (the Fp8LinearMethod instance) makes the design fragile. Although a new instance is currently created for each layer, this might change in the future (e.g., for optimization), which could lead to this flag persisting incorrectly across different layers.

To make this more robust, this state should be attached to the layer object, which is guaranteed to be unique. This change should be made in conjunction with the corresponding check in process_weights_after_loading.

Suggested change

self._already_called_process_weights_after_loading = True

layer._already_called_process_weights_after_loading = True

gemini-code-assist · 2025-11-21T19:01:30Z

vllm/model_executor/layers/quantization/fp8.py

@@ -487,6 +516,9 @@ def create_weights(
                layer.register_parameter("input_scale", None)

    def process_weights_after_loading(self, layer: Module) -> None:
+        if getattr(self, "_already_called_process_weights_after_loading", False):


To make the state management more robust and in conjunction with the suggested change for setting this flag, this check should be on the layer object instead of self.

Suggested change

if getattr(self, "_already_called_process_weights_after_loading", False):

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

kylesayrs

Thanks!

kylesayrs · 2025-12-03T16:53:50Z

vllm/model_executor/layers/quantization/fp8.py

+                res = weight_loader(param, loaded_weight, *args, **kwargs)  # type: ignore[misc]
+
+                # add a counter to track how many elements we have updated
+                if not hasattr(layer, "_target_loaded_numel"):


Are all of these hasattr checks necessary? It's a little hard to read right now, I think the attribute is guaranteed to exist if not self.quant_config.is_checkpoint_fp8_serialized

makes sense, sg, we can delete the _target_loaded_numel var and just compute it every time

The same goes for _loaded_numel. If this attribute didn't exist, then something has gone terribly wrong and it might be appropriate to error.

kylesayrs · 2025-12-03T16:55:08Z

vllm/model_executor/layers/quantization/fp8.py

+                # add a counter to track how many elements we have updated
+                if not hasattr(layer, "_target_loaded_numel"):
+                    # for linear, the only weight we need to load is `layer.weight`
+                    layer._target_loaded_numel = layer.weight.numel()


Nit: avoid extra flags by deleting _target_loaded_numel and just calling layer.weight.numel()

Summary: not for land, just a demo 1. during weight loading, keep track of how many elements we have loaded 2. when we have loaded all the elements, call post-processing can be used to call weight post-processing in a streaming fashion to minimize GPU memory usage. Will only work if we can assume we only load each weight chunk once. Test Plan: tested locally with facebook/opt-125m and `fp8` online quantization Reviewers: Subscribers: Tasks: Tags: Signed-off-by: vasiliy <vasiliy@fb.com>

Summary: When we added online fp8 quant with streaming weight post-processing in vllm-project#29196, a bug was introduced where TP>1 case was not always handled correctly. Specifically: * vllm-project#29196 assumed that `weight_loader` copies `loaded_weight` to `param` * this is not true, as `weight_loader` can call arbitrary logic on both `param` and `loaded_weight` before eventually calling `copy_`. An example is here: https://github.com/vllm-project/vllm/blob/e3a0f21e6ce78268865cafcdc3dc58c7a80dbc57/vllm/model_executor/parameter.py#L195 A fix is to track exactly how many number of elements were updated with `copy_`. The PR implements this fix with `TorchDispatchMode`. Test Plan: ```bash // tp 1 still works CUDA_VISIBLE_DEVICES=6,7 with-proxy python3 examples/offline_inference/basic/generate.py --model Qwen/Qwen1.5-MoE-A2.7B --enforce-eager --quantization fp8 -tp=1 // tp > 2 was broken before, now works CUDA_VISIBLE_DEVICES=6,7 with-proxy python3 examples/offline_inference/basic/generate.py --model Qwen/Qwen1.5-MoE-A2.7B --enforce-eager --quantization fp8 -tp=2 ``` Reviewers: Subscribers: Tasks: Tags: Signed-off-by: vasiliy <vasiliy@fb.com>

…29196) Signed-off-by: vasiliy <vasiliy@fb.com> Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>

Summary: vllm-project#29196 implemented streaming weight post-processing for online fp8 quant but did not actually reduce peak memory, because the linear|moe weights were created in bf16 and references to them were held for the entire `load_weights` loop in model loaders. this PR fixes it by changing fp8 online quant to create zero-sized weights in `create_weights`, and materialize them to the correct size just-in-time in `patched_weight_loader`. I would note that this PR is a bit hacky, and there are two more proper ways to fix this that I can think of, both with a much wider blast radius: - 1: change weight creation in vllm to be materialized just-in-time (same as this PR, just explicit instead of hacky callables) - 2: or, add an extension point for post-processing the weight before loading it (similar to vllm-project#27280) fixes vllm-project#31805 Test Plan: inspect memory usage inside of `load_weights` and verify that it increases ~monotonically as weights are loaded ```bash // dense python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --enforce-eager --dtype=bfloat16 --max_model_len=2048 --quantization=fp8 // moe CUDA_VISIBLE_DEVICES=7 python3 examples/offline_inference/basic/generate.py --model Qwen/Qwen3-30B-A3B --enforce-eager --dtype=bfloat16 --block-size=64 --max_model_len=2048 --gpu-memory-utilization=0.8 --trust-remote-code --quantization=fp8 ``` Reviewers: Subscribers: Tasks: Tags: Signed-off-by: vasiliy <vasiliy@fb.com>

vkuzo requested review from mgoin, pavanimajety, robertgshaw2-redhat, tlrmchlsmth and yewentao256 as code owners November 21, 2025 18:58

chatgpt-codex-connector bot reviewed Nov 21, 2025

View reviewed changes

vllm/model_executor/layers/quantization/fp8.py Show resolved Hide resolved

gemini-code-assist bot reviewed Nov 21, 2025

View reviewed changes

vkuzo force-pushed the 20251121_fp8_online_quant_hack branch from 5326892 to 9583e3b Compare November 21, 2025 19:07

kylesayrs mentioned this pull request Nov 22, 2025

[WIP] [Design] Online quantization #29247

Closed

vkuzo force-pushed the 20251121_fp8_online_quant_hack branch from 9583e3b to 8956364 Compare November 26, 2025 15:52

vkuzo changed the title ~~[not for land] online fp8 quant with streaming weight post-processing~~ online fp8 quant with streaming weight post-processing Nov 26, 2025

robertgshaw2-redhat added the ready ONLY add when PR is ready to merge/full CI is needed label Dec 3, 2025

kylesayrs approved these changes Dec 3, 2025

View reviewed changes

vkuzo force-pushed the 20251121_fp8_online_quant_hack branch from fda3623 to 427f05d Compare December 4, 2025 13:34

robertgshaw2-redhat approved these changes Dec 8, 2025

View reviewed changes

robertgshaw2-redhat enabled auto-merge (squash) December 8, 2025 19:44

robertgshaw2-redhat merged commit 0d402d2 into vllm-project:main Dec 8, 2025
54 checks passed

kylesayrs mentioned this pull request Dec 8, 2025

[Quantization] Support pre-load online quantization for compressed-tensors W8A8 channel-wise schema #27280

Open

5 tasks

vkuzo mentioned this pull request Dec 17, 2025

fix fp8 online quantization streaming with tp > 1 #30900

Merged

5 tasks

vkuzo mentioned this pull request Jan 7, 2026

fix memory for online fp8 quantization with streaming weight load #31914

Merged

5 tasks

dsuhinin pushed a commit to dsuhinin/vllm that referenced this pull request Jan 21, 2026

online fp8 quant with streaming weight post-processing (vllm-project#…

d6c50e2

…29196) Signed-off-by: vasiliy <vasiliy@fb.com> Signed-off-by: dsuhinin <suhinin.dmitriy@gmail.com>

CSWYF3634076 mentioned this pull request Jan 21, 2026

[Bugfix][Quant] fix online fp8 quantization oom #32773

Closed

5 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

online fp8 quant with streaming weight post-processing#29196

online fp8 quant with streaming weight post-processing#29196
robertgshaw2-redhat merged 1 commit intovllm-project:mainfrom
vkuzo:20251121_fp8_online_quant_hack

vkuzo commented Nov 21, 2025 •

edited by github-actions bot

Loading

Uh oh!

chatgpt-codex-connector bot left a comment

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Nov 21, 2025

Uh oh!

gemini-code-assist bot Nov 21, 2025

Uh oh!

kylesayrs left a comment

Uh oh!

kylesayrs Dec 3, 2025

Uh oh!

vkuzo Dec 3, 2025

Uh oh!

kylesayrs Dec 8, 2025

Uh oh!

kylesayrs Dec 3, 2025

Uh oh!

vkuzo Dec 3, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	self._already_called_process_weights_after_loading = True
	layer._already_called_process_weights_after_loading = True

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

Uh oh!

Conversation

vkuzo commented Nov 21, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 21, 2025

Choose a reason for hiding this comment

Uh oh!

kylesayrs left a comment

Choose a reason for hiding this comment

Uh oh!

kylesayrs Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

vkuzo Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

kylesayrs Dec 8, 2025

Choose a reason for hiding this comment

Uh oh!

kylesayrs Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

vkuzo Dec 3, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vkuzo commented Nov 21, 2025 •

edited by github-actions bot

Loading