[USAGE] Improve error handling for weight initialization in Unquantized…

[koiker]

…LinearMethod and Attention class.

When initializing vLLM the torch.OutOfMemoryError: CUDA out of memory. is opaque and does not indicate which step of the loading model it failed. Also it is not clear where the pre-allocated memory has been used.

Errors like:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 896.00 MiB. GPU 2 has a total capacity of 23.57 GiB of which 235.88 MiB is free. Including non-PyTorch memory, this process has 23.33 GiB memory in use. Of the allocated memory 23.03 GiB is allocated by PyTorch, and 13.48 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

This small change will include a more clear error message on which step of the model loading the error has occurred and allow users to change the settings to use a model that will fit the memory or adjust other settings in the vllm.

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 result's comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Improve the usage of vLLM when initializing a vLLM instance or starting a vLLM serve and loading the model throws an out-of-memory error, but it is not clear if there was a previous process using the memory or if the device set to load the model is the correct one. This change will include logging and debug information that will allow developers and engineers to better understand at which step in the model loading the error occurs. This will improve the troubleshooting process to resolve the issue of loading LLMs with vLLM

Test Plan

Load a model that is larger than the available memory and see the torch.OutOfMemoryError happens. Try to understand from the information displayed where is the issue.

Test Result

After the implementation of the code in this PR, this is an example of the output of an error loading a large LLM that does not fit into the GPU memory:

2025-07-01T15:52:01.113Z  INFO gpu_model_runner.load_model: Starting to load model meta-llama/Llama-3.3-70B-Instruct...   
2025-07-01T15:52:01.506Z  INFO gpu_model_runner.load_model: Loading model from scratch...   
2025-07-01T15:52:01.691Z  INFO cuda.get_attn_backend_cls: Using Flash Attention backend on V1 engine.   
2025-07-01T15:52:01.739Z ERROR linear.create_weights: Failed to create unquantized linear weights: CUDA out of memory. Tried to allocate 896.00 MiB. GPU 2 has a total capacity of 23.57 GiB of which 235.88 MiB is free. Including non-PyTorch memory, this process has 23.33 GiB memory in use. Of the allocated memory 23.03 GiB is allocated by PyTorch, and 13.48 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)   
2025-07-01T15:52:01.739Z DEBUG linear.create_weights: CUDA device: 2   
2025-07-01T15:52:01.739Z DEBUG linear.create_weights: Allocated: 23.03 GB   
2025-07-01T15:52:01.740Z DEBUG linear.create_weights: Reserved: 23.04 GB   
2025-07-01T15:52:02.595Z ERROR core.run_engine_core: EngineCore failed to start.
Traceback (most recent call last):
  File "/mnt/data/dynamo/.venv/lib/python3.12/site-packages/vllm/model_executor/layers/linear.py", line 196, in create_weights
    weight = Parameter(torch.empty(sum(output_partition_sizes),
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/data/dynamo/.venv/lib/python3.12/site-packages/torch/utils/_device.py", line 104, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 896.00 MiB. GPU 2 has a total capacity of 23.57 GiB of which 235.88 MiB is free. Including non-PyTorch memory, this process has 23.33 GiB memory in use. Of the allocated memory 23.03 GiB is allocated by PyTorch, and 13.48 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

Note: In this example, the environment variable VLLM_LOGGING_LEVEL was set to DEBUG to show extra information about the GPU ID and memory utilization.

(Optional) Documentation Update

[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 @koiker, 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 significantly improves the user experience when encountering torch.OutOfMemoryError during vLLM model initialization. It introduces more specific and informative error messages, along with debug details about CUDA memory usage, to help users pinpoint the exact stage of model loading where the memory exhaustion occurs, thereby streamlining the troubleshooting process.

Highlights

• Improved OOM Error Messages: Enhanced error handling in the UnquantizedLinearMethod and Attention classes to provide more specific and informative messages when torch.OutOfMemoryError occurs during weight or constant initialization.
• Detailed CUDA Memory Debugging: Added logging of CUDA device ID, allocated memory, and reserved memory (in GB) to the error messages, which will significantly aid in troubleshooting memory-related issues during model loading.
• Specific Failure Point Identification: The new error messages clarify whether the Out-of-Memory error happened during the initialization of attention layer constants (q_range, k_range, v_range) or unquantized linear layer weights, making it easier to pinpoint the exact stage of failure.

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 code changes improve error handling for out-of-memory situations during model weight initialization by logging CUDA memory statistics. The exception type caught in vllm/attention/layer.py and vllm/model_executor/layers/linear.py can be narrowed for better precision.

[hmellor]

I think there may be a git issue as the diff is unusually large

[koiker]

Hi @hmellor, I have two files changed.

• vllm/model_executor/layers/linear.py
• vllm/attention/layer.py

But for some reason, there are other commits in the same PR.
I see the lint and pre-commit tests are passing now. But there are some time-outs on the buildkite that may not be related to my changes.
Can you try to re-run the buildkite?

[mergify]

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

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

[koiker]

The Lint and Deploy Charts / lint-and-deploy (pull_request) failed due to a docker build error.
How can I re-run the checks?

[DarkLight1337]

I have rerun it for you

[koiker]

buildkite is failing to complete some tasks.
Can someone restart those tasks? If possible, can I have permission to run actions?

[koiker]

Hi @DarkLight1337, I have re-run the checks, all appears to be good.
Can do you do the review?

[DarkLight1337]

Maybe we can use vllm.utils.MemorySnapshot class to avoid code duplication

[hmellor]

Looks good overall, just some nits about GiB and @DarkLight1337's comment about using a memory util

[hmellor]

Please resolve the merge conflict and we can get this merged