[1/n] lora support - Auto detect lora target modules (#21439)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

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
+                        # Resolve by scanning the base model for all
-                        raise ValueError(
+                        # LoRA-compatible linear modules.
-                            f"LoRA adapter '{lora_name}' uses "
+                        adapter_target_modules = auto_detect_lora_target_modules(
-                            f"target_modules='{config.target_modules}' which cannot "
+                            self.base_model
-                            "be resolved automatically. Please explicitly specify "
-                            "--lora-target-modules during server startup. You can "
-                            "specify 'all' to enable all supported module types."
                         )
+                        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
                 )

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
+    {"all"} as a sentinel value.  Callers that need a concrete module set
-    back to the CLI --lora-target-modules to determine the 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
+    # Handle PEFT shorthand strings — return {"all"} as sentinel.
-    # module names without inspecting the base model, so we return {"all"}
+    # Callers can resolve to concrete names via auto_detect_lora_target_modules().
-    # and let the caller fall back to the CLI --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.

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()