From 6d48719e31aa42adfd718dd0da2becb1f9be1d2a Mon Sep 17 00:00:00 2001 From: "Ethan (Yusheng) Su" Date: Fri, 27 Mar 2026 16:08:36 -0700 Subject: [PATCH] [1/n] lora support - Auto detect lora target modules (#21439) Co-authored-by: Baizhou Zhang --- python/sglang/srt/lora/lora_manager.py | 26 ++- python/sglang/srt/lora/utils.py | 53 ++++- .../lora/test_lora_qwen3_8b_logprob_diff.py | 202 ++++++++++++++++++ 3 files changed, 266 insertions(+), 15 deletions(-) create mode 100644 test/registered/lora/test_lora_qwen3_8b_logprob_diff.py diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 73f6bc235..7b222161f 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -36,6 +36,7 @@ from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.lora.mem_pool import LoRAMemoryPool from sglang.srt.lora.utils import ( LoRAType, + auto_detect_lora_target_modules, get_normalized_target_modules, get_target_module_name, ) @@ -424,9 +425,6 @@ class LoRAManager: for lora_id, config in self.configs.items(): # Handle PEFT shorthand strings like "all-linear" or "all". - # These cannot be resolved to concrete module names without - # inspecting the base model, so we require the user to specify - # --lora-target-modules explicitly when such shorthands are used. if isinstance(config.target_modules, str): if config.target_modules in ("all-linear", "all"): if target_modules is not None: @@ -434,14 +432,20 @@ class LoRAManager: # per-adapter inference for this adapter. continue else: - lora_name = self.lora_refs[lora_id].lora_name - raise ValueError( - f"LoRA adapter '{lora_name}' uses " - f"target_modules='{config.target_modules}' which cannot " - "be resolved automatically. Please explicitly specify " - "--lora-target-modules during server startup. You can " - "specify 'all' to enable all supported module types." + # Resolve by scanning the base model for all + # LoRA-compatible linear modules. + adapter_target_modules = auto_detect_lora_target_modules( + self.base_model ) + logger.info( + "LoRA adapter '%s' uses target_modules='%s'. " + "Resolved to %s by inspecting the base model.", + self.lora_refs[lora_id].lora_name, + config.target_modules, + sorted(adapter_target_modules), + ) + self.target_modules.update(adapter_target_modules) + continue else: raise ValueError( f"SGLang does not recognize target_modules=" @@ -672,6 +676,8 @@ class LoRAManager: # The module should be converted if it is included in target_names if module_name.split(".")[-1] in self.target_modules: layer_id = get_layer_id(module_name) + if layer_id is None: + continue self.lora_modules[layer_id][module_name] = self.set_lora_module( module_name, module ) diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index 45987d736..1a12b1bd6 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -113,13 +113,17 @@ def get_normalized_target_modules( Handles both base module names (e.g., "gate_proj") and prefixed module names (e.g., "feed_forward.gate_proj"). Also handles PEFT shorthand strings like "all-linear" or "all" by returning - {"all"} as a sentinel value (the caller should check for "all" and fall - back to the CLI --lora-target-modules to determine the concrete module set). + {"all"} as a sentinel value. Callers that need a concrete module set + should use :func:`auto_detect_lora_target_modules` to resolve the shorthand + against the loaded base model. """ - # Handle PEFT shorthand strings — these cannot be resolved to concrete - # module names without inspecting the base model, so we return {"all"} - # and let the caller fall back to the CLI --lora-target-modules. + # Handle PEFT shorthand strings — return {"all"} as sentinel. + # Callers can resolve to concrete names via auto_detect_lora_target_modules(). if isinstance(target_modules, str): + if target_modules not in ["all", "all-linear"]: + raise ValueError( + "Only 'all' or 'all-linear' can be used as the string for target module" + ) return {"all"} params_mapping = { @@ -175,6 +179,45 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s EMBEDDING_NAMES = ["embed_tokens", "lm_head"] ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj", "down_proj_moe"] +# Normalized module names that the LoRA system fully supports +# (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them). +_KNOWN_LORA_TARGET_MODULES = frozenset( + { + "qkv_proj", + "o_proj", + "gate_up_proj", + "down_proj", + "embed_tokens", + "lm_head", + } +) + + +def auto_detect_lora_target_modules(model: "torch.nn.Module") -> set: + """Discover LoRA-compatible modules by inspecting the base model. + + Walks the model graph and returns the set of *normalized* target-module + names that (a) actually exist in the model and (b) the LoRA memory pool + can handle. This is used to resolve PEFT shorthands like ``"all-linear"`` + without requiring the user to enumerate modules on the CLI. + """ + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + raw_names: set = set() + for name, module in model.named_modules(): + if isinstance(module, FusedMoE): + raw_names.add("gate_up_proj") + raw_names.add("down_proj") + elif isinstance(module, ParallelLMHead): + raw_names.add("lm_head") + elif isinstance(module, LinearBase): + raw_names.add(name.split(".")[-1]) + + normalized = get_normalized_target_modules(raw_names) + return normalized & _KNOWN_LORA_TARGET_MODULES + def get_lm_head_lora_b_shard_size(output_dim: int, shard_indices=None) -> int: """Get the LoRA B output dimension for lm_head, accounting for TP. diff --git a/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py b/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py new file mode 100644 index 000000000..04f60cfdc --- /dev/null +++ b/test/registered/lora/test_lora_qwen3_8b_logprob_diff.py @@ -0,0 +1,202 @@ +# Copyright 2023-2025 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. +# ============================================================================== + +""" +Regression test for Qwen3-8B LoRA logprob accuracy. + +Compares SGLang LoRA logprobs against reference training logprobs from a +pre-computed dataset. The LoRA adapter and reference data are downloaded from: +https://huggingface.co/datasets/yushengsu/lora-diff-Qwen3-8B + +Usage: + python -m unittest test_lora_qwen3_8b_logprob_diff +""" + +import multiprocessing as mp +import os +import unittest +from unittest.mock import patch + +import torch +import torch.nn as nn +from huggingface_hub import snapshot_download + +import sglang as sgl +from sglang.srt.lora.utils import auto_detect_lora_target_modules +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci( + est_time=200, + suite="stage-b-test-1-gpu-large", +) + +BASE_MODEL = "Qwen/Qwen3-8B" +LORA_HF_REPO = "yushengsu/lora-diff-Qwen3-8B" +LORA_BACKEND = "triton" +MAX_LORA_RANK = 32 +TP_SIZE = 1 +DISABLE_CUDA_GRAPH = True +PREFILL_ATTENTION_BACKEND = "fa3" +DECODE_ATTENTION_BACKEND = "fa3" + +KL_THRESHOLD = 1e-2 + + +def kl_v2(a, b): + a = torch.tensor(a) if not torch.is_tensor(a) else a + b = torch.tensor(b) if not torch.is_tensor(b) else b + return (((a - b) ** 2) * 0.5).mean().item() + + +def get_prompt_logprobs(engine, input_ids, lora_path): + out = engine.generate( + input_ids=input_ids, + sampling_params={"max_new_tokens": 0, "temperature": 0.0}, + return_logprob=True, + logprob_start_len=0, + lora_path=lora_path, + ) + return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:] + + +class _MockLinearBase(nn.Module): + pass + + +class _MockFusedMoE(nn.Module): + pass + + +class _MockParallelLMHead(nn.Module): + pass + + +def _build_qwen3_mock(): + """Build a lightweight nn.Module tree that mirrors Qwen3-8B's named modules.""" + model = nn.Module() + inner = nn.Module() + layer = nn.Module() + + attn = nn.Module() + attn.qkv_proj = _MockLinearBase() + attn.o_proj = _MockLinearBase() + layer.self_attn = attn + + mlp = nn.Module() + mlp.gate_up_proj = _MockLinearBase() + mlp.down_proj = _MockLinearBase() + layer.mlp = mlp + + inner.layers = nn.ModuleList([layer]) + inner.embed_tokens = nn.Embedding(10, 8) # not a LinearBase — should be excluded + model.model = inner + model.lm_head = _MockParallelLMHead() + return model + + +class TestLoRAQwen3_8BLogprobDiff(CustomTestCase): + + def test_auto_detect_lora_target_modules(self): + """Verify auto_detect_lora_target_modules returns the expected module + set for a Qwen3-8B-like (dense) architecture. Catches silent renames + of internal param names that would break LoRA auto-detection.""" + model = _build_qwen3_mock() + + with patch("sglang.srt.layers.linear.LinearBase", _MockLinearBase), patch( + "sglang.srt.layers.moe.fused_moe_triton.layer.FusedMoE", _MockFusedMoE + ), patch( + "sglang.srt.layers.vocab_parallel_embedding.ParallelLMHead", + _MockParallelLMHead, + ): + detected = auto_detect_lora_target_modules(model) + + expected = {"qkv_proj", "o_proj", "gate_up_proj", "down_proj", "lm_head"} + self.assertEqual(detected, expected) + + def test_lora_qwen3_8b_logprob_accuracy(self): + adapter_path = snapshot_download( + LORA_HF_REPO, + repo_type="dataset", + ) + + engine = sgl.Engine( + model_path=BASE_MODEL, + tp_size=TP_SIZE, + enable_lora=True, + max_lora_rank=MAX_LORA_RANK, + lora_paths={"my_lora": adapter_path}, + lora_backend=LORA_BACKEND, + attention_backend="flashinfer", + disable_cuda_graph=DISABLE_CUDA_GRAPH, + prefill_attention_backend=PREFILL_ATTENTION_BACKEND, + decode_attention_backend=DECODE_ATTENTION_BACKEND, + ) + + try: + cdata = torch.load( + os.path.join(adapter_path, "compare_sample_train_data.pt"), + weights_only=False, + ) + + base_logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path=None) + logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path="my_lora") + + base_t = torch.tensor(base_logprobs) + lora_t = torch.tensor(logprobs) + diff = (base_t - lora_t).abs() + print( + f"[VERIFY] base vs lora: mean_diff={diff.mean().item():.6f}, " + f"max_diff={diff.max().item():.6f}, " + f"identical={torch.equal(base_t, lora_t)}" + ) + + self.assertFalse( + torch.equal(base_t, lora_t), + "LoRA logprobs should differ from base model logprobs", + ) + + kl_sglang_trainer = kl_v2(cdata["training_logprobs"], logprobs) + kl_orig_trainer = kl_v2( + cdata["training_logprobs"], cdata["sampling_logprobs"] + ) + kl_sglang_orig = kl_v2(logprobs, cdata["sampling_logprobs"]) + + print(f"KL(orig_sampler, trainer) = {kl_orig_trainer:.6e}") + print(f"KL(sglang, trainer) = {kl_sglang_trainer:.6e}") + print(f"KL(sglang, orig_sampler) = {kl_sglang_orig:.6e}") + + self.assertLessEqual( + kl_sglang_trainer, + KL_THRESHOLD, + f"KL(sglang, trainer) = {kl_sglang_trainer:.6e} exceeds " + f"threshold {KL_THRESHOLD}", + ) + + finally: + engine.shutdown() + + +if __name__ == "__main__": + try: + mp.set_start_method("spawn") + except RuntimeError: + pass + + try: + unittest.main(warnings="ignore", verbosity=2) + finally: + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.synchronize()