[TestFix] rewrite LoRA overlap loading tests (#18047)

This commit is contained in:
Glen Liu
2026-02-01 17:52:08 -05:00
committed by GitHub
parent 993ec178ef
commit 99dad105fd
2 changed files with 18 additions and 83 deletions

View File

@@ -640,6 +640,7 @@ def run_lora_multiple_batch_on_model_cases(
disable_cuda_graph: bool = True,
enable_deterministic_inference: bool = False,
disable_radix_cache: bool = True,
enable_lora_overlap_loading: Optional[bool] = None,
):
for model_case in model_cases:
for torch_dtype in TORCH_DTYPES:
@@ -673,6 +674,7 @@ def run_lora_multiple_batch_on_model_cases(
torch_dtype=torch_dtype,
model_type="generation",
lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]],
enable_lora_overlap_loading=enable_lora_overlap_loading,
max_loras_per_batch=len(lora_adapter_paths) + 1,
max_loaded_loras=model_case.max_loaded_loras,
sleep_on_idle=True, # Eliminate non-determinism by forcing all requests to be processed in one batch.
@@ -733,6 +735,7 @@ def run_lora_batch_splitting_equivalence_test(
attention_backend: str = "torch_native",
disable_cuda_graph: bool = True,
disable_radix_cache: bool = True,
enable_lora_overlap_loading: Optional[bool] = None,
):
"""
Test that SRT correctly handles batch splitting with multiple LoRA adapters.
@@ -801,6 +804,7 @@ def run_lora_batch_splitting_equivalence_test(
torch_dtype=torch_dtype,
model_type="generation",
lora_paths=lora_adapter_paths,
enable_lora_overlap_loading=enable_lora_overlap_loading,
max_loras_per_batch=max_loras_per_batch,
max_loaded_loras=model_case.max_loaded_loras,
sleep_on_idle=True,

View File

@@ -18,97 +18,28 @@ End-to-end tests for the --enable-lora-overlap-loading server argument.
import multiprocessing as mp
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.lora_utils import (
CI_MULTI_LORA_MODELS,
TEST_MULTIPLE_BATCH_PROMPTS,
TORCH_DTYPES,
LoRAModelCase,
ensure_reproducibility,
run_lora_batch_splitting_equivalence_test,
run_lora_multiple_batch_on_model_cases,
)
from sglang.test.runners import SRTRunner
from sglang.test.test_utils import CustomTestCase, calculate_rouge_l
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(
est_time=300,
suite="stage-b-test-small-1-gpu",
disabled="Flaky test - outputs differ between overlap/no-overlap loading modes. See https://github.com/sgl-project/sglang/actions/runs/21320657015/job/61370002606",
)
register_cuda_ci(est_time=100, suite="stage-b-test-large-1-gpu")
register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd")
class TestLoRAPipelineLoading(CustomTestCase):
def _run_mixed_batch_test(
self,
model_case: LoRAModelCase,
torch_dtype,
):
base_path = model_case.base
adaptor_paths = [a.name for a in model_case.adaptors]
print(
f"\n========== Testing mixed batch LoRA overlap loading on base '{base_path}' "
f"with dtype={torch_dtype} ==========\n"
)
ensure_reproducibility()
max_new_tokens = 32
prompts = TEST_MULTIPLE_BATCH_PROMPTS[:3]
configs = [
[None, adaptor_paths[0], adaptor_paths[1]],
[adaptor_paths[0], None, adaptor_paths[1]],
[adaptor_paths[0], adaptor_paths[1], None],
[adaptor_paths[1], adaptor_paths[0], adaptor_paths[1]],
]
common_args = dict(
torch_dtype=torch_dtype,
model_type="generation",
tp_size=model_case.tp_size,
lora_paths=adaptor_paths,
max_loras_per_batch=model_case.max_loras_per_batch,
max_loaded_loras=model_case.max_loaded_loras,
disable_cuda_graph=True,
disable_radix_cache=True,
mem_fraction_static=0.65,
sleep_on_idle=True,
class TestLoRAOverlapLoading(CustomTestCase):
def test_ci_lora_models_batch_splitting(self):
run_lora_batch_splitting_equivalence_test(
CI_MULTI_LORA_MODELS, enable_lora_overlap_loading=True
)
results_no_overlap_loading = []
with SRTRunner(
base_path, enable_lora_overlap_loading=False, **common_args
) as runner:
for lora_paths in configs:
results_no_overlap_loading.append(
runner.batch_forward(
prompts, max_new_tokens=max_new_tokens, lora_paths=lora_paths
).output_strs
)
results_overlap_loading = []
with SRTRunner(
base_path, enable_lora_overlap_loading=True, **common_args
) as runner:
for lora_paths in configs:
results_overlap_loading.append(
runner.batch_forward(
prompts, max_new_tokens=max_new_tokens, lora_paths=lora_paths
).output_strs
)
for i, (res_no_overlap_loading, res_overlap_loading) in enumerate(
zip(results_no_overlap_loading, results_overlap_loading)
):
scores = calculate_rouge_l(res_overlap_loading, res_no_overlap_loading)
for j, score in enumerate(scores):
assert score >= model_case.rouge_l_tolerance, (
f"Batch {i} prompt {j} mismatch: {score}\n"
f"Overlap loading: {res_overlap_loading[j]}\n"
f"No overlap loading: {res_no_overlap_loading[j]}"
)
def test_mixed_batch(self):
for model_case in CI_MULTI_LORA_MODELS:
for dtype in TORCH_DTYPES:
self._run_mixed_batch_test(model_case, dtype)
def test_ci_lora_models_multi_batch(self):
run_lora_multiple_batch_on_model_cases(
CI_MULTI_LORA_MODELS, enable_lora_overlap_loading=True
)
if __name__ == "__main__":