@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 157 lines 5.13 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """LM eval harness on model to compare vs HF baseline computed offline.
4 Configs are found in configs/$MODEL.yaml
5
6 pytest -s -v test_lm_eval_correctness.py \
7 --config-list-file=configs/models-small.txt \
8 --tp-size=1
9 """
10
11 import os
12 from contextlib import contextmanager
13
14 import lm_eval
15 import pytest
16 import yaml
17
18 from vllm.platforms import current_platform
19
20 DEFAULT_RTOL = 0.08
21
22
23 @contextmanager
24 def scoped_env_vars(new_env: dict[str, str]):
25 if not new_env:
26 # Fast path: nothing to do
27 yield
28 return
29
30 old_values = {}
31 new_keys = []
32
33 try:
34 for key, value in new_env.items():
35 if key in os.environ:
36 old_values[key] = os.environ[key]
37 else:
38 new_keys.append(key)
39 os.environ[key] = str(value)
40 yield
41 finally:
42 # Restore / clean up
43 for key, value in old_values.items():
44 os.environ[key] = value
45 for key in new_keys:
46 os.environ.pop(key, None)
47
48
49 def launch_lm_eval(eval_config, tp_size):
50 trust_remote_code = eval_config.get("trust_remote_code", False)
51 max_model_len = eval_config.get("max_model_len", 4096)
52 batch_size = eval_config.get("batch_size", "auto")
53 backend = eval_config.get("backend", "vllm")
54 enforce_eager = eval_config.get("enforce_eager", "true")
55 kv_cache_dtype = eval_config.get("kv_cache_dtype", "auto")
56 model_args = (
57 f"pretrained={eval_config['model_name']},"
58 f"tensor_parallel_size={tp_size},"
59 f"enforce_eager={enforce_eager},"
60 f"kv_cache_dtype={kv_cache_dtype},"
61 f"add_bos_token=true,"
62 f"trust_remote_code={trust_remote_code},"
63 f"max_model_len={max_model_len},"
64 "allow_deprecated_quantization=True,"
65 )
66
67 if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]:
68 model_args += "attention_backend=TRITON_ATTN"
69
70 moe_backend = eval_config.get("moe_backend", None)
71 if moe_backend is not None:
72 model_args += f"moe_backend={moe_backend},"
73
74 if current_platform.is_rocm():
75 rocm_load_strategy = eval_config.get("rocm_safetensors_load_strategy")
76 if rocm_load_strategy is not None:
77 model_args += f"safetensors_load_strategy={rocm_load_strategy},"
78
79 tokenizer_mode = eval_config.get("tokenizer_mode", None)
80 if tokenizer_mode is not None:
81 model_args += f"tokenizer_mode={tokenizer_mode},"
82
83 env_vars = eval_config.get("env_vars", None)
84 with scoped_env_vars(env_vars):
85 results = lm_eval.simple_evaluate(
86 model=backend,
87 model_args=model_args,
88 tasks=[task["name"] for task in eval_config["tasks"]],
89 num_fewshot=eval_config["num_fewshot"],
90 limit=eval_config["limit"],
91 # TODO(yeq): using chat template w/ fewshot_as_multiturn is supposed help
92 # text models. however, this is regressing measured strict-match for
93 # existing text models in CI, so only apply it for mm, or explicitly set
94 apply_chat_template=eval_config.get(
95 "apply_chat_template", backend == "vllm-vlm"
96 ),
97 fewshot_as_multiturn=eval_config.get("fewshot_as_multiturn", False),
98 # Forward decoding and early-stop controls (e.g., max_gen_toks, until=...)
99 gen_kwargs=eval_config.get("gen_kwargs"),
100 batch_size=batch_size,
101 )
102 return results
103
104
105 def _check_rocm_gpu_arch_requirement(eval_config):
106 """Skip the test if the model requires a ROCm GPU arch not present.
107
108 Model YAML configs can specify::
109
110 required_gpu_arch:
111 - gfx942
112 - gfx950
113
114 The check only applies on ROCm. On other platforms (e.g. CUDA) the
115 field is ignored so that shared config files work for both NVIDIA and
116 AMD CI pipelines.
117 """
118 required_archs = eval_config.get("required_gpu_arch")
119 if not required_archs:
120 return
121
122 if not current_platform.is_rocm():
123 return
124
125 from vllm.platforms.rocm import _GCN_ARCH # noqa: E402
126
127 if not any(arch in _GCN_ARCH for arch in required_archs):
128 pytest.skip(
129 f"Model requires GPU arch {required_archs}, "
130 f"but detected arch is '{_GCN_ARCH}'"
131 )
132
133
134 def test_lm_eval_correctness_param(config_filename, tp_size):
135 eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
136
137 _check_rocm_gpu_arch_requirement(eval_config)
138
139 results = launch_lm_eval(eval_config, tp_size)
140
141 rtol = eval_config.get("rtol", DEFAULT_RTOL)
142
143 success = True
144 for task in eval_config["tasks"]:
145 for metric in task["metrics"]:
146 ground_truth = metric["value"]
147 measured_value = results["results"][task["name"]][metric["name"]]
148 print(
149 f"{task['name']} | {metric['name']}: "
150 f"ground_truth={ground_truth:.3f} | "
151 f"measured={measured_value:.3f} | rtol={rtol}"
152 )
153
154 min_acceptable = ground_truth * (1 - rtol)
155 success = success and measured_value >= min_acceptable
156
157 assert success