@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 165 lines 5.28 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import fnmatch
5 import multiprocessing as mp
6 import os
7 import shutil
8 from tempfile import TemporaryDirectory
9
10 import pytest
11 import torch
12
13 from vllm import LLM, SamplingParams
14 from vllm.model_executor.model_loader import ShardedStateLoader
15 from vllm.transformers_utils.repo_utils import hf_api
16
17 prompts = [
18 "Hello, my name is",
19 "The president of the United States is",
20 "The capital of France is",
21 "The future of AI is",
22 ]
23
24 # Create a sampling params object.
25 sampling_params = SamplingParams(
26 temperature=0,
27 max_tokens=256,
28 ignore_eos=True,
29 )
30
31
32 def test_filter_subtensors():
33 state_dict = {
34 "a": torch.empty(2),
35 "b": torch.empty((2, 4)),
36 "c": torch.empty((2, 4, 8)),
37 }
38 state_dict.update(
39 {
40 "x": state_dict["b"],
41 "y": state_dict["c"][1, 2, :],
42 "z": state_dict["c"][1, :, 4],
43 }
44 )
45 filtered_state_dict = ShardedStateLoader._filter_subtensors(state_dict)
46 assert tuple(filtered_state_dict.keys()) == ("a", "b", "c")
47 for key, tensor in filtered_state_dict.items():
48 # NOTE: don't use `equal` here, as the tensor might contain NaNs
49 assert tensor is state_dict[key]
50
51
52 @pytest.fixture(scope="module")
53 def llama_3p2_1b_files():
54 input_dir = hf_api().snapshot_download(
55 "meta-llama/Llama-3.2-1B-Instruct", ignore_patterns=["*.bin*", "original/*"]
56 )
57
58 yield input_dir
59
60
61 def _run_writer(input_dir, output_dir, weights_patterns, **kwargs):
62 llm_sharded_writer = LLM(model=input_dir, **kwargs)
63
64 # Dump worker states to output directory
65 llm_sharded_writer.llm_engine.engine_core.save_sharded_state(path=output_dir)
66
67 # Copy metadata files to output directory
68 for file in os.listdir(input_dir):
69 if os.path.isdir(os.path.join(input_dir, file)):
70 shutil.copytree(
71 os.path.join(input_dir, file), os.path.join(output_dir, file)
72 )
73 elif not any(fnmatch.fnmatch(file, ext) for ext in weights_patterns):
74 shutil.copy(os.path.join(input_dir, file), output_dir)
75
76
77 def _run_generate(input_dir, queue: mp.Queue, **kwargs):
78 llm = LLM(model=input_dir, **kwargs)
79 gen = llm.generate(prompts, sampling_params)
80 queue.put([g.outputs[0].__dict__ for g in gen])
81 queue.close()
82 queue.join_thread()
83
84
85 @pytest.mark.parametrize("enable_lora", [False, True])
86 @pytest.mark.parametrize("tp_size", [1, 2])
87 def test_sharded_state_loader(
88 enable_lora, tp_size, num_gpus_available, llama_3p2_1b_files
89 ):
90 if num_gpus_available < tp_size:
91 pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")
92
93 weights_patterns = ("*.safetensors",)
94 gpu_memory_utilization = 0.8
95 input_dir = llama_3p2_1b_files
96 ctx = mp.get_context("spawn")
97
98 # Keep batching deterministic: this test compares exact greedy outputs
99 # across separate engine processes, whose schedulers may otherwise form
100 # different batches depending on prompt-processing timing.
101 platform_args = {"max_num_seqs": 1}
102
103 # Run in separate processes for memory & CUDA isolation
104 with TemporaryDirectory() as output_dir:
105 p = ctx.Process(
106 target=_run_writer,
107 args=(input_dir, output_dir, weights_patterns),
108 kwargs=dict(
109 tensor_parallel_size=tp_size,
110 gpu_memory_utilization=gpu_memory_utilization,
111 enforce_eager=True,
112 **platform_args,
113 ),
114 )
115 p.start()
116 p.join()
117
118 queue = ctx.Queue()
119
120 p = ctx.Process(
121 target=_run_generate,
122 args=(input_dir, queue),
123 kwargs=dict(
124 enable_lora=enable_lora,
125 gpu_memory_utilization=gpu_memory_utilization,
126 tensor_parallel_size=tp_size,
127 **platform_args,
128 ),
129 )
130 p.start()
131 # Call queue.get() before p.join() to prevent deadlock:
132 # If p.join() is called before queue.get() and the queue is full,
133 # the child process may block while writing to the queue and never
134 # terminate, causing the parent to wait indefinitely on p.join().
135 # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814
136 out_before = queue.get()
137 p.join()
138 queue.close()
139 queue.join_thread()
140
141 queue = ctx.Queue()
142
143 p = ctx.Process(
144 target=_run_generate,
145 args=(output_dir, queue),
146 kwargs=dict(
147 enable_lora=enable_lora,
148 gpu_memory_utilization=gpu_memory_utilization,
149 tensor_parallel_size=tp_size,
150 load_format="sharded_state",
151 **platform_args,
152 ),
153 )
154 p.start()
155 # Call queue.get() before p.join() to prevent deadlock:
156 # If p.join() is called before queue.get() and the queue is full,
157 # the child process may block while writing to the queue and never
158 # terminate, causing the parent to wait indefinitely on p.join().
159 # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814
160 out_after = queue.get()
161 p.join()
162 queue.close()
163 queue.join_thread()
164
165 assert out_before == out_after