|
| 1 | +# This file contains end-to-end tests for the RunAI Model Streamer loader. |
| 2 | +# |
| 3 | +# The RunAI Model Streamer is a high-performance model loader that serves as an |
| 4 | +# alternative to the default Hugging Face loader. Instead of downloading a model |
| 5 | +# to local disk, it streams the weights from object storage (like GCS) into |
| 6 | +# GPU memory. This streaming process is significantly faster than the |
| 7 | +# traditional disk-based loading method. |
| 8 | + |
| 9 | +# The tests in this file verify that loading model weights using the |
| 10 | +# streamer produces the same results as loading the same model using the |
| 11 | +# standard Hugging Face loader. This ensures the correctness of the streamer |
| 12 | +# integration. |
| 13 | + |
| 14 | +# The tests are performed by: |
| 15 | +# 1. Loading a model from Google Cloud Storage using the `runai_streamer` format. |
| 16 | +# 2. Generating output with this model. |
| 17 | +# 3. Loading the same model from Hugging Face using the default loader. |
| 18 | +# 4. Generating output with this second model. |
| 19 | +# 5. Asserting that the outputs from both models are identical. |
| 20 | + |
| 21 | +from __future__ import annotations |
| 22 | + |
| 23 | +import time |
| 24 | + |
| 25 | +import pytest |
| 26 | +from vllm import LLM, SamplingParams |
| 27 | + |
| 28 | + |
| 29 | +@pytest.fixture |
| 30 | +def sampling_config(): |
| 31 | + return SamplingParams(temperature=0, max_tokens=10, ignore_eos=True) |
| 32 | + |
| 33 | + |
| 34 | +@pytest.fixture |
| 35 | +# TODO(amacaskill): Replace with GKE owned GCS bucket. |
| 36 | +def gcs_model_name(): |
| 37 | + return "gs://vertex-model-garden-public-us/llama3/llama3-8b-hf" |
| 38 | + |
| 39 | + |
| 40 | +@pytest.fixture |
| 41 | +def hf_model_name(): |
| 42 | + return "meta-llama/Meta-Llama-3-8B" |
| 43 | + |
| 44 | + |
| 45 | +@pytest.fixture |
| 46 | +def prompt(): |
| 47 | + return "Hello, my name is" |
| 48 | + |
| 49 | + |
| 50 | +def test_correctness( |
| 51 | + sampling_config: SamplingParams, |
| 52 | + gcs_model_name: str, |
| 53 | + hf_model_name: str, |
| 54 | + prompt: str, |
| 55 | + monkeypatch: pytest.MonkeyPatch, |
| 56 | +): |
| 57 | + ''' |
| 58 | + Compare the outputs of a model loaded from GCS via runai_model_streamer |
| 59 | + and a model loaded from Hugging Face. The outputs should be the same. |
| 60 | + These tests attempt to use tensor_parallel_size=1. The model is 16GB, |
| 61 | + # and v6e has 32GB of HBM, so it will fit. |
| 62 | + ''' |
| 63 | + # Set ENV variables so that runai_model_streamer uses anonymous GCS access. |
| 64 | + monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "fake-project") |
| 65 | + monkeypatch.setenv("RUNAI_STREAMER_GCS_USE_ANONYMOUS_CREDENTIALS", "true") |
| 66 | + monkeypatch.setenv("CLOUD_STORAGE_EMULATOR_ENDPOINT", |
| 67 | + "https://storage.googleapis.com") |
| 68 | + gcs_llm = LLM(model=gcs_model_name, |
| 69 | + load_format="runai_streamer", |
| 70 | + max_model_len=128, |
| 71 | + max_num_seqs=16, |
| 72 | + max_num_batched_tokens=256) |
| 73 | + gcs_outputs = gcs_llm.generate([prompt], sampling_config) |
| 74 | + gcs_output_text = gcs_outputs[0].outputs[0].text |
| 75 | + del gcs_llm |
| 76 | + time.sleep(10) # Wait for TPUs to be released |
| 77 | + |
| 78 | + # Test with Hugging Face model |
| 79 | + hf_llm = LLM(model=hf_model_name, |
| 80 | + max_model_len=128, |
| 81 | + max_num_seqs=16, |
| 82 | + max_num_batched_tokens=256) |
| 83 | + hf_outputs = hf_llm.generate([prompt], sampling_config) |
| 84 | + hf_output_text = hf_outputs[0].outputs[0].text |
| 85 | + del hf_llm |
| 86 | + time.sleep(10) # Wait for TPUs to be released |
| 87 | + |
| 88 | + assert gcs_output_text == hf_output_text, ( |
| 89 | + f"Outputs do not match! " |
| 90 | + f"GCS output: {gcs_output_text}, HF output: {hf_output_text}") |
0 commit comments