[CI/Build] Consolidate model loader tests and requirements (#25765)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from vllm import SamplingParams
|
||||
|
||||
test_model = "openai-community/gpt2"
|
||||
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
# Create a sampling params object.
|
||||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
|
||||
|
||||
|
||||
def test_model_loader_download_files(vllm_runner):
|
||||
with vllm_runner(test_model, load_format="fastsafetensors") as llm:
|
||||
deserialized_outputs = llm.generate(prompts, sampling_params)
|
||||
assert deserialized_outputs
|
||||
@@ -0,0 +1,47 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import glob
|
||||
import tempfile
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
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):
|
||||
fastsafetensors_tensors[name] = tensor
|
||||
|
||||
for name, tensor in safetensors_weights_iterator(safetensors, True):
|
||||
hf_safetensors_tensors[name] = tensor
|
||||
|
||||
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]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_fastsafetensors_model_loader()
|
||||
Reference in New Issue
Block a user