nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/sparsity/sparsegpt.py

@@ -0,0 +1,276 @@
+# 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.
+
+"""Utility functions of SparseGPT."""
+
+import math
+import warnings
+
+import torch
+import torch.nn as nn
+
+from modelopt.torch.opt.searcher import SearchConfig
+from modelopt.torch.utils import print_rank_0
+
+from .magnitude import get_nmprune_info
+from .module import SparseModule
+from .searcher import BaseSparseSearcher
+
+
+def invert(hessian: torch.Tensor) -> torch.Tensor:
+    """Invert a Hessian matrix."""
+    try:
+        hessian_inv = torch.linalg.cholesky(hessian)
+        hessian_inv = torch.cholesky_inverse(hessian_inv)
+        hessian_inv = torch.linalg.cholesky(hessian_inv, upper=True)
+    except RuntimeError:
+        cols = hessian.size(0)
+        eps = 1e-6 * torch.eye(cols).to(hessian.device)
+        hessian_inv = torch.cholesky_inverse(torch.linalg.cholesky(hessian + eps))
+
+    return hessian_inv
+
+
+def prepare(
+    tensor: torch.Tensor, hessian: torch.Tensor, hessian_damp: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Prepare the inverse Hessian matrix."""
+    weight = tensor.detach().clone()
+    # move the hessian matrix from CPU to GPU for acceleration
+    hessian = hessian.to(weight.device)
+    if len(weight.size()) == 4:
+        weight = weight.flatten(1)
+
+    zero = torch.diag(hessian) == 0
+    hessian[zero, zero] = 1
+    weight[:, zero] = 0
+
+    damp = hessian_damp * torch.mean(torch.diag(hessian))
+    cols = weight.size(1)
+    diag = torch.arange(cols)
+    hessian[diag, diag] += damp
+
+    hessian_inv = invert(hessian)
+
+    # remove the Hessian matrix to save GPU memory
+    del hessian
+    torch.cuda.empty_cache()
+
+    return weight, hessian_inv
+
+
+@torch.no_grad()
+def create_sgpt_mask(
+    tensor: torch.Tensor, hessian: torch.Tensor, config: SearchConfig
+) -> torch.Tensor:
+    """Create a sparse mask for the given tensor."""
+    shape = tensor.size()
+    weight, hessian_inv = prepare(tensor, hessian, config["hessian_damp"])
+    rows, cols = weight.size()
+    hessian_inv_diag = torch.diagonal(hessian_inv, dim1=0, dim2=1)
+
+    is_nm_prune, n, m = get_nmprune_info(config["pattern"])
+    col_bs = config["col_block_size"]
+    row_bs = config["row_block_size"]
+    # if row_bs is not specified, prune the whole weight block
+    if row_bs == -1:
+        row_bs = rows
+
+    for r1 in range(0, rows, row_bs):
+        r2 = min(r1 + row_bs, rows)
+        # the mask of the weights not to be pruned
+        w_rows = weight[r1:r2].float()
+
+        # pruning the weight block W[row:row+row_bs, i1:i1+col_bs]
+        for i1 in range(0, cols, col_bs):
+            i2 = min(i1 + col_bs, cols)
+            w_blk = w_rows[:, i1:i2].clone()
+            q_blk = torch.zeros_like(w_blk)
+            # the error of the weights to be pruned
+            delta_blk = torch.zeros_like(w_blk)
+            hinv_blk = hessian_inv[i1:i2, i1:i2]
+            hinv_diag_blk = hessian_inv_diag[i1:i2]
+
+            errors_blk = (w_blk**2) / (hinv_diag_blk**2 + 1e-9)
+            if torch.isnan(errors_blk).any():
+                print("nan in errors_blk.")
+
+            mask_blk = torch.zeros_like(w_blk, dtype=torch.bool)
+
+            for j in range(i2 - i1):
+                # compute the error of the weights to be pruned
+                w = w_blk[:, j]
+                d = hinv_diag_blk[j]
+                if is_nm_prune and j % m == 0:
+                    errors_blk = (w_blk[:, j : j + m] ** 2) / (hinv_diag_blk[j : j + m] ** 2 + 1e-9)
+                    mask_blk.scatter_(
+                        1, j + torch.topk(errors_blk, n, dim=1, largest=False)[1], True
+                    )
+
+                q = w.clone()
+                q[mask_blk[:, j]] = 0
+                q_blk[:, j] = q
+
+                # update the remaining weights in the col_bs block to compensate the error caused by pruning W[:, j]
+                err = (w - q) / d
+                w_blk[:, j:] -= err.unsqueeze(1).matmul(hinv_blk[j, j:].unsqueeze(0))
+                delta_blk[:, j] = err
+
+            # compensate the error caused by pruning W[:, i: i + col_bs] with the weights update in W[:, i + col_bs:]
+            w_rows[:, i1:i2] = q_blk
+            w_rows[:, i2:] -= delta_blk.matmul(hessian_inv[i1:i2, i2:])
+            if torch.isnan(w_rows[:, i2:]).any():
+                print("nan")
+
+        weight[r1:r2] = w_rows
+
+    mask = weight != 0
+
+    return mask.view(shape)
+
+
+class SparseGPTSearcher(BaseSparseSearcher):
+    """SparseGPT-based sparse mask searching algorithm."""
+
+    @property
+    def default_search_config(self) -> SearchConfig:
+        """Get the default config for the searcher."""
+        return {
+            **super().default_search_config,
+            "col_block_size": 128,  # column block size in sparsegpt
+            "row_block_size": -1,  # row block size in sparsegpt
+            "hessian_damp": 0.1,  # hessian damp in sparsegpt
+            "calib_size": 256,  # calibration size for hessian matrix calculation
+            "device": "cuda",  # device of hessian matrix
+        }
+
+    def _check_weight_size(self, weight, mod_name) -> bool:
+        """Check if the weight size is supported by SparseGPT."""
+        _, _, m = get_nmprune_info(self.config["pattern"])
+
+        # the column size must be divisible by m
+        if weight.size(0) % m != 0 or weight.size(1) % m != 0:
+            warnings.warn(
+                f"Skipping pruning {mod_name} of size={weight.size()!s} and"
+                f" type={weight.dtype!s} for SparseGPT"
+            )
+            return False
+
+        return True
+
+    def _compute_mask(self, module: SparseModule) -> torch.BoolTensor:
+        """Compute the mask (and weight update) for the given module."""
+        return create_sgpt_mask(module.weight, module.hessian, self.config)
+
+    @torch.no_grad()
+    def before_search(self):
+        """Register the forward hook to collect the hessian matrix."""
+        super().before_search()
+
+        handles = []
+        for _, module in self._named_sparsifiable_modules():
+            # setup and register the forward hook
+            self._setup_forward_hook(module)
+            handles.append(module.register_forward_hook(self._hook_compute_hessian))
+
+        print_rank_0(f"Collecting Hessian statistics for {len(handles)} modules.")
+
+        # run a forward loop to collect the hessian matrix
+        assert self.forward_loop is not None, "Please provide `data_loader` or `forward_loop`!"
+        self.forward_loop(self.model)
+
+        # remove the forward hooks
+        for handle in handles:
+            handle.remove()
+
+    def after_search(self):
+        """Remove Hessian artifacts from network."""
+        super().after_search()
+        for _, module in self._named_sparsifiable_modules():
+            del module.hessian
+            del module.samples
+
+    @staticmethod
+    def _is_memory_sufficient(device_id, threshold):
+        """Check if the memory usage on the CUDA device is below the threshold."""
+        total_memory = torch.cuda.get_device_properties(device_id).total_memory
+        allocated_memory = torch.cuda.memory_allocated(device_id)
+        free_memory = total_memory - allocated_memory
+        return free_memory / total_memory > (1 - threshold)
+
+    @classmethod
+    def _setup_forward_hook(cls, mod: SparseModule) -> None:
+        """Setup the attributes we need for our forward hook during the SparseGPT search."""
+        # initialize the hessian matrix
+        if isinstance(mod, nn.Conv2d):
+            # For conv2d layers, the hessian matrix is calculated as X * X^T, where X is the
+            # flattened weight matrix.
+            cols = mod.weight.size(1) * mod.weight.size(2) * mod.weight.size(3)
+        else:
+            # For linear layers, the hessian matrix is calculated as X * X^T, where X is the
+            # weight matrix.
+            cols = mod.weight.size(1)
+
+        target_device = mod.weight.device
+        # Hessian matrix is stored in the GPU memory by default
+        if target_device.type == "cuda" and cls._is_memory_sufficient(target_device.index, 0.8):
+            hessian = torch.zeros((cols, cols), dtype=torch.float32).to(target_device)
+        else:
+            hessian = torch.zeros((cols, cols), dtype=torch.float32).to("cpu")
+
+        # store the hessian matrix and the number of samples
+        # TODO: this should probably be improved eventually!!
+        mod.hessian = hessian
+        mod.samples = 0
+
+    @classmethod
+    def _hook_compute_hessian(cls, mod: nn.Module, inp: torch.Tensor, out: torch.Tensor):
+        with torch.inference_mode():
+            # TODO: move the hessian matrix to GPU memory if possible
+            if isinstance(inp, tuple):
+                inp = inp[0]
+            # use torch.float32 to avoid overflow in Hessian
+            if len(inp.shape) == 2:
+                inp = inp.unsqueeze(0)
+            tmp = inp.shape[0]
+            # nn.Linear and *ParallelLinear in mcore
+            if "Linear" in type(mod).__name__:
+                if len(inp.shape) == 3:
+                    inp = inp.reshape((-1, inp.shape[-1]))
+                inp.t_()
+            if isinstance(mod, nn.Conv2d):
+                unfold = nn.Unfold(
+                    mod.kernel_size,
+                    dilation=mod.dilation,
+                    stride=mod.stride,
+                )
+                inp = unfold(inp)
+                inp = inp.permute([1, 0, 2])
+                inp = inp.flatten(1)
+            mod.hessian *= mod.samples / (mod.samples + tmp)
+            mod.samples += tmp
+            inp = math.sqrt(2 / mod.samples) * inp.float()
+
+            # the hessian matrix is calculated as X * X^T
+            target_device = mod.hessian.device
+            if mod.hessian.device.type == "cuda":
+                if cls._is_memory_sufficient(mod.hessian.device.index, 0.8):
+                    mod.hessian += inp.matmul(inp.t()).to(mod.hessian.device)
+                else:
+                    target_device = "cpu"
+                    mod.hessian = mod.hessian.to("cpu")
+            mod.hessian += inp.matmul(inp.t()).to(target_device)
+
+            assert not torch.isinf(mod.hessian).any(), "Hessian contains inf"

modelopt/torch/sparsity/sparsification.py

@@ -0,0 +1,123 @@
+# 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.
+
+"""High-level API to automatically sparsify your model with various algorithms."""
+
+from typing import Any
+
+from torch import nn
+
+from modelopt.torch.opt.conversion import apply_mode, get_mode
+from modelopt.torch.opt.mode import ModeLike
+from modelopt.torch.opt.searcher import BaseSearcher, SearchConfig
+from modelopt.torch.utils import unwrap_model
+
+from .mode import SparsityModeRegistry
+
+__all__ = ["export", "sparsify"]
+
+
+def sparsify(
+    model: nn.Module, mode: ModeLike, config: SearchConfig | None = None
+) -> tuple[nn.Module, dict[str, Any]]:
+    """Sparsify a given model and search for they optimal sparsified weights.
+
+    Args:
+        model: A standard model that contains standard building blocks to be sparsified in-place.
+        mode: A (list of) string(s) or Mode(s) or a list of tuples containing the mode and its
+            config indicating the desired mode(s) (and configurations) for the convert
+            process. Modes set up the model for different algorithms for model optimization. The
+            following modes are available:
+
+            *   :class:`"sparse_magnitude"<modelopt.torch.sparsity.mode.SparseMagnitudeModeDescriptor>`:
+                The ``model`` will be sparsified according to the magnitude of weights in each
+                layer. The mode's config is described in
+                :class:`SparseMagnitudeConfig<modelopt.torch.sparsity.config.SparseMagnitudeConfig>`.
+            *   :class:`"sparsegpt"<modelopt.torch.sparsity.mode.SparseGPTModeDescriptor>`:
+                The ``model`` will be sparsified and weights are updated optimally using an Hessian
+                approximation of the loss function (see SparseGPT paper for details). The mode's
+                config is described in
+                :class:`SparseGPTConfig<modelopt.torch.sparsity.config.SparseGPTConfig>`.
+
+            If the mode argument is specified as a dictionary, the keys should indicate the mode and
+            the values specify the per-mode configuration. If not provided, then default
+            configuration would be used.
+
+        config: Additional optional arguments to configure the search. Currently, we support:
+
+            * ``verbose``: Whether to print detailed search stats during search.
+            * ``forward_loop``: A ``Callable`` that takes a model as input and runs a forward loop
+                on it. It is recommended to choose the data loader used inside the forward loop
+                carefully to reduce the runtime. Cannot be provided at the same time as
+                ``data_loader`` and ``collect_func``.
+            * ``data_loader``: An iterator yielding batches of data for calibrating the
+              normalization layers in the model or compute gradient scores. It is recommended to use
+              the same data loader as for training but with significantly fewer iterations. Cannot
+              be provided at the same time as ``forward_loop``.
+            * ``collect_func``: A ``Callable`` that takes a batch of data from the data loader as
+              input and returns the input to ``model.forward()`` as described in
+              :meth:`run_forward_loop <modelopt.torch.utils.network.run_forward_loop>`. Cannot
+              be provided at the same time as ``forward_loop``.
+
+            .. note::
+
+                Additional configuration options may be added by individual algorithms. Please
+                refer to the documentation of the individual algorithms for more information.
+
+    Returns: A sparsified model
+
+    .. note::
+
+        The given model is sparsified in-place. The returned model is thus a reference to the same
+        model instance as the input model.
+    """
+    # apply sparsity to the model
+    model = apply_mode(model, mode, registry=SparsityModeRegistry)
+
+    # retrieve searcher class
+    searcher_cls: type[BaseSearcher] = getattr(get_mode(model), "search_algorithm")
+
+    # run search+sparsification algorithm
+    searcher = searcher_cls()
+    searcher.search(model, {}, (), config)
+
+    # return the sparsified model
+    return model
+
+
+def export(model: nn.Module) -> nn.Module:
+    """Export a sparse dynamic model to a regular model.
+
+    This should be done after the model is fine-tuned and the weights are fixed.
+
+    .. warning::
+
+        After the call to ``export()``, the sparsity mask will no longer be enforced. This means any
+        future weight updates would destroy the sparsity pattern. If you want to continue training,
+        call ``export()`` after training is finished.
+    """
+    # unwrap a DP/DDP model
+    model = unwrap_model(
+        model,
+        warn=True,
+        msg=(
+            f"Unwrapping a {type(model).__name__} model for export! Note that the export is"
+            " in-place and the model wrapper should be re-created after export since the wrapper"
+            " might not support changing parameters after initialization."
+        ),
+    )
+
+    # apply export mode and return model
+    return apply_mode(model, "export_sparse", registry=SparsityModeRegistry)

modelopt/torch/speculative/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+"""Speculative Decoding Optimizations."""
+
+from . import eagle, medusa, mode, plugins
+from .config import *
+from .speculative_decoding import *

modelopt/torch/speculative/config.py

@@ -0,0 +1,100 @@
+# 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.
+
+"""Configurations for speculative decoding modes."""
+
+from copy import deepcopy
+
+from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField
+
+from .eagle.default_config import default_eagle_config
+
+eagle3_default_config = deepcopy(default_eagle_config)
+eagle_mtp_default_config = deepcopy(default_eagle_config)
+
+eagle3_default_config.update({"use_aux_hidden_state": True, "use_last_layernorm": True})
+eagle_mtp_default_config.update({"use_last_layernorm": True, "use_mtp_layernorm": True})
+
+EAGLE1_DEFAULT_CFG = {
+    "algorithm": "eagle",
+    "config": {
+        "eagle_architecture_config": deepcopy(default_eagle_config),
+    },
+}
+
+EAGLE3_DEFAULT_CFG = {
+    "algorithm": "eagle",
+    "config": {
+        "eagle_architecture_config": eagle3_default_config,
+    },
+}
+
+EAGLE_MTP_DEFAULT_CFG = {
+    "algorithm": "eagle",
+    "config": {
+        "eagle_reuse_base_decoder": True,
+        "eagle_architecture_config": eagle_mtp_default_config,
+    },
+}
+
+
+class MedusaConfig(ModeloptBaseConfig):
+    """Medusa config."""
+
+    medusa_num_heads: int = ModeloptField(
+        default=2,
+        description=("The number of medusa heads added to the model."),
+    )
+
+    medusa_num_layers: int = ModeloptField(
+        default=1,
+        description=("The number of ResBlocks used in medusa head."),
+    )
+
+
+class EagleConfig(ModeloptBaseConfig):
+    """Eagle config."""
+
+    eagle_offline: bool = ModeloptField(
+        default=False, description=("Whether to use detached Eagle.")
+    )
+
+    eagle_hidden_state_distillation: bool = ModeloptField(
+        default=False, description=("Whether to use feature hidden states distillation.")
+    )
+
+    eagle_self_logit_distillation: bool = ModeloptField(
+        default=True, description=("Whether to use logit distillation.")
+    )
+
+    eagle_freeze_base_model: bool = ModeloptField(
+        default=True, description=("Whether to freeze base model during eagle module training.")
+    )
+
+    eagle_report_acc: bool = ModeloptField(
+        default=True, description=("Whether to report eval accuracy.")
+    )
+
+    eagle_reuse_base_decoder: bool = ModeloptField(
+        default=False, description=("Whether to reuse base model decoder in eagle module.")
+    )
+
+    eagle_loss_decay_factor: float = ModeloptField(
+        default=0.9, description=("The decay factor for multiple eagle_loss.")
+    )
+
+    eagle_architecture_config: dict = ModeloptField(
+        default={}, description=("The config for eagle module architecture.")
+    )

modelopt/torch/speculative/eagle/__init__.py

@@ -0,0 +1,20 @@
+# 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.
+
+"""Eagle Optimization Method."""
+
+from .conversion import *
+from .default_config import *
+from .eagle_model import *

modelopt/torch/speculative/eagle/conversion.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.
+
+"""Eagle conversion/restore utilities."""
+
+from torch import nn
+
+from modelopt.torch.opt.conversion import ModelLikeModule
+from modelopt.torch.opt.dynamic import _DMRegistryCls
+from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
+
+from ..config import EagleConfig
+
+EagleDMRegistry = _DMRegistryCls(prefix="Eagle")  # global instance for the registry
+OfflineEagleDMRegistry = _DMRegistryCls(prefix="DetachedEagle")  # global instance for the registry
+
+
+def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertReturnType:
+    """Convert the model to a eagle model as per `config`."""
+    # initialize the true module if necessary
+    model = model.init_modellike() if isinstance(model, ModelLikeModule) else model
+
+    registry = OfflineEagleDMRegistry if config.eagle_offline else EagleDMRegistry
+
+    original_cls = type(model)
+    if original_cls not in registry:
+        for cls in registry._registry:
+            if issubclass(original_cls, cls):
+                registry.register({original_cls: "base_model_class"})(registry[cls])
+                break
+
+    eagle_model = registry.convert(model)
+    eagle_model.modify(
+        eagle_offline=config.eagle_offline,
+        eagle_hidden_state_distillation=config.eagle_hidden_state_distillation,
+        eagle_self_logit_distillation=config.eagle_self_logit_distillation,
+        eagle_freeze_base_model=config.eagle_freeze_base_model,
+        eagle_report_acc=config.eagle_report_acc,
+        eagle_reuse_base_decoder=config.eagle_reuse_base_decoder,
+        eagle_loss_decay_factor=config.eagle_loss_decay_factor,
+        eagle_architecture_config=config.eagle_architecture_config,
+    )
+
+    # no metadata, all specified via config.
+    metadata = {}
+
+    return eagle_model, metadata
+
+
+def restore_eagle_model(model: nn.Module, config: EagleConfig, metadata: MetadataDict) -> nn.Module:
+    """Function for restoring a previously convert model to a eagle model."""
+    # the metadata should be empty
+    assert not metadata, "No metadata expected!"
+
+    return convert_to_eagle_model(model, config)[0]

modelopt/torch/speculative/eagle/default_config.py

@@ -0,0 +1,50 @@
+# 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.
+
+"""Default EAGLE architecture config."""
+
+default_eagle_config = {
+    "hidden_act": "silu",
+    "torch_dtype": "bfloat16",
+    "vocab_size": 128256,
+    "draft_vocab_size": 128256,
+    "max_position_embeddings": 8192,
+    "position_embedding_type": "rope",
+    "rope_scaling": {
+        "factor": 8.0,
+        "low_freq_factor": 1.0,
+        "high_freq_factor": 4.0,
+        "original_max_position_embeddings": 8192,
+        "rope_type": "llama3",
+    },
+    "rope_theta": 500000.0,
+    "num_hidden_layers": 1,
+    "hidden_size": 4096,
+    "intermediate_size": 14336,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "initializer_range": 0.01,
+    "rms_norm_eps": 1e-05,
+    "mlp_bias": False,
+    "attention_bias": False,
+    "attention_dropout": 0.0,
+    "use_input_layernorm_in_first_layer": True,
+    "use_last_layernorm": False,
+    "use_aux_hidden_state": False,
+    "eagle_aux_hidden_state_layer_ids": [],
+    "use_mtp_layernorm": False,
+    "parallel_draft_step": 1,
+    "has_lm_head": False,
+}