nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/export/model_config_utils.py

@@ -0,0 +1,392 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Common utils for the ModelConfig."""
+
+import dataclasses
+import math
+from types import UnionType
+from typing import Union, get_args, get_origin
+
+import numpy as np
+import torch
+
+from .model_config import (
+    QUANTIZATION_FP8_PC_PT,
+    QUANTIZATION_INT4_AWQ,
+    QUANTIZATION_W4A8_AWQ,
+    DecoderLayerConfig,
+    LayernormConfig,
+    LinearConfig,
+    MLPConfig,
+    ModelConfig,
+    MOEConfig,
+    QKVConfig,
+)
+from .quant_utils import to_quantized_weight
+
+# numpy doesn't know bfloat16, define abstract binary type instead
+np_bfloat16 = np.dtype("V2", metadata={"dtype": "bfloat16"})
+
+
+def _numpy_to_torch(x):
+    """Convert numpy array to torch tensor."""
+    if isinstance(x, torch.Tensor):
+        return x
+
+    if x.dtype != np_bfloat16:
+        return torch.tensor(x)
+    return torch.tensor(x.view(np.int16)).view(torch.bfloat16)
+
+
+def model_config_to_dict(model_config: ModelConfig) -> dict:
+    """Converts the instance to a python dict."""
+    assert model_config is not None, "model_config is None"
+
+    def _to_dict(obj):
+        if dataclasses.is_dataclass(obj):
+            return {k: _to_dict(v) for k, v in vars(obj).items()}
+        elif isinstance(obj, dict):
+            return {k: _to_dict(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [_to_dict(v) for v in obj]
+        return obj
+
+    return _to_dict(model_config)
+
+
+def split_config_and_weights(
+    config,
+    weights: dict[str, torch.tensor],
+    prefix: str = "transformer",
+    layer_config_dict: dict = {},
+):
+    """Util function to split the weights or any torch.Tensor in nested config to weights.
+
+    A weight id starts with transformers or lm_head will also be generated to link the original key to the weights dict.
+    The weights in the weights dict are contiguous.
+
+    layer_config_dict: A dictionary containing layerwise quantization format information and awq_block_size information
+    when relevant. It is used to export quantization.json for auto_quant checkpoint.
+    """
+    if isinstance(config, dict):
+        for k, v in config.items():
+            if k == "lm_head" and "medusa_heads" not in prefix:
+                # lm_head is not part of the transformer.
+                array_key = k
+            elif k == "experts":
+                # Omit the 'experts' key that is not in the model name
+                array_key = prefix
+            elif k == "medusa_heads":
+                # medusa_heads is not part of the transformer
+                array_key = k
+            elif "rel_attn_table" in prefix:
+                # rel_attn_table is treated as a weight for quantize loop whereas in TRTLLM it is a Tensor
+                array_key = prefix
+            else:
+                array_key = f"{prefix}.{k}"
+
+            # TensorRT-LLM 0.17 uses the kv_cache_scaling_factor,
+            # the following code will be deprecated after the 0.18 upgrade.
+            if "v_cache_scaling_factor" in array_key:
+                k_cache_scaling_factor_key = array_key.replace(
+                    "v_cache_scaling_factor", "k_cache_scaling_factor"
+                )
+                if k_cache_scaling_factor_key in weights:
+                    weights[
+                        array_key.replace("v_cache_scaling_factor", "kv_cache_scaling_factor")
+                    ] = torch.maximum(v, weights[k_cache_scaling_factor_key])
+
+            # Construct per_layer quantization dictionary, with block size information
+            if array_key != "transformer.quantization" and (
+                "quantization" in array_key or "awq_block_size" in array_key
+            ):
+                layer_config_dict[array_key] = v
+
+            if isinstance(v, torch.Tensor):
+                weights[array_key] = v
+                config[k] = f"{array_key}"
+            else:
+                split_config_and_weights(v, weights, array_key, layer_config_dict)
+    elif isinstance(config, list):
+        for i, v in enumerate(config):
+            array_key = f"{prefix}.{i}"
+            if isinstance(v, torch.Tensor):
+                weights[array_key] = v
+                config[i] = f"{array_key}"
+            else:
+                split_config_and_weights(v, weights, array_key, layer_config_dict)
+
+
+def _unified_weights_key(k: str) -> str:
+    """Try to unify the weights dict key between old npz and the new safetensors format."""
+    prefixes = ["transformer.", "_np:"]
+    for prefix in prefixes:
+        k = k.removeprefix(prefix)
+
+    k = k.replace("final_layernorm", "ln_f")
+
+    return k.replace(":", ".")
+
+
+def _restore_model_config(model_config, weights: dict[str, np.ndarray | torch.Tensor]):
+    def _is_tensor_key(k):
+        return isinstance(k, str) and _unified_weights_key(k) in weights
+
+    if isinstance(model_config, dict):
+        for k, v in model_config.items():
+            if _is_tensor_key(v):
+                model_config[k] = _numpy_to_torch(weights[_unified_weights_key(v)])
+            else:
+                _restore_model_config(v, weights)
+    if isinstance(model_config, list):
+        for i, v in enumerate(model_config):
+            if _is_tensor_key(v):
+                model_config[i] = _numpy_to_torch(weights[_unified_weights_key(v)])
+            else:
+                _restore_model_config(v, weights)
+
+
+def restore_model_config(model_config, weights: dict[str, np.ndarray | torch.Tensor]):
+    """Recursively restores the model_config from json and loads np.ndarray or torch.Tensor weights from weights."""
+    unified_key_weights = {}
+    for k, v in weights.items():
+        unified_key_weights[_unified_weights_key(k)] = v
+
+    _restore_model_config(model_config, unified_key_weights)
+
+
+def _from_dict(class_type, data):
+    """Helper function to load the data as a class_type. class_type must be a dataclass."""
+    if data is None:
+        return None
+
+    # class_type of quantization is str | None, which is catergrorized as Union
+    if class_type != str | None and get_origin(class_type) in [Union, UnionType]:
+        # Handle QKV
+        if all(key in data for key in ["q", "k", "v"]):
+            # splitted qkv case
+            class_type = QKVConfig
+        elif all(key in data for key in ["router", "experts"]):
+            # moe
+            class_type = MOEConfig
+        elif all(key in data for key in ["fc", "gate", "proj"]):
+            # mlp
+            class_type = MLPConfig
+        else:
+            # merged qkv case
+            assert "linear_type" in data, f"{data} is not a valid LinearConfig"
+            class_type = LinearConfig
+
+    if dataclasses.is_dataclass(class_type):
+        fieldtypes = {f.name: f.type for f in dataclasses.fields(class_type)}
+        fields_map = {}
+        for k, v in data.items():
+            if k in fieldtypes:
+                # We only handle keys available in the fields.
+                # Deprecated fields in the checkpoint will be ignored.
+                fields_map[k] = _from_dict(fieldtypes[k], v)
+        return class_type(**fields_map)
+    elif get_origin(class_type) is list and dataclasses.is_dataclass(get_args(class_type)[0]):
+        list_value = []
+        for child in data:
+            child_class_type = get_args(class_type)[0]
+            list_value.append(_from_dict(child_class_type, child))
+        return list_value
+    else:
+        return data
+
+
+def model_config_from_dict(d: dict) -> ModelConfig:
+    """Load a dict to a `ModelConfig` instance."""
+    config_type = ModelConfig
+
+    config_type_map = {}
+    for t in [ModelConfig, DecoderLayerConfig, LayernormConfig, LinearConfig]:
+        config_type_map[t.__name__] = t
+
+    if "__name__" in d:
+        config_name = d.pop("__name__")
+        try:
+            config_type = config_type_map[config_name]
+        except Exception as e:
+            raise NotImplementedError(f"{config_name} not supported") from e
+
+    return _from_dict(config_type, d)
+
+
+def pad_weights(weights, tp_size):
+    """Returns the padded weights to tp_size."""
+    assert len(weights.shape) > 1
+
+    def _pad_size(original_size, tp_size):
+        return int(math.ceil(original_size / tp_size) * tp_size)
+
+    original_size = weights.shape[0]
+    padded_size = _pad_size(original_size, tp_size)
+
+    if original_size != padded_size:
+        pad_width = padded_size - original_size
+        return torch.nn.functional.pad(weights, (0, 0, 0, pad_width), "constant", value=0)
+    return weights
+
+
+def merge_qkv(model_config):
+    """Merges the qkv fields in model_config from QKVConfig to a single LinearConfig."""
+    for decoder_config in model_config.layers:
+        for attention_key in ["attention", "self_attention", "cross_attention"]:
+            attention = getattr(decoder_config, attention_key, None)
+            if attention and isinstance(attention.qkv, QKVConfig):
+                splitted_qkv = attention.qkv
+                attention.qkv = LinearConfig()
+                attention.qkv.weight = splitted_qkv.weight
+                attention.qkv.bias = splitted_qkv.bias
+                attention.qkv.activation_scaling_factor = splitted_qkv.activation_scaling_factor
+                attention.qkv.weights_scaling_factor = splitted_qkv.weights_scaling_factor
+                attention.qkv.weights_scaling_factor_2 = splitted_qkv.weights_scaling_factor_2
+                attention.qkv.prequant_scaling_factor = splitted_qkv.prequant_scaling_factor
+                attention.qkv.awq_block_size = splitted_qkv.awq_block_size
+                # Assert q,k,v have same quantization formats before merging
+                assert (
+                    splitted_qkv.q.quantization
+                    == splitted_qkv.k.quantization
+                    == splitted_qkv.v.quantization
+                ), "Quantization formats of q,k,v must be the same."
+                attention.qkv.quantization = splitted_qkv.q.quantization
+
+                # Collect GPU memory from the deleted tensors
+                del splitted_qkv
+
+
+def merge_gate_fc(model_config):
+    """Postprocess the MLP config for TensorRT-LLM export."""
+    for decoder_config in model_config.layers:
+        mlp = None
+        if isinstance(decoder_config.mlp, MLPConfig):
+            mlp = decoder_config.mlp
+        elif (
+            isinstance(decoder_config.mlp, MOEConfig)
+            and decoder_config.mlp.shared_expert is not None
+        ):
+            mlp = decoder_config.mlp.shared_expert
+
+        if mlp is not None and mlp.merge_gate_fc and mlp.gate is not None and mlp.fc is not None:
+            mlp.fc.weight = torch.cat(
+                [
+                    mlp.gate.weight,
+                    mlp.fc.weight,
+                ],
+                dim=0,
+            )
+
+            if (
+                mlp.fc.weights_scaling_factor is not None
+                and mlp.fc.weights_scaling_factor.numel() > 1
+            ):
+                mlp.fc.weights_scaling_factor = torch.cat(
+                    [mlp.gate.weights_scaling_factor, mlp.fc.weights_scaling_factor], dim=0
+                )
+
+            mlp.gate = None
+
+
+def pack_linear_weights(model_config: ModelConfig):
+    """Packs the quantized linear weights in the model_config to the quantized format."""
+
+    def _linear_layer_to_quantized_weight(linear_layers):
+        for linear_layer in linear_layers:
+            # Check if quantization of the layer is None to support auto_quant
+            if isinstance(linear_layer, LinearConfig) and (
+                linear_layer.weights_scaling_factor is not None
+                and linear_layer.quantization is not None
+            ):
+                # Quantize on CPU if we are short of GPU memory.
+                # Using 2x of the tensor size as a threshold.
+                if linear_layer.weight.is_cuda:
+                    free_mem, _ = torch.cuda.mem_get_info(linear_layer.weight.device)
+                    if (
+                        free_mem
+                        < 2 * linear_layer.weight.element_size() * linear_layer.weight.nelement()
+                    ):
+                        linear_layer.weight = linear_layer.weight.cpu()
+
+                # Save the quantize layer weights to cpu and save gpu memory.
+                if linear_layer.weight.element_size() > 1:
+                    linear_layer.weight = to_quantized_weight(
+                        linear_layer.weight,
+                        linear_layer.weights_scaling_factor,
+                        linear_layer.quantization,
+                        linear_layer.weights_scaling_factor_2,
+                        linear_layer.awq_block_size,
+                    ).cpu()
+
+                # Convert to int8 if to make the checkpoint compatible with the latest TensorRT-LLM release.
+                # The future TensorRT-LLM release will use uint8 weights instead.
+                if linear_layer.quantization in [QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ]:
+                    linear_layer.weight = linear_layer.weight.view(torch.int8)
+
+                # TensorRT-LLM uses per_channel_scale for FP8 per channel weight quantization.
+                if linear_layer.quantization == QUANTIZATION_FP8_PC_PT:
+                    linear_layer.per_channel_scale = linear_layer.weights_scaling_factor.cpu()
+                    linear_layer.weights_scaling_factor = None
+                else:
+                    linear_layer.weights_scaling_factor = linear_layer.weights_scaling_factor.cpu()
+
+    if not model_config.quantization:
+        return
+
+    def _find_linear_configs_recursive(model_config):
+        linear_configs = []
+
+        # Base case - not a dataclass
+        if not dataclasses.is_dataclass(model_config):
+            return linear_configs
+
+        # Check if current object is a LinearConfig
+        if isinstance(model_config, LinearConfig):
+            linear_configs.append(model_config)
+            return linear_configs
+
+        # Recursively check all fields
+        for field in dataclasses.fields(model_config):
+            value = getattr(model_config, field.name)
+
+            if isinstance(value, list):
+                for item in value:
+                    linear_configs.extend(_find_linear_configs_recursive(item))
+
+            elif isinstance(value, dict):
+                for item in value.values():
+                    linear_configs.extend(_find_linear_configs_recursive(item))
+
+            # Handle nested dataclasses
+            elif dataclasses.is_dataclass(value):
+                linear_configs.extend(_find_linear_configs_recursive(value))
+
+        return linear_configs
+
+    linear_layers = _find_linear_configs_recursive(model_config)
+
+    _linear_layer_to_quantized_weight(linear_layers)
+
+    if model_config.medusa_heads is not None:
+        linear_layers = []
+
+        for head in model_config.medusa_heads:
+            linear_layers.append(head.lm_head)
+            for layer in head.medusa_layers:
+                linear_layers.append(layer.linear)
+
+        _linear_layer_to_quantized_weight(linear_layers)

modelopt/torch/export/model_utils.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""Utility functions for model type detection and classification."""
+
+MODEL_NAME_TO_TYPE = {
+    "GPT2": "gpt",
+    "Mllama": "mllama",
+    "Llama4": "llama4",
+    "Llama": "llama",
+    "Mistral": "llama",
+    "GPTJ": "gptj",
+    "FalconForCausalLM": "falcon",
+    "RWForCausalLM": "falcon",
+    "baichuan": "baichuan",
+    "MPT": "mpt",
+    "Bloom": "bloom",
+    "ChatGLM": "chatglm",
+    "QWen": "qwen",
+    "RecurrentGemma": "recurrentgemma",
+    "Gemma3": "gemma3",
+    "Gemma2": "gemma2",
+    "Gemma": "gemma",
+    "phi3small": "phi3small",
+    "phi3": "phi3",
+    "PhiMoEForCausalLM": "phi3",
+    "Phi4MMForCausalLM": "phi4mm",
+    "phi": "phi",
+    "TLGv4ForCausalLM": "phi",
+    "MixtralForCausalLM": "llama",
+    "ArcticForCausalLM": "llama",
+    "StarCoder": "gpt",
+    "Dbrx": "dbrx",
+    "T5": "t5",
+    "Bart": "bart",
+    "GLM": "glm",
+    "InternLM2ForCausalLM": "internlm",
+    "ExaoneForCausalLM": "exaone",
+    "Nemotron": "gpt",
+    "Deepseek": "deepseek",
+    "Whisper": "whisper",
+    "gptoss": "gptoss",
+}
+
+__doc__ = f"""Utility functions for model type detection and classification.
+
+    .. code-block:: python
+
+        {MODEL_NAME_TO_TYPE=}
+"""
+
+__all__ = ["get_model_type"]
+
+
+def get_model_type(model):
+    """Try get the model type from the model name. If not found, return None."""
+    for k, v in MODEL_NAME_TO_TYPE.items():
+        if k.lower() in type(model).__name__.lower():
+            return v
+    return None

modelopt/torch/export/plugins/__init__.py

@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Export package plugin."""
+
+from modelopt.torch.utils import import_plugin
+
+with import_plugin("megatron_importer"):
+    from .megatron_importer import *

modelopt/torch/export/plugins/mcore_common.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Custom mapping from Megatron Core models to their Hugging Face counter part."""
+
+from typing import Any
+
+from .mcore_deepseek import deepseek_causal_lm_export, deepseek_causal_lm_import
+from .mcore_gptoss import gptoss_causal_lm_export, gptoss_causal_lm_import
+from .mcore_llama import (
+    eagle3_llama_causal_lm_export,
+    eagle_llama_causal_lm_export,
+    llama4_causal_lm_export,
+    llama4_causal_lm_import,
+    llama_causal_lm_export,
+    llama_causal_lm_import,
+)
+from .mcore_nemotron import (
+    nemotron_causal_lm_export,
+    nemotron_h_causal_lm_export,
+    nemotron_h_causal_lm_import,
+)
+from .mcore_qwen import (
+    qwen3_causal_lm_export,
+    qwen3_causal_lm_import,
+    qwen25_causal_lm_export,
+    qwen25_causal_lm_import,
+)
+
+all_mcore_hf_export_mapping: dict[str, Any] = {
+    "DeepseekV2ForCausalLM": deepseek_causal_lm_export,
+    "DeepseekV3ForCausalLM": deepseek_causal_lm_export,
+    "LlamaForCausalLM": llama_causal_lm_export,
+    "Llama4ForConditionalGeneration": llama4_causal_lm_export,
+    "NemotronForCausalLM": nemotron_causal_lm_export,
+    "NemotronHForCausalLM": nemotron_h_causal_lm_export,
+    "LlamaForCausalLMEagle": eagle_llama_causal_lm_export,
+    "LlamaForCausalLMEagle3": eagle3_llama_causal_lm_export,
+    "Qwen3ForCausalLM": qwen3_causal_lm_export,
+    "Qwen3MoeForCausalLM": qwen3_causal_lm_export,
+    "Qwen2ForCausalLM": qwen25_causal_lm_export,
+    "GptOssForCausalLM": gptoss_causal_lm_export,
+}
+
+all_mcore_hf_import_mapping: dict[str, Any] = {
+    "LlamaForCausalLM": llama_causal_lm_import,
+    "Llama4ForConditionalGeneration": llama4_causal_lm_import,
+    "DeepseekV2ForCausalLM": deepseek_causal_lm_import,
+    "DeepseekV3ForCausalLM": deepseek_causal_lm_import,
+    "NemotronHForCausalLM": nemotron_h_causal_lm_import,
+    "Qwen3ForCausalLM": qwen3_causal_lm_import,
+    "Qwen3MoeForCausalLM": qwen3_causal_lm_import,
+    "Qwen2ForCausalLM": qwen25_causal_lm_import,
+    "GptOssForCausalLM": gptoss_causal_lm_import,
+}