feat: offload weights to cpu before fp8 online quant

[yma11]

We offload weights to cpu before online quantization so that device doesn't need to hold whole model weights, which fix the limitation mentioned in upstream fp8 quantization. Note that, this option can't work together with cpu_offload_gb to avoid obscure behaviors.

This feature controlled by env variable VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT, example command:

VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT=1 VLLM_MLA_DISABLE=1 python3 examples/offline_inference/basic/generate.py --model deepseek-ai/DeepSeek-v2-lite  --enforce-eager --dtype=float16 --block-size=64 --max_model_len=2048 --gpu-memory-utilization=0.8 --quantization fp8 --trust-remote-code

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Summary of Changes

Hello @yma11, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a crucial memory optimization for FP8 online quantization by enabling the temporary offloading of model weights to the CPU before the quantization process begins. This addresses a significant limitation by reducing the GPU memory footprint required during model loading and quantization. The feature is controlled via a new environment variable and includes a safeguard to prevent its concurrent use with existing CPU offload functionalities, ensuring system stability.

Highlights

• Enhanced FP8 Quantization Memory Efficiency: Implemented a mechanism to offload model weights to CPU memory prior to online FP8 quantization. This significantly reduces the peak GPU memory usage during the model loading and quantization process, addressing a previous limitation where the entire model had to reside on the GPU.
• New Configuration Option: Introduced VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT as an environment variable. Setting this to 1 enables the new weight offloading behavior, allowing for more memory-efficient FP8 quantization.
• Feature Incompatibility Enforcement: Added a validation check to prevent simultaneous use of VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT and cpu_offload_gb. This ensures stable operation and avoids potential conflicting behaviors between the two memory management strategies.
• Documentation Update: Updated the FP8 quantization documentation (docs/features/quantization/fp8.md) to reflect the new memory optimization feature and provide guidance on its usage.
• New Test Coverage: Added a dedicated test case (test_offload_weights_before_quant_fp8 in tests/quantization/test_cpu_offload.py) to verify the correct functionality of the VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT feature.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

The pull request introduces a new feature to offload model weights to the CPU before FP8 online quantization. This addresses memory limitations when quantizing large models. The changes include modifications to the documentation, a new test case, updates to the configuration and environment variable handling, and adjustments to the FP8 quantization process to accommodate CPU offloading. The code appears well-structured and addresses the intended problem. However, some comments could be improved for clarity.

[gemini-code-assist]

Consider rephrasing for clarity. Suggestion: "Currently, the model is loaded at its original precision before being quantized to 8-bits. This requires sufficient memory to hold the entire model. To avoid this, set VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT=1 to offload weights to CPU before quantization; the quantized weights will then be stored on the device."

Suggested change

	Currently, by default we load the model at original precision before quantizing down to 8-bits, so you need enough memory to load the whole model. To avoid this, adding `VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT=1` can allow offloading weights to cpu before quantization and quantized weights will be kept in device.
	Currently, by default we load the model at original precision before quantizing down to 8-bits, so you need enough memory to load the whole model. To avoid this, adding `VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT=1` can allow offloading weights to cpu before quantization and quantized weights will be kept in device.

[gemini-code-assist]

Consider adding a more descriptive error message, including the names of the conflicting variables. Suggestion: "cpu_offload_gb and VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT cannot be enabled simultaneously. Disable one to resolve the conflict."

            raise ValueError("CPU offload can't work together with"
                             "OFFLOAD_WEIGHTS_BEFORE_QUANT")

[gemini-code-assist]

Consider rephrasing for clarity. Suggestion: "If weights were loaded onto the CPU for FP8 online quantization, there is no need to move them back to the original device."

Suggested change

	if VLLM_OFFLOAD_WEIGHTS_BEFORE_QUANT:
	return
	# If weights were loaded on cpu for fp8 online quantization, won't need move
	# back to original device

[yma11]

@Mogin @comaniac Can you help take a look at this PR? Thanks.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yma11.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork