# 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" # LoRA layers class inheritance adapted from: # https://github.com/vllm-project/vllm/blob/4abf6336ec65c270343eb895e7b18786e9274176/vllm/lora/layers.py import logging from typing import Dict, List import torch from torch import nn from sglang.srt.configs.load_config import LoadConfig from sglang.srt.layers.utils import get_layer_id from sglang.srt.lora.backend.base_backend import BaseLoRABackend from sglang.srt.lora.backend.lora_registry import LORA_SUPPORTED_BACKENDS from sglang.srt.lora.lora_config import LoRAConfig from sglang.srt.model_loader.loader import DefaultModelLoader from sglang.srt.utils.hf_transformers_utils import AutoConfig logger = logging.getLogger(__name__) class LoRALayer(nn.Module): def __init__(self, config: LoRAConfig, base_hf_config: AutoConfig): super().__init__() self.config: LoRAConfig = config self.base_hf_config: AutoConfig = base_hf_config # lora weights in cpu. The weights are loaded from checkpoint. self.weights: Dict[str, torch.Tensor] = {} class LoRAAdapter(nn.Module): def __init__( self, uid: str, config: LoRAConfig, base_hf_config: AutoConfig, load_config: LoadConfig, lora_backend: BaseLoRABackend, ): super().__init__() self.uid: str = uid self.config: LoRAConfig = config assert self.config.hf_config["peft_type"].lower() == "lora" self.base_hf_config: AutoConfig = base_hf_config self.load_config: LoadConfig = load_config self.lora_backend: BaseLoRABackend = lora_backend self.scaling: float = self.config.lora_alpha / self.config.r self.layers: List[LoRALayer] = nn.ModuleList( [ LoRALayer(config, base_hf_config) for _ in range(base_hf_config.num_hidden_layers) ] ) self.embedding_layers: Dict[str, torch.Tensor] = {} self.added_tokens_embeddings: Dict[str, torch.Tensor] = {} def initialize_weights(self): model_path = self.config.path loader = DefaultModelLoader(self.load_config) revision = getattr(self.config.hf_config, "revision", None) # Get normalized target modules for filtering for name, loaded_weight in loader._get_weights_iterator( DefaultModelLoader.Source( model_path, revision=revision, fall_back_to_pt=True ) ): self._process_weight(name, loaded_weight) self._normalize_weights() def initialize_weights_from_tensors(self, tensors: Dict[str, torch.Tensor]): for name, tensor in tensors.items(): self._process_weight(name, tensor) self._normalize_weights() def _process_weight(self, name: str, loaded_weight: torch.Tensor): from sglang.srt.lora.utils import get_normalized_target_modules normalized_target_modules = get_normalized_target_modules( 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. # 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 ( "all" in normalized_target_modules or module_name in normalized_target_modules ): self.embedding_layers[name] = loaded_weight.cpu() else: logger.debug( f"Skipping {name} as '{module_name}' is not in adapter's target_modules: {self.config.target_modules}" ) elif "input_embeddings" in name or "output_embeddings" in name: # added/extra token emb self.added_tokens_embeddings[name] = loaded_weight.cpu() assert loaded_weight.shape[0] == self.config.lora_added_tokens_size, ( f"LoRA adapter {self.uid} has extra_vocab_size {self.config.extra_vocab_size} specified in the config, " f"but the loaded weight has {loaded_weight.shape[0]} extra vocab size" ) def _normalize_weights(self): # normalize kv_proj and gate_up_proj for layer in self.layers: weight_names = list(layer.weights.keys()) self.normalize_qkv_proj(weight_names, layer.weights) self.normalize_gate_up_proj(weight_names, layer.weights) def normalize_qkv_proj( self, weight_names: List[str], weights: Dict[str, torch.Tensor] ): # Collect target q/k/v modules. This process is necessary since there might be no lora attached to k_proj target_module = set() for weight_name in weight_names: if "k_proj" in weight_name: target_module.add("k_proj") if "q_proj" in weight_name: target_module.add("q_proj") if "v_proj" in weight_name: target_module.add("v_proj") if "qkv_proj" in weight_name: target_module.add("qkv_proj") if len(target_module) == 0: return for weight_name in weight_names: # We assume every lora adaptor should contain lora modules for q_proj if "q_proj" in weight_name: q_name = weight_name k_name = weight_name.replace("q_proj", "k_proj") v_name = weight_name.replace("q_proj", "v_proj") qkv_name = weight_name.replace("q_proj", "qkv_proj") # If k_proj doesn't have lora, initialize it to zero k_proj_weight = ( weights[k_name] if "k_proj" in target_module else torch.zeros_like(weights[v_name]) ) weights[qkv_name] = torch.cat( ( weights[q_name], k_proj_weight, weights[v_name], ), 0, ) weights.pop(q_name) if "k_proj" in target_module: weights.pop(k_name) weights.pop(v_name) elif "qkv_proj" in weight_name: # If qkv_proj is already stacked, we normalize it following the SGL convention. qkv_name = weight_name q_name = weight_name.replace("qkv_proj", "q_proj") k_name = weight_name.replace("qkv_proj", "k_proj") v_name = weight_name.replace("qkv_proj", "v_proj") if "lora_A" in weight_name: weights[qkv_name] = weights[qkv_name].repeat(3, 1) # else: no-op as LoRA B weight is already stacked. def normalize_gate_up_proj( self, weight_names: List[str], weights: Dict[str, torch.Tensor] ): for weight_name in weight_names: if "gate_proj" in weight_name: up_name = weight_name.replace("gate_proj", "up_proj") gate_up_name = weight_name.replace("gate_proj", "gate_up_proj") if up_name not in weights: weights[up_name] = torch.zeros_like(weights[weight_name]) assert self.lora_backend.name in LORA_SUPPORTED_BACKENDS, ( f"LoRA weight initialization currently only supported for LoRA backends: {', '.join(b for b in LORA_SUPPORTED_BACKENDS)}" f"Received backend: {self.lora_backend.name}. Please verify your backend configuration " f"or consider implementing custom initialization logic for other backends." ) weights[gate_up_name] = torch.cat( (weights[weight_name], weights[up_name]), 0 ) weights.pop(weight_name) if up_name in weights: weights.pop(up_name) elif "gate_up_proj" in weight_name: # If gate_up_proj is already stacked, we normalize it following the SGL convention gate_up_name = weight_name if "lora_A" in weight_name: weights[gate_up_name] = weights[gate_up_name].repeat(2, 1) # else: no-op as LoRA B weight is already stacked. def pin_weights_in_cpu(self): for layer in self.layers: for name, weight in layer.weights.items(): layer.weights[name] = weight.pin_memory() for name, weight in self.embedding_layers.items(): self.embedding_layers[name] = weight.pin_memory() for name, weight in self.added_tokens_embeddings.items(): self.added_tokens_embeddings[name] = weight.pin_memory()