optimum/neuron/configuration_utils.py

# coding=utf-8
# Copyright 2025 The HuggingFace Inc. team.
#
# 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.
"""Neuron configuration class and utilities."""

import copy
import enum
import json
import os
from dataclasses import dataclass, is_dataclass
from typing import Any, Dict, Optional, Union

import torch
from transformers.utils import (
    PushToHubMixin,
    cached_file,
    download_url,
    is_remote_url,
    logging,
)

from .utils.version_utils import get_neuronxcc_version
from .version import __version__


logger = logging.get_logger(__name__)

NEURON_CONFIG_NAME = "neuron_config.json"

_NEURON_CONFIG_FOR_KEY = {}
_KEY_FOR_NEURON_CONFIG = {}


def register_neuron_config(cls):
    key = cls.__name__
    assert issubclass(cls, NeuronConfig)
    _KEY_FOR_NEURON_CONFIG[cls] = key
    _NEURON_CONFIG_FOR_KEY[key] = cls
    return cls


@dataclass
class NeuronConfig(PushToHubMixin):
    # no-format
    """
    Class that holds a configuration for a neuron model.

    Must be subclassed.

    """

    def __hash__(self):
        return hash(self.to_json_string())

    def __eq__(self, other):
        if not isinstance(other, NeuronConfig):
            return False

        self_json = self.to_json_string()
        other_json = other.to_json_string()
        return self_json == other_json

    def __repr__(self):
        return f"{self.__class__.__name__} {self.to_json_string()}"

    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        config_file_name: Optional[Union[str, os.PathLike]] = None,
        push_to_hub: bool = False,
        **kwargs,
    ):
        r"""
        Save a neuron configuration object to the directory `save_directory`, so that it can be re-loaded using the
        [`~NeuronConfig.from_pretrained`] class method.

        Args:
            save_directory (`str` or `os.PathLike`):
                Directory where the configuration JSON file will be saved (will be created if it does not exist).
            config_file_name (`str` or `os.PathLike`, *optional*, defaults to `"neuron_config.json"`):
                Name of the generation configuration JSON file to be saved in `save_directory`.
            push_to_hub (`bool`, *optional*, defaults to `False`):
                Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
                namespace).
            kwargs (`Dict[str, Any]`, *optional*):
                Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
        """

        config_file_name = config_file_name if config_file_name is not None else NEURON_CONFIG_NAME

        if os.path.isfile(save_directory):
            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")

        os.makedirs(save_directory, exist_ok=True)

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
            repo_id = self._create_repo(repo_id, **kwargs)
            files_timestamps = self._get_files_timestamps(save_directory)

        output_config_file = os.path.join(save_directory, config_file_name)

        self.to_json_file(output_config_file)
        logger.info(f"Configuration saved in {output_config_file}")

        if push_to_hub:
            self._upload_modified_files(
                save_directory,
                repo_id,
                files_timestamps,
                commit_message=commit_message,
                token=kwargs.get("token"),
            )

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name: Union[str, os.PathLike],
        config_file_name: Optional[Union[str, os.PathLike]] = None,
        token: Optional[Union[str, bool]] = None,
        revision: str = "main",
    ) -> "NeuronConfig":
        r"""
        Instantiate a [`NeuronConfig`] from a neuron configuration file.

        Args:
            pretrained_model_name (`str` or `os.PathLike`):
                This can be either:

                - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
                  huggingface.co.
                - a path to a *directory* containing a configuration file saved using the
                  [`~NeuronConfig.save_pretrained`] method, e.g., `./my_model_directory/`.
            config_file_name (`str` or `os.PathLike`, *optional*, defaults to `"neuron_config.json"`):
                Name of the neuron configuration JSON file to be loaded from `pretrained_model_name`.
            token (`str` or `bool`, *optional*):
                The token to use as HTTP bearer authorization for remote files. If `True`, or not specified, will use
                the token generated when running `huggingface-cli login` (stored in `~/.huggingface`).
            revision (`str`, *optional*, defaults to `"main"`):
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
                identifier allowed by git.

        Returns:
            [`NeuronConfig`]: The configuration object instantiated from this pretrained model.

        ```"""
        config_file_name = config_file_name if config_file_name is not None else NEURON_CONFIG_NAME
        config_path = os.path.join(pretrained_model_name, config_file_name)
        config_path = str(config_path)

        is_local = os.path.exists(config_path)
        if os.path.isfile(config_path):
            # Special case when config_path is a local file
            resolved_config_file = config_path
            is_local = True
        elif is_remote_url(config_path):
            configuration_file = config_path
            resolved_config_file = download_url(config_path)
        else:
            configuration_file = config_file_name
            try:
                # Load from local folder or from cache or download from model Hub and cache
                resolved_config_file = cached_file(
                    pretrained_model_name,
                    configuration_file,
                    token=token,
                    revision=revision,
                )
            except EnvironmentError:
                # Raise any environment error raise by `cached_file`. It will have a helpful error message adapted to
                # the original exception.
                raise
            except Exception:
                # For any other exception, we throw a generic error.
                raise EnvironmentError(
                    f"Can't load the configuration of '{pretrained_model_name}'. If you were trying to load it"
                    " from 'https://huggingface.co/models', make sure you don't have a local directory with the same"
                    f" name. Otherwise, make sure '{pretrained_model_name}' is the correct path to a directory"
                    f" containing a {configuration_file} file"
                )

        try:
            # Load config dict
            with open(resolved_config_file, "r", encoding="utf-8") as reader:
                text = reader.read()
            config_dict = json.loads(text)
        except (json.JSONDecodeError, UnicodeDecodeError):
            raise EnvironmentError(
                f"It looks like the config file at '{resolved_config_file}' is not a valid JSON file."
            )

        if is_local:
            logger.info(f"loading configuration file {resolved_config_file}")
        else:
            logger.info(f"loading configuration file {configuration_file} from cache at {resolved_config_file}")

        return cls.from_dict(config_dict)

    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> "NeuronConfig":
        """
        Instantiates a [`NeuronConfig`] from a Python dictionary of parameters.

        Args:
            config_dict (`Dict[str, Any]`):
                Dictionary that will be used to instantiate the configuration object.

        Returns:
            [`NeuronConfig`]: The configuration object instantiated from those parameters.
        """
        for key in ["optimum_neuron_version", "_serialized_key", "neuronxcc_version"]:
            if key not in config_dict:
                raise ValueError(f"Invalid neuron configuration: missing mandatory {key} attribute.")
        # Check versions
        version = config_dict.pop("optimum_neuron_version")
        if version != __version__:
            raise ValueError(
                f"This neuron configuration file was generated with optimum-neuron version {version}, but this version is {__version__}."
            )
        cfg_neuronxcc_version = config_dict.pop("neuronxcc_version")
        neuronxcc_version = get_neuronxcc_version()
        if cfg_neuronxcc_version != neuronxcc_version:
            raise RuntimeError(
                "This neuron configuration corresponds to a model exported using version"
                f" {cfg_neuronxcc_version} of the neuronxcc compiler, but {neuronxcc_version} is installed."
            )
        # Retrieve neuron config class from serialized key
        _serialized_key = config_dict.pop("_serialized_key")
        if cls is NeuronConfig:
            # Force the registration of the neuron configuration classes
            from .models import neuron_config  # noqa F401

            # We need to identify the actual neuron configuration class for this serialized key
            if _serialized_key is None:
                raise ValueError("Neuron configuration is invalid: unable to identify the serialized key")
            cls = _NEURON_CONFIG_FOR_KEY.get(_serialized_key, None)
            if cls is None:
                raise ValueError(f"No neuron configuration registered for the {_serialized_key} serialized key")
        else:
            assert _KEY_FOR_NEURON_CONFIG[cls] == _serialized_key
        config = cls(**{**config_dict})
        logger.info(f"Neuron config {config}")
        return config

    def to_dict(self) -> Dict[str, Any]:
        """
        Serializes this instance to a Python dictionary.

        Returns:
            `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance.
        """
        output = copy.deepcopy(self)

        serializable_types = (str, int, float, bool)

        def _to_dict(obj):
            if obj is None or isinstance(obj, serializable_types):
                return obj
            elif isinstance(obj, enum.Enum):
                return obj.value
            elif isinstance(obj, torch.Tensor):
                return obj.tolist()
            elif isinstance(obj, list):
                return [_to_dict(e) for e in obj]
            elif isinstance(obj, dict):
                return {_to_dict(k): _to_dict(v) for k, v in obj.items()}
            elif isinstance(obj, tuple):
                return str(tuple(_to_dict(e) for e in obj))
            elif isinstance(obj, torch.dtype):
                return str(obj).split(".")[1]
            else:
                as_dict = obj.__dict__
                return _to_dict(as_dict)

        output = _to_dict(output)

        # Add serialized key as it is required to identify the NeuronConfig class when deserializing the file
        cls = self.__class__
        _serialized_key = _KEY_FOR_NEURON_CONFIG.get(cls, None)
        if _serialized_key is None:
            raise ValueError(f"Unable to identify the serialized key for {cls.__name__}. Did you register it ?")
        output["_serialized_key"] = _serialized_key
        # Add optimum-neuron version to check compatibility
        output["optimum_neuron_version"] = __version__
        # Also add compiler version
        output["neuronxcc_version"] = get_neuronxcc_version()

        return output

    def to_json_string(self) -> str:
        """
        Serializes this instance to a JSON string.

        Returns:
            `str`: String containing all the attributes that make up this configuration instance in JSON format.
        """
        config_dict = self.to_dict()

        def convert_keys_to_string(obj):
            if isinstance(obj, dict):
                return {str(key): convert_keys_to_string(value) for key, value in obj.items()}
            elif isinstance(obj, list):
                return [convert_keys_to_string(item) for item in obj]
            else:
                return obj

        def convert_dataclass_to_dict(obj):
            if isinstance(obj, dict):
                return {key: convert_dataclass_to_dict(value) for key, value in obj.items()}
            elif is_dataclass(obj):
                return obj.to_dict()
            else:
                return obj

        config_dict = convert_keys_to_string(config_dict)
        config_dict = convert_dataclass_to_dict(config_dict)

        return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"

    def to_json_file(self, json_file_path: Union[str, os.PathLike]):
        """
        Save this instance to a JSON file.

        Args:
            json_file_path (`str` or `os.PathLike`):
                Path to the JSON file in which this configuration instance's parameters will be saved.
        """
        with open(json_file_path, "w", encoding="utf-8") as writer:
            writer.write(self.to_json_string())