[LoRA] Remove tensorizer (#60024)

Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
Jee Jee Li authored and GitHub committed 2026-10-05 07:01:56 -07:00
1 parent d3547f9d03
commit f9c9e8ac24
9 files changed
+30 -375

No files matched your search

+3 -17
View File
@@ -32,25 +32,15 @@ python examples/features/tensorize_vllm_model.py \
--suffix v1
```
This saves the model tensors at `s3://my-bucket/vllm/facebook/opt-125m/v1`. If you intend on applying a LoRA adapter to your tensorized model, you can pass the HF id of the LoRA adapter in the above command, and the artifacts will be saved there too:
```bash
python examples/features/tensorize_vllm_model.py \
--model facebook/opt-125m \
--lora-path <lora_id> \
serialize \
--serialized-directory s3://my-bucket \
--suffix v1
```
This saves the model tensors at `s3://my-bucket/vllm/facebook/opt-125m/v1`.
## Serving the model using Tensorizer
Once the model is serialized where you want it, you can load the model using `vllm serve` or the `LLM` entrypoint. You can pass the directory where you saved the model to the `model` argument for `LLM()` and `vllm serve`. For example, to serve the tensorized model saved previously with the LoRA adapter, you'd do:
Once the model is serialized where you want it, you can load the model using `vllm serve` or the `LLM` entrypoint. You can pass the directory where you saved the model to the `model` argument for `LLM()` and `vllm serve`. For example, to serve the tensorized model saved previously, you'd do:
```bash
vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
--load-format tensorizer \
--enable-lora
--load-format tensorizer
```
Or, with `LLM()`:
@@ -60,7 +50,6 @@ from vllm import LLM
llm = LLM(
"s3://my-bucket/vllm/facebook/opt-125m/v1",
load_format="tensorizer",
enable_lora=True,
)
```
@@ -73,7 +62,6 @@ As an example, CPU concurrency can be limited when serializing with `tensorizer`
```bash
python examples/features/tensorize_vllm_model.py \
--model facebook/opt-125m \
--lora-path <lora_id> \
serialize \
--serialized-directory s3://my-bucket \
--serialization-kwargs '{"limit_cpu_concurrency": 2}' \
@@ -85,7 +73,6 @@ As an example when customizing the loading process via `TensorDeserializer`, you
```bash
vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
--load-format tensorizer \
--enable-lora \
--model-loader-extra-config '{"deserialization_kwargs": {"num_readers": 2}}'
```
@@ -96,7 +83,6 @@ from vllm import LLM
llm = LLM(
"s3://my-bucket/vllm/facebook/opt-125m/v1",
load_format="tensorizer",
enable_lora=True,
model_loader_extra_config={"deserialization_kwargs": {"num_readers": 2}},
)
```
+7 -70
View File
@@ -6,13 +6,11 @@ import logging
import os
import uuid
from vllm import LLM, SamplingParams
from vllm import LLM
from vllm.engine.arg_utils import EngineArgs
from vllm.lora.request import LoRARequest
from vllm.model_executor.model_loader.tensorizer import (
TensorizerArgs,
TensorizerConfig,
tensorize_lora_adapter,
tensorize_vllm_model,
tensorizer_kwargs_arg,
)
@@ -105,21 +103,6 @@ loading with tensorizer that are given to `TensorizerConfig`, run:
under the `tensorizer options` section. These can also be used for
deserialization in this example script, although `--tensorizer-uri` and
`--path-to-tensors` are functionally the same in this case.
Tensorizer can also be used to save and load LoRA adapters. A LoRA adapter
can be serialized directly with the path to the LoRA adapter on HF Hub and
a TensorizerConfig object. In this script, passing a HF id to a LoRA adapter
will serialize the LoRA adapter artifacts to `--serialized-directory`.
You can then use the LoRA adapter with `vllm serve`, for instance, by ensuring
the LoRA artifacts are in your model artifacts directory and specifying
`--enable-lora`. For instance:
```
vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
--load-format tensorizer \
--enable-lora
```
"""
@@ -134,18 +117,6 @@ def get_parser():
)
parser = EngineArgs.add_cli_args(parser)
parser.add_argument(
"--lora-path",
type=str,
required=False,
help="Path to a LoRA adapter to "
"serialize along with model tensors. This can then be deserialized "
"along with the model by instantiating a TensorizerConfig object, "
"creating a dict from it with TensorizerConfig.to_serializable(), "
"and passing it to LoRARequest's initializer with the kwarg "
"tensorizer_config_dict.",
)
subparsers = parser.add_subparsers(dest="command", required=True)
serialize_parser = subparsers.add_parser(
@@ -263,42 +234,12 @@ def merge_extra_config_with_tensorizer_config(extra_cfg: dict, cfg: TensorizerCo
def deserialize(args, tensorizer_config):
if args.lora_path:
tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
llm = LLM(
model=args.model,
load_format="tensorizer",
tensor_parallel_size=args.tensor_parallel_size,
model_loader_extra_config=tensorizer_config,
enable_lora=True,
)
sampling_params = SamplingParams(
temperature=0, max_tokens=256, stop=["[/assistant]"]
)
# Truncating this as the extra text isn't necessary
prompts = ["[user] Write a SQL query to answer the question based on ..."]
# Test LoRA load
print(
llm.generate(
prompts,
sampling_params,
lora_request=LoRARequest(
"sql-lora",
1,
args.lora_path,
tensorizer_config_dict=tensorizer_config.to_serializable(),
),
)
)
else:
llm = LLM(
model=args.model,
load_format="tensorizer",
tensor_parallel_size=args.tensor_parallel_size,
model_loader_extra_config=tensorizer_config,
)
llm = LLM(
model=args.model,
load_format="tensorizer",
tensor_parallel_size=args.tensor_parallel_size,
model_loader_extra_config=tensorizer_config,
)
return llm
@@ -366,10 +307,6 @@ def main():
**credentials,
)
if args.lora_path:
tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
tensorize_lora_adapter(args.lora_path, tensorizer_config)
merge_extra_config_with_tensorizer_config(extra_config, tensorizer_config)
tensorize_vllm_model(engine_args, tensorizer_config)
@@ -13,13 +13,11 @@ from tests.utils import RemoteOpenAIServer
from vllm.engine.arg_utils import EngineArgs
from vllm.model_executor.model_loader.tensorizer import (
TensorizerConfig,
tensorize_lora_adapter,
tensorize_vllm_model,
)
from vllm.platforms import current_platform
MODEL_NAME = "unsloth/llama-3.2-1b-Instruct"
LORA_PATH = "davzoku/finqa_adapter_1b"
def _cleanup():
@@ -44,11 +42,10 @@ def model_uri(tmp_dir):
@pytest.fixture(scope="module")
def tensorize_model_and_lora(tmp_dir, model_uri):
tensorizer_config = TensorizerConfig(tensorizer_uri=model_uri, lora_dir=tmp_dir)
def tensorize_model(model_uri):
tensorizer_config = TensorizerConfig(tensorizer_uri=model_uri)
args = EngineArgs(model=MODEL_NAME)
tensorize_lora_adapter(LORA_PATH, tensorizer_config)
tensorize_vllm_model(args, tensorizer_config)
# Manually invoke a _cleanup() here, as the cleanup()
@@ -59,7 +56,7 @@ def tensorize_model_and_lora(tmp_dir, model_uri):
@pytest.fixture(scope="module")
def server(model_uri, tensorize_model_and_lora):
def server(model_uri, tensorize_model):
# In this case, model_uri is a directory with a model.tensors
# file and all necessary model artifacts, particularly a
# HF `config.json` file. In this case, Tensorizer can infer the
@@ -72,7 +69,6 @@ def server(model_uri, tensorize_model_and_lora):
"tensorizer",
"--served-model-name",
MODEL_NAME,
"--enable-lora",
]
if current_platform.is_rocm():
args += ["--attention-backend", "TRITON_ATTN"]
+9 -124
View File
@@ -1,18 +1,14 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import subprocess
import sys
import pytest
import vllm
import vllm.config
from vllm import LLM
from vllm.lora.request import LoRARequest
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from vllm.platforms import current_platform
from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
from ..utils import create_new_process_for_each_test, multi_gpu_test
PROMPT_TEMPLATE = """<|eot_id|><|start_header_id|>user<|end_header_id|>
I want you to act as a SQL terminal in front of an example database, you need only to return the sql command to me.Below is an instruction that describes a task, Write a response that appropriately completes the request.
@@ -40,7 +36,6 @@ def do_sample(
llm: vllm.LLM,
lora_path: str,
lora_id: int,
tensorizer_config_dict: dict | None = None,
) -> list[str]:
prompts = [
PROMPT_TEMPLATE.format(context="How many candidates are there?"),
@@ -56,27 +51,11 @@ def do_sample(
sampling_params = vllm.SamplingParams(
temperature=0, max_tokens=64, stop=["<|im_end|>"]
)
if tensorizer_config_dict is not None:
outputs = llm.generate(
prompts,
sampling_params,
lora_request=LoRARequest(
str(lora_id),
lora_id,
lora_path,
tensorizer_config_dict=tensorizer_config_dict,
)
if lora_id
else None,
)
else:
outputs = llm.generate(
prompts,
sampling_params,
lora_request=LoRARequest(str(lora_id), lora_id, lora_path)
if lora_id
else None,
)
outputs = llm.generate(
prompts,
sampling_params,
lora_request=LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None,
)
lora_request = LoRARequest(str(lora_id), lora_id, lora_path) if lora_id else None
generated_texts: list[str] = []
for output in outputs:
@@ -94,31 +73,13 @@ def do_sample(
return generated_texts
def generate_and_test(
llm, llama32_lora_files, tensorizer_config_dict: dict | None = None
):
def generate_and_test(llm, llama32_lora_files):
print("lora adapter created")
print("lora 1")
assert (
do_sample(
llm,
llama32_lora_files,
tensorizer_config_dict=tensorizer_config_dict,
lora_id=1,
)
== EXPECTED_LORA_OUTPUT
)
assert do_sample(llm, llama32_lora_files, lora_id=1) == EXPECTED_LORA_OUTPUT
print("lora 2")
assert (
do_sample(
llm,
llama32_lora_files,
tensorizer_config_dict=tensorizer_config_dict,
lora_id=2,
)
== EXPECTED_LORA_OUTPUT
)
assert do_sample(llm, llama32_lora_files, lora_id=2) == EXPECTED_LORA_OUTPUT
print("removing lora")
@@ -168,79 +129,3 @@ def test_llama_lora_tp4_fully_sharded_loras(llama32_lora_files):
fully_sharded_loras=True,
)
generate_and_test(llm, llama32_lora_files)
@multi_gpu_test(num_gpus=2)
def test_tp2_serialize_and_deserialize_lora(
tmp_path,
llama32_lora_files,
):
# Run the tensorizing of the LoRA adapter and the model in a subprocess
# to guarantee cleanup
tp_size = 2
model_name = "model-rank-%03d.tensors"
model_ref = MODEL_PATH
lora_path = llama32_lora_files
suffix = "test"
try:
result = subprocess.run(
[
sys.executable,
f"{VLLM_PATH}/examples/features/tensorize_vllm_model.py",
"--model",
MODEL_PATH,
"--lora-path",
lora_path,
"--tensor-parallel-size",
str(tp_size),
"serialize",
"--serialized-directory",
str(tmp_path),
"--suffix",
suffix,
"--serialization-kwargs",
'{"limit_cpu_concurrency": 4}',
],
check=True,
capture_output=True,
text=True,
)
except subprocess.CalledProcessError as e:
print("Tensorizing failed.")
print("STDOUT:\n", e.stdout)
print("STDERR:\n", e.stderr)
raise
print("STDOUT:\n", result.stdout)
model_uri = tmp_path / "vllm" / model_ref / suffix / model_name
tensorizer_config = TensorizerConfig(tensorizer_uri=str(model_uri))
loaded_llm = LLM(
model=model_ref,
load_format="tensorizer",
enable_lora=True,
enforce_eager=True,
model_loader_extra_config=tensorizer_config,
max_num_seqs=7,
max_model_len=1024,
tensor_parallel_size=2,
max_loras=2,
# Leave headroom for LoRA adapter loading and Triton JIT
# compilation, which can allocate GPU memory concurrently
# during the first inference step.
gpu_memory_utilization=0.85,
)
tc_as_dict = tensorizer_config.to_serializable()
print("lora adapter created")
print("lora 1")
assert (
do_sample(
loaded_llm, llama32_lora_files, tensorizer_config_dict=tc_as_dict, lora_id=1
)
== EXPECTED_LORA_OUTPUT
)
+1 -22
View File
@@ -15,7 +15,6 @@ from vllm.lora.utils import (
is_base_embedding_weights,
parse_fine_tuned_lora_name,
)
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
from vllm.model_executor.models.utils import WeightsMapper
from vllm.utils.torch_utils import PIN_MEMORY
@@ -204,7 +203,6 @@ class LoRAModel:
dtype: torch.dtype | None = None,
model_vocab_size: int | None = None,
weights_mapper: WeightsMapper | None = None,
tensorizer_config_dict: dict | None = None,
skip_prefixes: list[str] | None = None,
moe_ep_spec: MoEEPLoadSpec | None = None,
) -> "LoRAModel":
@@ -223,8 +221,6 @@ class LoRAModel:
embedding deltas.
weights_mapper: Optional mapper rewriting checkpoint weight names
to vLLM names.
tensorizer_config_dict: Optional tensorizer config used to load
the checkpoint via tensorizer instead of from disk.
skip_prefixes: List of module name prefixes to skip during loading.
Models can define this to skip modules not used in inference
(e.g., MTP layers). Format: ["mtp."]
@@ -280,24 +276,7 @@ class LoRAModel:
f" Please verify that the loaded LoRA module is correct"
)
if tensorizer_config_dict:
from tensorizer import TensorDeserializer
tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
tensorizer_dir = tensorizer_config.tensorizer_dir
if tensorizer_dir is None:
raise ValueError("tensorizer_dir must be set in tensorizer config.")
lora_tensor_path = os.path.join(tensorizer_dir, "adapter_model.tensors")
tensorizer_args = tensorizer_config._construct_tensorizer_args()
tensors = TensorDeserializer(
lora_tensor_path,
dtype=tensorizer_config.dtype,
device=device,
**tensorizer_args.deserialization_kwargs,
)
check_unexpected_modules(tensors)
elif os.path.isfile(lora_tensor_path):
if os.path.isfile(lora_tensor_path):
# Find unexpected modules.
# Use safetensor key as a source of truth to find expected modules.
# in peft if you have target_modules A, B, C and C does not exist
+2 -26
View File
@@ -11,7 +11,6 @@ from typing import Literal
from vllm.config.lora import LoRAConfig
from vllm.logger import init_logger
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
logger = init_logger(__name__)
@@ -89,33 +88,10 @@ class PEFTHelper:
cls,
lora_path: str,
max_position_embeddings: int | None,
tensorizer_config_dict: dict | None = None,
) -> "PEFTHelper":
lora_config_path = os.path.join(lora_path, "adapter_config.json")
if tensorizer_config_dict:
tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
tensorizer_args = tensorizer_config._construct_tensorizer_args()
from tensorizer.stream_io import open_stream
tensorizer_dir = tensorizer_config.tensorizer_dir
if tensorizer_dir is None:
raise ValueError("tensorizer_dir must be set in tensorizer config.")
lora_config_path = os.path.join(tensorizer_dir, "adapter_config.json")
with open_stream(
lora_config_path, mode="rb", **tensorizer_args.stream_kwargs
) as f:
config = json.load(f)
logger.info(
"Successfully deserialized LoRA config from %s",
tensorizer_config.tensorizer_dir,
)
else:
with open(lora_config_path) as f:
config = json.load(f)
with open(lora_config_path) as f:
config = json.load(f)
config["vllm_max_position_embeddings"] = max_position_embeddings
return cls.from_dict(config)
-1
View File
@@ -25,7 +25,6 @@ class LoRARequest(
lora_int_id: int
lora_path: str = ""
base_model_name: str | None = msgspec.field(default=None)
tensorizer_config_dict: dict | None = None
load_inplace: bool = False
is_3d_lora_weight: bool = False
"""Whether this adapter's MoE weights are stored in the 3D fused
+1 -4
View File
@@ -119,9 +119,7 @@ class WorkerLoRAManager:
lora_path = get_adapter_absolute_path(lora_request.lora_path)
peft_helper = PEFTHelper.from_local_dir(
lora_path,
self.max_position_embeddings,
lora_request.tensorizer_config_dict,
lora_path, self.max_position_embeddings
)
# Validates the LoRA configuration against requirements before
@@ -147,7 +145,6 @@ class WorkerLoRAManager:
device="cpu",
dtype=self.lora_config.lora_dtype,
model_vocab_size=self.vocab_size,
tensorizer_config_dict=lora_request.tensorizer_config_dict,
weights_mapper=hf_to_vllm_mapper,
skip_prefixes=lora_skip_prefixes,
moe_ep_spec=self._adapter_manager.moe_ep_load_spec,
+4 -104
View File
@@ -23,7 +23,6 @@ from transformers import PreTrainedConfig
import vllm.envs as envs
from vllm.config import ModelConfig, ParallelConfig, VllmConfig, set_current_vllm_config
from vllm.logger import init_logger
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
from vllm.platforms import current_platform
from vllm.transformers_utils.repo_utils import hf_api
from vllm.utils.argparse_utils import FlexibleArgumentParser
@@ -167,7 +166,6 @@ class TensorizerConfig(MutableMapping):
s3_access_key_id: str | None = None
s3_secret_access_key: str | None = None
s3_endpoint: str | None = None
lora_dir: str | None = None
stream_kwargs: dict[str, Any] | None = None
serialization_kwargs: dict[str, Any] | None = None
deserialization_kwargs: dict[str, Any] | None = None
@@ -185,11 +183,8 @@ class TensorizerConfig(MutableMapping):
Attributes:
tensorizer_uri: Path to serialized model tensors. Can be a local file
path or a S3 URI. This is a required field unless lora_dir is
provided and the config is meant to be used for the
`tensorize_lora_adapter` function. Unless a `tensorizer_dir` or
`lora_dir` is passed to this object's initializer, this is
a required argument.
path or a S3 URI. Unless a `tensorizer_dir` is passed to this
object's initializer, this is a required argument.
tensorizer_dir: Path to a directory containing serialized model tensors,
and all other potential model artifacts to load the model, such as
configs and tokenizer files. Can be passed instead of
@@ -219,10 +214,6 @@ class TensorizerConfig(MutableMapping):
be set via the S3_SECRET_ACCESS_KEY environment variable.
s3_endpoint: The endpoint for the S3 bucket. Can also be set via the
S3_ENDPOINT_URL environment variable.
lora_dir: Path to a directory containing LoRA adapter artifacts for
serialization or deserialization. When serializing LoRA adapters
this is the only necessary parameter to pass to this object's
initializer.
"""
def __post_init__(self):
@@ -232,12 +223,6 @@ class TensorizerConfig(MutableMapping):
and re.search(r"%0\dd", self.tensorizer_uri) is not None
)
if self.tensorizer_dir and self.lora_dir:
raise ValueError(
"Only one of tensorizer_dir or lora_dir may be specified. "
"Use lora_dir exclusively when serializing LoRA adapters, "
"and tensorizer_dir or tensorizer_uri otherwise."
)
if self.tensorizer_dir and self.tensorizer_uri:
logger.warning_once(
"Provided both tensorizer_dir and tensorizer_uri. "
@@ -246,17 +231,12 @@ class TensorizerConfig(MutableMapping):
)
self.tensorizer_dir = os.path.dirname(self.tensorizer_uri)
if not self.tensorizer_uri:
if self.lora_dir:
self.tensorizer_uri = f"{self.lora_dir}/adapter_model.tensors"
elif self.tensorizer_dir:
if self.tensorizer_dir:
self.tensorizer_uri = f"{self.tensorizer_dir}/model.tensors"
else:
raise ValueError(
"Unable to resolve tensorizer_uri. "
"A valid tensorizer_uri or tensorizer_dir "
"must be provided for deserialization, and a "
"valid tensorizer_uri, tensorizer_uri, or "
"lora_dir for serialization."
"A valid tensorizer_uri or tensorizer_dir must be provided."
)
else:
self.tensorizer_dir = os.path.dirname(self.tensorizer_uri)
@@ -288,9 +268,6 @@ class TensorizerConfig(MutableMapping):
if "tensorizer_uri" in raw_tc_dict and "tensorizer_dir" in raw_tc_dict:
blacklisted.append("tensorizer_dir")
if "tensorizer_dir" in raw_tc_dict and "lora_dir" in raw_tc_dict:
blacklisted.append("tensorizer_dir")
tc_dict = {}
for k, v in raw_tc_dict.items():
if (
@@ -495,25 +472,6 @@ def _check_tensors_on_meta_device(model: nn.Module) -> None:
)
def _resize_lora_embeddings(model: nn.Module):
"""Modify LoRA embedding layers to use bigger tensors
to allow for adapter added tokens."""
for child in model.modules():
if (
isinstance(child, VocabParallelEmbedding)
and child.weight.shape[0] < child.num_embeddings_per_partition
):
new_weight = torch.empty(
child.num_embeddings_per_partition,
child.embedding_dim,
dtype=child.weight.dtype,
device=child.weight.device,
)
new_weight[: child.weight.shape[0]].copy_(child.weight.data)
new_weight[child.weight.shape[0] :].fill_(0)
child.weight.data = new_weight
def init_tensorizer_model(
tensorizer_config: TensorizerConfig, vllm_config: VllmConfig
) -> nn.Module:
@@ -567,7 +525,6 @@ def deserialize_tensorizer_model(
logger.info("Memory usage after: %s", after_mem)
_check_tensors_on_meta_device(model)
_resize_lora_embeddings(model)
del model.vllm_tensorized_marker
@@ -769,60 +726,3 @@ def tensorize_vllm_model(
if error is not None:
raise error
def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig):
"""Uses tensorizer to serialize a LoRA adapter. Assumes that the files
needed to load a LoRA adapter are a safetensors-format file called
adapter_model.safetensors and a json config file called adapter_config.json.
Serializes the files in the tensorizer_config.tensorizer_dir
"""
import safetensors
from vllm.lora.utils import get_adapter_absolute_path
lora_dir = get_adapter_absolute_path(lora_path)
tensor_path = config_path = ""
for file in os.listdir(lora_dir):
if file.startswith("adapter_model"):
tensor_path = lora_dir + "/" + file
if file.startswith("adapter_config"):
config_path = lora_dir + "/" + file
if tensor_path and config_path:
break
if tensor_path.endswith(".safetensors"):
tensors = safetensors.torch.load_file(tensor_path)
elif tensor_path.endswith(".bin"):
tensors = torch.load(tensor_path, weights_only=True)
else:
raise ValueError(
f"Unsupported adapter model file: {tensor_path}. "
f"Must be a .safetensors or .bin file."
)
with open(config_path) as f:
config = json.load(f)
tensorizer_args = tensorizer_config._construct_tensorizer_args()
with open_stream(
f"{tensorizer_config.tensorizer_dir}/adapter_config.json",
mode="wb+",
**tensorizer_args.stream_kwargs,
) as f:
f.write(json.dumps(config).encode("utf-8"))
lora_uri = f"{tensorizer_config.tensorizer_dir}/adapter_model.tensors"
with open_stream(lora_uri, mode="wb+", **tensorizer_args.stream_kwargs) as f:
serializer = TensorSerializer(f)
serializer.write_state_dict(tensors)
serializer.close()
logger.info(
"Successfully serialized LoRA files to %s",
str(tensorizer_config.tensorizer_dir),
)