# 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. # ============================================================================== # Integrates "S-LoRA: Serving Thousands of Concurrent LoRA Adapters" # and "Punica: Multi-Tenant LoRA Serving" import logging from typing import Dict, Iterable, List, Optional import torch from sglang.srt.configs.load_config import LoadConfig from sglang.srt.environ import envs from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.utils import get_layer_id from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.backend.lora_registry import get_backend_from_name from sglang.srt.lora.layers import BaseLayerWithLoRA, FusedMoEWithLoRA, get_lora_layer from sglang.srt.lora.lora import LoRAAdapter from sglang.srt.lora.lora_config import LoRAConfig from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.lora.mem_pool import LoRAMemoryPool from sglang.srt.lora.utils import ( DSA_INDEXER_LORA_NAMES, EMBEDDING_NAMES, LoRAType, auto_detect_lora_target_modules, get_normalized_target_modules, get_target_module_name, ) from sglang.srt.managers.io_struct import LoRAUpdateOutput from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import ServerArgs from sglang.srt.utils import replace_submodule from sglang.srt.utils.hf_transformers_utils import AutoConfig _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() logger = logging.getLogger(__name__) class LoRAManager: def __init__( self, base_model: torch.nn.Module, base_hf_config: AutoConfig, max_loras_per_batch: int, load_config: LoadConfig, dtype: torch.dtype, server_args: ServerArgs, lora_backend: str = "triton", tp_size: int = 1, tp_rank: int = 0, max_lora_rank: Optional[int] = None, target_modules: Optional[Iterable[str]] = None, lora_paths: Optional[List[LoRARef]] = None, ): self.base_model: torch.nn.Module = base_model if hasattr(base_hf_config, "get_text_config"): self.base_hf_config: AutoConfig = base_hf_config.get_text_config() else: self.base_hf_config: AutoConfig = base_hf_config self.max_loras_per_batch: int = max_loras_per_batch self.load_config: LoadConfig = load_config self.dtype: torch.dtype = dtype self.device: torch.device = next(self.base_model.parameters()).device self.tp_size: int = tp_size self.tp_rank: int = tp_rank self.lora_added_tokens_size: Optional[int] = None self.enable_lora_overlap_loading: Optional[bool] = ( server_args.enable_lora_overlap_loading ) self.eviction_policy = server_args.lora_eviction_policy self._experts_shared_outer_override: Optional[bool] = ( server_args.experts_shared_outer_loras ) self.lora_use_virtual_experts: bool = server_args.lora_use_virtual_experts self.lora_strict_loading: bool = getattr( server_args, "lora_strict_loading", False ) # LoRA backend for running sgemm kernels logger.info(f"Using {lora_backend} as backend of LoRA kernels.") backend_type = get_backend_from_name(lora_backend) self.lora_backend: BaseLoRABackend = backend_type( max_loras_per_batch=max_loras_per_batch, device=self.device, server_args=server_args, ) # Initialize mutable internal state of the LoRAManager. self.init_state( max_lora_rank=max_lora_rank, target_modules=target_modules, lora_paths=lora_paths, ) def init_cuda_graph_batch_info( self, max_bs_in_cuda_graph: int, num_tokens_per_bs: int ): """Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata. Called during CudaGraphRunner.__init__(), after init_memory_pool(). Phase 1 (MoE buffers) is handled earlier via init_cuda_graph_moe_buffers(). """ self.max_bs_in_cuda_graph = max_bs_in_cuda_graph self.lora_backend.init_cuda_graph_batch_info( max_bs_in_cuda_graph=max_bs_in_cuda_graph, num_tokens_per_bs=num_tokens_per_bs, ) # ===== TO BE REFACTORED ==== # Pre-create the experimental LoRA two-stream side stream now (gated) so the # torch.cuda.Stream() call never lands inside a cuda-graph capture region. if _SGLANG_EXPERIMENTAL_LORA_OPTI: from sglang.srt.lora.trtllm_lora_temp import ( init_lora_two_stream_resources, ) init_lora_two_stream_resources(self.device) # ===== END TO BE REFACTORED ==== def init_cuda_graph_moe_buffers( self, max_bs: int, max_loras: int, compute_dtype, moe_layer ): """Phase 1 of LoRA CUDA graph init: MoE intermediate buffers. Called before init_memory_pool() so memory profiling accounts for them. Phase 2 (dense batch metadata) is handled later via init_cuda_graph_batch_info(). """ self.lora_backend.init_cuda_graph_moe_buffers( max_bs=max_bs, max_loras=max_loras, compute_dtype=compute_dtype, moe_layer=moe_layer, ) def create_lora_update_result( self, success: bool, error_message: str = "" ) -> LoRAUpdateOutput: return LoRAUpdateOutput( success=success, error_message=error_message, loaded_adapters={ lora_ref.lora_name: lora_ref.lora_path for lora_ref in self.lora_refs.values() }, ) def load_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput: """ Load a single LoRA adapter from the specified path. Args: lora_ref (LoRARef): The LoRARef object containing the LoRA name, path, and ID. """ assert ( lora_ref.lora_name is not None and lora_ref.lora_path is not None ), "LoRARef must have both lora_name and lora_path set for loading." assert ( lora_ref.lora_id not in self.loras ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." try: # load configs new_adapter = LoRAConfig( lora_ref.lora_path, base_vocab_size=self.base_hf_config.vocab_size, ) self.validate_new_adapter(new_adapter, lora_ref) self.configs[lora_ref.lora_id] = new_adapter # load weights self.load_lora_weights(lora_ref) # keep metadata for displayed messages self.lora_refs[lora_ref.lora_id] = lora_ref self.num_pinned_loras += int(lora_ref.pinned) except Exception as e: return self.create_lora_update_result( success=False, error_message=str(e), ) return self.create_lora_update_result(success=True) def validate_new_adapter(self, lora_config: LoRAConfig, lora_ref: LoRARef): """ Validate if an adapter can be loaded into the current LoRA memory pool and generate error if it is incompatible. """ if lora_config.lora_added_tokens_size > 0: raise ValueError( f"Failed to load {lora_ref.lora_name} because LoRA serving currently doesn't support adapters that add tokens to the vocabulary" ) if lora_config.use_dora: raise ValueError( f"Failed to load {lora_ref.lora_name} because LoRA serving currently doesn't support DoRA adapters" ) # Check if this LoRA adapter is already loaded for existing_lora_ref in self.lora_refs.values(): if lora_ref.lora_name == existing_lora_ref.lora_name: raise ValueError( f"Failed to load LoRA adapter {lora_ref.lora_name} because it is already loaded" ) if lora_ref.lora_path == existing_lora_ref.lora_path: logger.warning( f"{lora_ref.lora_path} is already loaded with name: {existing_lora_ref.lora_name}, " f"but another copy is being loaded with name: {lora_ref.lora_name}" ) # Check if the LoRA adapter shape is compatible with the current LoRA memory pool configuration. memory_pool = getattr(self, "memory_pool", None) incompatible = memory_pool and not memory_pool.can_support(lora_config) if incompatible: raise ValueError( f"LoRA adapter {lora_ref.lora_name} with rank {lora_config.r} is incompatible with the current " "LoRA memory pool configuration. Please ensure that the LoRA adapter's rank is within the configured " "`--max-lora-rank` and that the target modules are included in `--lora-target-modules`." ) # Ensure pinned LoRA adapters does not exceed maximal limit or cause starvation. if lora_ref.pinned and self.num_pinned_loras >= self.max_loras_per_batch - 1: raise ValueError( f"Failed to load LoRA adapter {lora_ref.lora_name} as a pinned adapter. It is not allowed to pin all slots " "in the LoRA memory pool to avoid starvation for unpinned adapters and base models. Please increase your " "`--max-loras-per-batch` or load it as unpinned LoRA adapters." ) def unload_lora_adapter(self, lora_ref: LoRARef) -> LoRAUpdateOutput: """ Unload LoRA adapters by their names. This will remove the adapters from the memory pool and delete the corresponding LoRA modules. """ adapter = self.configs.get(lora_ref.lora_id) lora_ref = self.lora_refs.get(lora_ref.lora_id) assert ( adapter is not None and lora_ref is not None ), f"LoRA adapter with ID {lora_ref.lora_id} is not loaded. This should have been verified before request is sent to the backend." try: del self.configs[lora_ref.lora_id] del self.loras[lora_ref.lora_id] del self.lora_refs[lora_ref.lora_id] self.num_pinned_loras -= int(lora_ref.pinned) except Exception as e: return self.create_lora_update_result( success=False, error_message=str(e), ) return self.create_lora_update_result(success=True) def validate_lora_batch(self, lora_ids: set[Optional[str]]) -> bool: """ Validate if the LoRA IDs in the batch can be loaded into the current LoRA memory pool. """ if len(lora_ids) > self.max_loras_per_batch: return False # skip pinned LoRA check if no pinned LoRA adapters are loaded. if self.num_pinned_loras == 0: return True # counting the number of pinned LoRA adapters in the batch. pinned_loras_in_batch = 0 for lora_id in lora_ids: if lora_id is not None: lora_ref = self.lora_refs.get(lora_id) assert ( lora_ref is not None ), f"LoRA ID {lora_id} not found in lora_refs." pinned_loras_in_batch += int(lora_ref.pinned) assert pinned_loras_in_batch <= self.num_pinned_loras, ( f"Number of pinned LoRA adapters in the batch ({pinned_loras_in_batch}) exceeds the total number of pinned adapters " f"({self.num_pinned_loras}). This indicates a bug in the LoRA loading logic." ) required_slots = len(lora_ids) - pinned_loras_in_batch mem_pool_vacancy = self.memory_pool.max_loras_per_batch - self.num_pinned_loras return required_slots <= mem_pool_vacancy def fetch_new_loras( self, new_loras: set[Optional[str]], running_loras: set[Optional[str]] = set() ): # Load active loras into lora memory pool cur_uids = new_loras | running_loras assert len(cur_uids) <= self.max_loras_per_batch self.memory_pool.prepare_lora_batch( cur_uids=cur_uids, lora_adapters=self.loras, lora_modules=self.lora_modules, lora_refs=self.lora_refs.copy(), # copy snapshot of current lora_refs to avoid mutation during the batch preparation. lora_embed_tokens_module=self.embed_tokens_module, # merge into embedding or lora module lora_lm_head_module=self.lm_head_module, # merge into embedding or lora module ) def prepare_lora_batch(self, forward_batch: ForwardBatch): # set up batch info shared by all lora modules bs = forward_batch.batch_size use_cuda_graph = ( hasattr(self, "max_bs_in_cuda_graph") and bs <= self.max_bs_in_cuda_graph and forward_batch.forward_mode.is_cuda_graph() ) weight_indices = [0] * len(forward_batch.lora_ids) lora_ranks = [0] * self.max_loras_per_batch scalings = [0] * self.max_loras_per_batch for i, uid in enumerate(forward_batch.lora_ids): if uid not in self.memory_pool.uid_to_buffer_id: continue weight_indices[i] = self.memory_pool.get_buffer_id(uid) if uid is not None: lora = self.loras[uid] lora_ranks[weight_indices[i]] = lora.config.r scalings[weight_indices[i]] = lora.scaling # Do in-place updates when CUDA graph is enabled and the batch forward mode # could use CUDA graph. self.lora_backend.prepare_lora_batch( forward_batch=forward_batch, weight_indices=weight_indices, lora_ranks=lora_ranks, scalings=scalings, use_cuda_graph=use_cuda_graph, ) self.lora_backend.batch_info.has_active_lora = any( lora_ranks[wi] > 0 for wi in weight_indices ) def update_lora_info(self): """ Update all LoRA modules to associate them with the latest memory buffer. """ for layer_id, layer_modules in enumerate(self.lora_modules): for module_name, module in layer_modules.items(): # Hack for FusedMoE layer if isinstance(module, FusedMoEWithLoRA) and all( x in self.target_modules for x in ["gate_up_proj", "down_proj"] ): gate_up_key = ( "gate_up_proj_moe" if "gate_up_proj_moe" in self.memory_pool.A_buffer else "gate_up_proj" ) down_key = ( "down_proj_moe" if "down_proj_moe" in self.memory_pool.A_buffer else "down_proj" ) gate_up_a = self.memory_pool.get_tensor( target_module=gate_up_key, layer_id=layer_id, lora_type=LoRAType.LORA_A, ) gate_up_b = self.memory_pool.get_tensor( target_module=gate_up_key, layer_id=layer_id, lora_type=LoRAType.LORA_B, ) down_a = self.memory_pool.get_tensor( target_module=down_key, layer_id=layer_id, lora_type=LoRAType.LORA_A, ) down_b = self.memory_pool.get_tensor( target_module=down_key, layer_id=layer_id, lora_type=LoRAType.LORA_B, ) module.set_lora_info( gate_up_lora_a_weights=gate_up_a, gate_up_lora_b_weights=gate_up_b, down_lora_a_weights=down_a, down_lora_b_weights=down_b, ) continue target_module = get_target_module_name( module_name, self.memory_pool.target_modules ) module.set_lora_info( self.memory_pool.get_tensor( target_module=target_module, layer_id=layer_id, lora_type=LoRAType.LORA_A, ), self.memory_pool.get_tensor( target_module=target_module, layer_id=layer_id, lora_type=LoRAType.LORA_B, ), ) # Update embedding layer if present - gotta merge (refer to PR codebase) if self.embed_tokens_module is not None: self.embed_tokens_module.set_lora_info( self.memory_pool.get_embedding_tensor("added_tokens", LoRAType.LORA_A), self.memory_pool.get_embedding_tensor("embed_tokens", LoRAType.LORA_A), self.memory_pool.get_embedding_tensor("embed_tokens", LoRAType.LORA_B), ) # Update lm_head layer if present if self.lm_head_module is not None: self.lm_head_module.set_lora_info( self.memory_pool.get_embedding_tensor("lm_head", LoRAType.LORA_A), self.memory_pool.get_embedding_tensor("lm_head", LoRAType.LORA_B), ) def init_state( self, max_lora_rank: Optional[int] = None, target_modules: Optional[Iterable[str]] = None, lora_paths: Optional[List[LoRARef]] = None, ): """ Initialize the internal (mutable) state of the LoRAManager. When `lora_paths` is provided and not empty, it might be used for inferring LoRA shape info such as the target modules and max_lora_rank. """ assert lora_paths or ( max_lora_rank is not None and target_modules is not None ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." self.init_lora_adapters(lora_paths) self.init_lora_shapes( max_lora_rank=max_lora_rank, target_modules=target_modules, ) if self._experts_shared_outer_override is not None: self.experts_shared_outer_loras = self._experts_shared_outer_override else: self.experts_shared_outer_loras = self._detect_shared_outer_loras() if self.experts_shared_outer_loras: logger.info( "Shared outer LoRA mode enabled: gate_up lora_A and " "down lora_B will be shared across experts (expert_dim=1)." ) self.init_lora_modules() self.init_memory_pool() self.update_lora_info() def init_lora_adapters(self, lora_paths: Optional[List[LoRARef]] = None): # Configs of all active LoRA adapters, indexed by LoRA ID. self.configs: Dict[str, LoRAConfig] = {} # LoRA adapter weights cached in CPU memory, indexed by LoRA ID. self.loras: Dict[str, LoRAAdapter] = {} # Mapping from LoRA ID to LoRARef object. self.lora_refs: Dict[str, LoRARef] = {} # Count of pinned LoRA adapters. self.num_pinned_loras: int = 0 if lora_paths: for lora_ref in lora_paths: result = self.load_lora_adapter(lora_ref) if not result.success: raise RuntimeError( f"Failed to load LoRA adapter {lora_ref.lora_name}: {result.error_message}" ) def _detect_shared_outer_loras(self) -> bool: """Auto-detect shared outer LoRA format from loaded adapter weights. MoE adapters with shared outer experts store 3D tensors where dim[0]=1 indicates weights shared across all experts, while dim[0]=num_experts indicates per-expert weights. Returns True if gate_up lora_A has expert_dim=1 (shared). All loaded adapters that expose a 3D gate_up lora_A must agree; mixed formats raise RuntimeError. """ shared_outer: Optional[bool] = None for adapter_id, adapter in self.loras.items(): found = False for layer in adapter.layers: for name, weight in layer.weights.items(): if ( "gate_up_proj" in name and "lora_A" in name and weight.dim() == 3 ): is_shared = weight.shape[0] == 1 if shared_outer is None: shared_outer = is_shared elif shared_outer != is_shared: raise RuntimeError( "Mixed shared-outer LoRA formats detected across " f"loaded adapters (conflict in adapter '{adapter_id}'). " "All MoE adapters must either all use shared outer " "experts (expert_dim=1) or all use per-expert weights." ) found = True break if found: break return bool(shared_outer) if shared_outer is not None else False def init_lora_shapes( self, max_lora_rank: Optional[int] = None, target_modules: Optional[Iterable[str]] = None, ): """Infer LoRA target modules and max_lora_rank from loaded adapters if not provided.""" if target_modules and target_modules == {"all"}: self.target_modules = auto_detect_lora_target_modules(self.base_model) self.target_modules.update(EMBEDDING_NAMES) logger.info( "CLI --lora-target-modules='all' resolved to %s " "by inspecting the base model.", sorted(self.target_modules), ) target_modules = self.target_modules elif target_modules: self.target_modules = get_normalized_target_modules(target_modules) else: self.target_modules = set() for lora_id, config in self.configs.items(): # Handle PEFT shorthand strings like "all-linear" or "all". 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: # 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=" 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 " "suffixes is provided in `target_modules` field of PEFT config. Please explicitly " "specify `--lora-target-modules` during server startup. You can specify `all` to " "enable all support modules types. " ) adapter_target_modules = get_normalized_target_modules( config.target_modules ) if target_modules is not None: # When `--lora-target-modules` is provided, validate adapter target modules is a subset of the specified target modules. if not adapter_target_modules.issubset(self.target_modules): unsupported_modules = adapter_target_modules - self.target_modules lora_name = self.lora_refs[lora_id].lora_name raise ValueError( f"LoRA adapter '{lora_name}' contains target modules {sorted(unsupported_modules)} " f"that are not included in the specified --lora-target-modules {sorted(self.target_modules)}. " f"Please update --lora-target-modules to include all required modules: " f"{sorted(self.target_modules | adapter_target_modules)}, or use 'all' to enable all supported modules." ) else: # Otherwise, infer target_modules from adapter configs. self.target_modules.update(adapter_target_modules) # Fusion folds wk + weights_proj into wk_weights_proj, so the modules # LoRA wraps are absent and an indexer-targeted adapter is silently dropped. indexer_targets = self.target_modules & DSA_INDEXER_LORA_NAMES if indexer_targets: from sglang.srt.layers.attention.dsa.dsa_indexer import ( _use_dsa_indexer_fusion, ) if _use_dsa_indexer_fusion: raise ValueError( f"LoRA targets the DSA indexer ({sorted(indexer_targets)}), which is " "incompatible with DSA indexer Q/K fusion. Set " "SGLANG_DISABLE_DSA_INDEXER_FUSION=1 to disable fusion and use indexer LoRA." ) if max_lora_rank is not None: self.max_lora_rank = max_lora_rank else: self.max_lora_rank = max( [x.r for x in self.configs.values()], default=0, ) # Auto-infer self.lora_added_vocab_size from loaded LoRA configs # This happens automatically without requiring user input # if self.lora_added_vocab_size is None: if self.lora_added_tokens_size is None: inferred_extra_vocab_size = next( ( x.lora_added_tokens_size for x in self.configs.values() if x.lora_added_tokens_size > 0 ), 0, ) if inferred_extra_vocab_size > 0: logger.info( f"self.lora_added_tokens_size={inferred_extra_vocab_size} from LoRA adapters." ) self.lora_added_tokens_size = inferred_extra_vocab_size def load_lora_weights(self, lora_ref: LoRARef): """ Load the weights of a LoRA adapter to CPU memory and conducts post-loading validation. """ lora_adapter = LoRAAdapter( lora_ref.lora_id, self.configs[lora_ref.lora_id], self.base_hf_config, self.load_config, self.lora_backend, base_model=self.base_model, ) lora_adapter.initialize_weights() self.loras[lora_ref.lora_id] = lora_adapter def load_lora_weights_from_tensors( self, lora_ref: LoRARef, tensors: Dict[str, torch.Tensor] ): """ Load the weights of a LoRA adapter from tensors to CPU memory. """ lora_adapter = LoRAAdapter( lora_ref.lora_id, self.configs[lora_ref.lora_id], self.base_hf_config, self.load_config, self.lora_backend, base_model=self.base_model, ) lora_adapter.initialize_weights_from_tensors(tensors) self.loras[lora_ref.lora_id] = lora_adapter def load_lora_adapter_from_tensors( self, lora_ref: LoRARef, tensors: Dict[str, torch.Tensor], config_dict: Dict, added_tokens_config: Optional[Dict] = None, ) -> LoRAUpdateOutput: """ Load a single LoRA adapter from tensors and config dict. """ assert ( lora_ref.lora_name is not None and lora_ref.lora_path is not None ), "LoRARef must have both lora_name and lora_path set for loading." assert ( lora_ref.lora_id not in self.loras ), f"LoRA adapter with ID {lora_ref.lora_id} is already loaded. This should have been verified before request is sent to the backend." try: new_adapter = LoRAConfig.from_dict( config_dict, added_tokens_config, base_vocab_size=self.base_hf_config.vocab_size, ) self.validate_new_adapter(new_adapter, lora_ref) self.configs[lora_ref.lora_id] = new_adapter self.load_lora_weights_from_tensors(lora_ref, tensors) self.lora_refs[lora_ref.lora_id] = lora_ref self.num_pinned_loras += int(lora_ref.pinned) except Exception as e: return self.create_lora_update_result( success=False, error_message=str(e), ) return self.create_lora_update_result(success=True) def init_memory_pool(self): """(Re)initialize the LoRA memory pool based on the current configurations.""" self.memory_pool = LoRAMemoryPool( base_hf_config=self.base_hf_config, max_loras_per_batch=self.max_loras_per_batch, dtype=self.dtype, tp_size=self.tp_size, tp_rank=self.tp_rank, max_lora_rank=self.max_lora_rank, target_modules=self.target_modules, base_model=self.base_model, eviction_policy=self.eviction_policy, lora_added_tokens_size=self.lora_added_tokens_size, experts_shared_outer_loras=self.experts_shared_outer_loras, strict_loading=self.lora_strict_loading, enable_lora_overlap_loading=self.enable_lora_overlap_loading, ) # Initializing memory pool with base model self.fetch_new_loras({None}) def set_lora_module(self, module_name, module): """Wrap any module (standard or MoE) with LoRA support.""" lora_module = get_lora_layer(module, self.lora_backend) replace_submodule(self.base_model, module_name, lora_module) return lora_module def init_lora_modules(self): # Look-up table that essentially maps (layer_index, module_name) to the corresponding LoRA module. self.lora_modules: List[Dict[str, BaseLayerWithLoRA]] = [ {} for _ in range(self.base_hf_config.num_hidden_layers) ] 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(): # Handle embed_tokens and lm_head before the should_apply_lora gate, # since VL models' should_apply_lora patterns only match language # model layers and would incorrectly skip these. # Handle embed_tokens if "embed_tokens" in module_name and "embed_tokens" in self.target_modules: if isinstance(module, VocabParallelEmbedding) and not isinstance( module, BaseLayerWithLoRA ): lora_module = self.set_lora_module(module_name, module) self.embed_tokens_module = lora_module continue # Handle lm_head if "lm_head" in module_name and "lm_head" in self.target_modules: if isinstance(module, ParallelLMHead) and not isinstance( module, BaseLayerWithLoRA ): lora_module = self.set_lora_module(module_name, module) self.lm_head_module = lora_module continue # Handle DeepSeek MLA fused projection: set the boundary # between q_a and kv_a output partitions so the LoRA layer # can apply separate B projections for each. if ( "fused_qkv_a_proj_with_mqa" in self.target_modules and module_name.endswith("fused_qkv_a_proj_with_mqa") ): from sglang.srt.lora.layers import ReplicatedLinearWithLoRA layer_id = get_layer_id(module_name) if layer_id is None: continue lora_module = self.set_lora_module(module_name, module) if isinstance(lora_module, ReplicatedLinearWithLoRA): q_lora_rank = getattr(self.base_hf_config, "q_lora_rank", None) or 0 lora_module.first_output_dim = q_lora_rank self.lora_modules[layer_id][module_name] = lora_module continue # The module should be converted if it is included in target_names parts = module_name.split(".") if ( parts[-1] in self.target_modules or ".".join(parts[-2:]) 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 ) continue if isinstance(module, FusedMoE) and all( x in self.target_modules for x in ["gate_up_proj", "down_proj"] ): layer_id = get_layer_id(module_name) if layer_id is None: # FusedMoE submodules outside the decoder layer hierarchy # (e.g. nested helpers under non-".layers." prefixes) have # no resolvable layer id; skip them so we don't index # `self.lora_modules` with `None`. continue lora_module = self.set_lora_module(module_name, module) lora_module.experts_shared_outer_loras = self.experts_shared_outer_loras lora_module.lora_use_virtual_experts = self.lora_use_virtual_experts self.lora_modules[layer_id][module_name] = lora_module