minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Lora] Lora quat info re-factor and support deepseekv3 mla lora (#22323)

Browse Source
This commit is contained in:
Ethan (Yusheng) Su
2026-04-09 14:19:58 -07:00
committed by GitHub
parent 60acdc31f2
commit 28ef6de091
16 changed files with 458 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/base_config.py
+14
View File
@@ -10,6 +10,7 @@ from torch import nn
if TYPE_CHECKING:
from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
from sglang.srt.models.utils import WeightsMapper
@@ -106,6 +107,19 @@ class FusedMoEMethodBase(QuantizeMethodBase):
) -> CombineInput:
raise NotImplementedError
def get_triton_quant_info(self, layer: torch.nn.Module) -> "TritonMoeQuantInfo":
"""Return a ``TritonMoeQuantInfo`` describing the quantisation state
stored on *layer*.
The LoRA MoE runner calls this so that ``invoke_fused_moe_kernel``
receives the correct flags / scales / block-shape for the base
weights. Each quantisation method must override this with the
same construction it already uses inside ``apply()``.
"""
raise NotImplementedError(
f"{type(self).__name__} must implement get_triton_quant_info()"
)
class QuantizationConfig(ABC):
"""Base class for quantization configs."""
python/sglang/srt/layers/quantization/blockwise_int8.py
+10 -7
View File
@@ -360,13 +360,8 @@ class BlockInt8MoEMethod(FusedMoEMethodBase):
self.moe_runner_config = moe_runner_config
self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
quant_info = TritonMoeQuantInfo(
def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
return TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
use_int8_w8a8=True,
@@ -377,4 +372,12 @@ class BlockInt8MoEMethod(FusedMoEMethodBase):
block_shape=self.quant_config.weight_block_size,
)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
quant_info = self.get_triton_quant_info(layer)
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/layers/quantization/fp8.py
+21 -20
View File
@@ -1487,6 +1487,26 @@ class Fp8MoEMethod(FusedMoEMethodBase):
# TODO(cwan): refactor other backends
pass
def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
return TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
use_fp8_w8a8=True,
w13_scale=(
layer.w13_weight_scale_inv
if self.block_quant
else layer.w13_weight_scale
),
w2_scale=(
layer.w2_weight_scale_inv if self.block_quant else layer.w2_weight_scale
),
a13_scale=layer.w13_input_scale,
a2_scale=layer.w2_input_scale,
block_shape=self.quant_config.weight_block_size,
)
def apply(
self,
layer: torch.nn.Module,
@@ -1663,26 +1683,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
),
)
elif self.runner.runner_backend.is_triton():
quant_info = TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
use_fp8_w8a8=True,
w13_scale=(
layer.w13_weight_scale_inv
if self.block_quant
else layer.w13_weight_scale
),
w2_scale=(
layer.w2_weight_scale_inv
if self.block_quant
else layer.w2_weight_scale
),
a13_scale=layer.w13_input_scale,
a2_scale=layer.w2_input_scale,
block_shape=self.quant_config.weight_block_size,
)
quant_info = self.get_triton_quant_info(layer)
else:
raise NotImplementedError(
"Unsupported runner backend: %s" % self.runner.runner_backend
python/sglang/srt/layers/quantization/moe_wna16.py
+13 -11
View File
@@ -364,19 +364,10 @@ class MoeWNA16Method(FusedMoEMethodBase):
self.moe_runner_config = moe_runner_config
self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
assert (
self.moe_runner_config.activation == "silu"
), "Only SiLU activation is supported."
def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
weight_bits = self.quant_config.weight_bits
has_zp = self.quant_config.has_zp
quant_info = TritonMoeQuantInfo(
return TritonMoeQuantInfo(
w13_weight=layer.w13_qweight,
w2_weight=layer.w2_qweight,
use_int4_w4a16=weight_bits == 4,
@@ -387,6 +378,17 @@ class MoeWNA16Method(FusedMoEMethodBase):
w2_zp=layer.w2_qzeros if has_zp else None,
block_shape=[0, layer.group_size],
)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
assert (
self.moe_runner_config.activation == "silu"
), "Only SiLU activation is supported."
quant_info = self.get_triton_quant_info(layer)
return self.runner.run(dispatch_output, quant_info)
@staticmethod
python/sglang/srt/layers/quantization/unquant.py
+9 -6
View File
@@ -554,6 +554,14 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
)
return StandardCombineInput(hidden_states=output)
def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
return TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
)
def forward_xpu(
self,
layer: torch.nn.Module,
@@ -594,12 +602,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
), f"activation = {moe_runner_config.activation} is not supported \
for Triton PATH, please set ENV SGLANG_USE_SGL_XPU=1."
quant_info = TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
)
quant_info = self.get_triton_quant_info(layer)
return self.runner.run(dispatch_output, quant_info)
def forward_npu(
python/sglang/srt/layers/quantization/w8a8_fp8.py
+10 -7
View File
@@ -286,13 +286,8 @@ class W8A8FP8MoEMethod(FusedMoEMethodBase):
self.moe_runner_config = moe_runner_config
self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
quant_info = TritonMoeQuantInfo(
def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
return TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
use_fp8_w8a8=True,
@@ -302,4 +297,12 @@ class W8A8FP8MoEMethod(FusedMoEMethodBase):
a13_scale=layer.w13_input_scale,
a2_scale=layer.w2_input_scale,
)
def apply(
self,
layer: torch.nn.Module,
dispatch_output: StandardDispatchOutput,
) -> CombineInput:
quant_info = self.get_triton_quant_info(layer)
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/layers/quantization/w8a8_int8.py
+13 -10
View File
@@ -331,6 +331,18 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
self.moe_runner_config = moe_runner_config
self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
def get_triton_quant_info(self, layer: torch.nn.Module) -> TritonMoeQuantInfo:
return TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
use_int8_w8a8=True,
per_channel_quant=True,
w13_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
a13_scale=layer.w13_input_scale,
a2_scale=layer.w2_input_scale,
)
def apply(
self,
layer: torch.nn.Module,
@@ -365,14 +377,5 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
)
return StandardCombineInput(hidden_states=output)
quant_info = TritonMoeQuantInfo(
w13_weight=layer.w13_weight,
w2_weight=layer.w2_weight,
use_int8_w8a8=True,
per_channel_quant=True,
w13_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
a13_scale=layer.w13_input_scale,
a2_scale=layer.w2_input_scale,
)
quant_info = self.get_triton_quant_info(layer)
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/lora/backend/chunked_backend.py
+2 -1
View File
@@ -66,6 +66,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
x: torch.Tensor,
weights: torch.Tensor,
pruned_batch_info: LoRABatchInfo = None,
stack_num: int = 1,
*args,
**kwargs,
) -> torch.Tensor:
@@ -76,7 +77,7 @@ class ChunkedSgmvLoRABackend(BaseLoRABackend):
x=x,
weights=weights,
batch_info=batch_info,
num_slices=1,
num_slices=stack_num,
)
def run_lora_b_sgemm(
python/sglang/srt/lora/backend/torch_backend.py
+7 -2
View File
@@ -39,7 +39,12 @@ class TorchNativeLoRABackend(BaseLoRABackend):
super().__init__(max_loras_per_batch, device)
def run_lora_a_sgemm(
self, x: torch.Tensor, weights: torch.Tensor, *args, **kwargs
self,
x: torch.Tensor,
weights: torch.Tensor,
stack_num: int = 1,
*args,
**kwargs,
) -> torch.Tensor:
output_tensor = sgemm_lora_a_fwd(
inputs=x,
@@ -48,7 +53,7 @@ class TorchNativeLoRABackend(BaseLoRABackend):
seg_len_tensor=self.batch_info.seg_lens_cpu,
lora_ranks=self.batch_info.lora_ranks_cpu,
scaling_tensor=self.batch_info.scalings_cpu,
num_slices=1,
num_slices=stack_num,
)
return output_tensor
python/sglang/srt/lora/backend/triton_backend.py
+2 -1
View File
@@ -53,13 +53,14 @@ class TritonLoRABackend(BaseLoRABackend):
x: torch.Tensor,
weights: torch.Tensor,
pruned_batch_info: LoRABatchInfo = None,
stack_num: int = 1,
*args,
**kwargs,
) -> torch.Tensor:
batch_info = (
pruned_batch_info if pruned_batch_info is not None else self.batch_info
)
return sgemm_lora_a_fwd(x, weights, batch_info)
return sgemm_lora_a_fwd(x, weights, batch_info, stack_num=stack_num)
def run_lora_b_sgemm(
self,
python/sglang/srt/lora/layers.py
+91 -7
View File
@@ -14,6 +14,7 @@ from sglang.srt.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
QKVParallelLinear,
ReplicatedLinear,
RowParallelLinear,
)
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
@@ -687,6 +688,94 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
return B
class ReplicatedLinearWithLoRA(BaseLayerWithLoRA):
"""LoRA wrapper for ReplicatedLinear (no TP sharding).
Used for DeepSeek MLA's fused_qkv_a_proj_with_mqa, which fuses
q_a_proj and kv_a_proj_with_mqa into a single replicated linear.
The two sub-projections have unequal output dimensions, so LoRA B
is applied via two separate sgemm calls, one per partition.
"""
first_output_dim: int = 0
def __init__(
self,
base_layer: ReplicatedLinear,
lora_backend: BaseLoRABackend,
) -> None:
super().__init__(base_layer, lora_backend)
self.output_size = base_layer.output_size
def set_lora_info(self, A_buffer: torch.Tensor, B_buffer: torch.Tensor):
self.set_lora = True
self.A_buffer = A_buffer
self.B_buffer = B_buffer
first = self.first_output_dim
if first > 0 and first < B_buffer.shape[-2]:
self.B_first = B_buffer[:, :first, :].contiguous()
self.B_second = B_buffer[:, first:, :].contiguous()
output_size = B_buffer.shape[-2]
self.first_offset = torch.tensor(
[0, first], dtype=torch.int32, device=B_buffer.device
)
self.second_offset = torch.tensor(
[0, output_size - first], dtype=torch.int32, device=B_buffer.device
)
else:
self.B_first = None
self.B_second = None
self.output_offset = torch.tensor(
[0, self.output_size],
dtype=torch.int32,
device=B_buffer.device,
)
def apply_lora(self, base_output: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
if self.B_first is not None:
rank = self.B_buffer.shape[-1]
lora_a_output = self.lora_backend.run_lora_a_sgemm(
x, self.A_buffer, stack_num=2
)
first_out = base_output[:, : self.first_output_dim]
second_out = base_output[:, self.first_output_dim :]
self.lora_backend.run_lora_b_sgemm(
x=lora_a_output[:, :rank].contiguous(),
weights=self.B_first,
output_offset=self.first_offset,
base_output=first_out,
)
self.lora_backend.run_lora_b_sgemm(
x=lora_a_output[:, rank:].contiguous(),
weights=self.B_second,
output_offset=self.second_offset,
base_output=second_out,
)
return base_output
else:
lora_a_output = self.lora_backend.run_lora_a_sgemm(x, self.A_buffer)
return self.lora_backend.run_lora_b_sgemm(
x=lora_a_output,
weights=self.B_buffer,
output_offset=self.output_offset,
base_output=base_output,
)
def forward(self, x: torch.Tensor):
bias = self.base_layer.bias if not self.base_layer.skip_bias_add else None
output = self.base_layer.quant_method.apply(self.base_layer, x, bias)
if self.set_lora:
output = self.apply_lora(output, x)
output_bias = self.base_layer.bias if self.base_layer.skip_bias_add else None
return output, output_bias
def slice_lora_a_weights(self, A: torch.Tensor, tp_rank: int):
return A
def slice_lora_b_weights(self, B: torch.Tensor, tp_rank: int):
return B
class FusedMoEWithLoRA(BaseLayerWithLoRA):
"""
Wrapper around FusedMoE that integrates LoRA into the MoE computation.
@@ -721,7 +810,6 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
# initialize triton_lora moe runner for batches with lora enabled
from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
self._lora_runner = MoeRunner(
base_layer.quant_method.runner.runner_backend,
@@ -730,12 +818,7 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
)
# Pre-compute quant info for efficiency (weights don't change during inference)
self._quant_info = TritonMoeQuantInfo(
w13_weight=base_layer.w13_weight,
w2_weight=base_layer.w2_weight,
b13=getattr(base_layer, "w13_weight_bias", None),
b2=getattr(base_layer, "w2_weight_bias", None),
)
self._quant_info = base_layer.quant_method.get_triton_quant_info(base_layer)
def set_lora_info(
self,
@@ -941,6 +1024,7 @@ def get_lora_layer(
FusedMoE: FusedMoEWithLoRA,
ParallelLMHead: ParallelLMHeadWithLoRA,
VocabParallelEmbedding: VocabParallelEmbeddingWithLoRA,
ReplicatedLinear: ReplicatedLinearWithLoRA,
QKVParallelLinear: QKVParallelLinearWithLoRA,
MergedColumnParallelLinear: MergedColumnParallelLinearWithLoRA,
ColumnParallelLinear: ColumnParallelLinearWithLoRA,
python/sglang/srt/lora/lora.py
+29 -1
View File
@@ -133,13 +133,14 @@ class LoRAAdapter(nn.Module):
)
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._rename_expert_w_to_proj(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_qkv_proj(
self, weight_names: List[str], weights: Dict[str, torch.Tensor]
@@ -242,6 +243,33 @@ class LoRAAdapter(nn.Module):
weights[gate_up_name] = weights[gate_up_name].repeat(*repeat_dims)
# else: no-op as LoRA B weight is already stacked.
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():
python/sglang/srt/lora/lora_manager.py
+19
View File
@@ -763,6 +763,25 @@ class LoRAManager:
self.lm_head_module = lora_module
continue
# Handle DeepSeek MLA fused projection: set the boundary
# between q_a and kv_a output partitions so the LoRA layer
# can apply separate B projections for each.
if (
"fused_qkv_a_proj_with_mqa" in self.target_modules
and module_name.endswith("fused_qkv_a_proj_with_mqa")
):
from sglang.srt.lora.layers import ReplicatedLinearWithLoRA
layer_id = get_layer_id(module_name)
if layer_id is None:
continue
lora_module = self.set_lora_module(module_name, module)
if isinstance(lora_module, ReplicatedLinearWithLoRA):
q_lora_rank = getattr(self.base_hf_config, "q_lora_rank", None) or 0
lora_module.first_output_dim = q_lora_rank
self.lora_modules[layer_id][module_name] = lora_module
continue
# The module should be converted if it is included in target_names
if module_name.split(".")[-1] in self.target_modules:
layer_id = get_layer_id(module_name)
python/sglang/srt/lora/mem_pool.py
+21 -4
View File
@@ -12,6 +12,7 @@ from sglang.srt.lora.lora_config import LoRAConfig
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.lora.utils import (
EMBEDDING_NAMES,
REPLICATED_LINEAR_LORA_NAMES,
ROW_PARALLELISM_LINEAR_LORA_NAMES,
LoRAType,
get_hidden_dim,
@@ -188,7 +189,11 @@ class LoRAMemoryPool:
module_name, self.base_hf_config, base_model, layer_idx
)
c = get_stacked_multiply(module_name)
if self.tp_size > 1 and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES:
if (
self.tp_size > 1
and module_name in ROW_PARALLELISM_LINEAR_LORA_NAMES
and module_name not in REPLICATED_LINEAR_LORA_NAMES
):
input_dim = divide(input_dim, self.tp_size)
if self.is_moe_module(module_name):
@@ -240,7 +245,11 @@ class LoRAMemoryPool:
_, output_dim = get_hidden_dim(
module_name, self.base_hf_config, base_model, layer_idx
)
if self.tp_size > 1 and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES:
if (
self.tp_size > 1
and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES
and module_name not in REPLICATED_LINEAR_LORA_NAMES
):
output_dim = divide(output_dim, self.tp_size)
# Check if MoE module and return appropriate shape
@@ -619,7 +628,12 @@ class LoRAMemoryPool:
if name in ["gate_up_proj_moe", "down_proj_moe"]:
if self.experts_shared_outer_loras and name == "gate_up_proj_moe":
if isinstance(weights, torch.Tensor) and weights.dim() == 3:
if weights is None:
buffer_view = target_buffer[
buffer_id, 0, : lora_rank * c, :
]
load_lora_weight_tensor(buffer_view, None)
elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
if weights.shape[0] != 1:
raise ValueError(
f"experts_shared_outer_loras is enabled but "
@@ -669,7 +683,10 @@ class LoRAMemoryPool:
if name in ["gate_up_proj_moe", "down_proj_moe"]:
if self.experts_shared_outer_loras and name == "down_proj_moe":
if isinstance(weights, torch.Tensor) and weights.dim() == 3:
if weights is None:
buffer_view = target_buffer[buffer_id, 0, :, :lora_rank]
load_lora_weight_tensor(buffer_view, None)
elif isinstance(weights, torch.Tensor) and weights.dim() == 3:
if weights.shape[0] != 1:
raise ValueError(
f"experts_shared_outer_loras is enabled but "
python/sglang/srt/lora/utils.py
+41 -3
View File
@@ -83,14 +83,47 @@ def get_hidden_dim(
config.num_attention_heads + config.num_key_value_heads * 2
)
elif module_name == "o_proj":
o_head_dim = getattr(config, "v_head_dim", None) or head_dim
return (
head_dim * config.num_attention_heads,
o_head_dim * config.num_attention_heads,
config.hidden_size,
)
elif module_name == "gate_up_proj":
return config.hidden_size, config.intermediate_size * 2
inter = config.intermediate_size
first_k = getattr(config, "first_k_dense_replace", None)
moe_freq = getattr(config, "moe_layer_freq", 1)
if (
first_k is not None
and layer_idx >= first_k
and layer_idx % moe_freq == 0
):
moe_inter = getattr(config, "moe_intermediate_size", None)
n_shared = getattr(config, "n_shared_experts", None)
if moe_inter is not None and n_shared is not None:
inter = moe_inter * n_shared
return config.hidden_size, inter * 2
elif module_name == "down_proj":
return config.intermediate_size, config.hidden_size
inter = config.intermediate_size
first_k = getattr(config, "first_k_dense_replace", None)
moe_freq = getattr(config, "moe_layer_freq", 1)
if (
first_k is not None
and layer_idx >= first_k
and layer_idx % moe_freq == 0
):
moe_inter = getattr(config, "moe_intermediate_size", None)
n_shared = getattr(config, "n_shared_experts", None)
if moe_inter is not None and n_shared is not None:
inter = moe_inter * n_shared
return inter, config.hidden_size
elif module_name == "fused_qkv_a_proj_with_mqa":
q_lora_rank = getattr(config, "q_lora_rank", None) or 0
kv_lora_rank = config.kv_lora_rank
qk_rope_head_dim = config.qk_rope_head_dim
return (
config.hidden_size,
q_lora_rank + kv_lora_rank + qk_rope_head_dim,
)
elif module_name == "gate_up_proj_moe":
moe_inter = (
getattr(config, "moe_intermediate_size", None)
@@ -151,6 +184,8 @@ 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",
}
result = set()
@@ -169,6 +204,7 @@ def get_stacked_multiply(module_name: str) -> int:
"qkv_proj": 3,
"gate_up_proj": 2,
"gate_up_proj_moe": 2,
"fused_qkv_a_proj_with_mqa": 2,
}
return stacked_rank[module_name] if module_name in stacked_rank else 1
@@ -190,6 +226,7 @@ def get_target_module_name(full_module_name: str, target_modules: Set[str]) -> s
EMBEDDING_NAMES = ["embed_tokens", "lm_head"]
ROW_PARALLELISM_LINEAR_LORA_NAMES = ["o_proj", "down_proj", "down_proj_moe"]
REPLICATED_LINEAR_LORA_NAMES = ["fused_qkv_a_proj_with_mqa"]
# Normalized module names that the LoRA system fully supports
# (i.e. get_hidden_dim, init_buffers, and init_lora_modules can handle them).
@@ -201,6 +238,7 @@ _KNOWN_LORA_TARGET_MODULES = frozenset(
"down_proj",
"embed_tokens",
"lm_head",
"fused_qkv_a_proj_with_mqa",
}
)