79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import re
|
|
import torch
|
|
from .qwen2 import convert_qwen2_to_hf
|
|
|
|
|
|
def convert_mimo_to_hf(args, name, param):
|
|
"""
|
|
Convert MiMo model parameters from Megatron to HuggingFace format.
|
|
|
|
MiMo extends Qwen2 with MTP (Multi-Token Prediction) layers.
|
|
"""
|
|
|
|
if "mtp" in name:
|
|
return convert_mimo_mtp_param(args, name, param)
|
|
|
|
return convert_qwen2_to_hf(args, name, param)
|
|
|
|
|
|
def convert_mimo_mtp_param(args, name, param):
|
|
"""
|
|
Convert MTP layer parameters from Megatron to HuggingFace format.
|
|
|
|
MTP layers in MiMo contain:
|
|
- LayerNorms (token_layernorm, hidden_layernorm, final_layernorm)
|
|
- Input projection (input_proj)
|
|
- Self attention (reuses Qwen2 attention structure)
|
|
- MLP (reuses Qwen2 MLP structure)
|
|
|
|
Based on MimoBridge._convert_mtp_param logic (reverse mapping)
|
|
"""
|
|
mtp_pattern = r"module\.module\.mtp\.layers\.(\d+)\.(.+)"
|
|
match = re.match(mtp_pattern, name)
|
|
|
|
if not match:
|
|
raise ValueError(f"Invalid MTP parameter name: {name}")
|
|
|
|
layer_idx, component = match.groups()
|
|
|
|
# Direct mappings for MTP-specific components (Megatron -> HF)
|
|
# Based on MimoBridge direct_name_mapping (reversed)
|
|
direct_mappings = {
|
|
"enorm.weight": f"model.mtp_layers.{layer_idx}.token_layernorm.weight",
|
|
"hnorm.weight": f"model.mtp_layers.{layer_idx}.hidden_layernorm.weight",
|
|
"eh_proj.weight": f"model.mtp_layers.{layer_idx}.input_proj.weight",
|
|
"final_layernorm.weight": f"model.mtp_layers.{layer_idx}.final_layernorm.weight",
|
|
}
|
|
if component == "eh_proj.weight":
|
|
first_half, second_half = param.chunk(2, dim=1)
|
|
param = torch.cat([second_half, first_half], dim=1)
|
|
|
|
# Check direct mappings first
|
|
if component in direct_mappings:
|
|
return [(direct_mappings[component], param)]
|
|
|
|
# Handle transformer_layer components
|
|
if component.startswith("transformer_layer."):
|
|
# Remove "transformer_layer." prefix
|
|
transformer_component = component[len("transformer_layer.") :]
|
|
|
|
# Create proxy name for reusing existing Qwen2 conversion functions
|
|
proxy_name = f"module.module.decoder.layers.{layer_idx}.{transformer_component}"
|
|
|
|
# Use existing convert_qwen2_to_hf function for transformer components
|
|
results = convert_qwen2_to_hf(args, proxy_name, param)
|
|
|
|
# Replace model.layers with mtp_layers in results
|
|
converted_results = []
|
|
for hf_name, hf_param in results:
|
|
# Replace model.layers.{idx} with mtp_layers.{idx}
|
|
hf_name = hf_name.replace(f"model.layers.{layer_idx}", f"model.mtp_layers.{layer_idx}")
|
|
converted_results.append((hf_name, hf_param))
|
|
|
|
return converted_results
|
|
|
|
raise ValueError(f"Unknown MTP component: {component} in {name}")
|