[Fix] Add lora tied lm head support (for Qwen2.5, Gemma, etc model need) (#18634)

python/sglang/srt/lora/lora.py

@@ -102,13 +102,24 @@ class LoRAAdapter(nn.Module):
             self.config.target_modules
         )
 
+        # Remap PEFT "unembed_tokens" key to "lm_head" so the weight is
+        # recognized and loaded into the correct buffer.
+        if "unembed_tokens" in name:
+            name = name.replace("unembed_tokens", "lm_head")
+
         layer_id = get_layer_id(name)
         if layer_id is not None:
             self.layers[layer_id].weights[name] = loaded_weight.cpu()
         elif "embed_tokens" in name or "lm_head" in name:
-            # Check if this module is declared in target_modules before loading
+            # Check if this module is declared in target_modules before loading.
+            # When normalized_target_modules is {"all"} (e.g. target_modules was
+            # "all-linear"), we allow loading since the server-level
+            # --lora-target-modules will govern which modules are active.
             module_name = "embed_tokens" if "embed_tokens" in name else "lm_head"
-            if module_name in normalized_target_modules:
+            if (
+                "all" in normalized_target_modules
+                or module_name in normalized_target_modules
+            ):
                 self.embedding_layers[name] = loaded_weight.cpu()
             else:
                 logger.debug(

python/sglang/srt/lora/lora_manager.py

@@ -387,6 +387,33 @@ 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:
+                        # CLI --lora-target-modules already provided; skip
+                        # 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."
+                        )
+                else:
+                    raise ValueError(
+                        f"SGLang does not recognize target_modules="
+                        f"'{config.target_modules}'. Please use a list of module "
+                        "name suffixes in the adapter's PEFT config, or explicitly "
+                        "specify --lora-target-modules during server startup."
+                    )
+
             if not isinstance(config.target_modules, list):
                 raise ValueError(
                     f"SGLang currently only supports inferring LoRA target modules when a list of "
@@ -541,6 +568,40 @@ class LoRAManager:
         self.embed_tokens_module: Optional[BaseLayerWithLoRA] = None
         self.lm_head_module: Optional[BaseLayerWithLoRA] = None
 
+        # When tie_word_embeddings=True, lm_head is the same Python object as
+        # embed_tokens. PyTorch's named_modules() deduplicates by object identity,
+        # so lm_head will not appear as a separate entry in the scan below,
+        # preventing LoRA from wrapping it. To fix this, we create a new
+        # ParallelLMHead that shares the same base weight tensor (no extra GPU
+        # memory) so that named_modules() yields it as an independent module.
+        if "lm_head" in self.target_modules:
+            lm_head = getattr(self.base_model, "lm_head", None)
+            embed_tokens = None
+            for name, mod in self.base_model.named_modules():
+                if name.endswith("embed_tokens"):
+                    embed_tokens = mod
+                    break
+            if (
+                lm_head is not None
+                and embed_tokens is not None
+                and lm_head is embed_tokens
+            ):
+                logger.info(
+                    "lm_head is tied with embed_tokens. Creating a separate "
+                    "ParallelLMHead that shares the base weight for LoRA support."
+                )
+                untied_lm_head = ParallelLMHead(
+                    num_embeddings=embed_tokens.org_vocab_size,
+                    embedding_dim=embed_tokens.embedding_dim,
+                    params_dtype=embed_tokens.weight.dtype,
+                    org_num_embeddings=embed_tokens.org_vocab_size,
+                )
+                # Share the base weight tensor — no additional GPU memory.
+                untied_lm_head.weight = embed_tokens.weight
+                # Replace the model attribute so named_modules() sees it
+                # independently.
+                self.base_model.lm_head = untied_lm_head
+
         for module_name, module in self.base_model.named_modules():
             # TODO (lifuhuang): in the future, we should consider generalizing the
             # should_apply_lora function to support mapping by full module name instead

python/sglang/srt/lora/mem_pool.py

@@ -115,6 +115,8 @@ class LoRAMemoryPool:
             if config.lora_added_tokens_size > self.lora_added_tokens_size:
                 return False
             target_module_names = get_normalized_target_modules(config.target_modules)
+            if "all" in target_module_names:
+                return True
             return target_module_names.issubset(self.target_modules)
 
         if isinstance(config, LoRAConfig):

python/sglang/srt/lora/utils.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Iterable, Optional, Set, Tuple
+from typing import Iterable, Optional, Set, Tuple, Union
 
 import torch
 
@@ -98,12 +98,22 @@ def get_hidden_dim(
 
 
 def get_normalized_target_modules(
-    target_modules: Iterable[str],
+    target_modules: Union[str, Iterable[str]],
 ) -> set[str]:
     """
     Mapping a list of target module name to names of the normalized LoRA weights.
     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).
     """
+    # 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.
+    if isinstance(target_modules, str):
+        return {"all"}
+
     params_mapping = {
         "q_proj": "qkv_proj",
         "k_proj": "qkv_proj",
@@ -116,6 +126,7 @@ def get_normalized_target_modules(
         "word_embeddings": "embed_tokens",
         "lm_head": "lm_head",
         "output": "lm_head",
+        "unembed_tokens": "lm_head",
     }
 
     result = set()