fix fp8 online quantization streaming with tp > 1

[vkuzo]

Summary:

Fix for #30830

When we added online fp8 quant with streaming weight post-processing in #29196, a bug was introduced where TP>1 case was not always handled correctly. Specifically:

• online fp8 quant with streaming weight post-processing #29196 assumed that weight_loader copies loaded_weight to param directly
• this is not always true, as weight_loader can call arbitrary logic on both param and loaded_weight before eventually calling copy_. An example is here:

  vllm/vllm/model_executor/parameter.py

  Line 195 in e3a0f21

  param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)

A fix is to track exactly how many number of elements were updated with copy_. The PR implements this fix with TorchDispatchMode.

Test Plan:

// 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:

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 usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request addresses a bug in the online FP8 quantization streaming logic when using tensor parallelism (TP > 1). The issue stemmed from an incorrect assumption about the number of elements loaded by the weight_loader. The fix introduces a CopyNumelCounter class using TorchDispatchMode to accurately track the number of elements modified by copy_ operations, even when weights are transformed (e.g., with narrow) before being copied. This is a clever and robust solution. The changes in Fp8LinearMethod and Fp8OnlineMoEMethod correctly apply this new counter. The implementation is clean and effectively resolves the described problem. I have no further comments.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[kylesayrs]

Totally correct, thanks for the fix

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[mergify]

Hi @vkuzo, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint