mirror of
https://github.com/vllm-project/vllm.git
synced 2026-10-28 09:45:26 +08:00
[LoRA] Remove tensorizer (#60024)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
1 parent
d3547f9d03
commit
f9c9e8ac24
9 files changed
+30
-375
No files matched your search
@@ -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}},
|
||||
)
|
||||
```
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
Reference in new issue
Block a user