[Feature] Add LoRA support for embedding layers (#14177)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Beichen-Ma <bm685@cornell.edu>

python/sglang/srt/lora/utils.py

@@ -46,7 +46,11 @@ class LoRAType(Enum):
 
 
 def get_hidden_dim(
-    module_name: str, config: AutoConfig, base_model: torch.nn.Module, layer_idx: int
+    module_name: str,
+    config: AutoConfig,
+    base_model: torch.nn.Module,
+    layer_idx: int,
+    lora_added_vocab_size: int = 0,
 ) -> Tuple[int]:
     """
     Given a module_name (might be a stacked name), return the hidden dims of modules' input and output.
@@ -78,6 +82,14 @@ def get_hidden_dim(
             return config.hidden_size, config.intermediate_size * 2
         elif module_name == "down_proj":
             return config.intermediate_size, config.hidden_size
+        elif module_name == "embed_tokens":
+            # For embedding: input is vocab_size (as embedding lookup), output is hidden_size
+            # if contain extra tokens will be added; otherwise is 0.
+            return config.vocab_size + lora_added_vocab_size, config.hidden_size
+        elif module_name == "lm_head":
+            # For lm_head: input is hidden_size, output is vocab_size
+            # if contain extra tokens will be added; otherwise is 0.
+            return config.hidden_size, config.vocab_size + lora_added_vocab_size
         else:
             raise NotImplementedError()
 
@@ -95,6 +107,12 @@ def get_normalized_target_modules(
         "v_proj": "qkv_proj",
         "gate_proj": "gate_up_proj",
         "up_proj": "gate_up_proj",
+        "embed_tokens": "embed_tokens",
+        "vocab_emb": "embed_tokens",
+        "embeddings": "embed_tokens",
+        "word_embeddings": "embed_tokens",
+        "lm_head": "lm_head",
+        "output": "lm_head",
     }
 
     result = set()
@@ -131,4 +149,5 @@ 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"]