[sergeykochetkov]

FIX #4186 #2816

Support lm_head and embed_tokens fully trained in LoRA.

We found that quality of our adapters significantly drops without fully-trained lm_head or lm_head trained in LoRA style.
This is functionality of peft modules_to_save=[lm_head, mebed_tokens] https://huggingface.co/docs/peft/v0.12.0/en/package_reference/#peft.LoraConfig.modules_to_save

The idea is to replace base_model VocabParallelEmbedding and ParallelLMHead by layers loaded from modules_to_save at inferencing LoRA

• dirty implementation
• tests for new functionality
• checking old functionality is working
• inference with fully trained lm_head performance measurement
• implement embed_tokens fully trained as well

This PR is clean recreation of #8082

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
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 do one of these:

• Add ready label to the PR
• Enable auto-merge.

🚀

[sergeykochetkov]

My PR with modules_to_save in LoRA moved here, please review
@WoosukKwon @youkaichao

[sergeykochetkov]

/ready

[youkaichao]

for lora code, please ping @jeejeelee for review.

[mergify]

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

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

[jeejeelee]

This feature isn't really very compatible with multi-LoRA.

[sergeykochetkov]

This feature isn't really very compatible with multi-LoRA.

Why?
I implemented punica style triton-kernels which allows to select lora weights (lm_head or embed_tokens) according to multi-lora indexes. Please check files
vllm/lora/ops/bgmv_embed.py
vllm/lora/ops/bgmv_sample.py

[jeejeelee]

Cases like this - have you considered them?

[sergeykochetkov]

Cases like this - have you considered them?

modules_to_save ['resample' ] is not implemented, it will raise exception. This is a task for another PR.

modules_to_save:['token_embeds'] - will work

[jeejeelee]

If we support module_to_save, we need to consider all possible scenarios, which would make LoRA difficult to maintain.

We also appreciate your willingness to contribute this. Give me some time to think through this issue.

[jeejeelee]

Can your provided test script run successfully? I can't get it to run locally.

[balachandarsv]

This will be useful to us, we have a model where we have trained lora on lm_head.
@sergeykochetkov @jeejeelee

[nole70]

any updates on this? Would greatly appreciate this feature getting through.

[hixulei]

same issues

[devansh-e2e]

I have similar issue of lm_head, is this one fixed or still in development?

[JEM-Mosig]

Also have the same issue

[yanghailong-git]

@sergeykochetkov hello
I install the code from the sergeykochetkov:lora_with_modules_to_save_s_2 branch and get the following error when running tests/lora/test_lora_with_modules_to_save.py:

However, when combining t-tech/T-lite-it-1.0 with SergeyKochetkovT/qwen2.5-lora-with-modules-to-save for testing, all the tests passed successfully.

Is this the only combination of base model and LoRA model that can pass the test?

My python environment is as follows:

accelerate 1.9.0
aiohappyeyeballs 2.6.1
aiohttp 3.12.15
aiohttp-cors 0.8.1
aiosignal 1.4.0
airportsdata 20250706
annotated-types 0.7.0
anyio 4.9.0
astor 0.8.1
attrs 25.3.0
blake3 1.0.5
cachetools 5.5.2
certifi 2025.7.14
charset-normalizer 3.4.2
click 8.2.1
cloudpickle 3.1.1
colorful 0.5.7
compressed-tensors 0.9.0
depyf 0.18.0
dill 0.4.0
diskcache 5.6.3
distlib 0.4.0
distro 1.9.0
einops 0.8.1
fastapi 0.116.1
filelock 3.18.0
frozenlist 1.7.0
fsspec 2025.7.0
gguf 0.10.0
google-api-core 2.25.1
google-auth 2.40.3
googleapis-common-protos 1.70.0
grpcio 1.74.0
h11 0.16.0
hf-xet 1.1.5
httpcore 1.0.9
httptools 0.6.4
httpx 0.28.1
huggingface-hub 0.34.3
idna 3.10
importlib_metadata 8.7.0
iniconfig 2.1.0
interegular 0.3.3
Jinja2 3.1.6
jiter 0.10.0
jsonschema 4.25.0
jsonschema-specifications 2025.4.1
lark 1.2.2
lm-format-enforcer 0.10.11
MarkupSafe 3.0.2
mistral_common 1.8.3
mpmath 1.3.0
msgpack 1.1.1
msgspec 0.19.0
multidict 6.6.3
nest-asyncio 1.6.0
networkx 3.5
ninja 1.11.1.4
numpy 1.26.4
nvidia-cublas-cu12 12.4.5.8
nvidia-cuda-cupti-cu12 12.4.127
nvidia-cuda-nvrtc-cu12 12.4.127
nvidia-cuda-runtime-cu12 12.4.127
nvidia-cudnn-cu12 9.1.0.70
nvidia-cufft-cu12 11.2.1.3
nvidia-curand-cu12 10.3.5.147
nvidia-cusolver-cu12 11.6.1.9
nvidia-cusparse-cu12 12.3.1.170
nvidia-ml-py 12.575.51
nvidia-nccl-cu12 2.21.5
nvidia-nvjitlink-cu12 12.4.127
nvidia-nvtx-cu12 12.4.127
openai 1.97.1
opencensus 0.11.4
opencensus-context 0.1.3
opencv-python-headless 4.11.0.86
opentelemetry-api 1.36.0
opentelemetry-exporter-prometheus 0.57b0
opentelemetry-proto 1.36.0
opentelemetry-sdk 1.36.0
opentelemetry-semantic-conventions 0.57b0
outlines 0.1.11
outlines_core 0.1.26
packaging 25.0
partial-json-parser 0.2.1.1.post6
peft 0.16.0
pillow 11.3.0
pip 25.1
platformdirs 4.3.8
pluggy 1.6.0
prometheus_client 0.22.1
prometheus-fastapi-instrumentator 7.1.0
propcache 0.3.2
proto-plus 1.26.1
protobuf 6.31.1
psutil 7.0.0
py-cpuinfo 9.0.0
py-spy 0.4.0
pyasn1 0.6.1
pyasn1_modules 0.4.2
pycountry 24.6.1
pydantic 2.11.7
pydantic_core 2.33.2
pydantic-extra-types 2.10.5
Pygments 2.19.2
pytest 8.4.1
python-dotenv 1.1.1
PyYAML 6.0.2
pyzmq 27.0.0
ray 2.48.0
referencing 0.36.2
regex 2025.7.31
requests 2.32.4
rpds-py 0.26.0
rsa 4.9.1
safetensors 0.5.3
sentencepiece 0.2.0
setuptools 78.1.1
six 1.17.0
smart_open 7.3.0.post1
sniffio 1.3.1
starlette 0.47.2
sympy 1.13.1
tiktoken 0.9.0
tokenizers 0.21.4
torch 2.5.1
torchaudio 2.5.1
torchvision 0.20.1
tqdm 4.67.1
transformers 4.54.1
triton 3.1.0
typing_extensions 4.14.1
typing-inspection 0.4.1
urllib3 2.5.0
uvicorn 0.35.0
uvloop 0.21.0
virtualenv 20.32.0
vllm 0.1.dev4430+g0402980.precompiled /mnt/yanghl/github/vllm_for_lora/vllm_spec_decoding
watchfiles 1.1.0
websockets 15.0.1
wheel 0.45.1
wrapt 1.17.2
xformers 0.0.28.post3
xgrammar 0.1.22
yarl 1.20.1
zipp 3.23.0

[yanghailong-git]

@jeejeelee I’m following the implementation of this feature closely. Why hasn't this PR been merged yet? Are there any plans to merge it in the future?

If cases like this can be solved，has the PR implementing the modules_to_save feature been considered for merging?

Is there anything else that needs to be taken into account? I’m currently thinking about doing some work and optimizations on top of it.

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[AlongWY]

any update?

[mergify]

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

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

[Serzhanov]

Looking forward for this feature!

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[mergify]

Hi @sergeykochetkov, 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

[mergify]

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

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

[github-actions]

This pull request has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this pull request should remain open. Thank you!

[github-actions]

This pull request has been automatically closed due to inactivity. Please feel free to reopen if you intend to continue working on it. Thank you!