fix: fix NVFP4 Kimi-K2.5 weight mapping and exclude list (#18370)

This commit is contained in:
Mohammad Miadh Angkad
2026-02-08 10:23:48 +08:00
committed by GitHub
parent 64950d8f97
commit 7b83659310
2 changed files with 30 additions and 1 deletions

View File

@@ -64,6 +64,7 @@ if TYPE_CHECKING:
CombineInput,
StandardDispatchOutput,
)
from sglang.srt.models.utils import WeightsMapper
fp4_quantize = None
try:
@@ -304,6 +305,22 @@ class ModelOptQuantConfig(QuantizationConfig):
def get_scaled_act_names(self) -> List[str]:
return []
def apply_weight_name_mapper(
self, hf_to_sglang_mapper: "WeightsMapper"
): # noqa: B027
# Map excluded module patterns from HF layout to sglang layout.
# Ref: HF hf_quant_config.json for nvidia/Kimi-K2.5-NVFP4
# https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/hf_quant_config.json
if self.exclude_modules:
mapped = hf_to_sglang_mapper.apply_list(self.exclude_modules)
expanded: List[str] = []
for name in mapped:
expanded.append(name)
if name.startswith("language_model."):
expanded.append(name.removeprefix("language_model."))
# Preserve order, drop duplicates.
self.exclude_modules = list(dict.fromkeys(expanded))
class ModelOptFp8Config(ModelOptQuantConfig):
"""Configuration for ModelOpt FP8 quantization, including serialization and compatibility checks."""

View File

@@ -34,6 +34,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.deepseek_v2 import DeepseekV3ForCausalLM
from sglang.srt.models.kimi_vl_moonvit import MLP2
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.utils import add_prefix
KIMIV_VT_INFER_MAX_PATCH_NUM = 16328
@@ -643,6 +644,15 @@ def vision_tower_forward_auto(
class KimiK25ForConditionalGeneration(nn.Module):
# Support nvidia/Kimi-K2.5-NVFP4 naming: language_model.layers.*.
# Ref: HF config.json for nvidia/Kimi-K2.5-NVFP4
# https://huggingface.co/nvidia/Kimi-K2.5-NVFP4/blob/main/config.json
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_prefix={
"language_model.layers.": "language_model.model.layers.",
}
)
def __init__(
self,
config: KimiK25Config,
@@ -710,7 +720,9 @@ class KimiK25ForConditionalGeneration(nn.Module):
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
"""Load weights for the model, separating vision and language weights"""
weights = list(weights)
mapper = getattr(self, "hf_to_sglang_mapper", None)
if mapper is not None:
weights = mapper.apply(weights)
# Separate vision tower weights and language model weights
vision_weights = []