From 6327dff24264955fc828906cd85c72bb8e7648fc Mon Sep 17 00:00:00 2001 From: Glen Liu <62917497+glenliu21@users.noreply.github.com> Date: Sat, 10 Jan 2026 03:49:00 -0500 Subject: [PATCH] enhance LoRA tests and fix base model LoRA eviction in Scheduler (#16333) --- python/sglang/srt/managers/scheduler.py | 13 +- python/sglang/test/lora_utils.py | 130 ++++++++++++++++++ test/manual/lora/test_lora_spec_decoding.py | 20 +-- test/registered/lora/test_lora.py | 54 -------- test/registered/lora/test_lora_qwen3.py | 20 +-- .../lora/test_multi_lora_backend.py | 24 +--- 6 files changed, 144 insertions(+), 117 deletions(-) delete mode 100644 test/registered/lora/test_lora.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 35cb7a379..107988d56 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1954,14 +1954,11 @@ class Scheduler( | set([req.lora_id]) ) if not self.tp_worker.can_run_lora_batch(new_lora_set): - # If this is a LoRA request that would exceed the LoRA slot limit, - # skip it and continue to try scheduling non-LoRA requests. - # Non-LoRA requests (lora_id=None) share a single reserved slot - # and should never cause this check to fail. - if req.lora_id is not None: - # Skip this LoRA request - it would trigger adapter eviction/loading - # which is slow. We'll try to schedule it in a future iteration. - continue + # Batch would exceed the LoRA slot limit. + # Skip this request and try scheduling it in a future iteration. + # Note: When eviction is needed, the eviction policy prefers to + # evict LoRA adapters over base model (None) - see mem_pool.py. + continue running_bs = len(self.running_batch.reqs) if len(adder.can_run_list) >= self.get_num_allocatable_reqs(running_bs): diff --git a/python/sglang/test/lora_utils.py b/python/sglang/test/lora_utils.py index e2cb1f17c..171dc9fc1 100644 --- a/python/sglang/test/lora_utils.py +++ b/python/sglang/test/lora_utils.py @@ -25,6 +25,7 @@ class LoRAModelCase: decode_tolerance: float = 1e-1 rouge_l_tolerance: float = 1.0 max_loras_per_batch: int = 1 + max_loaded_loras: Optional[int] = None skip_long_prompt: bool = False def __post_init__(self): @@ -115,6 +116,23 @@ ALL_OTHER_MULTI_LORA_MODELS = [ ), ] +LORA_MODELS_QWEN3 = [ + LoRAModelCase( + base="Qwen/Qwen3-4B", + adaptors=[ + LoRAAdaptor( + name="nissenj/Qwen3-4B-lora-v2", + prefill_tolerance=3e-1, + ), + LoRAAdaptor( + name="y9760210/Qwen3-4B-lora_model", + prefill_tolerance=3e-1, + ), + ], + max_loras_per_batch=2, + ), +] + def safe_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Matrix multiplication with mixed precision handling for float16""" @@ -314,6 +332,7 @@ def run_lora_test_one_by_one( adaptor.name for adaptor in model_case.adaptors if adaptor.name is not None ], max_loras_per_batch=model_case.max_loras_per_batch, + max_loaded_loras=model_case.max_loaded_loras, lora_backend=backend, disable_cuda_graph=disable_cuda_graph, disable_radix_cache=disable_radix_cache, @@ -461,6 +480,7 @@ def run_lora_test_by_batch( adaptor.name for adaptor in model_case.adaptors if adaptor.name is not None ], max_loras_per_batch=model_case.max_loras_per_batch, + max_loaded_loras=model_case.max_loaded_loras, lora_backend=backend, disable_cuda_graph=disable_cuda_graph, disable_radix_cache=disable_radix_cache, @@ -651,6 +671,7 @@ def run_lora_multiple_batch_on_model_cases( model_type="generation", lora_paths=[lora_adapter_paths[0], lora_adapter_paths[1]], 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. attention_backend=attention_backend, enable_deterministic_inference=enable_deterministic_inference, @@ -702,3 +723,112 @@ def run_lora_multiple_batch_on_model_cases( ) print(f"--- Batch {i+1} Comparison Passed --- ") + + +def run_lora_batch_splitting_equivalence_test( + model_cases: List[LoRAModelCase], + attention_backend: str = "torch_native", + disable_cuda_graph: bool = True, + disable_radix_cache: bool = True, +): + """ + Test that SRT correctly handles batch splitting with multiple LoRA adapters. + + When the number of distinct adapters (including None for base model) exceeds + max_loras_per_batch, SRT internally splits requests into microbatches. + + This test validates: + 1. SRT can process batches that trigger internal splitting without errors + 2. Different adapters don't produce all identical outputs (i.e., at least one + output differs, indicating adapters are being applied correctly) + + Args: + model_cases: List of LoRAModelCase configurations to test + attention_backend: Attention backend to use + disable_cuda_graph: Whether to disable CUDA graph + disable_radix_cache: Whether to disable radix cache + """ + max_loras_per_batch = 2 + + def _run_test(model_case: LoRAModelCase, torch_dtype: torch.dtype): + lora_adapter_paths = [a.name for a in model_case.adaptors] + assert ( + len(lora_adapter_paths) >= max_loras_per_batch + ), f"Need at least {max_loras_per_batch} adapters for this test" + + max_new_tokens = 64 + base_path = model_case.base + + print( + f"\n========== Testing batch splitting on base '{base_path}', " + f"dtype={torch_dtype} ==========" + ) + + prompts = [TEST_MULTIPLE_BATCH_PROMPTS[0]] * 3 + test_cases = [ + ( + prompts, + [None, lora_adapter_paths[0], lora_adapter_paths[1]], + ), + ( + prompts, + [lora_adapter_paths[0], None, lora_adapter_paths[1]], + ), + ( + prompts, + [lora_adapter_paths[0], lora_adapter_paths[1], None], + ), + ( + prompts, + [None, lora_adapter_paths[1], None], + ), + ( + prompts, + [lora_adapter_paths[0], lora_adapter_paths[1], lora_adapter_paths[0]], + ), + ( + prompts, + [None, None, None], + ), + ] + + ensure_reproducibility() + with SRTRunner( + base_path, + torch_dtype=torch_dtype, + model_type="generation", + lora_paths=lora_adapter_paths, + max_loras_per_batch=max_loras_per_batch, + max_loaded_loras=model_case.max_loaded_loras, + sleep_on_idle=True, + attention_backend=attention_backend, + disable_cuda_graph=disable_cuda_graph, + disable_radix_cache=disable_radix_cache, + ) as srt_runner: + for batch_idx, (batch_prompts, lora_paths) in enumerate(test_cases): + print(f"\n--- Batch {batch_idx + 1} ---") + print(f" Adapters: {lora_paths}") + + srt_outputs = srt_runner.batch_forward( + batch_prompts, + max_new_tokens=max_new_tokens, + lora_paths=lora_paths, + ) + + # If different adapters are used in this batch, verify that not every + # output is identical (at least one should differ) + unique_adapters = set(lora_paths) + if len(unique_adapters) >= 2: + all_outputs = [s.strip() for s in srt_outputs.output_strs] + all_identical = all(out == all_outputs[0] for out in all_outputs) + assert not all_identical, ( + f"Every output was identical despite using different adapters for " + f"base '{base_path}', batch {batch_idx + 1}: " + f"adapters={lora_paths}. Expected at least one output to differ." + ) + + print(f"--- Batch {batch_idx + 1} passed ---") + + for model_case in model_cases: + for torch_dtype in TORCH_DTYPES: + _run_test(model_case, torch_dtype) diff --git a/test/manual/lora/test_lora_spec_decoding.py b/test/manual/lora/test_lora_spec_decoding.py index e00fafa75..e247cdddd 100644 --- a/test/manual/lora/test_lora_spec_decoding.py +++ b/test/manual/lora/test_lora_spec_decoding.py @@ -17,29 +17,11 @@ import unittest from sglang.test.lora_utils import ( CI_MULTI_LORA_MODELS, - LoRAAdaptor, - LoRAModelCase, + LORA_MODELS_QWEN3, run_lora_multiple_batch_on_model_cases, ) from sglang.test.test_utils import CustomTestCase -LORA_MODELS_QWEN3 = [ - LoRAModelCase( - base="Qwen/Qwen3-4B", - adaptors=[ - LoRAAdaptor( - name="nissenj/Qwen3-4B-lora-v2", - prefill_tolerance=3e-1, - ), - LoRAAdaptor( - name="y9760210/Qwen3-4B-lora_model", - prefill_tolerance=3e-1, - ), - ], - max_loras_per_batch=2, - ), -] - class TestLoRASpecDecoding(CustomTestCase): def test_qwen(self): diff --git a/test/registered/lora/test_lora.py b/test/registered/lora/test_lora.py deleted file mode 100644 index b9e1559fa..000000000 --- a/test/registered/lora/test_lora.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright 2023-2024 SGLang Team -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== - -import multiprocessing as mp -import os -import unittest - -from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci -from sglang.test.lora_utils import ( - ALL_OTHER_MULTI_LORA_MODELS, - CI_MULTI_LORA_MODELS, - run_lora_multiple_batch_on_model_cases, -) -from sglang.test.test_utils import CustomTestCase, is_in_ci - -register_cuda_ci(est_time=82, suite="stage-b-test-small-1-gpu") -register_amd_ci(est_time=82, suite="stage-b-test-small-1-gpu-amd") - - -class TestLoRA(CustomTestCase): - def test_ci_lora_models(self): - run_lora_multiple_batch_on_model_cases(CI_MULTI_LORA_MODELS) - - def test_all_lora_models(self): - if is_in_ci(): - return - - filtered_models = [] - for model_case in ALL_OTHER_MULTI_LORA_MODELS: - if "ONLY_RUN" in os.environ and os.environ["ONLY_RUN"] != model_case.base: - continue - filtered_models.append(model_case) - - run_lora_multiple_batch_on_model_cases(filtered_models) - - -if __name__ == "__main__": - try: - mp.set_start_method("spawn") - except RuntimeError: - pass - - unittest.main(warnings="ignore") diff --git a/test/registered/lora/test_lora_qwen3.py b/test/registered/lora/test_lora_qwen3.py index da6999f42..912c60cf1 100644 --- a/test/registered/lora/test_lora_qwen3.py +++ b/test/registered/lora/test_lora_qwen3.py @@ -17,8 +17,7 @@ import unittest from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.lora_utils import ( - LoRAAdaptor, - LoRAModelCase, + LORA_MODELS_QWEN3, run_lora_multiple_batch_on_model_cases, ) @@ -26,23 +25,6 @@ register_cuda_ci(est_time=97, suite="nightly-1-gpu", nightly=True) from sglang.test.test_utils import CustomTestCase -LORA_MODELS_QWEN3 = [ - LoRAModelCase( - base="Qwen/Qwen3-4B", - adaptors=[ - LoRAAdaptor( - name="nissenj/Qwen3-4B-lora-v2", - prefill_tolerance=3e-1, - ), - LoRAAdaptor( - name="y9760210/Qwen3-4B-lora_model", - prefill_tolerance=3e-1, - ), - ], - max_loras_per_batch=2, - ), -] - class TestLoRAQwen3(CustomTestCase): def test_ci_lora_models(self): diff --git a/test/registered/lora/test_multi_lora_backend.py b/test/registered/lora/test_multi_lora_backend.py index 7180bebfe..62c7e2c64 100644 --- a/test/registered/lora/test_multi_lora_backend.py +++ b/test/registered/lora/test_multi_lora_backend.py @@ -20,30 +20,20 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.lora_utils import ( ALL_OTHER_MULTI_LORA_MODELS, CI_MULTI_LORA_MODELS, + run_lora_batch_splitting_equivalence_test, run_lora_multiple_batch_on_model_cases, ) from sglang.test.test_utils import CustomTestCase, is_in_ci -register_cuda_ci(est_time=60, suite="stage-b-test-small-1-gpu") -register_amd_ci(est_time=60, suite="stage-b-test-small-1-gpu-amd") - -# All prompts are used at once in a batch. -PROMPTS = [ - "AI is a field of computer science focused on", - """ - ### Instruction: - Tell me about llamas and alpacas - ### Response: - Llamas are large, long-necked animals with a woolly coat. They have two toes on each foot instead of three like other camelids. - ### Question: - What do you know about llamas? - ### Answer: - """, -] +register_cuda_ci(est_time=100, suite="stage-b-test-small-1-gpu") +register_amd_ci(est_time=100, suite="stage-b-test-small-1-gpu-amd") class TestMultiLoRABackend(CustomTestCase): - def test_ci_lora_models(self): + def test_ci_lora_models_batch_splitting(self): + run_lora_batch_splitting_equivalence_test(CI_MULTI_LORA_MODELS) + + def test_ci_lora_models_multi_batch(self): run_lora_multiple_batch_on_model_cases(CI_MULTI_LORA_MODELS) def test_all_lora_models(self):