[CI] Fix test_lora_with_spec_decode on V2 model runner (#43314)

Signed-off-by: haosdent <haosdent@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
This commit is contained in:
haosdent
2026-05-22 14:24:18 +08:00
committed by GitHub
co-authored by Wentao Ye
parent fa1ff88b31
commit 5ea76fa89a
2 changed files with 51 additions and 34 deletions
@@ -67,6 +67,12 @@ def test_batch_inference_correctness(
method, model_name, spec_model_name, lora_path, tp_size = model_setup
prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
lora_request = LoRARequest("adapter", 1, lora_path)
sampling_params = SamplingParams(
temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
)
# without speculative decoding
ref_llm = LLM(
model=model_name,
@@ -79,19 +85,14 @@ def test_batch_inference_correctness(
max_cpu_loras=1,
max_lora_rank=16,
)
prompts = [LORA_TEST_PROMPT_MAP[lora_path]] * 100
lora_request = LoRARequest("adapter", 1, lora_path)
sampling_params = SamplingParams(
temperature=0.0, top_p=1.0, top_k=-1, seed=SEED, max_tokens=128
)
ref_outputs = ref_llm.generate(
prompts, sampling_params, lora_request=lora_request
)
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
try:
ref_outputs = ref_llm.generate(
prompts, sampling_params, lora_request=lora_request
)
finally:
del ref_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
lora_spec_llm = LLM(
model=model_name,
@@ -110,25 +111,29 @@ def test_batch_inference_correctness(
max_cpu_loras=1,
max_lora_rank=16,
)
try:
lora_spec_outputs = lora_spec_llm.generate(
prompts, sampling_params, lora_request=lora_request
)
lora_spec_outputs = lora_spec_llm.generate(
prompts, sampling_params, lora_request=lora_request
)
matches = 0
for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
matches = 0
misses = 0
for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs):
if ref_output.outputs[0].text == spec_output.outputs[0].text:
matches += 1
else:
misses += 1
print(f"ref_output: {ref_output.outputs[0].text}")
print(f"spec_output: {spec_output.outputs[0].text}")
# Heuristic: expect at least 90% of the prompts to match exactly
# Upon failure, inspect the outputs to check for inaccuracy.
print(f"match ratio: {matches}/{len(ref_outputs)}")
assert matches > int(0.90 * len(ref_outputs))
del lora_spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()
# Heuristic threshold: under greedy verification, the spec-decode
# output should equal the non-spec output (modulo FP noise from the
# target's verify-path matmul running at seqlen
# num_speculative_tokens+1 vs 1). 90% leaves slack for that noise.
threshold = int(0.90 * len(ref_outputs))
print(f"match ratio: {matches}/{len(ref_outputs)}")
assert matches > threshold, (
f"match ratio {matches}/{len(ref_outputs)} <= {threshold}"
)
finally:
del lora_spec_llm
torch.accelerator.empty_cache()
cleanup_dist_env_and_memory()