Add ModelOpt Qwen3 nvfp4 support (#20101)
Signed-off-by: Zhiyu Cheng <zhiyuc@nvidia.com>
This commit is contained in:
@@ -764,39 +764,41 @@ def maybe_remap_kv_scale_name(name: str, params_dict: dict) -> Optional[str]:
|
||||
return None
|
||||
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"
|
||||
# Define scale name mapping patterns in order of precedence
|
||||
scale_mapping_patterns = [
|
||||
# ModelOpt format: .self_attn.{k,v}_proj.{k,v}_scale ->
|
||||
# .self_attn.attn.{k,v}_scale
|
||||
(r"\.self_attn\.([kv])_proj\.([kv])_scale$",
|
||||
r".self_attn.attn.\2_scale"),
|
||||
# QKV proj format: .self_attn.qkv_proj.{k,v}_scale ->
|
||||
# .self_attn.attn.{k,v}_scale
|
||||
(r"\.self_attn\.qkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"),
|
||||
# Qwen3 MoE format: .self_attn.qkqkv_proj.{k,v}_scale ->
|
||||
# .self_attn.attn.{k,v}_scale
|
||||
(r"\.self_attn\.qkqkv_proj\.([kv])_scale$", r".self_attn.attn.\1_scale"
|
||||
),
|
||||
# Default format: .{k,v}_scale -> .attn.{k,v}_scale
|
||||
(r"\.([kv])_scale$", r".attn.\1_scale"),
|
||||
]
|
||||
# Also support qkv_proj scale parameters (from stacked parameter processing)
|
||||
qkv_proj_scale_names = [
|
||||
".self_attn.qkv_proj.k_scale", ".self_attn.qkv_proj.v_scale"
|
||||
]
|
||||
for scale_name in possible_scale_names:
|
||||
if name.endswith(scale_name):
|
||||
if any(mo_scale_name in name
|
||||
for mo_scale_name in modelopt_scale_names):
|
||||
remapped_name = name.replace(
|
||||
f".self_attn.{scale_name[1]}_proj{scale_name}",
|
||||
f".self_attn.attn{scale_name}")
|
||||
elif any(qkv_scale_name in name
|
||||
for qkv_scale_name in qkv_proj_scale_names):
|
||||
# Handle qkv_proj scale parameters
|
||||
remapped_name = name.replace(
|
||||
f".self_attn.qkv_proj{scale_name}",
|
||||
f".self_attn.attn{scale_name}")
|
||||
else:
|
||||
remapped_name = name.replace(scale_name, f".attn{scale_name}")
|
||||
if remapped_name not in params_dict:
|
||||
logger.warning_once(
|
||||
"Found %s in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). %s is not loaded.", # noqa: E501
|
||||
scale_name,
|
||||
name,
|
||||
remapped_name,
|
||||
scale_name,
|
||||
)
|
||||
return None
|
||||
return remapped_name
|
||||
|
||||
# Check if name ends with k_scale or v_scale
|
||||
if name.endswith((".k_scale", ".v_scale")):
|
||||
import regex as re
|
||||
|
||||
for pattern, replacement in scale_mapping_patterns:
|
||||
if re.search(pattern, name):
|
||||
remapped_name = re.sub(pattern, replacement, name)
|
||||
if remapped_name not in params_dict:
|
||||
scale_type = name.split(".")[-1]
|
||||
logger.warning_once(
|
||||
"Found %s in the checkpoint (e.g. %s), but not found the expected name in the model (e.g. %s). %s is not loaded.", # noqa: E501
|
||||
scale_type,
|
||||
name,
|
||||
remapped_name,
|
||||
scale_type,
|
||||
)
|
||||
return None
|
||||
return remapped_name
|
||||
|
||||
# If there were no matches, return the untouched param name
|
||||
return name
|
||||
|
||||
Reference in New Issue
Block a user