Fix V4 tensor naming: .scale companions, w1/w3 expert pairs, ffn.gate, hc_* preserve
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user