# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import re
from dataclasses import asdict, dataclass, field
from functools import cached_property, partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
import torch.nn.functional as F
import yaml
from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
from megatron.core.transformer.identity_op import IdentityOp
from megatron.core.transformer.transformer_config import MLATransformerConfig
from safetensors.torch import load_file
from torch import nn
from transformers import AutoConfig

from nemo.collections.llm.gpt.model.base import (
    HAVE_TE,
    GPTConfig,
    GPTModel,
    gpt_data_step,
    torch_dtype_from_dict_config,
)
from nemo.collections.llm.utils import is_safe_repo
from nemo.export.trt_llm.nemo_ckpt_loader.nemo_file import load_distributed_model_weights
from nemo.lightning import io, teardown
from nemo.lightning.io.state import TransformFns, _ModelState
from nemo.lightning.pytorch.optim import OptimizerModule
from nemo.lightning.pytorch.utils import dtype_from_hf
from nemo.utils import logging

if TYPE_CHECKING:
    from megatron.core.transformer import ModuleSpec
    from transformers import AutoModelForCausalLM
    from transformers import DeepseekV3Config as HFDeepseekV3Config

    from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
    from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec

if HAVE_TE:
    from megatron.core.utils import is_te_min_version


@dataclass
class DeepSeekConfig(MLATransformerConfig, GPTConfig):
    """
    Base config for DeepSeek V2 and V3 models.
    """

    transformer_layer_spec: Union['ModuleSpec', Callable[["GPTConfig"], 'ModuleSpec']] = partial(
        get_gpt_decoder_block_spec, use_transformer_engine=HAVE_TE
    )

    # Model
    normalization: str = "RMSNorm"
    activation_func: Callable = F.silu
    gated_linear_unit: bool = True  # swiglu
    position_embedding_type: str = "rope"
    add_bias_linear: bool = False
    share_embeddings_and_output_weights: bool = False
    num_attention_heads: int = 128
    kv_channels: int = 128
    max_position_embeddings: int = 4096
    seq_length: int = 4096
    rotary_base: float = 10000.0
    make_vocab_size_divisible_by: int = 3200
    mtp_num_layers: Optional[int] = None
    mtp_loss_scaling_factor: Optional[float] = None

    # Regularization
    attention_dropout: float = 0.0
    hidden_dropout: float = 0.0
    qk_layernorm: bool = True

    # MoE
    moe_grouped_gemm: bool = True
    moe_router_pre_softmax: bool = True
    moe_token_dispatcher_type: str = "alltoall"
    moe_router_load_balancing_type: str = 'seq_aux_loss'
    moe_shared_expert_overlap: bool = True
    moe_router_dtype: Optional[str] = 'fp32'

    # MLA
    q_lora_rank: int = 1536
    kv_lora_rank: int = 512
    qk_head_dim: int = 128
    qk_pos_emb_head_dim: int = 64
    v_head_dim: int = 128
    rotary_scaling_factor: float = 40
    mscale: float = 1.0
    mscale_all_dim: float = 1.0

    # Miscellaneous
    init_method_std: float = 0.006
    layernorm_epsilon: float = 1e-6
    bf16: bool = True
    params_dtype: torch.dtype = torch.bfloat16
    async_tensor_model_parallel_allreduce: bool = True
    attention_softmax_in_fp32: bool = False
    persist_layer_norm: bool = True
    num_layers_in_first_pipeline_stage: Optional[int] = None
    num_layers_in_last_pipeline_stage: Optional[int] = None
    account_for_embedding_in_pipeline_split: bool = False
    account_for_loss_in_pipeline_split: bool = False

    # fusions
    apply_rope_fusion: bool = False
    bias_activation_fusion: bool = True
    bias_dropout_fusion: bool = True
    masked_softmax_fusion: bool = True
    gradient_accumulation_fusion: bool = True
    cross_entropy_loss_fusion: bool = True
    cross_entropy_fusion_impl: str = "te"
    moe_permute_fusion: bool = is_te_min_version("2.1.0") if HAVE_TE else False

    def __post_init__(self):
        super().__post_init__()
        if self.mtp_num_layers is not None:
            self.data_step_fn = partial(gpt_data_step, use_mtp=True)


@dataclass
class DeepSeekV2Config(DeepSeekConfig):
    """
    DeepSeek-V2 Model: https://github.com/deepseek-ai/DeepSeek-V2
    """

    num_layers: int = 60
    hidden_size: int = 5120
    ffn_hidden_size: int = 12288
    num_moe_experts: int = 160
    moe_ffn_hidden_size: int = 1536
    moe_shared_expert_intermediate_size: int = 3072  # 1536 * 2 shared experts
    moe_layer_freq: Union[int, List[int]] = field(default_factory=lambda: [0] + [1] * 59)  # first layer is dense
    moe_router_topk: int = 6
    moe_router_num_groups: int = 8
    moe_router_group_topk: int = 3
    moe_router_topk_scaling_factor: float = 16.0
    moe_aux_loss_coeff: float = 1e-3
    mscale: float = 0.707
    mscale_all_dim: float = 0.707


@dataclass
class DeepSeekV2LiteConfig(DeepSeekV2Config):
    """
    DeepSeek-V2-Lite Model: https://github.com/deepseek-ai/DeepSeek-V2
    HuggingFace: https://huggingface.co/deepseek-ai/DeepSeek-V2-Lite
    """

    num_layers: int = 27
    hidden_size: int = 2048
    ffn_hidden_size: int = 10944
    num_attention_heads: int = 16
    kv_channels: int = 16
    q_lora_rank: int = None
    num_moe_experts: int = 64
    moe_ffn_hidden_size: int = 1408
    moe_shared_expert_intermediate_size: int = 2816  # 1408 * 2 shared experts
    moe_layer_freq: Union[int, List[int]] = field(default_factory=lambda: [0] + [1] * 26)  # first layer is dense
    moe_router_topk: int = 6
    moe_router_num_groups: int = 1
    moe_router_group_topk: int = 1
    moe_router_topk_scaling_factor: float = 1.0


@dataclass
class DeepSeekV3Config(DeepSeekConfig):
    """
    DeepSeek-V3 Model: https://github.com/deepseek-ai/DeepSeek-V3
    """

    num_layers: int = 61
    hidden_size: int = 7168
    ffn_hidden_size: int = 18432
    num_moe_experts: int = 256
    moe_ffn_hidden_size: int = 2048
    moe_shared_expert_intermediate_size: int = 2048  # 2048 * 1 shared expert
    moe_layer_freq: Union[int, List[int]] = field(
        default_factory=lambda: [0] * 3 + [1] * 58
    )  # first three layers are dense
    moe_router_topk: int = 8
    moe_router_num_groups: int = 8
    moe_router_group_topk: int = 4
    moe_router_topk_scaling_factor: float = 2.5
    moe_aux_loss_coeff: float = 1e-4
    make_vocab_size_divisible_by: int = 1280
    moe_router_score_function: str = "sigmoid"
    moe_router_enable_expert_bias: bool = True
    moe_router_bias_update_rate: float = 1e-3
    mscale: float = 1.0
    mscale_all_dim: float = 1.0


class DeepSeekModel(GPTModel):
    # pylint: disable=C0115,C0116
    def __init__(
        self,
        config: Optional[DeepSeekConfig] = None,
        optim: Optional[OptimizerModule] = None,
        tokenizer: Optional["TokenizerSpec"] = None,
        model_transform: Optional[Callable[[nn.Module], nn.Module]] = None,
    ):
        super().__init__(
            config or DeepSeekV2Config(), optim=optim, tokenizer=tokenizer, model_transform=model_transform
        )


@io.model_importer(DeepSeekModel, ext="hf")
class HFDeepSeekImporter(io.ModelConnector["AutoModelForCausalLM", DeepSeekModel]):
    # pylint: disable=C0115,C0116
    def init(self) -> DeepSeekModel:
        return DeepSeekModel(self.config, tokenizer=self.tokenizer)

    def apply(self, output_path: Path, convert_mtp: bool = False, trust_remote_code: bool | None = None) -> Path:
        from transformers import AutoModelForCausalLM

        self.trust_remote_code = trust_remote_code
        self.convert_mtp = convert_mtp
        self._verify_source()
        source = AutoModelForCausalLM.from_pretrained(
            str(self),
            trust_remote_code=is_safe_repo(
                trust_remote_code=self.trust_remote_code,
                hf_path=str(self),
            ),
            torch_dtype='auto',
        )
        target = self.init()
        trainer = self.nemo_setup(target)
        self.convert_state(source, target)
        self.nemo_save(output_path, trainer)

        logging.info(f"Converted DeepSeek model to Nemo, model saved to {output_path}")

        teardown(trainer, target)
        del trainer, target

        return output_path

    def _verify_source(self):
        source_config = AutoConfig.from_pretrained(
            str(self),
            trust_remote_code=is_safe_repo(
                trust_remote_code=self.trust_remote_code,
                hf_path=str(self),
            ),
        )
        assert 'quantization_config' not in source_config, (
            "HuggingFace cannot load DeepSeek V3's FP8 checkpoint directly. You must convert the checkpoint "
            "to BF16. See NeMo documentation for more details: "
            "https://nemo-framework-tme.gitlab-master-pages.nvidia.com/documentation/user-guide/latest/llms/"
            "deepseek_v3.html#nemo-2-0-finetuning-recipes "
        )

    def _modify_source_state(self, source: nn.Module) -> _ModelState:
        """
        In deepseek, HF weight `model.layers.*.post_attention_layernorm.weight` is mapped to mcore weight
        a) `decoder.layers.*.mlp.linear_fc1.layer_norm_weight`, if the layer is dense
        b) `decoder.layers.*.pre_mlp_layernorm.weight`, if the layer is MoE

        We rename model.layers.*.post_attention_layernorm.weight in the first case to prevent a one-to-many mapping
        """

        state_dict = source.state_dict()

        for layer_i, use_moe in enumerate(self.config.moe_layer_freq):
            if use_moe == 0:
                weight = state_dict.pop(f"model.layers.{layer_i}.post_attention_layernorm.weight")
                state_dict[f"model.layers.{layer_i}.dense-post_attention_layernorm.weight"] = weight

        source = _ModelState(state_dict)
        return source

    def _add_mtp_to_source(self, source: nn.Module | _ModelState) -> None:
        # Load MTP weights from disk, since it is not in HF model
        mtp_hf_layer_low = self.config.num_layers  # 61 if DeepSeek V3
        mtp_hf_layer_high = self.config.num_layers + self.config.mtp_num_layers - 1  # 61 if DeepSeek V3
        # Identify which file to load
        with open(self / "model.safetensors.index.json", 'r') as file:
            manifest = json.load(file)

        safetensor_files_to_load = set()
        mtp_hf_keys = set()
        for k, fname in manifest['weight_map'].items():
            if match := re.match(r".*\.layers\.(\d+)\.", k):
                if mtp_hf_layer_low <= int(match.group(1)) <= mtp_hf_layer_high:
                    safetensor_files_to_load.add(fname)
                    mtp_hf_keys.add(k)

        mtp_state_dict = {}
        for safetensor_file in safetensor_files_to_load:
            for k, v in load_file(self / safetensor_file).items():
                if k in mtp_hf_keys:
                    # ensure HF keys "mtp" are alphabetically after "layers",
                    # since mcore keys "mtp" are after "decoder"
                    # This allows us to reuse the mapping and transforms for MTP
                    mtp_state_dict[k.replace(".layers.", ".mtp.")] = v
        source.state_dict().update(mtp_state_dict)

    def convert_state(self, source, target):
        # pylint: disable=C0301
        mapping = {
            # Embed
            "model.embed_tokens.weight": "embedding.word_embeddings.weight",
            # Attention
            "**.input_layernorm.weight": "**.input_layernorm.weight",
            "**.self_attn.o_proj.weight": "**.self_attention.linear_proj.weight",
            "**.self_attn.q_a_proj.weight": "**.self_attention.linear_q_down_proj.weight",
            "**.self_attn.q_b_proj.weight": "**.self_attention.linear_q_up_proj.weight",
            "**.self_attn.kv_a_proj_with_mqa.weight": "**.self_attention.linear_kv_down_proj.weight",
            "**.self_attn.kv_b_proj.weight": "**.self_attention.linear_kv_up_proj.weight",
            "**.self_attn.q_a_layernorm.weight": "**.self_attention.linear_q_up_proj.layer_norm_weight",
            "**.self_attn.kv_a_layernorm.weight": "**.self_attention.linear_kv_up_proj.layer_norm_weight",
            "**.dense-post_attention_layernorm.weight": "**.mlp.linear_fc1.layer_norm_weight",
            "**.post_attention_layernorm.weight": "**.pre_mlp_layernorm.weight",
            # Dense MLP
            # **.mlp.{gate|up}_proj.weight: **.mlp.linear_fc1.weight
            "**.mlp.down_proj.weight": "**.mlp.linear_fc2.weight",
            # MoE
            "**.mlp.gate.weight": "**.mlp.router.weight",
            # **.mlp.experts.*.{gate|up}_proj.weight: **.mlp.experts.linear_fc1.weight*
            "**.mlp.experts.*.down_proj.weight": "**.mlp.experts.linear_fc2.weight*",
            # **.mlp.shared_experts.{gate|up}_proj.weight： **.mlp.shared_experts.linear_fc1.weight
            "**.mlp.shared_experts.down_proj.weight": "**.mlp.shared_experts.linear_fc2.weight",
            # LM Head
            "model.norm.weight": "decoder.final_layernorm.weight",
            "lm_head.weight": "output_layer.weight",
        }
        # For lite model
        if self.config.q_lora_rank is None:
            del mapping["**.self_attn.q_a_proj.weight"]
            del mapping["**.self_attn.q_b_proj.weight"]
            mapping["**.self_attn.q_proj.weight"] = "**.self_attention.linear_q_proj.weight"
        # Account for Mcore local spec
        if self.config.q_lora_rank is not None and not isinstance(
            target.module.decoder.layers[0].self_attention.q_layernorm, IdentityOp
        ):
            mapping["**.self_attn.q_a_layernorm.weight"] = "**.self_attention.q_layernorm.weight"

        if not isinstance(target.module.decoder.layers[0].self_attention.kv_layernorm, IdentityOp):
            mapping["**.self_attn.kv_a_layernorm.weight"] = "**.self_attention.kv_layernorm.weight"

        if not isinstance(target.module.decoder.layers[0].pre_mlp_layernorm, IdentityOp):
            del mapping["**.dense-post_attention_layernorm.weight"]
            source = _ModelState(source.state_dict)
        else:
            source = self._modify_source_state(source)

        if hasattr(self.config, "moe_router_enable_expert_bias") and self.config.moe_router_enable_expert_bias:
            mapping.update(
                {
                    "**.mlp.gate.e_score_correction_bias": "**.mlp.router.expert_bias",
                }
            )

        transforms = [
            io.state_transform(
                source_key=("**.mlp.gate_proj.weight", "**.mlp.up_proj.weight"),
                target_key="**.mlp.linear_fc1.weight",
                fn=TransformFns.merge_fc1,
            ),
            io.state_transform(
                source_key=(
                    "**.mlp.experts.*.gate_proj.weight",
                    "**.mlp.experts.*.up_proj.weight",
                ),
                target_key="**.mlp.experts.linear_fc1.weight*",
                fn=TransformFns.merge_fc1,
            ),
            io.state_transform(
                source_key=(
                    "**.mlp.shared_experts.gate_proj.weight",
                    "**.mlp.shared_experts.up_proj.weight",
                ),
                target_key="**.mlp.shared_experts.linear_fc1.weight",
                fn=TransformFns.merge_fc1,
            ),
        ]

        # Convert MTP weights
        if getattr(self.config, "mtp_num_layers", None) and self.convert_mtp:
            self._add_mtp_to_source(source)
            mapping.update(
                {
                    'model.mtp.*.eh_proj.weight': "mtp.layers.*.eh_proj.weight",
                    'model.mtp.*.enorm.weight': "mtp.layers.*.enorm.weight",
                    'model.mtp.*.hnorm.weight': "mtp.layers.*.hnorm.weight",
                    'model.mtp.*.shared_head.norm.weight': "mtp.layers.*.shared_head_norm.weight",
                }
            )

        return io.apply_transforms(
            source,
            target,
            mapping=mapping,
            transforms=transforms,
        )

    @cached_property
    def tokenizer(self) -> "AutoTokenizer":
        from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer

        return AutoTokenizer(self.save_hf_tokenizer_assets(str(self)), use_fast=True)

    @cached_property
    def config(self) -> DeepSeekConfig:
        from transformers import AutoConfig as HFAutoConfig
        from transformers import GenerationConfig

        source = HFAutoConfig.from_pretrained(
            str(self),
            trust_remote_code=is_safe_repo(
                trust_remote_code=self.trust_remote_code,
                hf_path=str(self),
            ),
        )
        try:
            generation_config = GenerationConfig.from_pretrained(str(self))
        except OSError:
            generation_config = None

        n_moe_layers = source.num_hidden_layers - source.first_k_dense_replace
        is_v3 = source.scoring_func == "sigmoid"
        if is_v3:
            v3_kwargs = {
                "moe_router_score_function": "sigmoid",
                "moe_router_enable_expert_bias": True,
                "mtp_num_layers": source.num_nextn_predict_layers if self.convert_mtp else None,
            }
        else:
            v3_kwargs = {}
        return DeepSeekConfig(
            num_layers=source.num_hidden_layers,
            hidden_size=source.hidden_size,
            ffn_hidden_size=source.intermediate_size,
            num_attention_heads=source.num_attention_heads,
            kv_channels=source.num_key_value_heads,
            q_lora_rank=source.q_lora_rank,
            num_moe_experts=source.n_routed_experts,
            moe_ffn_hidden_size=source.moe_intermediate_size,
            moe_shared_expert_intermediate_size=source.moe_intermediate_size * source.n_shared_experts,
            moe_layer_freq=[0] * source.first_k_dense_replace + [1] * n_moe_layers,
            moe_router_topk=source.num_experts_per_tok,
            moe_router_num_groups=source.n_group,
            moe_router_group_topk=source.topk_group,
            moe_router_topk_scaling_factor=source.routed_scaling_factor,
            moe_aux_loss_coeff=getattr(source, "aux_loss_alpha", 0.001),
            kv_lora_rank=source.kv_lora_rank,
            qk_head_dim=source.qk_nope_head_dim,
            qk_pos_emb_head_dim=source.qk_rope_head_dim,
            v_head_dim=source.v_head_dim,
            make_vocab_size_divisible_by=1280 if is_v3 else 3200,
            fp16=(dtype_from_hf(source) == torch.float16),
            bf16=(dtype_from_hf(source) == torch.bfloat16),
            params_dtype=dtype_from_hf(source),
            generation_config=generation_config,
            **v3_kwargs,
        )


@io.model_exporter(DeepSeekModel, "hf")
class HFDeepSeekExporter(io.ModelConnector[DeepSeekModel, "AutoModelForCausalLM"]):
    # pylint: disable=C0115,C0116
    def init(self, dtype=torch.bfloat16, model_name="deepseek-ai/DeepSeek-V3") -> "AutoModelForCausalLM":
        from transformers import AutoConfig, AutoModelForCausalLM
        from transformers.modeling_utils import no_init_weights

        with no_init_weights():
            # Since DeepSeek is not importable from transformers, we can only initialize the HF model
            # from a known checkpoint folder containing the config file and modeling files.
            # The model_name will need to be passed in.
            config = AutoConfig.from_pretrained(
                model_name,
                trust_remote_code=is_safe_repo(
                    trust_remote_code=self.trust_remote_code,
                    hf_path=model_name,
                ),
            )
            hf_model = AutoModelForCausalLM.from_config(
                config,
                trust_remote_code=is_safe_repo(
                    trust_remote_code=self.trust_remote_code,
                    hf_path=model_name,
                ),
                torch_dtype=dtype,
            )
            # Register the AutoModel Hook so that the custom modeling files are saved during save_pretrained()
            type(hf_model).register_for_auto_class("AutoModelForCausalLM")
            return hf_model

    def _detect_hf_deepseek_version(self, source_config: Dict[str, Any]) -> str:
        """
        Detect the HF DeepSeek version based on the source NeMo config.

        Args:
            source_config (Dict[str, Any]): The source NeMo model config.

        Returns:
            str: The DeepSeek version in the Hugging Face Hub convention.
        """
        if source_config['moe_router_enable_expert_bias']:
            target_model_name = "deepseek-ai/DeepSeek-V3"
        elif source_config['q_lora_rank'] is not None:
            target_model_name = "deepseek-ai/DeepSeek-V2"
        else:
            target_model_name = "deepseek-ai/DeepSeek-V2-Lite"
        logging.info(
            f"Your model is determined to be {target_model_name} based on the config. If this is not correct, "
            f"please pass in a local HF checkpoint."
        )
        return target_model_name

    def ckpt_load(self, path: Path) -> Tuple[Dict, Dict]:
        """
        This function loads the state dict directly from a distributed checkpoint, and modify the state dict
        so that it is consistent with the key names you would get from loading the checkpoint into a model.
        This is a more memory-efficient method to obtain a state dict without initializing the nemo model.

        Args:
            path (Path): The path from which the model will be loaded.

        Returns
        -------
            Tuple[Dict, Dict]: The loaded state dict and the yaml config dict.
        """
        model_yaml = path / "context" / "model.yaml"
        if not model_yaml.exists():
            raise FileNotFoundError("model.yaml is not found in the context folder of the checkpoint.")
        with open(model_yaml, 'r') as stream:
            config = yaml.safe_load(stream)

        dist_ckpt_folder = path / "weights"
        state_dict = {}
        for k, v in load_distributed_model_weights(dist_ckpt_folder, True).items():
            if '_extra_state' in k:
                continue
            new_k = k.replace("module.", "")
            if '.experts.experts.' in k:
                # split experts into multiple tensors
                for i in range(v.size(0)):
                    state_dict[new_k.replace(".experts.experts.", ".experts.") + str(i)] = v[i]
            else:
                state_dict[new_k] = v
        return state_dict, config['config']

    def apply(self, output_path: Path, target_model_name=None, trust_remote_code: bool | None = None) -> Path:
        logging.info("Loading DeepSeek NeMo checkpoint. This may take a while...")
        self.trust_remote_code = trust_remote_code
        source, source_config = self.ckpt_load(self)
        logging.info("DeepSeek NeMo checkpoint loaded.")
        if target_model_name is None:
            # Before DeepSeek is fully supported by HF, it is necessary to pass in a local HF checkpoint that
            # is used to initialize the HF model.
            logging.warning(
                "Before DeepSeek is officially supported in HF, you should pass in a local HF "
                "checkpoint using llm.export_ckpt(..., target_model_name=<local hf path>)"
            )
            target_model_name = self._detect_hf_deepseek_version(source_config)

        target = self.init(torch_dtype_from_dict_config(source_config), model_name=target_model_name)
        target = self.convert_state(source, target, source_config)

        target = target.cpu()
        logging.info(f"Converted DeepSeek model to HF, saving model to {output_path}...")
        target.save_pretrained(output_path, safe_serialization=False)
        self.tokenizer.save_pretrained(output_path)

        return output_path

    def convert_state(self, source, target, source_config):
        # pylint: disable=C0301
        mapping = {
            # Embed
            "embedding.word_embeddings.weight": "model.embed_tokens.weight",
            # Attention
            "decoder.layers.*.input_layernorm.weight": "model.layers.*.input_layernorm.weight",
            "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight",
            "decoder.layers.*.self_attention.linear_q_down_proj.weight": "model.layers.*.self_attn.q_a_proj.weight",
            "decoder.layers.*.self_attention.linear_q_up_proj.weight": "model.layers.*.self_attn.q_b_proj.weight",
            "decoder.layers.*.self_attention.linear_kv_down_proj.weight": "model.layers.*.self_attn.kv_a_proj_with_mqa.weight",
            "decoder.layers.*.self_attention.linear_kv_up_proj.weight": "model.layers.*.self_attn.kv_b_proj.weight",
            "decoder.layers.*.self_attention.linear_q_up_proj.layer_norm_weight": "model.layers.*.self_attn.q_a_layernorm.weight",
            "decoder.layers.*.self_attention.linear_kv_up_proj.layer_norm_weight": "model.layers.*.self_attn.kv_a_layernorm.weight",
            "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.post_attention_layernorm.weight",
            # Dense MLP
            # decoder.layers.*.mlp.linear_fc1.weight: model.layers.*.mlp.{gate|up}_proj.weight
            "decoder.layers.*.mlp.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight",
            # MoE
            "decoder.layers.*.mlp.router.weight": "model.layers.*.mlp.gate.weight",
            # decoder.layers.*.mlp.experts.linear_fc1.weight*: model.layers.*.mlp.experts.*.{gate|up}_proj.weight
            "decoder.layers.*.mlp.experts.linear_fc2.weight*": "model.layers.*.mlp.experts.*.down_proj.weight",
            # decoder.layers.*.mlp.shared_experts.linear_fc1.weight: model.layers.*.mlp.shared_experts.{gate|up}_proj.weight
            "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.shared_experts.down_proj.weight",
            # LM Head
            "decoder.final_layernorm.weight": "model.norm.weight",
            "output_layer.weight": "lm_head.weight",
        }
        # For lite model
        if source_config['q_lora_rank'] is None:
            del mapping["decoder.layers.*.self_attention.linear_q_down_proj.weight"]
            del mapping["decoder.layers.*.self_attention.linear_q_up_proj.weight"]
            mapping["decoder.layers.*.self_attention.linear_q_proj.weight"] = "model.layers.*.self_attn.q_proj.weight"
        # Account for Mcore local spec
        if source_config['q_lora_rank'] is not None and 'decoder.layers.0.self_attention.q_layernorm.weight' in source:
            mapping["decoder.layers.*.self_attention.q_layernorm.weight"] = mapping.pop(
                "decoder.layers.*.self_attention.linear_q_up_proj.layer_norm_weight"
            )

        if 'decoder.layers.0.self_attention.kv_layernorm.weight' in source:
            mapping["decoder.layers.*.self_attention.kv_layernorm.weight"] = mapping.pop(
                "decoder.layers.*.self_attention.linear_kv_up_proj.layer_norm_weight"
            )

        if source_config.get('moe_router_enable_expert_bias', False):
            mapping.update(
                {
                    "decoder.layers.*.mlp.router.expert_bias": "model.layers.*.mlp.gate.e_score_correction_bias",
                }
            )

        transforms = [
            io.state_transform(
                source_key="decoder.layers.*.mlp.linear_fc1.weight",
                target_key=("model.layers.*.mlp.gate_proj.weight", "model.layers.*.mlp.up_proj.weight"),
                fn=TransformFns.split_fc1,
            ),
            io.state_transform(
                source_key="decoder.layers.*.mlp.experts.linear_fc1.weight*",
                target_key=(
                    "model.layers.*.mlp.experts.*.gate_proj.weight",
                    "model.layers.*.mlp.experts.*.up_proj.weight",
                ),
                fn=TransformFns.split_fc1,
            ),
            io.state_transform(
                source_key="decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
                target_key=(
                    "model.layers.*.mlp.shared_experts.gate_proj.weight",
                    "model.layers.*.mlp.shared_experts.up_proj.weight",
                ),
                fn=TransformFns.split_fc1,
            ),
        ]
        source = self._modify_source_state(source, source_config)

        return io.apply_transforms(
            source,
            target,
            mapping=mapping,
            transforms=transforms,
            cast_dtype=torch.bfloat16,
        )

    def _modify_source_state(self, source: Dict[str, Any], source_config: Dict[str, Any]) -> _ModelState:
        """
        In deepseek, HF weight `model.layers.*.post_attention_layernorm.weight` is mapped to mcore weight
        a) `decoder.layers.*.mlp.linear_fc1.layer_norm_weight`, if the layer is dense
        b) `decoder.layers.*.pre_mlp_layernorm.weight`, if the layer is MoE

        We rename decoder.layers.*.mlp.linear_fc1.layer_norm_weight in the first case to unify key names
        """
        for layer_i in range(source_config['num_layers']):
            if f"decoder.layers.{layer_i}.mlp.linear_fc1.layer_norm_weight" in source:
                weight = source.pop(f"decoder.layers.{layer_i}.mlp.linear_fc1.layer_norm_weight")
                source[f"decoder.layers.{layer_i}.pre_mlp_layernorm.weight"] = weight
        modified_source = _ModelState(source)
        return modified_source

    @property
    def tokenizer(self) -> 'AutoTokenizer':
        return io.load_context(self, subpath="model").tokenizer

    @property
    def config(self) -> "HFDeepseekV3Config":
        """Create a HF DeepseekV3Config from the NeMo model config.

        Translates the NeMo configuration parameters to the equivalent HF
        configuration.

        Currently only supports DeepseekV3Config based on availability
        in the Transformers library.

        Returns:
            HFDeepseekV3Config: HF configuration for DeepSeekV3 models
        """
        # TODO: Get config for all DeepSeek model variants once available in transformers

        from transformers import DeepseekV3Config as HFDeepseekV3Config

        source: DeepSeekV3Config = io.load_context(str(self)).model.config

        target_model_name = self._detect_hf_deepseek_version(asdict(source))
        if target_model_name != "deepseek-ai/DeepSeek-V3":
            raise ValueError(f"Getting config for model other than {target_model_name} is not supported.")

        # Figure out the number of zeros in the prefix of moe_layer_freq array
        # for the HF first_k_dense_replace parameter and validate the reminder:
        k = 0
        while k < len(source.moe_layer_freq) and source.moe_layer_freq[k] == 0:
            k += 1
        assert all(x == 1 for x in source.moe_layer_freq[k:])

        return HFDeepseekV3Config(
            architectures=["DeepseekV3ForCausalLM"],
            num_hidden_layers=source.num_layers,
            hidden_size=source.hidden_size,
            intermediate_size=source.ffn_hidden_size,
            num_attention_heads=source.num_attention_heads,
            q_lora_rank=source.q_lora_rank,
            qk_nope_head_dim=source.qk_head_dim,
            qk_rope_head_dim=source.qk_pos_emb_head_dim,
            v_head_dim=source.v_head_dim,
            kv_lora_rank=source.kv_lora_rank,
            num_key_value_heads=source.kv_channels,
            n_routed_experts=source.num_moe_experts,
            moe_intermediate_size=source.moe_ffn_hidden_size,
            first_k_dense_replace=k,
            num_experts_per_tok=source.moe_router_topk,
            n_group=source.moe_router_num_groups,
            topk_group=source.moe_router_group_topk,
            routed_scaling_factor=source.moe_router_topk_scaling_factor,
            aux_loss_alpha=source.moe_aux_loss_coeff,
            max_position_embeddings=source.max_position_embeddings,
            vocab_size=self.tokenizer.vocab_size,
        )


__all__ = [
    "DeepSeekConfig",
    "DeepSeekV2Config",
    "DeepSeekV2LiteConfig",
    "DeepSeekV3Config",
    "DeepSeekModel",
]
