Setting the file. One moment. scripts/inspect_vllm_uv.py
scripts/inspect_vllm_uv.py
Python·306 lines·9 KB
17
18Usage (standalone):
19 uv run scripts/inspect_vllm_uv.py --model "meta-llama/Llama-3.2-1B" --task "mmlu"
20
21Model backends:
22 - vllm: Fast inference with vLLM (recommended for large models)
23 - hf: HuggingFace Transformers backend (broader model compatibility)
24"""
25
26from __future__ import annotations
27
28import argparse
29import os
30import subprocess
31import sys
32from typing import Optional
33
34
35def setup_environment() -> None:
36 """Configure environment variables for HuggingFace authentication."""
37 hf_token = os.getenv("HF_TOKEN")
38 if hf_token:
39 os.environ.setdefault("HUGGING_FACE_HUB_TOKEN", hf_token)
40 os.environ.setdefault("HF_HUB_TOKEN", hf_token)
41
42
43def run_inspect_vllm(
44 model_id: str,
45 task: str,
46 limit: Optional[int] = None,
47 max_connections: int = 4,
48 temperature: float = 0.0,
49 tensor_parallel_size: int = 1,
50 gpu_memory_utilization: float = 0.8,
51 dtype: str = "auto",
52 trust_remote_code: bool = False,
53 log_level: str = "info",
54) -> None:
55 """
56 Run inspect-ai evaluation with vLLM backend.
57
58 Args:
59 model_id: HuggingFace model ID
60 task: inspect-ai task to execute (e.g., "mmlu", "gsm8k")
61 limit: Limit number of samples to evaluate
62 max_connections: Maximum concurrent connections
63 temperature: Sampling temperature
64 tensor_parallel_size: Number of GPUs for tensor parallelism
65 gpu_memory_utilization: GPU memory fraction
66 dtype: Data type (auto, float16, bfloat16)
67 trust_remote_code: Allow remote code execution
68 log_level: Logging level
69 """
70 setup_environment()
71
72 model_spec = f"vllm/{model_id}"
73 cmd = [
74 "inspect",
75 "eval",
76 task,
77 "--model",
78 model_spec,
79 "--log-level",
80 log_level,
81 "--max-connections",
82 str(max_connections),
83 ]
84
85 # vLLM supports temperature=0 unlike HF inference providers
86 cmd.extend(["--temperature", str(temperature)])
87
88 # Older inspect-ai CLI versions do not support --model-args; rely on defaults
89 # and let vLLM choose sensible settings for small models.
90 if tensor_parallel_size != 1:
91 cmd.extend(["--tensor-parallel-size", str(tensor_parallel_size)])
92 if gpu_memory_utilization != 0.8:
93 cmd.extend(["--gpu-memory-utilization", str(gpu_memory_utilization)])
94 if dtype != "auto":
95 cmd.extend(["--dtype", dtype])
96 if trust_remote_code:
97 cmd.append("--trust-remote-code")
98
99 if limit:
100 cmd.extend(["--limit", str(limit)])
101
102 print(f"Running: {' '.join(cmd)}")
103
104 try:
105 subprocess.run(cmd, check=True)
106 print("Evaluation complete.")
107 except subprocess.CalledProcessError as exc:
108 print(f"Evaluation failed with exit code {exc.returncode}", file=sys.stderr)
109 sys.exit(exc.returncode)
110
111
112def run_inspect_hf(
113 model_id: str,
114 task: str,
115 limit: Optional[int] = None,
116 max_connections: int = 1,
117 temperature: float = 0.001,
118 device: str = "auto",
119 dtype: str = "auto",
120 trust_remote_code: bool = False,
121 log_level: str = "info",
122) -> None:
123 """
124 Run inspect-ai evaluation with HuggingFace Transformers backend.
125
126 Use this when vLLM doesn't support the model architecture.
127
128 Args:
129 model_id: HuggingFace model ID
130 task: inspect-ai task to execute
131 limit: Limit number of samples
132 max_connections: Maximum concurrent connections (keep low for memory)
133 temperature: Sampling temperature
134 device: Device to use (auto, cuda, cpu)
135 dtype: Data type
136 trust_remote_code: Allow remote code execution
137 log_level: Logging level
138 """
139 setup_environment()
140
141 model_spec = f"hf/{model_id}"
142
143 cmd = [
144 "inspect",
145 "eval",
146 task,
147 "--model",
148 model_spec,
149 "--log-level",
150 log_level,
151 "--max-connections",
152 str(max_connections),
153 "--temperature",
154 str(temperature),
155 ]
156
157 if device != "auto":
158 cmd.extend(["--device", device])
159 if dtype != "auto":
160 cmd.extend(["--dtype", dtype])
161 if trust_remote_code:
162 cmd.append("--trust-remote-code")
163
164 if limit:
165 cmd.extend(["--limit", str(limit)])
166
167 print(f"Running: {' '.join(cmd)}")
168
169 try:
170 subprocess.run(cmd, check=True)
171 print("Evaluation complete.")
172 except subprocess.CalledProcessError as exc:
173 print(f"Evaluation failed with exit code {exc.returncode}", file=sys.stderr)
174 sys.exit(exc.returncode)
175
176
177def main() -> None:
178 parser = argparse.ArgumentParser(
179 description="Run inspect-ai evaluations with vLLM or HuggingFace Transformers on custom models",
180 formatter_class=argparse.RawDescriptionHelpFormatter,
181 epilog="""
182Examples:
183 # Run MMLU with vLLM backend
184 uv run scripts/inspect_vllm_uv.py --model meta-llama/Llama-3.2-1B --task mmlu
185
186 # Run with HuggingFace Transformers backend
187 uv run scripts/inspect_vllm_uv.py --model meta-llama/Llama-3.2-1B --task mmlu --backend hf
188
189 # Run with limited samples for testing
190 uv run scripts/inspect_vllm_uv.py --model meta-llama/Llama-3.2-1B --task mmlu --limit 10
191
192 # Run on multiple GPUs with tensor parallelism
193 uv run scripts/inspect_vllm_uv.py --model meta-llama/Llama-3.2-70B --task mmlu --tensor-parallel-size 4
194
195Available tasks (from inspect-evals):
196 - mmlu: Massive Multitask Language Understanding
197 - gsm8k: Grade School Math
198 - hellaswag: Common sense reasoning
199 - arc_challenge: AI2 Reasoning Challenge
200 - truthfulqa: TruthfulQA benchmark
201 - winogrande: Winograd Schema Challenge
202 - humaneval: Code generation (HumanEval)
203
204 """,
205 )
206
207 parser.add_argument(
208 "--model",
209 required=True,
210 help="HuggingFace model ID (e.g., meta-llama/Llama-3.2-1B)",
211 )
212 parser.add_argument(
213 "--task",
214 required=True,
215 help="inspect-ai task to execute (e.g., mmlu, gsm8k)",
216 )
217 parser.add_argument(
218 "--backend",
219 choices=["vllm", "hf"],
220 default="vllm",
221 help="Model backend (default: vllm)",
222 )
223 parser.add_argument(
224 "--limit",
225 type=int,
226 default=None,
227 help="Limit number of samples to evaluate",
228 )
229 parser.add_argument(
230 "--max-connections",
231 type=int,
232 default=None,
233 help="Maximum concurrent connections (default: 4 for vllm, 1 for hf)",
234 )
235 parser.add_argument(
236 "--temperature",
237 type=float,
238 default=None,
239 help="Sampling temperature (default: 0.0 for vllm, 0.001 for hf)",
240 )
241 parser.add_argument(
242 "--tensor-parallel-size",
243 type=int,
244 default=1,
245 help="Number of GPUs for tensor parallelism (vLLM only, default: 1)",
246 )
247 parser.add_argument(
248 "--gpu-memory-utilization",
249 type=float,
250 default=0.8,
251 help="GPU memory fraction to use (vLLM only, default: 0.8)",
252 )
253 parser.add_argument(
254 "--dtype",
255 default="auto",
256 choices=["auto", "float16", "bfloat16", "float32"],
257 help="Data type for model weights (default: auto)",
258 )
259 parser.add_argument(
260 "--device",
261 default="auto",
262 help="Device for HF backend (auto, cuda, cpu)",
263 )
264 parser.add_argument(
265 "--trust-remote-code",
266 action="store_true",
267 help="Allow executing remote code from model repository",
268 )
269 parser.add_argument(
270 "--log-level",
271 default="info",
272 choices=["debug", "info", "warning", "error"],
273 help="Logging level (default: info)",
274 )
275
276 args = parser.parse_args()
277
278 if args.backend == "vllm":
279 run_inspect_vllm(
280 model_id=args.model,
281 task=args.task,
282 limit=args.limit,
283 max_connections=args.max_connections or 4,
284 temperature=args.temperature if args.temperature is not None else 0.0,
285 tensor_parallel_size=args.tensor_parallel_size,
286 gpu_memory_utilization=args.gpu_memory_utilization,
287 dtype=args.dtype,
288 trust_remote_code=args.trust_remote_code,
289 log_level=args.log_level,
290 )
291 else:
292 run_inspect_hf(
293 model_id=args.model,
294 task=args.task,
295 limit=args.limit,
296 max_connections=args.max_connections or 1,
297 temperature=args.temperature if args.temperature is not None else 0.001,
298 device=args.device,
299 dtype=args.dtype,
300 trust_remote_code=args.trust_remote_code,
301 log_level=args.log_level,
302 )
303
304
305if __name__ == "__main__":
306 main()