[LoRA] feat: Support LoRA for DeepSeek V4 (#53361)
Signed-off-by: Hollow Man <hollowman@opensuse.org> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed
Aug 24, 2026 at 06:57 UTC
702e1d718646b5290f17533c04932d58bf03dad6
3 files changed
+70
-5
vllm/lora/layers/base.py
+18
@@ -15,6 +15,24 @@ if TYPE_CHECKING:
15
16
17
class BaseLayerWithLoRA(nn.Module):
18
+ def __getattr__(self, name):
19
+ d = self.__dict__
20
+ if name in d.get("_parameters", ()):
21
+ return d["_parameters"][name]
22
+ if name in d.get("_buffers", ()):
23
+ return d["_buffers"][name]
24
+ if name in d.get("_modules", ()):
25
+ return d["_modules"][name]
26
+ # Forward public misses to ``base_layer``; private names are framework
27
+ # bookkeeping and must stay local.
28
+ if not name.startswith("_"):
29
+ base_layer = d.get("_modules", {}).get("base_layer")
30
+ if base_layer is not None:
31
+ return getattr(base_layer, name)
32
+ raise AttributeError(
33
+ f"{type(self).__name__!r} object has no attribute {name!r}"
34
+ )
35
+
36
def load_weights(
37
self, weights: Iterable[tuple[str, torch.Tensor]]
38
) -> Iterable[str]:
vllm/model_executor/layers/fused_moe/routed_experts.py
+17
-4
@@ -1002,6 +1002,8 @@ class RoutedExperts(PluggableLayer):
1002
See `build_expert_params_mapping` for the returned tuple format.
1003
"""
1004
has_base_layer = any(".base_layer." in n for n, _ in model.named_parameters())
1005
+ prefix = "base_layer." if has_base_layer else ""
1006
+ # These loaders index ``params_dict[full_name]``, so both sides get it.
1007
return RoutedExperts.build_expert_params_mapping(
1008
ckpt_gate_proj_name,
1009
ckpt_down_proj_name,
@@ -1009,7 +1011,8 @@ class RoutedExperts(PluggableLayer):
1011
num_experts,
1012
num_redundant_experts,
1013
routed_experts_prefix,
1012
- "base_layer." if has_base_layer else "",
1014
+ lora_base_layer_prefix=prefix,
1015
+ lora_base_layer_prefix_on_param_name=prefix,
1016
)
1017
1018
@staticmethod
@@ -1021,6 +1024,7 @@ class RoutedExperts(PluggableLayer):
1024
num_redundant_experts: int = 0,
1025
routed_experts_prefix: str = "routed_experts",
1026
lora_base_layer_prefix: str = "",
1027
+ lora_base_layer_prefix_on_param_name: str = "",
1028
include_fused: bool = False,
1029
) -> list[tuple[str, str, int, str]]:
1030
"""
@@ -1035,7 +1039,14 @@ class RoutedExperts(PluggableLayer):
1039
ckpt_up_proj_name: Name of up projection in checkpoint
1040
num_experts: Number of logical (non-redundant) experts
1041
num_redundant_experts: Number of redundant experts
1038
- lora_base_layer_prefix: Prefix to add if this layer is a LoRA base layer
1042
+ lora_base_layer_prefix: LoRA ``base_layer.`` prefix for the
1043
+ ``weight_name`` (checkpoint) side
1044
+ lora_base_layer_prefix_on_param_name: same, for the ``param_name``
1045
+ side. Independent because ``get_expert_mapping`` resolves
1046
+ ``param_name`` via ``getattr`` against this layer's bare
1047
+ ``w13_weight``/``w2_weight`` (no prefix), while
1048
+ ``make_expert_params_mapping`` indexes the model-wide
1049
+ ``params_dict`` (prefix included).
1050
include_fused: Prepend the fused pre-fused-checkpoint entries
1051
1052
Returns:
@@ -1061,8 +1072,10 @@ class RoutedExperts(PluggableLayer):
1072
if routed_experts_prefix != "":
1073
routed_experts_prefix = f"{routed_experts_prefix}."
1074
1064
- w13 = f"experts.{lora_base_layer_prefix}{routed_experts_prefix}w13_"
1065
- w2 = f"experts.{lora_base_layer_prefix}{routed_experts_prefix}w2_"
1075
+ w13 = (
1076
+ f"experts.{lora_base_layer_prefix_on_param_name}{routed_experts_prefix}w13_"
1077
+ )
1078
+ w2 = f"experts.{lora_base_layer_prefix_on_param_name}{routed_experts_prefix}w2_"
1079
1080
fused_mapping = []
1081
if include_fused:
vllm/models/deepseek_v4/nvidia/model.py
+35
-1
@@ -55,6 +55,7 @@ from vllm.model_executor.models.interfaces import (
55
EagleModelMixin,
56
MixtureOfExperts,
57
SupportsEagle3,
58
+ SupportsLoRA,
59
SupportsPP,
60
)
61
from vllm.model_executor.models.utils import (
@@ -1468,6 +1469,11 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1469
1470
if is_pp_missing_parameter(name, self):
1471
break
1472
+ if name not in params_dict:
1473
+ head, _, leaf = name.rpartition(".")
1474
+ suffixed = f"{head}.base_layer.{leaf}"
1475
+ if suffixed in params_dict:
1476
+ name = suffixed
1477
param = params_dict[name]
1478
weight_loader = param.weight_loader
1479
weight_loader(param, loaded_weight, shard_id)
@@ -1522,6 +1528,13 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1528
else:
1529
if is_pp_missing_parameter(name, self):
1530
continue
1531
+ # Non-LoRA params on a LoRA-wrapped module live at
1532
+ # ``<head>.base_layer.<leaf>``; the checkpoint is plain.
1533
+ if name not in params_dict:
1534
+ head, _, leaf = name.rpartition(".")
1535
+ suffixed = f"{head}.base_layer.{leaf}"
1536
+ if suffixed in params_dict:
1537
+ name = suffixed
1538
param = params_dict[name]
1539
weight_loader = getattr(
1540
param, "weight_loader", default_weight_loader
@@ -1598,6 +1611,10 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1611
# shared experts use Fp8LinearMethod's block scales, which
1612
# register as ``weight_scale_inv``.
1613
scale_regex = {
1614
+ # ``.base_layer.``-namespace variant (LoRA-wrapped experts).
1615
+ re.compile(
1616
+ r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$"
1617
+ ): r"\1.weight_scale",
1618
re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale",
1619
re.compile(r"\.scale$"): ".weight_scale_inv",
1620
}
@@ -1606,6 +1623,10 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper:
1623
# scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys
1624
# there.
1625
scale_regex = {
1626
+ # ``.base_layer.``-namespace variant of the above.
1627
+ re.compile(
1628
+ r"(\.experts\.\d+\.w[123]\.base_layer)\.scale$"
1629
+ ): r"\1.weight_scale_inv",
1630
re.compile(r"\.scale$"): ".weight_scale_inv",
1631
}
1632
return WeightsMapper(
@@ -1667,7 +1688,11 @@ class DeepseekV4MixtureOfExperts(MixtureOfExperts):
1688
1689
1690
class DeepseekV4ForCausalLM(
1670
- nn.Module, SupportsPP, SupportsEagle3, DeepseekV4MixtureOfExperts
1691
+ nn.Module,
1692
+ SupportsPP,
1693
+ SupportsEagle3,
1694
+ SupportsLoRA,
1695
+ DeepseekV4MixtureOfExperts,
1696
):
1697
model_cls = DeepseekV4Model
1698
@@ -1675,6 +1700,15 @@ class DeepseekV4ForCausalLM(
1700
# Overridden per-instance in __init__ when expert_dtype != "fp4".
1701
hf_to_vllm_mapper = _make_deepseek_v4_weights_mapper("fp4")
1702
1703
+ packed_modules_mapping = {
1704
+ "gate_up_proj": ["w1", "w3"],
1705
+ "fused_wqa_wkv": ["wq_a", "wkv"],
1706
+ "fused_wkv_wgate": ["wkv", "wgate"],
1707
+ }
1708
+
1709
+ # The MTP draft head is not LoRA-adapted.
1710
+ lora_skip_prefixes = ["mtp."]
1711
+
1712
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1713
super().__init__()
1714