vllm-project/vllm

[Feature]: Composite model loading using `AutoWeightsLoader` for all models

Open

#15,697 opened on Mar 28, 2025

View on GitHub
 (39 comments) (0 reactions) (0 assignees)Python (80,034 stars) (16,816 forks)batch import
feature requestgood first issuekeep-open

Description

🚀 The feature, motivation and pitch

#9160 first introduced AutoWeightsLoader to recursively call load_weights on sub-modules. This lets composite models (most notably multi-modal models) use language backbones (*Model classes such as LlamaModel) without having to repeat their weight loading logic.

Currently, load_weights is only implemented in a few language backbones. It would be great to standardize this approach and apply it to all language backbones in vLLM. The steps to do this are pretty straightforward:

  1. Move the existing load_weights function from *ForCausalLM to *Model.
  2. Create a new load_weights function in *ForCausalLM that loads the weights using AutoWeightsLoader.
  3. Move any logic in *Model.load_weights that only applies to *ForCausalLM back to *ForCausalLM.load_weights. Usually, this involves lm_head.

For reference, you can look at the implementation for models such as Llama, Gemma2/3, Qwen2 and ChatGLM.

To avoid scope creep, I suggest opening a PR for updating only a few models at a time

Alternatives

No response

Additional context

No response

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Contributor guide

Tech stack
pythonpytorch
Domain
backendperformance
Issue type
feature
DifficultyEstimated implementation difficulty for a new contributor, from 1 for very small changes to 5 for expert-level work.
3
Estimated timeA rough time range for an experienced contributor to investigate, implement, test, and prepare a pull request.
half day
Activity statusHow available the issue appears right now: fresh, active, stale, blocked, or waiting on maintainer input.
active
ClarityHow clearly the issue explains the expected change, acceptance criteria, and next step.
mostly clear
Prerequisites
PythonPyTorchvLLM codebase understanding
Newbie friendlinessA 1-100 score estimating how approachable this issue is for first-time contributors.
65
Research direction
Review the existing implementation of AutoWeightsLoader in the codebase, particularly for models like Llama, Gemma2/3, Qwen2, and ChatGLM as referenced in the issue. Understand the current load weights functions in *ForCausalLM classes and *Model classes. Then, apply the three step process to a target model: move load weights from *ForCausalLM to *Model, create new load weights in *ForCausalLM using AutoWeightsLoader, and move any lm head specific logic back. Create a PR updating one or a few models at a time.