Fix V4 tensor naming: .scale companions, w1/w3 expert pairs, ffn.gate, hc_* preserve

This commit is contained in:
2026-05-07 00:03:20 +00:00
parent 2b7f063e39
commit 4470653e15

View File

@@ -57,18 +57,24 @@ from tqdm import tqdm
PRESERVE_REGEXES = [
r".*lm_head.*",
r".*embed_tokens.*",
r".*\.mlp\.gate(\.weight)?$", # MoE router (NOT gate_proj)
r".*\.(mlp|ffn)\.gate(\.weight)?$", # MoE router (NOT gate_proj)
r".*norm.*",
r".*indexer.*", # V3.2 DSA / V4 CSA indexer
r".*hyper_conn.*", # V4 mHC
r".*\.mhc.*",
r".*hc_attn.*", # V4 hyper-connection attn
r".*hc_ffn.*", # V4 hyper-connection ffn
r".*hc_head.*", # V4 hyper-connection head
r".*scoring.*",
r".*attn_sink.*", # V4 attention sink
r".*compressor\.ape.*", # V4 compressor absolute pos encoding
r".*tid2eid.*", # V4 MoE token-to-expert mapping
r".*\.bias$", # any biases
]
PRESERVE_RE = re.compile("|".join(f"(?:{p})" for p in PRESERVE_REGEXES))
# Identify expert pairs that need joint global scale
EXPERT_PAIR_RE = re.compile(r"(.*experts\.\d+)\.(gate_proj|up_proj)\.weight$")
EXPERT_PAIR_RE = re.compile(r"(.*experts\.\d+)\.(w1|w3)\.weight$")
def is_preserve(name: str) -> bool:
@@ -293,10 +299,10 @@ def build_plan(src_dir: Path):
# Classify everything else
solo_quantize = []
preserve = []
scale_companions = [] # weight_scale_inv tensors that get consumed during dequant
scale_companions = [] # .scale tensors that get consumed during dequant
for n in weight_map:
if n.endswith(".weight_scale_inv"):
if n.endswith(".scale") and n.replace(".scale", ".weight") in weight_map:
scale_companions.append(n)
continue
if n in paired_names:
@@ -354,7 +360,7 @@ class ShardCache:
def load_weight_and_scale(cache: ShardCache, weight_map, name):
"""Load an FP8 weight with its weight_scale_inv companion (if any)."""
weight = cache.get(weight_map[name]).get_tensor(name)
scale_name = name.replace(".weight", ".weight_scale_inv")
scale_name = name.replace(".weight", ".scale")
scale = None
if scale_name in weight_map:
scale = cache.get(weight_map[scale_name]).get_tensor(scale_name)