[Quantization] Support config.json quantization_config format, fix exclude_modules matching, and fix KV cache scale loading for Nemotron (#18546)

Signed-off-by: root <dafrimi@nvidia.com>
This commit is contained in:
danielafrimi
2026-02-21 16:14:29 +08:00
committed by GitHub
parent e239f8aa85
commit 33c33a7de9
3 changed files with 100 additions and 71 deletions
+13 -8
View File
@@ -1214,17 +1214,22 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
return remapped_name
possible_scale_names = [".k_scale", ".v_scale"]
modelopt_scale_names = [".self_attn.k_proj.k_scale", ".self_attn.v_proj.v_scale"]
# Patterns where modelopt stores scales under k_proj/v_proj
# but the model expects them under attn (RadixAttention)
modelopt_attn_prefixes = [".self_attn.", ".mixer."]
for scale_name in possible_scale_names:
if name.endswith(scale_name):
# Check and remap the name based on modelopt scale names
if any(
modelopt_scale_name in name
for modelopt_scale_name in modelopt_scale_names
):
# Check if this is a modelopt-style scale under k_proj/v_proj
matched_prefix = None
for attn_prefix in modelopt_attn_prefixes:
if f"{attn_prefix}{scale_name[1]}_proj{scale_name}" in name:
matched_prefix = attn_prefix
break
if matched_prefix is not None:
remapped_name = name.replace(
f".self_attn.{scale_name[1]}_proj{scale_name}",
f".self_attn.attn{scale_name}",
f"{matched_prefix}{scale_name[1]}_proj{scale_name}",
f"{matched_prefix}attn{scale_name}",
)
else:
remapped_name = name.replace(scale_name, f".attn{scale_name}")