chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity
Files
9c5aae4df55236214a72d9b60db2250c85583e27
sglang/python/sglang/srt/lora/lora_manager.py

640 lines
27 KiB
Python
Raw Blame History

# 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.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, 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 (
LoRAType,
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
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
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
)
# Store eviction policy from server args
self.eviction_policy = server_args.lora_eviction_policy
# 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
):
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,
)
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)
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"LoRA serving currently doesn't support adapters that add tokens to the vocabulary"
)
# 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):
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,
)
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():
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,
)
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 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."""
self.target_modules = (
get_normalized_target_modules(target_modules) if target_modules else set()
)
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 "
"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)
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,
)
lora_adapter.initialize_weights()
# If we want to overlap loading LoRA adapters with compute, they must be pinned in CPU memory
if self.enable_lora_overlap_loading:
lora_adapter.pin_weights_in_cpu()
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,
)
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)
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,
)
# Initializing memory pool with base model
self.fetch_new_loras({None})
def set_lora_module(self, module_name, module):
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():
# TODO (lifuhuang): in the future, we should consider generalizing the
# should_apply_lora function to support mapping by full module name instead
# of just the last part (e.g., "qkv_proj") to support scenarios with multiple
# attention stacks (e.g., multimodal models).
# See: https://github.com/sgl-project/sglang/issues/6608
if getattr(
self.base_model, "should_apply_lora", None
) and not self.base_model.should_apply_lora(module_name):
continue
# 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
# 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)
self.lora_modules[layer_id][module_name] = self.set_lora_module(
module_name, module
)
Reference in New Issue View Git Blame Copy Permalink