Thanaji Rao Thakkalapelli
1788d76eb3
Signed-off-by: Thanaji Rao Thakkalapelli <thanaji.rao.thakkalapelli@intel.com>
734 lines
31 KiB
Python
734 lines
31 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import math
|
|
import time
|
|
from collections import OrderedDict
|
|
from typing import get_args
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from vllm.config.lora import MaxLoRARanks
|
|
from vllm.logger import init_logger
|
|
from vllm.lora.layers import BaseLayerWithLoRA
|
|
from vllm.lora.lora_model import LoRAModel
|
|
from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
|
|
from vllm.lora.peft_helper import PEFTHelper
|
|
from vllm.lora.request import LoRARequest
|
|
from vllm.lora.utils import (
|
|
get_adapter_absolute_path,
|
|
get_supported_lora_modules,
|
|
replace_submodule,
|
|
)
|
|
from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear
|
|
|
|
from vllm_omni.config.lora import LoRAConfig
|
|
from vllm_omni.diffusion.lora.utils import (
|
|
_expand_expected_modules_for_packed_layers,
|
|
_match_target_modules,
|
|
from_layer_diffusion,
|
|
)
|
|
from vllm_omni.lora.utils import stable_lora_int_id
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class DiffusionLoRAManager:
|
|
"""Manager for LoRA adapters in diffusion models.
|
|
|
|
Reuses vLLM's LoRA infrastructure, adapted for diffusion pipelines.
|
|
Uses LRU cache management similar to LRUCacheLoRAModelManager.
|
|
"""
|
|
|
|
# Valid max allowed ranks for LoRA in vLLM
|
|
_VALID_MAX_RANKS: list[int] = sorted(get_args(MaxLoRARanks))
|
|
|
|
def __init__(
|
|
self,
|
|
pipeline: nn.Module,
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
max_cached_adapters: int = 1,
|
|
lora_path: str | None = None,
|
|
lora_scale: float = 1.0,
|
|
):
|
|
"""
|
|
Initialize the DiffusionLoRAManager.
|
|
|
|
Args:
|
|
max_cached_adapters: Maximum number of LoRA adapters to keep in the
|
|
CPU-side cache (LRU). This mirrors vLLM's `max_cpu_loras` and is
|
|
exposed to users via `OmniDiffusionConfig.max_cpu_loras`.
|
|
"""
|
|
self.pipeline = pipeline
|
|
self.device = device
|
|
self.dtype = dtype
|
|
|
|
# Cache supported/expected module suffixes once, before any layer
|
|
# replacement happens. After LoRA layers are injected, the original
|
|
# LinearBase layers become submodules named "*.base_layer", and calling
|
|
# vLLM's get_supported_lora_modules() again would incorrectly yield
|
|
# "base_layer" instead of the real target module suffixes.
|
|
self._supported_lora_modules = self._compute_supported_lora_modules()
|
|
self._packed_modules_mapping = self._compute_packed_modules_mapping()
|
|
self._expected_lora_modules = _expand_expected_modules_for_packed_layers(
|
|
self._supported_lora_modules,
|
|
self._packed_modules_mapping,
|
|
)
|
|
|
|
# LRU-style cache management
|
|
self.max_cached_adapters = max_cached_adapters # max_cpu_loras
|
|
self._registered_adapters: dict[int, LoRAModel] = {} # adapter_id -> LoRAModel
|
|
self._active_adapter_id: int | None = None
|
|
self._adapter_scales: dict[int, float] = {} # adapter_id -> external scale
|
|
|
|
# LRU cache tracking (adapter_id -> last_used_time)
|
|
self._adapter_access_order: OrderedDict[int, float] = OrderedDict()
|
|
# Pinned adapters are not evicted
|
|
self._pinned_adapters: set[int] = set()
|
|
|
|
# track replaced modules
|
|
# key: full module name (component.module.path); value: LoRA layer
|
|
self._lora_modules: dict[str, BaseLayerWithLoRA] = {}
|
|
# Track the maximum LoRA rank we've allocated buffers for.
|
|
self._max_lora_rank: int = 0
|
|
|
|
logger.info(
|
|
"Initializing DiffusionLoRAManager: device=%s, dtype=%s, max_cached_adapters=%d, static_lora_path=%s",
|
|
device,
|
|
dtype,
|
|
max_cached_adapters,
|
|
lora_path,
|
|
)
|
|
|
|
if lora_path is not None:
|
|
logger.info("Loading LoRA during initialization from %s with scale %.2f", lora_path, lora_scale)
|
|
init_request = LoRARequest(
|
|
lora_name="static",
|
|
lora_int_id=stable_lora_int_id(lora_path),
|
|
lora_path=lora_path,
|
|
)
|
|
self.set_active_adapter(init_request, lora_scale)
|
|
|
|
def _compute_supported_lora_modules(self) -> set[str]:
|
|
"""Compute supported LoRA module suffixes for this pipeline.
|
|
|
|
vLLM's get_supported_lora_modules() returns suffixes for LinearBase
|
|
modules. After this manager replaces layers with BaseLayerWithLoRA
|
|
wrappers, those LinearBase modules become nested under ".base_layer",
|
|
which would cause get_supported_lora_modules() to return "base_layer".
|
|
To make adapter loading stable across multiple adapters, we also accept
|
|
suffixes from existing BaseLayerWithLoRA wrappers and drop "base_layer"
|
|
when appropriate.
|
|
"""
|
|
supported = set(get_supported_lora_modules(self.pipeline))
|
|
|
|
has_lora_wrappers = False
|
|
for name, module in self.pipeline.named_modules():
|
|
if isinstance(module, BaseLayerWithLoRA):
|
|
has_lora_wrappers = True
|
|
supported.add(name.split(".")[-1])
|
|
|
|
if has_lora_wrappers:
|
|
supported.discard("base_layer")
|
|
|
|
return supported
|
|
|
|
def _compute_packed_modules_mapping(self) -> dict[str, list[str]]:
|
|
"""Collect packed->sublayer mappings from the diffusion model.
|
|
|
|
Diffusion models often use packed (fused) projections like `to_qkv` or
|
|
`w13`, while LoRA checkpoints are typically saved against the logical
|
|
sub-projections (e.g. `to_q`/`to_k`/`to_v`, `w1`/`w3`). Many diffusion
|
|
model implementations already define these relationships in
|
|
`load_weights()` via `stacked_params_mapping`. To avoid duplicating the
|
|
mapping in multiple places, we derive packed→sublayer mappings from the
|
|
model's `stacked_params_mapping`.
|
|
"""
|
|
|
|
def _derive_from_stacked_params_mapping(stacked: object) -> dict[str, list[str]]:
|
|
if not isinstance(stacked, (list, tuple)):
|
|
return {}
|
|
derived: dict[str, list[str]] = {}
|
|
for item in stacked:
|
|
if not isinstance(item, (list, tuple)) or len(item) < 2:
|
|
continue
|
|
packed_suffix, sub_suffix = item[0], item[1]
|
|
if not isinstance(packed_suffix, str) or not packed_suffix:
|
|
continue
|
|
if not isinstance(sub_suffix, str) or not sub_suffix:
|
|
continue
|
|
# The mapping strings are usually suffix patterns (e.g. ".to_qkv"),
|
|
# but some models scope them under submodules (e.g. ".attn1.to_qkv").
|
|
# For LoRA we only care about the leaf module names.
|
|
packed_name = packed_suffix.strip(".").split(".")[-1]
|
|
sub_name = sub_suffix.strip(".").split(".")[-1]
|
|
existing = derived.get(packed_name)
|
|
if existing is None:
|
|
derived[packed_name] = [sub_name]
|
|
elif sub_name not in existing:
|
|
existing.append(sub_name)
|
|
return derived
|
|
|
|
mapping: dict[str, list[str]] = {}
|
|
for module in self.pipeline.modules():
|
|
derived = _derive_from_stacked_params_mapping(getattr(module, "stacked_params_mapping", None))
|
|
for packed_name, sub_names in derived.items():
|
|
if not isinstance(packed_name, str) or not packed_name:
|
|
continue
|
|
if not isinstance(sub_names, (list, tuple)) or not all(isinstance(s, str) for s in sub_names):
|
|
continue
|
|
sub_names_list = list(sub_names)
|
|
if not sub_names_list:
|
|
continue
|
|
|
|
existing = mapping.get(packed_name)
|
|
if existing is None:
|
|
mapping[packed_name] = sub_names_list
|
|
elif existing != sub_names_list:
|
|
logger.warning(
|
|
"Conflicting packed module mapping for %s: %s vs %s; using %s",
|
|
packed_name,
|
|
existing,
|
|
sub_names_list,
|
|
existing,
|
|
)
|
|
|
|
return mapping
|
|
|
|
def _get_packed_sublayer_suffixes(self, packed_module_suffix: str, n_slices: int) -> list[str] | None:
|
|
sub_suffixes = self._packed_modules_mapping.get(packed_module_suffix)
|
|
if not sub_suffixes:
|
|
return None
|
|
if len(sub_suffixes) != n_slices:
|
|
logger.warning(
|
|
"Packed module mapping[%s] has %d slices but layer expects %d; skipping sublayer lookup",
|
|
packed_module_suffix,
|
|
len(sub_suffixes),
|
|
n_slices,
|
|
)
|
|
return None
|
|
return sub_suffixes
|
|
|
|
def set_active_adapter(self, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> None:
|
|
"""Set the active LoRA adapter for the pipeline.
|
|
|
|
Args:
|
|
lora_request: The LoRA request, or None to deactivate all adapters.
|
|
lora_scale: The external scale for the LoRA adapter.
|
|
"""
|
|
if lora_request is None:
|
|
if self._active_adapter_id is None:
|
|
logger.debug("No lora_request provided and adapters are already inactive")
|
|
return
|
|
logger.debug("No lora_request provided, deactivating all LoRA adapters")
|
|
self._deactivate_all_adapters()
|
|
return
|
|
elif math.isclose(0.0, lora_scale):
|
|
if self._active_adapter_id is None:
|
|
logger.debug("Received LoRA scale 0 with adapters already inactive")
|
|
return
|
|
logger.warning("Received a request with LoRA scale 0; deactivating all LoRA adapters")
|
|
self._deactivate_all_adapters()
|
|
return
|
|
|
|
adapter_id = lora_request.lora_int_id
|
|
logger.debug(
|
|
"Setting active adapter: id=%d, name=%s, path=%s, scale=%.2f, cache_size=%d/%d",
|
|
adapter_id,
|
|
lora_request.lora_name,
|
|
lora_request.lora_path,
|
|
lora_scale,
|
|
len(self._registered_adapters),
|
|
self.max_cached_adapters,
|
|
)
|
|
if adapter_id not in self._registered_adapters:
|
|
logger.info("Loading new adapter: id=%d, name=%s", adapter_id, lora_request.lora_name)
|
|
# Add the adapter + add to the cache
|
|
self.add_adapter(lora_request)
|
|
else:
|
|
# Just touch the cache access order
|
|
self._touch_adapter_info(adapter_id)
|
|
|
|
self._activate_adapter(adapter_id, lora_scale)
|
|
|
|
def _touch_adapter_info(self, adapter_id):
|
|
"""Update the current caching ordering info."""
|
|
self._adapter_access_order[adapter_id] = time.time()
|
|
self._adapter_access_order.move_to_end(adapter_id)
|
|
|
|
def _update_adapter_scale(self, adapter_id: int, lora_scale: float):
|
|
"""Update the adapter scale for a given adapter ID. To avoid potential
|
|
issues with using Floats as keys, for now, we round float values to
|
|
3 decimal points.
|
|
"""
|
|
scale = DiffusionLoRAManager._get_rounded_scale(lora_scale)
|
|
self._adapter_scales[adapter_id] = scale
|
|
|
|
@staticmethod
|
|
def _get_rounded_scale(lora_scale: float):
|
|
"""Normalizes a lora scale for use as a key in the _adapter_scales
|
|
dict; for now we just round scales to 3 decimal places.
|
|
"""
|
|
return round(lora_scale, 3)
|
|
|
|
def _load_adapter(
|
|
self,
|
|
lora_request: LoRARequest,
|
|
) -> tuple[LoRAModel, PEFTHelper]:
|
|
if not self._expected_lora_modules:
|
|
raise ValueError("No supported LoRA modules found in the diffusion pipeline.")
|
|
|
|
logger.debug("Supported LoRA modules: %s", self._expected_lora_modules)
|
|
|
|
lora_path = get_adapter_absolute_path(lora_request.lora_path)
|
|
logger.debug("Resolved LoRA path: %s", lora_path)
|
|
|
|
peft_helper = PEFTHelper.from_local_dir(
|
|
lora_path,
|
|
max_position_embeddings=None, # no need in diffusion
|
|
tensorizer_config_dict=lora_request.tensorizer_config_dict,
|
|
)
|
|
|
|
logger.info(
|
|
"Loaded PEFT config: r=%d, lora_alpha=%d, target_modules=%s",
|
|
peft_helper.r,
|
|
peft_helper.lora_alpha,
|
|
peft_helper.target_modules,
|
|
)
|
|
|
|
lora_model = LoRAModel.from_local_checkpoint(
|
|
lora_path,
|
|
expected_lora_modules=self._expected_lora_modules,
|
|
peft_helper=peft_helper,
|
|
lora_model_id=lora_request.lora_int_id,
|
|
device="cpu", # consistent w/ vllm's behavior
|
|
dtype=self.dtype,
|
|
model_vocab_size=None,
|
|
tensorizer_config_dict=lora_request.tensorizer_config_dict,
|
|
weights_mapper=None,
|
|
)
|
|
|
|
logger.info(
|
|
"Loaded LoRA model: id=%d, num_modules=%d, modules=%s",
|
|
lora_model.id,
|
|
len(lora_model.loras),
|
|
list(lora_model.loras.keys()),
|
|
)
|
|
|
|
for lora in lora_model.loras.values():
|
|
lora.optimize() # ref: _create_merged_loras_inplace, internal scaling
|
|
|
|
return lora_model, peft_helper
|
|
|
|
def _get_packed_modules_list(self, module: nn.Module) -> list[str]:
|
|
"""Return a packed_modules_list suitable for vLLM LoRA can_replace_layer().
|
|
|
|
Diffusion transformers frequently use packed projection layers like
|
|
QKVParallelLinear (fused QKV). vLLM's LoRA replacement logic relies on
|
|
`packed_modules_list` length to decide between single-slice vs packed
|
|
LoRA layer implementations.
|
|
"""
|
|
if isinstance(module, QKVParallelLinear):
|
|
# Treat diffusion QKV as a 3-slice packed projection by default.
|
|
return ["q", "k", "v"]
|
|
if isinstance(module, MergedColumnParallelLinear):
|
|
# 2-slice packed projection (e.g. fused MLP projections).
|
|
return ["0", "1"]
|
|
return []
|
|
|
|
def _replace_layers_with_lora(self, peft_helper: PEFTHelper) -> None:
|
|
self._ensure_max_lora_rank(peft_helper.r)
|
|
|
|
target_modules = getattr(peft_helper, "target_modules", None)
|
|
target_modules_list: list[str] | None = None
|
|
target_modules_pattern: str | None = None
|
|
if isinstance(target_modules, str) and target_modules:
|
|
target_modules_pattern = target_modules
|
|
elif isinstance(target_modules, list) and target_modules:
|
|
target_modules_list = target_modules
|
|
|
|
def _matches_target(module_name: str) -> bool:
|
|
if target_modules_pattern is not None:
|
|
import regex as re
|
|
|
|
return re.search(target_modules_pattern, module_name) is not None
|
|
if target_modules_list is None:
|
|
return True
|
|
return _match_target_modules(module_name, target_modules_list)
|
|
|
|
# dummy lora config
|
|
lora_config = LoRAConfig(
|
|
max_lora_rank=self._max_lora_rank,
|
|
max_loras=1,
|
|
max_cpu_loras=self.max_cached_adapters,
|
|
lora_dtype=self.dtype,
|
|
fully_sharded_loras=False,
|
|
)
|
|
|
|
# Default denoising components plus any a pipeline opts into via
|
|
# ``_lora_components`` (opt-in; other pipelines unchanged).
|
|
#
|
|
# NOTE: SDXL-style pipelines expose the denoiser as ``unet``.
|
|
# Without scanning this component, adapters can load/activate while
|
|
# effectively applying to zero layers, producing base-identical output.
|
|
default_components = (
|
|
"transformer",
|
|
"transformer_2",
|
|
"dit",
|
|
"bagel",
|
|
"unet",
|
|
)
|
|
extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ())
|
|
for component_name in (*default_components, *extra_components):
|
|
if not hasattr(self.pipeline, component_name):
|
|
continue
|
|
component = getattr(self.pipeline, component_name)
|
|
if not isinstance(component, nn.Module):
|
|
continue
|
|
|
|
# Collect replacements first to avoid mutating the module tree
|
|
# while iterating over named_modules().
|
|
pending_replacements: list[tuple[str, str, nn.Module, list[str]]] = []
|
|
|
|
for module_name, module in component.named_modules(remove_duplicate=False):
|
|
# Don't recurse into already-replaced LoRA wrappers. Their
|
|
# original LinearBase lives under "base_layer", and replacing
|
|
# that again would nest LoRA wrappers and break execution.
|
|
if isinstance(module, BaseLayerWithLoRA) or "base_layer" in module_name.split("."):
|
|
continue
|
|
|
|
full_module_name = f"{component_name}.{module_name}"
|
|
if full_module_name in self._lora_modules:
|
|
logger.debug("Layer %s already replaced, skipping", full_module_name)
|
|
continue
|
|
|
|
packed_modules_list = self._get_packed_modules_list(module)
|
|
if target_modules_pattern is not None or target_modules_list is not None:
|
|
should_replace = _matches_target(full_module_name)
|
|
if not should_replace and len(packed_modules_list) > 1:
|
|
prefix, _, packed_suffix = full_module_name.rpartition(".")
|
|
sub_suffixes = self._get_packed_sublayer_suffixes(packed_suffix, len(packed_modules_list))
|
|
if sub_suffixes is not None:
|
|
for sub_suffix in sub_suffixes:
|
|
sub_full_name = f"{prefix}.{sub_suffix}" if prefix else sub_suffix
|
|
if _matches_target(sub_full_name):
|
|
should_replace = True
|
|
break
|
|
|
|
if not should_replace:
|
|
continue
|
|
|
|
pending_replacements.append((module_name, full_module_name, module, packed_modules_list))
|
|
|
|
for module_name, full_module_name, module, packed_modules_list in pending_replacements:
|
|
lora_layer = from_layer_diffusion(
|
|
layer=module,
|
|
max_loras=1,
|
|
lora_config=lora_config,
|
|
packed_modules_list=packed_modules_list,
|
|
model_config=None,
|
|
)
|
|
|
|
if lora_layer is not module and isinstance(lora_layer, BaseLayerWithLoRA):
|
|
replace_submodule(component, module_name, lora_layer)
|
|
self._lora_modules[full_module_name] = lora_layer
|
|
logger.debug("Replaced layer: %s -> %s", full_module_name, type(lora_layer).__name__)
|
|
|
|
def _ensure_max_lora_rank(self, min_rank: int) -> None:
|
|
"""Ensure LoRA buffers can accommodate adapters up to `min_rank`.
|
|
|
|
We allocate per-layer LoRA buffers once when we first replace layers.
|
|
If a later adapter has a larger rank, we need to reinitialize those
|
|
buffers and re-apply the currently active adapter.
|
|
"""
|
|
if min_rank <= self._max_lora_rank:
|
|
return
|
|
|
|
valid_max_rank = self._get_smallest_valid_max_rank(min_rank)
|
|
|
|
logger.info("Increasing max LoRA rank: %d -> %d", self._max_lora_rank, valid_max_rank)
|
|
self._max_lora_rank = valid_max_rank
|
|
|
|
if not self._lora_modules:
|
|
return
|
|
|
|
lora_config = LoRAConfig(
|
|
max_lora_rank=self._max_lora_rank,
|
|
max_loras=1,
|
|
max_cpu_loras=self.max_cached_adapters,
|
|
lora_dtype=self.dtype,
|
|
fully_sharded_loras=False,
|
|
)
|
|
|
|
# Recreate per-layer buffers with the new maximum rank.
|
|
for lora_layer in self._lora_modules.values():
|
|
lora_layer.create_lora_weights(max_loras=1, lora_config=lora_config, model_config=None)
|
|
|
|
# Re-apply active adapter if needed (buffers were reset).
|
|
if self._active_adapter_id is not None:
|
|
active_id = self._active_adapter_id
|
|
active_scale = self._adapter_scales[active_id]
|
|
self._active_adapter_id = None
|
|
self._activate_adapter(active_id, active_scale)
|
|
|
|
@classmethod
|
|
def _get_smallest_valid_max_rank(cls, min_rank: int) -> int:
|
|
"""Given a LoRA rank, get the smallest max rank that can support it."""
|
|
if min_rank <= 0:
|
|
raise ValueError(f"Invalid LoRA rank: {min_rank}")
|
|
|
|
allowed_ranks = [rank for rank in cls._VALID_MAX_RANKS if rank >= min_rank]
|
|
if not allowed_ranks:
|
|
raise ValueError(f"LoRA rank of {min_rank} exceeds max allowed rank of {max(cls._VALID_MAX_RANKS)}")
|
|
|
|
return min(allowed_ranks)
|
|
|
|
def _get_lora_weights(
|
|
self,
|
|
lora_model: LoRAModel,
|
|
full_module_name: str,
|
|
) -> LoRALayerWeights | PackedLoRALayerWeights | None:
|
|
"""Best-effort lookup for LoRA weights by name.
|
|
|
|
Tries:
|
|
- Full module name (e.g. transformer.blocks.0.attn.to_qkv)
|
|
- Relative name without the top-level component (e.g. blocks.0.attn.to_qkv)
|
|
- Suffix-only name (e.g. to_qkv)
|
|
"""
|
|
lora_weights = lora_model.get_lora(full_module_name)
|
|
if lora_weights is not None:
|
|
return lora_weights
|
|
|
|
component_relative_name = full_module_name.split(".", 1)[-1] if "." in full_module_name else full_module_name
|
|
lora_weights = lora_model.get_lora(component_relative_name)
|
|
if lora_weights is not None:
|
|
return lora_weights
|
|
|
|
module_suffix = full_module_name.split(".")[-1]
|
|
return lora_model.get_lora(module_suffix)
|
|
|
|
def _is_active_at_scale(self, adapter_id: int, scale: float) -> bool:
|
|
"""True if the adapter_id is active and the current scale matches."""
|
|
rounded_scale = DiffusionLoRAManager._get_rounded_scale(scale)
|
|
is_active = self._active_adapter_id == adapter_id
|
|
matches_scale = self._adapter_scales.get(adapter_id) == rounded_scale
|
|
return is_active and matches_scale
|
|
|
|
def _activate_adapter(self, adapter_id: int, scale: float) -> None:
|
|
if self._is_active_at_scale(adapter_id, scale):
|
|
logger.debug("Adapter %d already active at scale %.3f skipping", adapter_id, scale)
|
|
return
|
|
|
|
logger.info("Activating adapter: id=%d", adapter_id)
|
|
lora_model = self._registered_adapters[adapter_id]
|
|
|
|
# activate weights in each LoRA layer
|
|
for full_module_name, lora_layer in self._lora_modules.items():
|
|
lora_weights = self._get_lora_weights(lora_model, full_module_name)
|
|
|
|
if lora_weights is None:
|
|
n_slices = getattr(lora_layer, "n_slices", 1)
|
|
if n_slices > 1:
|
|
prefix, _, packed_suffix = full_module_name.rpartition(".")
|
|
sub_suffixes = self._get_packed_sublayer_suffixes(packed_suffix, n_slices)
|
|
if sub_suffixes is None:
|
|
lora_layer.reset_lora(0)
|
|
continue
|
|
|
|
sub_loras: list[LoRALayerWeights | None] = []
|
|
any_found = False
|
|
for sub_suffix in sub_suffixes:
|
|
sub_full_name = f"{prefix}.{sub_suffix}" if prefix else sub_suffix
|
|
sub_lora = self._get_lora_weights(lora_model, sub_full_name)
|
|
if sub_lora is not None:
|
|
any_found = True
|
|
# Packed layers expect plain (non-packed) subloras.
|
|
if isinstance(sub_lora, PackedLoRALayerWeights):
|
|
sub_lora = None
|
|
sub_loras.append(sub_lora if isinstance(sub_lora, LoRALayerWeights) else None)
|
|
|
|
if not any_found:
|
|
lora_layer.reset_lora(0)
|
|
continue
|
|
|
|
lora_a_list: list[torch.Tensor | None] = []
|
|
lora_b_list: list[torch.Tensor | None] = []
|
|
for sub_lora in sub_loras:
|
|
if sub_lora is None:
|
|
lora_a_list.append(None)
|
|
lora_b_list.append(None)
|
|
continue
|
|
lora_a_list.append(sub_lora.lora_a)
|
|
lora_b_list.append(sub_lora.lora_b * scale)
|
|
|
|
lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
|
|
logger.debug(
|
|
"Activated packed LoRA for %s via submodules=%s (scale=%.2f)",
|
|
full_module_name,
|
|
sub_suffixes,
|
|
scale,
|
|
)
|
|
else:
|
|
lora_layer.reset_lora(0)
|
|
continue
|
|
|
|
# Packed LoRA weights already provide per-slice tensors.
|
|
if isinstance(lora_weights, PackedLoRALayerWeights):
|
|
lora_a_list = lora_weights.lora_a
|
|
lora_b_list = [
|
|
None if b is None else b * scale # type: ignore[operator]
|
|
for b in lora_weights.lora_b
|
|
]
|
|
lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
|
|
logger.debug(
|
|
"Activated packed LoRA for %s (scale=%.2f)",
|
|
full_module_name,
|
|
scale,
|
|
)
|
|
continue
|
|
|
|
# Fused (non-packed) weights: if the layer is multi-slice, split B.
|
|
n_slices = getattr(lora_layer, "n_slices", 1)
|
|
if n_slices > 1:
|
|
output_slices = getattr(lora_layer, "output_slices", None)
|
|
if output_slices is None:
|
|
lora_layer.reset_lora(0)
|
|
continue
|
|
|
|
total = sum(output_slices)
|
|
if lora_weights.lora_b.shape[0] != total:
|
|
logger.warning(
|
|
"Skipping LoRA for %s due to shape mismatch: lora_b[0]=%d != sum(output_slices)=%d",
|
|
full_module_name,
|
|
lora_weights.lora_b.shape[0],
|
|
total,
|
|
)
|
|
lora_layer.reset_lora(0)
|
|
continue
|
|
|
|
b_splits = list(torch.split(lora_weights.lora_b, list(output_slices), dim=0))
|
|
lora_a_list = [lora_weights.lora_a] * n_slices
|
|
lora_b_list = [b * scale for b in b_splits]
|
|
lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)
|
|
logger.debug(
|
|
"Activated fused LoRA for packed layer %s (scale=%.2f)",
|
|
full_module_name,
|
|
scale,
|
|
)
|
|
continue
|
|
|
|
scaled_lora_b = lora_weights.lora_b * scale
|
|
lora_layer.set_lora(index=0, lora_a=lora_weights.lora_a, lora_b=scaled_lora_b)
|
|
logger.debug(
|
|
"Activated LoRA for %s: lora_a shape=%s, lora_b shape=%s, scale=%.2f",
|
|
full_module_name,
|
|
lora_weights.lora_a.shape,
|
|
lora_weights.lora_b.shape,
|
|
scale,
|
|
)
|
|
|
|
self._active_adapter_id = adapter_id
|
|
self._update_adapter_scale(adapter_id, scale)
|
|
|
|
def _deactivate_all_adapters(self) -> None:
|
|
if self._active_adapter_id is None:
|
|
logger.debug("All adapters already inactive")
|
|
return
|
|
logger.info("Deactivating all adapters: %d layers", len(self._lora_modules))
|
|
for lora_layer in self._lora_modules.values():
|
|
lora_layer.reset_lora(0)
|
|
self._active_adapter_id = None
|
|
logger.debug("All adapters deactivated")
|
|
|
|
def _evict_for_new_adapter(self) -> None:
|
|
"""Evict unpinned registered adapters until we have room for a new
|
|
adapter to be loaded."""
|
|
while len(self._registered_adapters) > (self.max_cached_adapters - 1):
|
|
# Pick LRU among non-pinned adapters
|
|
evict_candidates = [aid for aid in self._adapter_access_order.keys() if aid not in self._pinned_adapters]
|
|
if not evict_candidates:
|
|
logger.warning(
|
|
"Cache full (%d) but all adapters are pinned; cannot evict. "
|
|
"Increase max_cached_adapters or unpin adapters.",
|
|
self.max_cached_adapters,
|
|
)
|
|
break
|
|
|
|
lru_adapter_id = evict_candidates[0]
|
|
logger.info(
|
|
"Evicting LRU adapter: id=%d (cache: %d/%d)",
|
|
lru_adapter_id,
|
|
len(self._registered_adapters),
|
|
self.max_cached_adapters,
|
|
)
|
|
self.remove_adapter(lru_adapter_id)
|
|
|
|
def add_adapter(self, lora_request: LoRARequest) -> bool:
|
|
"""
|
|
Add a new adapter to the cache without activating it.
|
|
"""
|
|
adapter_id = lora_request.lora_int_id
|
|
|
|
if adapter_id in self._registered_adapters:
|
|
logger.debug("Adapter %d already registered, skipping", adapter_id)
|
|
return False
|
|
|
|
logger.info("Adding new adapter: id=%d, name=%s", adapter_id, lora_request.lora_name)
|
|
|
|
# evict if cache full before adding the new adapter
|
|
# so that we don't go over capacity on the new load
|
|
self._evict_for_new_adapter()
|
|
|
|
lora_model, peft_helper = self._load_adapter(lora_request)
|
|
self._touch_adapter_info(adapter_id)
|
|
|
|
self._registered_adapters[adapter_id] = lora_model
|
|
|
|
self._replace_layers_with_lora(peft_helper)
|
|
|
|
logger.debug(
|
|
"Adapter %d added, cache size: %d/%d", adapter_id, len(self._registered_adapters), self.max_cached_adapters
|
|
)
|
|
return True
|
|
|
|
def remove_adapter(self, adapter_id: int) -> bool:
|
|
"""
|
|
Remove an adapter from the cache.
|
|
"""
|
|
if adapter_id not in self._registered_adapters:
|
|
logger.debug("Adapter %d not found, cannot remove", adapter_id)
|
|
return False
|
|
|
|
logger.info("Removing adapter: id=%d", adapter_id)
|
|
if self._active_adapter_id == adapter_id:
|
|
self._deactivate_all_adapters()
|
|
|
|
del self._registered_adapters[adapter_id]
|
|
self._adapter_scales.pop(adapter_id, None)
|
|
self._adapter_access_order.pop(adapter_id, None)
|
|
self._pinned_adapters.discard(adapter_id)
|
|
logger.debug(
|
|
"Adapter %d removed, cache size: %d/%d",
|
|
adapter_id,
|
|
len(self._registered_adapters),
|
|
self.max_cached_adapters,
|
|
)
|
|
return True
|
|
|
|
def list_adapters(self) -> list[int]:
|
|
"""Return list of registered adapter ids."""
|
|
return list(self._registered_adapters.keys())
|
|
|
|
def pin_adapter(self, adapter_id: int) -> bool:
|
|
"""Mark an adapter as pinned so it will not be evicted."""
|
|
if adapter_id not in self._registered_adapters:
|
|
logger.debug("Adapter %d not found, cannot pin", adapter_id)
|
|
return False
|
|
self._pinned_adapters.add(adapter_id)
|
|
# Touch access order so it is most recently used
|
|
self._adapter_access_order[adapter_id] = time.time()
|
|
self._adapter_access_order.move_to_end(adapter_id)
|
|
logger.info("Pinned adapter id=%d (won't be evicted)", adapter_id)
|
|
return True
|