diff --git a/fp8_to_nvfp4_streaming.py b/fp8_to_nvfp4_streaming.py index e364184..16a56c6 100644 --- a/fp8_to_nvfp4_streaming.py +++ b/fp8_to_nvfp4_streaming.py @@ -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)