biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Convert formatting to use ruff instead of yapf + isort (#26247)

Browse Source 
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
This commit is contained in:
Harry Mellor
2025-10-05 15:06:22 +01:00
committed by GitHub
parent 17edd8a807
commit d6953beb91
1508 changed files with 115244 additions and 94146 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py
View File

@@ -8,24 +8,25 @@ import huggingface_hub.constants
import torch
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf, fastsafetensors_weights_iterator,
safetensors_weights_iterator)
download_weights_from_hf,
fastsafetensors_weights_iterator,
safetensors_weights_iterator,
)
def test_fastsafetensors_model_loader():
with tempfile.TemporaryDirectory() as tmpdir:
huggingface_hub.constants.HF_HUB_OFFLINE = False
download_weights_from_hf("openai-community/gpt2",
allow_patterns=["*.safetensors"],
cache_dir=tmpdir)
download_weights_from_hf(
"openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
)
safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
assert len(safetensors) > 0
fastsafetensors_tensors = {}
hf_safetensors_tensors = {}
for name, tensor in fastsafetensors_weights_iterator(
safetensors, True):
for name, tensor in fastsafetensors_weights_iterator(safetensors, True):
fastsafetensors_tensors[name] = tensor
for name, tensor in safetensors_weights_iterator(safetensors, True):
@@ -34,13 +35,10 @@ def test_fastsafetensors_model_loader():
assert len(fastsafetensors_tensors) == len(hf_safetensors_tensors)
for name, fastsafetensors_tensor in fastsafetensors_tensors.items():
fastsafetensors_tensor = fastsafetensors_tensor.to('cpu')
assert fastsafetensors_tensor.dtype == hf_safetensors_tensors[
name].dtype
assert fastsafetensors_tensor.shape == hf_safetensors_tensors[
name].shape
assert torch.all(
fastsafetensors_tensor.eq(hf_safetensors_tensors[name]))
fastsafetensors_tensor = fastsafetensors_tensor.to("cpu")
assert fastsafetensors_tensor.dtype == hf_safetensors_tensors[name].dtype
assert fastsafetensors_tensor.shape == hf_safetensors_tensors[name].shape
assert torch.all(fastsafetensors_tensor.eq(hf_safetensors_tensors[name]))
if __name__ == "__main__":

27
tests/model_executor/model_loader/runai_model_streamer/test_runai_utils.py
View File

@@ -8,11 +8,12 @@ import tempfile
import huggingface_hub.constants
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf)
from vllm.transformers_utils.runai_utils import (ObjectStorageModel,
is_runai_obj_uri,
list_safetensors)
from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
from vllm.transformers_utils.runai_utils import (
ObjectStorageModel,
is_runai_obj_uri,
list_safetensors,
)
def test_is_runai_obj_uri():
@@ -24,14 +25,14 @@ def test_is_runai_obj_uri():
def test_runai_list_safetensors_local():
with tempfile.TemporaryDirectory() as tmpdir:
huggingface_hub.constants.HF_HUB_OFFLINE = False
download_weights_from_hf("openai-community/gpt2",
allow_patterns=["*.safetensors", "*.json"],
cache_dir=tmpdir)
download_weights_from_hf(
"openai-community/gpt2",
allow_patterns=["*.safetensors", "*.json"],
cache_dir=tmpdir,
)
safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
assert len(safetensors) > 0
parentdir = [
os.path.dirname(safetensor) for safetensor in safetensors
][0]
parentdir = [os.path.dirname(safetensor) for safetensor in safetensors][0]
files = list_safetensors(parentdir)
assert len(safetensors) == len(files)
@@ -50,9 +51,9 @@ def test_runai_pull_files_gcs(monkeypatch):
# | cut -d":" -f2 | base64 -d | xxd -p
expected_checksum = "f60dea775da1392434275b311b31a431"
hasher = hashlib.new("md5")
with open(os.path.join(model.dir, filename), 'rb') as f:
with open(os.path.join(model.dir, filename), "rb") as f:
# Read the file in chunks to handle large files efficiently
for chunk in iter(lambda: f.read(4096), b''):
for chunk in iter(lambda: f.read(4096), b""):
hasher.update(chunk)
actual_checksum = hasher.hexdigest()
assert actual_checksum == expected_checksum

15
tests/model_executor/model_loader/runai_model_streamer/test_weight_utils.py
View File

@@ -8,24 +8,25 @@ import huggingface_hub.constants
import torch
from vllm.model_executor.model_loader.weight_utils import (
download_weights_from_hf, runai_safetensors_weights_iterator,
safetensors_weights_iterator)
download_weights_from_hf,
runai_safetensors_weights_iterator,
safetensors_weights_iterator,
)
def test_runai_model_loader():
with tempfile.TemporaryDirectory() as tmpdir:
huggingface_hub.constants.HF_HUB_OFFLINE = False
download_weights_from_hf("openai-community/gpt2",
allow_patterns=["*.safetensors"],
cache_dir=tmpdir)
download_weights_from_hf(
"openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
)
safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
assert len(safetensors) > 0
runai_model_streamer_tensors = {}
hf_safetensors_tensors = {}
for name, tensor in runai_safetensors_weights_iterator(
safetensors, True):
for name, tensor in runai_safetensors_weights_iterator(safetensors, True):
runai_model_streamer_tensors[name] = tensor
for name, tensor in safetensors_weights_iterator(safetensors, True):

21
tests/model_executor/model_loader/tensorizer_loader/conftest.py
View File

@@ -32,7 +32,6 @@ def cleanup():
@pytest.fixture()
def just_serialize_model_tensors(model_ref, monkeypatch, tmp_path):
def noop(*args, **kwargs):
return None
@@ -56,8 +55,7 @@ def model_path(model_ref, tmp_path):
yield tmp_path / model_ref / "model.tensors"
def assert_from_collective_rpc(engine: LLM, closure: Callable,
closure_kwargs: dict):
def assert_from_collective_rpc(engine: LLM, closure: Callable, closure_kwargs: dict):
res = engine.collective_rpc(method=closure, kwargs=closure_kwargs)
return all(res)
@@ -67,18 +65,13 @@ def assert_from_collective_rpc(engine: LLM, closure: Callable,
# method. It's purely used as a dummy utility to run methods that test
# Tensorizer functionality
class DummyExecutor(UniProcExecutor):
def _init_executor(self) -> None:
"""Initialize the worker and load the model.
"""
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
rpc_rank=0)
distributed_init_method = get_distributed_init_method(
get_ip(), get_open_port())
"""Initialize the worker and load the model."""
self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config, rpc_rank=0)
distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())
local_rank = 0
# set local rank as the device index if specified
device_info = self.vllm_config.device_config.device.__str__().split(
":")
device_info = self.vllm_config.device_config.device.__str__().split(":")
if len(device_info) > 1:
local_rank = int(device_info[1])
rank = 0
@@ -91,7 +84,7 @@ class DummyExecutor(UniProcExecutor):
is_driver_worker=is_driver_worker,
)
self.mm_receiver_cache = None
self.collective_rpc("init_worker", args=([kwargs], ))
self.collective_rpc("init_worker", args=([kwargs],))
self.collective_rpc("init_device")
@property
@@ -99,5 +92,5 @@ class DummyExecutor(UniProcExecutor):
return 2
def shutdown(self):
if hasattr(self, 'thread_pool'):
if hasattr(self, "thread_pool"):
self.thread_pool.shutdown(wait=False)

244
tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py
View File

@@ -17,14 +17,19 @@ import vllm.model_executor.model_loader.tensorizer
from tests.utils import VLLM_PATH, RemoteOpenAIServer
from vllm import LLM, SamplingParams
from vllm.engine.arg_utils import EngineArgs
# yapf: disable
from vllm.model_executor.model_loader.tensorizer import (TensorizerConfig,
TensorSerializer,
is_vllm_tensorized,
open_stream,
tensorize_vllm_model)
from vllm.model_executor.model_loader.tensorizer import (
TensorizerConfig,
TensorSerializer,
is_vllm_tensorized,
open_stream,
tensorize_vllm_model,
)
from vllm.model_executor.model_loader.tensorizer_loader import (
BLACKLISTED_TENSORIZER_ARGS)
BLACKLISTED_TENSORIZER_ARGS,
)
# yapf: enable
from vllm.utils import PlaceholderModule
@@ -44,7 +49,7 @@ class TensorizerCaughtError(Exception):
EXAMPLES_PATH = VLLM_PATH / "examples"
pytest_plugins = "pytest_asyncio",
pytest_plugins = ("pytest_asyncio",)
prompts = [
"Hello, my name is",
@@ -56,8 +61,7 @@ prompts = [
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
def patch_init_and_catch_error(self, obj, method_name,
expected_error: type[Exception]):
def patch_init_and_catch_error(self, obj, method_name, expected_error: type[Exception]):
original = getattr(obj, method_name, None)
if original is None:
raise ValueError("Method '{}' not found.".format(method_name))
@@ -80,17 +84,19 @@ def assert_specific_tensorizer_error_is_raised(
expected_error: type[Exception],
):
with pytest.raises(TensorizerCaughtError):
executor.collective_rpc(patch_init_and_catch_error,
args=(
obj,
method_name,
expected_error,
))
executor.collective_rpc(
patch_init_and_catch_error,
args=(
obj,
method_name,
expected_error,
),
)
def is_curl_installed():
try:
subprocess.check_call(['curl', '--version'])
subprocess.check_call(["curl", "--version"])
return True
except (subprocess.CalledProcessError, FileNotFoundError):
return False
@@ -99,13 +105,14 @@ def is_curl_installed():
def write_keyfile(keyfile_path: str):
encryption_params = EncryptionParams.random()
pathlib.Path(keyfile_path).parent.mkdir(parents=True, exist_ok=True)
with open(keyfile_path, 'wb') as f:
with open(keyfile_path, "wb") as f:
f.write(encryption_params.key)
@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
def test_deserialized_encrypted_vllm_model_has_same_outputs(
model_ref, vllm_runner, tmp_path, model_path):
model_ref, vllm_runner, tmp_path, model_path
):
args = EngineArgs(model=model_ref)
with vllm_runner(model_ref) as vllm_model:
key_path = tmp_path / model_ref / "model.key"
@@ -113,29 +120,30 @@ def test_deserialized_encrypted_vllm_model_has_same_outputs(
outputs = vllm_model.generate(prompts, sampling_params)
config_for_serializing = TensorizerConfig(tensorizer_uri=str(model_path),
encryption_keyfile=str(key_path))
config_for_serializing = TensorizerConfig(
tensorizer_uri=str(model_path), encryption_keyfile=str(key_path)
)
tensorize_vllm_model(args, config_for_serializing)
config_for_deserializing = TensorizerConfig(
tensorizer_uri=str(model_path), encryption_keyfile=str(key_path))
tensorizer_uri=str(model_path), encryption_keyfile=str(key_path)
)
with vllm_runner(model_ref,
load_format="tensorizer",
model_loader_extra_config=config_for_deserializing
) as loaded_vllm_model: # noqa: E501
deserialized_outputs = loaded_vllm_model.generate(
prompts, sampling_params)
with vllm_runner(
model_ref,
load_format="tensorizer",
model_loader_extra_config=config_for_deserializing,
) as loaded_vllm_model: # noqa: E501
deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
# noqa: E501
assert outputs == deserialized_outputs
def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
tmp_path, model_ref,
model_path):
def test_deserialized_hf_model_has_same_outputs(
hf_runner, vllm_runner, tmp_path, model_ref, model_path
):
with hf_runner(model_ref) as hf_model:
max_tokens = 50
outputs = hf_model.generate_greedy(prompts, max_tokens=max_tokens)
@@ -143,14 +151,17 @@ def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
serializer = TensorSerializer(stream)
serializer.write_module(hf_model.model)
with vllm_runner(model_ref,
load_format="tensorizer",
model_loader_extra_config=TensorizerConfig(
tensorizer_uri=str(model_path),
num_readers=1,
)) as loaded_hf_model:
with vllm_runner(
model_ref,
load_format="tensorizer",
model_loader_extra_config=TensorizerConfig(
tensorizer_uri=str(model_path),
num_readers=1,
),
) as loaded_hf_model:
deserialized_outputs = loaded_hf_model.generate_greedy(
prompts, max_tokens=max_tokens)
prompts, max_tokens=max_tokens
)
assert outputs == deserialized_outputs
@@ -159,35 +170,37 @@ def test_load_without_tensorizer_load_format(vllm_runner, capfd, model_ref):
model = None
try:
model = vllm_runner(
model_ref,
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
model_ref, model_loader_extra_config=TensorizerConfig(tensorizer_uri="test")
)
pytest.fail("Expected RuntimeError for extra config keys")
except RuntimeError:
out, err = capfd.readouterr()
combined_output = out + err
assert ("ValueError: Unexpected extra config keys for load "
"format auto") in combined_output
assert (
"ValueError: Unexpected extra config keys for load format auto"
) in combined_output
finally:
del model
gc.collect()
torch.cuda.empty_cache()
def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd,
model_ref):
def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd, model_ref):
model = None
try:
model = vllm_runner(
model_ref,
load_format="safetensors",
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"),
)
pytest.fail("Expected RuntimeError for extra config keys")
except RuntimeError:
out, err = capfd.readouterr()
combined_output = out + err
assert ("ValueError: Unexpected extra config keys "
"for load format safetensors") in combined_output
assert (
"ValueError: Unexpected extra config keys for load format safetensors"
) in combined_output
finally:
del model
gc.collect()
@@ -214,21 +227,24 @@ def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
except RuntimeError:
out, err = capfd.readouterr()
combined_output = out + err
assert ("ValueError: For a sharded model, tensorizer_uri "
"should include a string format template like '%04d' "
"to be formatted with the rank "
"of the shard") in combined_output
assert (
"ValueError: For a sharded model, tensorizer_uri "
"should include a string format template like '%04d' "
"to be formatted with the rank "
"of the shard"
) in combined_output
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
vllm_runner, tmp_path):
vllm_runner, tmp_path
):
model_ref = "EleutherAI/pythia-1.4b"
# record outputs from un-sharded un-tensorized model
with vllm_runner(
model_ref,
disable_custom_all_reduce=True,
enforce_eager=True,
model_ref,
disable_custom_all_reduce=True,
enforce_eager=True,
) as base_model:
outputs = base_model.generate(prompts, sampling_params)
@@ -254,21 +270,22 @@ def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
assert os.path.isfile(model_path % 1), "Serialization subprocess failed"
with vllm_runner(
model_ref,
tensor_parallel_size=2,
load_format="tensorizer",
disable_custom_all_reduce=True,
enforce_eager=True,
model_loader_extra_config=tensorizer_config) as loaded_vllm_model:
deserialized_outputs = loaded_vllm_model.generate(
prompts, sampling_params)
model_ref,
tensor_parallel_size=2,
load_format="tensorizer",
disable_custom_all_reduce=True,
enforce_eager=True,
model_loader_extra_config=tensorizer_config,
) as loaded_vllm_model:
deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
assert outputs == deserialized_outputs
@pytest.mark.flaky(reruns=3)
def test_vllm_tensorized_model_has_same_outputs(model_ref, vllm_runner,
tmp_path, model_path):
def test_vllm_tensorized_model_has_same_outputs(
model_ref, vllm_runner, tmp_path, model_path
):
gc.collect()
torch.cuda.empty_cache()
config = TensorizerConfig(tensorizer_uri=str(model_path))
@@ -280,11 +297,10 @@ def test_vllm_tensorized_model_has_same_outputs(model_ref, vllm_runner,
tensorize_vllm_model(args, config)
assert is_vllm_tensorized(config)
with vllm_runner(model_ref,
load_format="tensorizer",
model_loader_extra_config=config) as loaded_vllm_model:
deserialized_outputs = loaded_vllm_model.generate(
prompts, sampling_params)
with vllm_runner(
model_ref, load_format="tensorizer", model_loader_extra_config=config
) as loaded_vllm_model:
deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
# noqa: E501
assert outputs == deserialized_outputs
@@ -314,15 +330,17 @@ def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
serialization_params = {
"limit_cpu_concurrency": 2,
}
model_ref = "facebook/opt-125m"
model_path = tmp_path / (model_ref + ".tensors")
config = TensorizerConfig(tensorizer_uri=str(model_path),
serialization_kwargs=serialization_params)
llm = LLM(model=model_ref, )
config = TensorizerConfig(
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
)
llm = LLM(
model=model_ref,
)
def serialization_test(self, *args, **kwargs):
# This is performed in the ephemeral worker process, so monkey-patching
@@ -340,10 +358,13 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
return original(self, *args, **kwargs)
tensorizer.serialization.TensorSerializer.__init__ = (
tensorizer_serializer_wrapper)
tensorizer_serializer_wrapper
)
tensorizer_config = TensorizerConfig(**kwargs["tensorizer_config"])
self.save_tensorized_model(tensorizer_config=tensorizer_config, )
self.save_tensorized_model(
tensorizer_config=tensorizer_config,
)
return to_compare | original_dict == to_compare
kwargs = {"tensorizer_config": config.to_serializable()}
@@ -351,9 +372,7 @@ def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
assert assert_from_collective_rpc(llm, serialization_test, kwargs)
def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
tmp_path, capfd):
def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
deserialization_kwargs = {
"num_readers": "bar", # illegal value
}
@@ -364,8 +383,9 @@ def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
model_ref = "facebook/opt-125m"
model_path = tmp_path / (model_ref + ".tensors")
config = TensorizerConfig(tensorizer_uri=str(model_path),
serialization_kwargs=serialization_params)
config = TensorizerConfig(
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
)
args = EngineArgs(model=model_ref)
tensorize_vllm_model(args, config)
@@ -393,7 +413,6 @@ def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
def test_assert_stream_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
deserialization_kwargs = {
"num_readers": 1,
}
@@ -404,8 +423,9 @@ def test_assert_stream_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
model_ref = "facebook/opt-125m"
model_path = tmp_path / (model_ref + ".tensors")
config = TensorizerConfig(tensorizer_uri=str(model_path),
serialization_kwargs=serialization_params)
config = TensorizerConfig(
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
)
args = EngineArgs(model=model_ref)
tensorize_vllm_model(args, config)
@@ -441,16 +461,24 @@ async def test_serialize_and_serve_entrypoints(tmp_path):
suffix = "test"
try:
result = subprocess.run([
sys.executable,
f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py", "--model",
model_ref, "serialize", "--serialized-directory",
str(tmp_path), "--suffix", suffix, "--serialization-kwargs",
'{"limit_cpu_concurrency": 4}'
],
check=True,
capture_output=True,
text=True)
result = subprocess.run(
[
sys.executable,
f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py",
"--model",
model_ref,
"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)
@@ -470,14 +498,20 @@ async def test_serialize_and_serve_entrypoints(tmp_path):
"deserialization_kwargs": {
"verify_hash": True,
"num_readers": 8,
}
},
}
cmd = [
"-m", "vllm.entrypoints.cli.main", "serve", "--host", "localhost",
"--load-format", "tensorizer", model_ref,
"-m",
"vllm.entrypoints.cli.main",
"serve",
"--host",
"localhost",
"--load-format",
"tensorizer",
model_ref,
"--model-loader-extra-config",
json.dumps(model_loader_extra_config, indent=2)
json.dumps(model_loader_extra_config, indent=2),
]
proc = await asyncio.create_subprocess_exec(
@@ -500,17 +534,16 @@ async def test_serialize_and_serve_entrypoints(tmp_path):
@pytest.mark.parametrize("illegal_value", BLACKLISTED_TENSORIZER_ARGS)
def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd,
illegal_value):
def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd, illegal_value):
serialization_params = {
"limit_cpu_concurrency": 2,
}
model_ref = "facebook/opt-125m"
model_path = tmp_path / (model_ref + ".tensors")
config = TensorizerConfig(tensorizer_uri=str(model_path),
serialization_kwargs=serialization_params)
config = TensorizerConfig(
tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
)
args = EngineArgs(model=model_ref)
tensorize_vllm_model(args, config)
@@ -526,5 +559,6 @@ def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd,
except RuntimeError:
out, err = capfd.readouterr()
combined_output = out + err
assert (f"ValueError: {illegal_value} is not an allowed "
f"Tensorizer argument.") in combined_output
assert (
f"ValueError: {illegal_value} is not an allowed Tensorizer argument."
) in combined_output

7
tests/model_executor/model_loader/test_registry.py
View File

@@ -6,22 +6,19 @@ from torch import nn
from vllm.config import ModelConfig
from vllm.config.load import LoadConfig
from vllm.model_executor.model_loader import (get_model_loader,
register_model_loader)
from vllm.model_executor.model_loader import get_model_loader, register_model_loader
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
@register_model_loader("custom_load_format")
class CustomModelLoader(BaseModelLoader):
def __init__(self, load_config: LoadConfig) -> None:
super().__init__(load_config)
def download_model(self, model_config: ModelConfig) -> None:
pass
def load_weights(self, model: nn.Module,
model_config: ModelConfig) -> None:
def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
pass

82
tests/model_executor/model_loader/test_sharded_state_loader.py
View File

@@ -35,11 +35,13 @@ def test_filter_subtensors():
"b": torch.empty((2, 4)),
"c": torch.empty((2, 4, 8)),
}
state_dict.update({
"x": state_dict["b"],
"y": state_dict["c"][1, 2, :],
"z": state_dict["c"][1, :, 4],
})
state_dict.update(
{
"x": state_dict["b"],
"y": state_dict["c"][1, 2, :],
"z": state_dict["c"][1, :, 4],
}
)
filtered_state_dict = ShardedStateLoader._filter_subtensors(state_dict)
assert tuple(filtered_state_dict.keys()) == ("a", "b", "c")
for key, tensor in filtered_state_dict.items():
@@ -49,8 +51,9 @@ def test_filter_subtensors():
@pytest.fixture(scope="module")
def llama_3p2_1b_files():
input_dir = snapshot_download("meta-llama/Llama-3.2-1B-Instruct",
ignore_patterns=["*.bin*", "original/*"])
input_dir = snapshot_download(
"meta-llama/Llama-3.2-1B-Instruct", ignore_patterns=["*.bin*", "original/*"]
)
yield input_dir
@@ -63,8 +66,7 @@ def _run_writer(input_dir, output_dir, weights_patterns, **kwargs):
if is_v1_engine:
# For V1 engine, we need to use engine_core.save_sharded_state
print("Using V1 engine save path")
llm_sharded_writer.llm_engine.engine_core.save_sharded_state(
path=output_dir)
llm_sharded_writer.llm_engine.engine_core.save_sharded_state(path=output_dir)
else:
# For V0 engine
print("Using V0 engine save path")
@@ -74,8 +76,9 @@ def _run_writer(input_dir, output_dir, weights_patterns, **kwargs):
# Copy metadata files to output directory
for file in os.listdir(input_dir):
if os.path.isdir(os.path.join(input_dir, file)):
shutil.copytree(os.path.join(input_dir, file),
os.path.join(output_dir, file))
shutil.copytree(
os.path.join(input_dir, file), os.path.join(output_dir, file)
)
elif not any(fnmatch.fnmatch(file, ext) for ext in weights_patterns):
shutil.copy(os.path.join(input_dir, file), output_dir)
@@ -90,37 +93,42 @@ def _run_generate(input_dir, queue: mp.Queue, **kwargs):
@pytest.mark.parametrize("enable_lora", [False, True])
@pytest.mark.parametrize("tp_size", [1, 2])
def test_sharded_state_loader(enable_lora, tp_size, num_gpus_available,
llama_3p2_1b_files):
def test_sharded_state_loader(
enable_lora, tp_size, num_gpus_available, llama_3p2_1b_files
):
if num_gpus_available < tp_size:
pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")
weights_patterns = ("*.safetensors", )
weights_patterns = ("*.safetensors",)
gpu_memory_utilization = 0.8
input_dir = llama_3p2_1b_files
ctx = mp.get_context("spawn")
# Run in separate processes for memory & CUDA isolation
with TemporaryDirectory() as output_dir:
p = ctx.Process(target=_run_writer,
args=(input_dir, output_dir, weights_patterns),
kwargs=dict(
tensor_parallel_size=tp_size,
gpu_memory_utilization=gpu_memory_utilization,
enforce_eager=True,
))
p = ctx.Process(
target=_run_writer,
args=(input_dir, output_dir, weights_patterns),
kwargs=dict(
tensor_parallel_size=tp_size,
gpu_memory_utilization=gpu_memory_utilization,
enforce_eager=True,
),
)
p.start()
p.join()
queue = ctx.Queue()
p = ctx.Process(target=_run_generate,
args=(input_dir, queue),
kwargs=dict(
enable_lora=enable_lora,
gpu_memory_utilization=gpu_memory_utilization,
tensor_parallel_size=tp_size,
))
p = ctx.Process(
target=_run_generate,
args=(input_dir, queue),
kwargs=dict(
enable_lora=enable_lora,
gpu_memory_utilization=gpu_memory_utilization,
tensor_parallel_size=tp_size,
),
)
p.start()
# Call queue.get() before p.join() to prevent deadlock:
# If p.join() is called before queue.get() and the queue is full,
@@ -134,14 +142,16 @@ def test_sharded_state_loader(enable_lora, tp_size, num_gpus_available,
queue = ctx.Queue()
p = ctx.Process(target=_run_generate,
args=(output_dir, queue),
kwargs=dict(
enable_lora=enable_lora,
gpu_memory_utilization=gpu_memory_utilization,
tensor_parallel_size=tp_size,
load_format="sharded_state",
))
p = ctx.Process(
target=_run_generate,
args=(output_dir, queue),
kwargs=dict(
enable_lora=enable_lora,
gpu_memory_utilization=gpu_memory_utilization,
tensor_parallel_size=tp_size,
load_format="sharded_state",
),
)
p.start()
# Call queue.get() before p.join() to prevent deadlock:
# If p.join() is called before queue.get() and the queue is full,