diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py index 8d406aa2fb49..e913ae0e8384 100644 --- a/python/sglang/srt/lora/layers.py +++ b/python/sglang/srt/lora/layers.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Union +from typing import Dict, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -630,18 +630,32 @@ def slice_lora_a_weights(self, A: torch.Tensor): return A def slice_lora_b_weights(self, B: torch.Tensor): - local_tp_rank = self.base_layer.tp_rank - partition_sizes = self.base_layer.output_partition_sizes - output_sizes = self.base_layer.output_sizes slices = [] offset = 0 - for full_size, part_size in zip(output_sizes, partition_sizes): - start_idx = local_tp_rank * part_size - end_idx = start_idx + part_size - slices.append(B[offset + start_idx : offset + end_idx, :]) + for member, full_size in enumerate(self.base_layer.output_sizes): + slices.append( + self.slice_lora_b_member(member, B[offset : offset + full_size, :]) + ) offset += full_size return torch.concat(slices, dim=0) + def lora_b_member_full_rows(self) -> Tuple[int, ...]: + """Adapter-side lora_B rows of each member (partition), in slot order.""" + return tuple(self.base_layer.output_sizes) + + def lora_b_member_local_offsets(self) -> Tuple[int, ...]: + """Where each member's rows start in this rank's lora_B.""" + offsets = [0] + for part_size in self.base_layer.output_partition_sizes[:-1]: + offsets.append(offsets[-1] + part_size) + return tuple(offsets) + + def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor: + """This rank's rows of one member's full lora_B.""" + part_size = self.base_layer.output_partition_sizes[member] + start_idx = self.base_layer.tp_rank * part_size + return B[start_idx : start_idx + part_size, :] + class InklingQKVRLinearWithLoRA(MergedColumnParallelLinearWithLoRA): """LoRA wrapper for Inkling's fused q/k/v/r projection. @@ -654,31 +668,43 @@ class InklingQKVRLinearWithLoRA(MergedColumnParallelLinearWithLoRA): """ def slice_lora_b_weights(self, B: torch.Tensor): + members = torch.split(B, self.lora_b_member_full_rows(), dim=0) + return torch.concat( + [ + self.slice_lora_b_member(i, member_b) + for i, member_b in enumerate(members) + ], + dim=0, + ) + + def lora_b_member_full_rows(self) -> Tuple[int, ...]: + bl = self.base_layer + q_size = bl.inkling_head_dim * bl.inkling_num_heads + kv_size = bl.inkling_head_dim * bl.inkling_num_kv_heads + return q_size, kv_size, kv_size, bl.inkling_d_rel * bl.inkling_num_heads + + def _lora_b_member_local_rows(self) -> Tuple[int, ...]: bl = self.base_layer - tp_rank = bl.tp_rank - hd, nkv, nh, dr, tp = ( - bl.inkling_head_dim, - bl.inkling_num_kv_heads, - bl.inkling_num_heads, - bl.inkling_d_rel, - bl.inkling_tp_size, + q_size, kv_size, _, r_size = self.lora_b_member_full_rows() + tp = bl.inkling_tp_size + kv_local = ( + bl.inkling_head_dim if tp > bl.inkling_num_kv_heads else kv_size // tp ) - q_size, kv_size, r_size = hd * nh, hd * nkv, dr * nh - q_off, k_off, v_off, r_off = 0, q_size, q_size + kv_size, q_size + 2 * kv_size - q_per, r_per = q_size // tp, r_size // tp - q = B[q_off + tp_rank * q_per : q_off + (tp_rank + 1) * q_per, :] - r = B[r_off + tp_rank * r_per : r_off + (tp_rank + 1) * r_per, :] - if tp > nkv: + return q_size // tp, kv_local, kv_local, r_size // tp + + def lora_b_member_local_offsets(self) -> Tuple[int, ...]: + q_local, kv_local, _, _ = self._lora_b_member_local_rows() + return 0, q_local, q_local + kv_local, q_local + 2 * kv_local + + def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor: + bl = self.base_layer + local_rows = self._lora_b_member_local_rows()[member] + if member in (1, 2) and bl.inkling_tp_size > bl.inkling_num_kv_heads: # Replicate: each rank takes the single kv-head it shares (mirrors #15). - replicas = tp // nkv - kv_idx = tp_rank // replicas - k = B[k_off + kv_idx * hd : k_off + (kv_idx + 1) * hd, :] - v = B[v_off + kv_idx * hd : v_off + (kv_idx + 1) * hd, :] + shard = bl.tp_rank // (bl.inkling_tp_size // bl.inkling_num_kv_heads) else: - kv_per = kv_size // tp - k = B[k_off + tp_rank * kv_per : k_off + (tp_rank + 1) * kv_per, :] - v = B[v_off + tp_rank * kv_per : v_off + (tp_rank + 1) * kv_per, :] - return torch.concat([q, k, v, r], dim=0) + shard = bl.tp_rank + return B[shard * local_rows : (shard + 1) * local_rows, :] class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): @@ -741,34 +767,37 @@ def slice_lora_a_weights(self, A: torch.Tensor): return A def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor: - base_layer = self.base_layer - q_proj_shard_size = base_layer.q_proj_shard_size - kv_proj_shard_size = base_layer.kv_proj_shard_size - num_kv_head_replicas = base_layer.num_kv_head_replicas - local_tp_rank = base_layer.tp_rank - - q_start_idx = q_proj_shard_size * local_tp_rank - q_end_idx = q_start_idx + q_proj_shard_size - - kv_shard_id = local_tp_rank // num_kv_head_replicas - kv_start_idx = kv_proj_shard_size * kv_shard_id - kv_end_idx = kv_start_idx + kv_proj_shard_size - - q_size = base_layer.output_sizes[0] - k_size = base_layer.output_sizes[1] // num_kv_head_replicas - B_q_shard = B[q_start_idx:q_end_idx, :] - B_k_shard = B[q_size + kv_start_idx : q_size + kv_end_idx, :] - B_v_shard = B[q_size + k_size + kv_start_idx : q_size + k_size + kv_end_idx, :] - + members = torch.split(B, self.lora_b_member_full_rows(), dim=0) return torch.concat( - ( - B_q_shard, - B_k_shard, - B_v_shard, - ), + [ + self.slice_lora_b_member(i, member_b) + for i, member_b in enumerate(members) + ], dim=0, ) + def lora_b_member_full_rows(self) -> Tuple[int, ...]: + """Adapter-side q/k/v lora_B rows; output_sizes counts replicated kv heads.""" + base_layer = self.base_layer + q_size = base_layer.output_sizes[0] + kv_size = base_layer.output_sizes[1] // base_layer.num_kv_head_replicas + return q_size, kv_size, kv_size + + def lora_b_member_local_offsets(self) -> Tuple[int, ...]: + q_shard = self.base_layer.q_proj_shard_size + kv_shard = self.base_layer.kv_proj_shard_size + return 0, q_shard, q_shard + kv_shard + + def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor: + base_layer = self.base_layer + if member == 0: + shard_size = base_layer.q_proj_shard_size + shard = base_layer.tp_rank + else: + shard_size = base_layer.kv_proj_shard_size + shard = base_layer.tp_rank // base_layer.num_kv_head_replicas + return B[shard * shard_size : (shard + 1) * shard_size, :] + class RowParallelLinearWithLoRA(BaseLayerWithLoRA): def __init__( @@ -976,6 +1005,15 @@ def slice_lora_a_weights(self, A: torch.Tensor): def slice_lora_b_weights(self, B: torch.Tensor): return B + def lora_b_member_full_rows(self) -> Tuple[int, ...]: + return self.first_output_dim, self.output_size - self.first_output_dim + + def lora_b_member_local_offsets(self) -> Tuple[int, ...]: + return 0, self.first_output_dim + + def slice_lora_b_member(self, member: int, B: torch.Tensor) -> torch.Tensor: + return B + class FusedMoEWithLoRA(BaseLayerWithLoRA): """ diff --git a/python/sglang/srt/lora/lora.py b/python/sglang/srt/lora/lora.py index e0230377067e..a24461f787c5 100644 --- a/python/sglang/srt/lora/lora.py +++ b/python/sglang/srt/lora/lora.py @@ -37,6 +37,7 @@ # keys ("...experts...."), while excluding "shared_experts." (where the # preceding char is "_", not "."). _ROUTED_EXPERT_PATTERN = re.compile(r"\.experts\.") +_MOE_EXPERT_PATTERN = re.compile(r"\.(?:experts|shared_experts?)\.") logger = logging.getLogger(__name__) @@ -216,10 +217,9 @@ def _process_weight(self, name: str, loaded_weight: torch.Tensor): ) def _normalize_weights(self): + # qkv / qkvr / MLA qkv_a / dense gate_up stay per member: the memory pool writes each + # member into its slot of the fused buffer (LORA_FUSED_MEMBERS). for layer in self.layers: - weight_names = list(layer.weights.keys()) - self.normalize_qkv_proj(weight_names, layer.weights) - self.normalize_inkling_qkvr_proj(weight_names, layer.weights) self._normalize_shared_expert_moe(layer.weights) self._rename_expert_w_to_proj(layer.weights) # Stack gate_proj + x_proj → in_proj for Mamba layers (before gate_up normalization) @@ -230,107 +230,6 @@ def _normalize_weights(self): self._normalize_in_proj_ba(layer.weights) weight_names = list(layer.weights.keys()) self.normalize_gate_up_proj(weight_names, layer.weights) - weight_names = list(layer.weights.keys()) - self.normalize_fused_qkv_a_proj(weight_names, layer.weights) - - def normalize_inkling_qkvr_proj( - self, weight_names: List[str], weights: Dict[str, torch.Tensor] - ): - """Normalize Inkling split attention LoRA keys into the runtime qkvr layer. - - Inkling checkpoints expose separate attention projections - (wq_du/wk_dv/wv_dv/wr_du), while the serving model uses a single - MergedColumnParallelLinear named qkvr with those slices in that order. - """ - for weight_name in weight_names: - if "wq_du" not in weight_name: - continue - - q_name = weight_name - k_name = weight_name.replace("wq_du", "wk_dv") - v_name = weight_name.replace("wq_du", "wv_dv") - r_name = weight_name.replace("wq_du", "wr_du") - qkvr_name = weight_name.replace("wq_du", "qkvr") - - if any(name not in weights for name in (k_name, v_name, r_name)): - continue - - cat_dim = weights[q_name].dim() - 2 - weights[qkvr_name] = torch.cat( - ( - weights[q_name], - weights[k_name], - weights[v_name], - weights[r_name], - ), - cat_dim, - ) - weights.pop(q_name) - weights.pop(k_name) - weights.pop(v_name) - weights.pop(r_name) - - for weight_name in list(weights.keys()): - if "qkvr" not in weight_name or "lora_A" not in weight_name: - continue - if weights[weight_name].shape[-2] != self.config.r: - continue - repeat_dims = [1] * weights[weight_name].dim() - repeat_dims[-2] = 4 - weights[weight_name] = weights[weight_name].repeat(*repeat_dims) - - 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_shared_expert_moe(self, weights: Dict[str, torch.Tensor]): """Reshape flat Inkling shared-sink factors to shared-outer 3D form.""" @@ -538,7 +437,10 @@ def _normalize_in_proj_ba(self, weights: Dict[str, torch.Tensor]): def normalize_gate_up_proj( self, weight_names: List[str], weights: Dict[str, torch.Tensor] ): + """Stack routed- and shared-expert gate/up factors; dense gate/up load per member.""" for weight_name in weight_names: + if not _MOE_EXPERT_PATTERN.search(weight_name): + continue 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") @@ -576,33 +478,6 @@ def normalize_gate_up_proj( # Models with non-gated MLP/shared-experts declare up_proj in # supported_lora_modules so they get their own buffer and wrapping. - def normalize_fused_qkv_a_proj( - self, weight_names: List[str], weights: Dict[str, torch.Tensor] - ): - """Fuse separate q_a_proj and kv_a_proj_with_mqa LoRA weights into - a single fused_qkv_a_proj_with_mqa entry (concat along dim 0 for - both A and B), matching the DeepSeek MLA fused projection layout.""" - for weight_name in weight_names: - if "q_a_proj" not in weight_name: - continue - if "fused_qkv_a_proj_with_mqa" in weight_name: - continue - - q_a_name = weight_name - kv_a_name = weight_name.replace("q_a_proj", "kv_a_proj_with_mqa") - fused_name = weight_name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") - - kv_a_weight = ( - weights[kv_a_name] - if kv_a_name in weights - else torch.zeros_like(weights[q_a_name]) - ) - - weights[fused_name] = torch.cat((weights[q_a_name], kv_a_weight), dim=0) - weights.pop(q_a_name) - if kv_a_name in weights: - weights.pop(kv_a_name) - def pin_weights_in_cpu(self): for layer in self.layers: for name, weight in layer.weights.items(): diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index f1d61a7b521b..0fe8391ef65f 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -28,6 +28,8 @@ from sglang.srt.lora.utils import ( ATTN_TP_LORA_MODULE_NAMES, EMBEDDING_NAMES, + LORA_FUSED_MEMBER_SLOTS, + LORA_FUSED_MEMBERS, REPLICATED_LINEAR_LORA_NAMES, ROW_PARALLELISM_LINEAR_LORA_NAMES, LoRAType, @@ -90,6 +92,14 @@ def append_cache_key_suffix( return f"{cache_keys}#{suffix}" +def _fused_member_slot(weight_name: str) -> Optional[Tuple[str, int]]: + """(fused module, slot) of a dense per-member adapter weight, e.g. ``...self_attn.k_proj.lora_A.weight``.""" + if re.search(r"\.(?:experts|shared_experts?)\.", weight_name): + return None + projection = weight_name.split(".lora_", 1)[0].rsplit(".", 1)[-1] + return LORA_FUSED_MEMBER_SLOTS.get(projection) + + def _get_moe_ep_context() -> Tuple[int, int]: """Return `(moe_ep_size, moe_ep_rank)`, or `(1, 0)` if the MoE EP group is not initialized (hermetic tests or pure-TP launches).""" @@ -1052,6 +1062,10 @@ def load_lora_weight_tensor( if lora_adapter.embedding_layers: all_weight_names.extend(lora_adapter.embedding_layers.keys()) for name in all_weight_names: + fused_member = _fused_member_slot(name) + if fused_member is not None and fused_member[0] in self.target_modules: + matched_modules.add(fused_member[0]) + continue try: target_module = get_target_module_name(name, self.target_modules) matched_modules.add(target_module) @@ -1103,6 +1117,12 @@ def load_lora_weight_tensor( target_module: None for target_module in self.B_buffer } + fused_member_modules = self._fused_member_modules(cur_layer_modules) + # fused module -> {(side, slot): (member weight, cache key)} + temp_member_weights: Dict[ + str, Dict[Tuple[str, int], Tuple[torch.Tensor, str]] + ] = {target_module: {} for target_module in fused_member_modules} + fused_moe = self._get_fused_shared_moe(self.base_model, layer_id) for name, weights in layer_weights.items(): cache_name = name @@ -1114,6 +1134,24 @@ def load_lora_weight_tensor( f".experts.{fused_moe._num_global_routed}.", name, ) + fused_member = _fused_member_slot(name) + if fused_member is not None and fused_member[0] in fused_member_modules: + # PEFT may ship an untrained member as a zero-element placeholder. + if weights.numel() == 0: + continue + fused_module, slot = fused_member + side = "A" if "lora_A" in name else "B" + temp_member_weights[fused_module][(side, slot)] = ( + weights, + cache_name, + ) + continue + if fused_member is not None and not any( + target in name for target in self.target_modules + ): + # A member of a fused module this layer lacks (e.g. q_proj on a linear-attention + # layer) is not loaded, as a weight for any absent module is not. + continue target_module = get_target_module_name(name, self.target_modules) # Check if this is an MoE weight (has expert index in name) @@ -1233,6 +1271,8 @@ def load_lora_weight_tensor( # redundant zero-fills on slots no `update_lora_info` ever points # a forward-time module at. active_target_modules: Set[str] = set() + # fused module -> (A writes as (slot, weight, cache key), B writes as (row offset, weight, cache key)) + fused_member_writes: Dict[str, Tuple[list, list]] = {} from sglang.srt.lora.layers import FusedMoEWithLoRA for module_name, module in cur_layer_modules.items(): @@ -1288,6 +1328,23 @@ def load_lora_weight_tensor( # the new slot) and the wrapped layer module will read it. active_target_modules.add(target_module) + if target_module in fused_member_modules: + fused_member_writes[target_module] = self._fused_member_writes( + module, + target_module, + member_weights=temp_member_weights[target_module], + fused_weights=( + temp_A_buffer[target_module], + temp_B_buffer[target_module], + ), + fused_cache_keys=( + temp_A_cache_keys[target_module], + temp_B_cache_keys[target_module], + ), + lora_rank=lora_rank, + ) + continue + if temp_A_buffer[target_module] is None: # Skip weight slicing if the weight is not present in the adapter continue @@ -1313,8 +1370,23 @@ def load_lora_weight_tensor( f"tp{self.tp_rank}", ) + for name, (a_writes, b_writes) in fused_member_writes.items(): + self._load_fused_member_writes( + a_view=self.A_buffer[name][layer_id][ + buffer_id, + : lora_rank * get_stacked_multiply(name, self.base_model), + :, + ], + b_buffer=self.B_buffer[name][layer_id][buffer_id], + a_writes=a_writes, + b_writes=b_writes, + lora_rank=lora_rank, + pinned_layer_weights=pinned_layer_weights, + load_lora_weight_tensor=load_lora_weight_tensor, + ) + for name, weights in temp_A_buffer.items(): - if name not in active_target_modules: + if name not in active_target_modules or name in fused_member_writes: continue c = get_stacked_multiply(name, self.base_model) max_r = self.max_lora_rank @@ -1444,7 +1516,7 @@ def load_lora_weight_tensor( load_lora_weight_tensor(buffer_view, weights) for name, weights in temp_B_buffer.items(): - if name not in active_target_modules: + if name not in active_target_modules or name in fused_member_writes: continue target_buffer = self.B_buffer[name][layer_id] weights_cache_key = temp_B_cache_keys[name] @@ -1703,6 +1775,115 @@ def load_lora_weight_tensor( self.new_embeddings_buffer["input_embeddings"][buffer_id], None ) + def _fused_member_modules( + self, layer_modules: Dict[str, torch.nn.Module] + ) -> Dict[str, torch.nn.Module]: + """This layer's fused modules that load member by member (LORA_FUSED_MEMBERS).""" + from sglang.srt.lora.layers import FusedMoEWithLoRA + + modules = {} + for module_name, module in layer_modules.items(): + if isinstance(module, FusedMoEWithLoRA) or getattr( + module, "is_shared_fused_moe", False + ): + continue + target_module = get_target_module_name(module_name, self.target_modules) + if target_module in LORA_FUSED_MEMBERS: + modules[target_module] = module + return modules + + def _fused_member_writes( + self, + module: torch.nn.Module, + fused_module: str, + *, + member_weights: Dict[Tuple[str, int], Tuple[torch.Tensor, str]], + fused_weights: Tuple[Optional[torch.Tensor], Optional[torch.Tensor]], + fused_cache_keys: Tuple[Optional[str], Optional[str]], + lora_rank: int, + ) -> Tuple[list, list]: + """Each member's A slot and this rank's B rows, from per-member weights or one fused tensor.""" + num_slots = len(LORA_FUSED_MEMBERS[fused_module]) + fused_a, fused_b = fused_weights + if fused_a is not None or fused_b is not None: + assert fused_a is not None and fused_b is not None, ( + f"LoRA adapter carries only one of {fused_module}'s lora_A / lora_B" + ) + assert not member_weights, ( + f"LoRA adapter carries {fused_module} both fused and as separate projections" + ) + fused_a_key, fused_b_key = fused_cache_keys + # One A shared by every member, or one stacked per member. + if fused_a.shape[0] == lora_rank: + a_parts = [fused_a] * num_slots + else: + assert fused_a.shape[0] == num_slots * lora_rank, ( + f"{fused_module} lora_A has {fused_a.shape[0]} rows; expected {lora_rank} or " + f"{num_slots * lora_rank}" + ) + a_parts = list(torch.split(fused_a, lora_rank, dim=0)) + b_parts = torch.split(fused_b, module.lora_b_member_full_rows(), dim=0) + member_weights = {} + for slot in range(num_slots): + member_weights[("A", slot)] = ( + a_parts[slot], + append_cache_key_suffix(fused_a_key, f"member{slot}"), + ) + member_weights[("B", slot)] = ( + b_parts[slot], + append_cache_key_suffix(fused_b_key, f"member{slot}"), + ) + + b_offsets = module.lora_b_member_local_offsets() + a_writes, b_writes = [], [] + for (side, slot), (weight, cache_key) in member_weights.items(): + cache_key = append_cache_key_suffix(cache_key, f"tp{self.tp_rank}") + if side == "A": + a_writes.append((slot, module.slice_lora_a_weights(weight), cache_key)) + else: + b_writes.append( + ( + b_offsets[slot], + module.slice_lora_b_member(slot, weight), + cache_key, + ) + ) + return a_writes, b_writes + + def _load_fused_member_writes( + self, + *, + a_view: torch.Tensor, + b_buffer: torch.Tensor, + a_writes: list, + b_writes: list, + lora_rank: int, + pinned_layer_weights: Dict[str, torch.Tensor], + load_lora_weight_tensor, + ) -> None: + """Zero the adapter's fused slots, then write each present member; absent members stay zero.""" + load_lora_weight_tensor(a_view, None) + for slot, weight, cache_key in a_writes: + weight = self._get_maybe_cached_weight_for_transfer( + pinned_layer_weights, cache_key, weight + ) + load_lora_weight_tensor( + a_view[slot * lora_rank : (slot + 1) * lora_rank], weight + ) + b_view = b_buffer[:, :lora_rank] + load_lora_weight_tensor(b_view, None) + for row_offset, weight, cache_key in b_writes: + weight = self._get_maybe_cached_weight_for_transfer( + pinned_layer_weights, cache_key, weight + ) + load_lora_weight_tensor( + b_view[row_offset : row_offset + weight.shape[0]], weight + ) + if _SGLANG_EXPERIMENTAL_LORA_OPTI: + # Zero beyond loaded rank: the experimental dense LoRA-B kernel + # contracts over the full padded max_rank, so the tail must be clean. + load_lora_weight_tensor(b_buffer[:, lora_rank:], None) + def get_embedding_tensor( self, target_module: str, lora_type: LoRAType ) -> Optional[torch.Tensor]: diff --git a/python/sglang/srt/lora/utils.py b/python/sglang/srt/lora/utils.py index e982d1d01123..7524ffc42cf3 100644 --- a/python/sglang/srt/lora/utils.py +++ b/python/sglang/srt/lora/utils.py @@ -272,6 +272,21 @@ def get_default_hidden_dim( raise NotImplementedError("get_hidden_dim not implemented for " + module_name) +# Fused LoRA modules loaded member by member: each adapter projection fills its own slot of +# the fused buffer, sliced by the wrapper of the fused layer; untrained members stay zero. +LORA_FUSED_MEMBERS = { + "qkv_proj": ("q_proj", "k_proj", "v_proj"), + "gate_up_proj": ("gate_proj", "up_proj"), + "fused_qkv_a_proj_with_mqa": ("q_a_proj", "kv_a_proj_with_mqa"), + "qkvr": ("wq_du", "wk_dv", "wv_dv", "wr_du"), +} +LORA_FUSED_MEMBER_SLOTS = { + member: (fused, slot) + for fused, members in LORA_FUSED_MEMBERS.items() + for slot, member in enumerate(members) +} + + def get_normalized_target_modules( target_modules: Union[str, Iterable[str]], base_model: Optional[torch.nn.Module] = None, @@ -295,11 +310,7 @@ def get_normalized_target_modules( return {"all"} params_mapping = { - "q_proj": "qkv_proj", - "k_proj": "qkv_proj", - "v_proj": "qkv_proj", - "gate_proj": "gate_up_proj", - "up_proj": "gate_up_proj", + **{member: fused for member, (fused, _) in LORA_FUSED_MEMBER_SLOTS.items()}, "in_proj_q": "in_proj_qkvz", "in_proj_k": "in_proj_qkvz", "in_proj_v": "in_proj_qkvz", @@ -315,8 +326,6 @@ def get_normalized_target_modules( "lm_head": "lm_head", "output": "lm_head", "unembed_tokens": "lm_head", - "q_a_proj": "fused_qkv_a_proj_with_mqa", - "kv_a_proj_with_mqa": "fused_qkv_a_proj_with_mqa", # DSA indexer projections are qualified with their parent module name # because the bare leaf names collide with unrelated modules in other # models (e.g. DeepSeek-V4 attention `wq_b`, Pixtral vision `wk`). @@ -348,14 +357,12 @@ def get_stacked_multiply( if base_model is not None and hasattr(base_model, "get_stacked_multiply"): return base_model.get_stacked_multiply(module_name) stacked_rank = { - "qkv_proj": 3, + **{fused: len(members) for fused, members in LORA_FUSED_MEMBERS.items()}, "in_proj_qkvz": 4, # GDN packed input projection "in_proj_ba": 2, # GDN packed b/a input projection - "gate_up_proj": 2, "gate_up_proj_moe": 2, "gate_up_proj_shared_moe": 2, "in_proj": 2, - "fused_qkv_a_proj_with_mqa": 2, } return stacked_rank[module_name] if module_name in stacked_rank else 1 diff --git a/test/registered/unit/lora/test_lora_fused_member_load.py b/test/registered/unit/lora/test_lora_fused_member_load.py new file mode 100644 index 000000000000..62e962a9b89b --- /dev/null +++ b/test/registered/unit/lora/test_lora_fused_member_load.py @@ -0,0 +1,381 @@ +"""CPU-only tests: dense fused LoRA modules (qkv, gate_up, MLA qkv_a, Inkling qkvr) load each adapter +projection into its own slot of the fused buffer. + +A complete adapter must land byte-identical to the previous concatenate-then-slice loader, and any +subset of a fused module's projections must load with the absent projections left zero.""" + +from __future__ import annotations + +import itertools +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.lora.layers import ( + InklingQKVRLinearWithLoRA, + MergedColumnParallelLinearWithLoRA, + QKVParallelLinearWithLoRA, + ReplicatedLinearWithLoRA, +) +from sglang.srt.lora.lora import LoRAAdapter +from sglang.srt.lora.mem_pool import LoRAMemoryPool +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_HIDDEN = 4 +_RANK = 2 +_PREFIX = "base_model.model.model.layers.0." + + +# --- the previous loader's fused-B slicing, the oracle for complete adapters --------------------- +# copied from sglang-miles@4076a68579 (QKVParallelLinearWithLoRA / InklingQKVRLinearWithLoRA / +# MergedColumnParallelLinearWithLoRA.slice_lora_b_weights) + + +def _reference_qkv_slice(bl, B): + q_start = bl.q_proj_shard_size * bl.tp_rank + kv_shard_id = bl.tp_rank // bl.num_kv_head_replicas + kv_start = bl.kv_proj_shard_size * kv_shard_id + kv_end = kv_start + bl.kv_proj_shard_size + q_size = bl.output_sizes[0] + k_size = bl.output_sizes[1] // bl.num_kv_head_replicas + return torch.concat( + ( + B[q_start : q_start + bl.q_proj_shard_size], + B[q_size + kv_start : q_size + kv_end], + B[q_size + k_size + kv_start : q_size + k_size + kv_end], + ) + ) + + +def _reference_inkling_slice(bl, B): + hd, nkv, nh, dr, tp, tp_rank = ( + bl.inkling_head_dim, + bl.inkling_num_kv_heads, + bl.inkling_num_heads, + bl.inkling_d_rel, + bl.inkling_tp_size, + bl.tp_rank, + ) + q_size, kv_size, r_size = hd * nh, hd * nkv, dr * nh + k_off, v_off, r_off = q_size, q_size + kv_size, q_size + 2 * kv_size + q_per, r_per = q_size // tp, r_size // tp + q = B[tp_rank * q_per : (tp_rank + 1) * q_per] + r = B[r_off + tp_rank * r_per : r_off + (tp_rank + 1) * r_per] + if tp > nkv: + kv_idx = tp_rank // (tp // nkv) + k = B[k_off + kv_idx * hd : k_off + (kv_idx + 1) * hd] + v = B[v_off + kv_idx * hd : v_off + (kv_idx + 1) * hd] + else: + kv_per = kv_size // tp + k = B[k_off + tp_rank * kv_per : k_off + (tp_rank + 1) * kv_per] + v = B[v_off + tp_rank * kv_per : v_off + (tp_rank + 1) * kv_per] + return torch.concat([q, k, v, r]) + + +def _reference_merged_slice(bl, B): + slices, offset = [], 0 + for full, part in zip(bl.output_sizes, bl.output_partition_sizes): + slices.append(B[offset + bl.tp_rank * part : offset + bl.tp_rank * part + part]) + offset += full + return torch.concat(slices) + + +# --- one fused layer of each kind --------------------------------------------------------------- + + +def _wrap(cls, base_layer, **attrs): + wrapper = cls.__new__(cls) + object.__setattr__(wrapper, "base_layer", base_layer) + for name, value in attrs.items(): + object.__setattr__(wrapper, name, value) + return wrapper + + +def _qkv(tp_rank, tp_size=4): + # 4 q heads, 2 kv heads, head_dim 2: at TP4 each kv head is replicated on two ranks + replicas = tp_size // 2 + bl = SimpleNamespace( + tp_rank=tp_rank, + q_proj_shard_size=8 // tp_size, + kv_proj_shard_size=2, + num_kv_head_replicas=replicas, + output_sizes=[8, 2 * tp_size, 2 * tp_size], + ) + return ( + "self_attn.qkv_proj", + ("q_proj", "k_proj", "v_proj"), + (8, 4, 4), + _wrap(QKVParallelLinearWithLoRA, bl), + lambda B: _reference_qkv_slice(bl, B), + ) + + +def _gate_up(tp_rank, tp_size=2): + bl = SimpleNamespace( + tp_rank=tp_rank, output_sizes=[6, 6], output_partition_sizes=[6 // tp_size] * 2 + ) + return ( + "mlp.gate_up_proj", + ("gate_proj", "up_proj"), + (6, 6), + _wrap(MergedColumnParallelLinearWithLoRA, bl), + lambda B: _reference_merged_slice(bl, B), + ) + + +def _qkvr(tp_rank, tp_size=4): + bl = SimpleNamespace( + tp_rank=tp_rank, + inkling_head_dim=2, + inkling_num_heads=4, + inkling_num_kv_heads=2, + inkling_d_rel=1, + inkling_tp_size=tp_size, + ) + return ( + "attn.qkvr", + ("wq_du", "wk_dv", "wv_dv", "wr_du"), + (8, 4, 4, 4), + _wrap(InklingQKVRLinearWithLoRA, bl), + lambda B: _reference_inkling_slice(bl, B), + ) + + +def _mla_qkv_a(tp_rank, tp_size=1): + del tp_rank, tp_size + wrapper = _wrap( + ReplicatedLinearWithLoRA, SimpleNamespace(), first_output_dim=3, output_size=5 + ) + return ( + "self_attn.fused_qkv_a_proj_with_mqa", + ("q_a_proj", "kv_a_proj_with_mqa"), + (3, 2), + wrapper, + lambda B: B, + ) + + +_LAYERS = { + "qkv": (_qkv, 4), + "gate_up": (_gate_up, 2), + "qkvr": (_qkvr, 4), + "mla_qkv_a": (_mla_qkv_a, 1), +} + + +# --- the memory pool under test ------------------------------------------------------------------- + + +def _pool(module_path, wrapper, full_rows, max_rank): + fused = module_path.rsplit(".", 1)[-1] + pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.num_layer = 1 + pool.tp_rank = ( + wrapper.base_layer.tp_rank if hasattr(wrapper.base_layer, "tp_rank") else 0 + ) + pool.max_lora_rank = max_rank + pool.target_modules = {fused} + pool.strict_loading = True + pool.pin_memory_available = False + pool.enable_lora_overlap_loading = False + pool.experts_shared_outer_loras = False + pool.base_model = torch.nn.Identity() + pool.lora_added_tokens_size = 0 + pool.embedding_A_buffer = {} + pool.embedding_B_buffer = {} + pool.lm_head_A_buffer = {} + pool.lm_head_B_buffer = {} + pool.new_embeddings_buffer = {} + local_rows = sum( + wrapper.slice_lora_b_member(i, torch.zeros(rows, 1)).shape[0] + for i, rows in enumerate(full_rows) + ) + pool.A_buffer = {fused: [torch.full((1, max_rank * len(full_rows), _HIDDEN), -9.0)]} + pool.B_buffer = {fused: [torch.full((1, local_rows, max_rank), -9.0)]} + return pool + + +def _load(pool, module_path, wrapper, weights): + adapter = SimpleNamespace( + config=SimpleNamespace(r=_RANK), + scaling=1.0, + layers=[SimpleNamespace(weights=weights, pinned_weights={})], + embedding_layers={}, + added_tokens_embeddings={}, + ) + pool.load_lora_weight_to_buffer( + uid="adapter", + buffer_id=0, + lora_adapter=adapter, + lora_modules=[{f"model.layers.0.{module_path}": wrapper}], + lora_embed_tokens_module=None, + lora_lm_head_module=None, + ) + fused = module_path.rsplit(".", 1)[-1] + return pool.A_buffer[fused][0][0], pool.B_buffer[fused][0][0] + + +def _member_weights(members, full_rows, present): + torch.manual_seed(0) + a = {m: torch.randn(_RANK, _HIDDEN) for m in members} + b = {m: torch.randn(rows, _RANK) for m, rows in zip(members, full_rows)} + block = _PREFIX + ( + "attn." + if members[0] == "wq_du" + else "mlp." + if members[0] == "gate_proj" + else "self_attn." + ) + weights = {} + for m in present: + weights[f"{block}{m}.lora_A.weight"] = a[m] + weights[f"{block}{m}.lora_B.weight"] = b[m] + return weights, a, b + + +def _expected(members, full_rows, a, b, present, reference_slice): + """The previous loader's buffer for the adapter with every absent member written as zeros.""" + fused_a = torch.cat( + [a[m] if m in present else torch.zeros_like(a[m]) for m in members] + ) + fused_b = torch.cat( + [ + b[m] if m in present else torch.zeros(rows, _RANK) + for m, rows in zip(members, full_rows) + ] + ) + return fused_a, reference_slice(fused_b) + + +def _cases(): + for kind, (factory, tp_size) in _LAYERS.items(): + for tp_rank in range(tp_size): + yield pytest.param(kind, tp_rank, id=f"{kind}-tp{tp_rank}") + + +@pytest.mark.parametrize("kind, tp_rank", list(_cases())) +@pytest.mark.parametrize("max_rank", [_RANK, _RANK + 1], ids=["max_r=r", "max_r>r"]) +def test_a_complete_adapter_loads_like_the_concatenating_loader( + kind, tp_rank, max_rank +): + module_path, members, full_rows, wrapper, reference_slice = _LAYERS[kind][0]( + tp_rank + ) + pool = _pool(module_path, wrapper, full_rows, max_rank) + weights, a, b = _member_weights(members, full_rows, members) + + a_buffer, b_buffer = _load(pool, module_path, wrapper, weights) + + expected_a, expected_b = _expected( + members, full_rows, a, b, members, reference_slice + ) + assert torch.equal(a_buffer[: len(members) * _RANK], expected_a) + assert torch.equal(b_buffer[:, :_RANK], expected_b) + + +@pytest.mark.parametrize("kind, tp_rank", list(_cases())) +def test_any_member_subset_loads_with_the_absent_members_zero(kind, tp_rank): + """A q-only adapter raised KeyError, a k/v-only or up-only adapter raised ValueError, and a + q_a-only MLA adapter got kv_a rows of q_a's width.""" + module_path, members, full_rows, wrapper, reference_slice = _LAYERS[kind][0]( + tp_rank + ) + for size in range(1, len(members)): + for present in itertools.combinations(members, size): + pool = _pool(module_path, wrapper, full_rows, _RANK) + weights, a, b = _member_weights(members, full_rows, present) + + a_buffer, b_buffer = _load(pool, module_path, wrapper, weights) + + expected_a, expected_b = _expected( + members, full_rows, a, b, present, reference_slice + ) + assert torch.equal(a_buffer[: len(members) * _RANK], expected_a), present + assert torch.equal(b_buffer[:, :_RANK], expected_b), present + + +def test_zero_element_placeholders_load_as_absent_members(): + """PEFT can ship an untrained up_proj as zero-element tensors; they must load as zeros.""" + module_path, members, full_rows, wrapper, reference_slice = _gate_up(tp_rank=0) + weights, a, b = _member_weights(members, full_rows, ["gate_proj"]) + weights[f"{_PREFIX}mlp.up_proj.lora_A.weight"] = torch.empty(0) + weights[f"{_PREFIX}mlp.up_proj.lora_B.weight"] = torch.empty(0) + + a_buffer, b_buffer = _load( + _pool(module_path, wrapper, full_rows, _RANK), module_path, wrapper, weights + ) + + expected_a, expected_b = _expected( + members, full_rows, a, b, ["gate_proj"], reference_slice + ) + assert torch.equal(a_buffer[: len(members) * _RANK], expected_a) + assert torch.equal(b_buffer[:, :_RANK], expected_b) + + +@pytest.mark.parametrize("stacked_a", [False, True], ids=["shared-A", "stacked-A"]) +def test_an_already_fused_qkv_adapter_matches_its_members(stacked_a): + module_path, members, full_rows, wrapper, reference_slice = _qkv(tp_rank=3) + torch.manual_seed(1) + fused_a = torch.randn(_RANK * (3 if stacked_a else 1), _HIDDEN) + fused_b = torch.randn(sum(full_rows), _RANK) + weights = { + f"{_PREFIX}self_attn.qkv_proj.lora_A.weight": fused_a, + f"{_PREFIX}self_attn.qkv_proj.lora_B.weight": fused_b, + } + + a_buffer, b_buffer = _load( + _pool(module_path, wrapper, full_rows, _RANK), module_path, wrapper, weights + ) + + expected_a = fused_a if stacked_a else fused_a.repeat(3, 1) + assert torch.equal(a_buffer[: 3 * _RANK], expected_a) + assert torch.equal(b_buffer[:, :_RANK], reference_slice(fused_b)) + + +def test_an_adapter_with_fused_and_separate_projections_is_rejected(): + module_path, members, full_rows, wrapper, _ = _qkv(tp_rank=0) + weights, _, _ = _member_weights(members, full_rows, ("q_proj",)) + weights[f"{_PREFIX}self_attn.qkv_proj.lora_A.weight"] = torch.randn(_RANK, _HIDDEN) + weights[f"{_PREFIX}self_attn.qkv_proj.lora_B.weight"] = torch.randn( + sum(full_rows), _RANK + ) + + with pytest.raises(AssertionError, match="both fused and as separate projections"): + _load( + _pool(module_path, wrapper, full_rows, _RANK), module_path, wrapper, weights + ) + + +def test_the_adapter_keeps_dense_projections_per_member(): + """Normalization no longer stacks dense projections; routed-expert gate/up still are.""" + adapter = LoRAAdapter.__new__(LoRAAdapter) + adapter.base_hf_config = SimpleNamespace() + object.__setattr__(adapter, "base_model", None) + object.__setattr__(adapter, "_moe_is_gated_by_layer", {}) + names = [ + f"{_PREFIX}self_attn.q_proj.lora_A.weight", + f"{_PREFIX}mlp.gate_proj.lora_A.weight", + f"{_PREFIX}self_attn.q_a_proj.lora_A.weight", + f"{_PREFIX}mlp.experts.0.gate_proj.lora_A.weight", + f"{_PREFIX}mlp.experts.0.up_proj.lora_A.weight", + ] + object.__setattr__( + adapter, + "layers", + [SimpleNamespace(weights={n: torch.ones(_RANK, _HIDDEN) for n in names})], + ) + + adapter._normalize_weights() + + assert sorted(adapter.layers[0].weights) == sorted( + names[:3] + [f"{_PREFIX}mlp.experts.0.gate_up_proj.lora_A.weight"] + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))