[Fix] Load HunyuanV3 NextN final_layernorm into the draft head's output norm (#30331)

Co-authored-by: Auroter <auroter@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-07-13 20:37:36 +08:00
parent a74bee2261
commit 9fec359a60
2 changed files with 97 additions and 0 deletions
@@ -186,6 +186,10 @@ class HYV3ForCausalLMNextN(nn.Module):
                subname = name[len(nextn_prefix) :]
                if any(subname.startswith(s) for s in spec_weight_names):
                    name = f"model.{subname}"
+                elif subname.startswith("final_layernorm"):
+                    # Released checkpoints store the draft head's output norm
+                    # as model.layers.<N>.final_layernorm.weight.
+                    name = "model.shared_head.norm.weight"
                else:
                    name = f"model.decoder.{subname}"
            elif name == "model.shared_head.norm.weight":