Break weight tying when quantizing input embedding (#37905)

<|original_code|>

    def post_init(self):
        r"""
        Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
        """
        if self.bits not in [2, 3, 4]:
            raise ValueError("bits must be 2, 3, or 4")
        if self.p not in [1, 2]:
            raise ValueError("p must be 1 or 2. 2 is always better in practice")
        if self.group_size not in [64, 128, 256]:
            raise ValueError("group_size must be 64, 128, or 256")
        if self.hadamard_size % self.group_size != 0:
            raise ValueError("hadamard_size must be divisible by group_size")

@dataclass
class TorchAoConfig(QuantizationConfigMixin):
    quant_method: QuantizationMethod
    quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
    modules_to_not_convert: Optional[List]
    quant_type_kwargs: Dict[str, Any]
    include_embedding: bool

    """This is a config class for torchao quantization/sparsity techniques.

    Args:
        quant_type (`Union[str, AOBaseConfig]`):
            The type of quantization we want to use. Can be either:
            - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
            - An AOBaseConfig instance: for more advanced configuration options.
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision.
        inlcude_embedding (`bool`, default to `False`):
            Whether to include embedding in quantization or not, input embedding will be removed from
            the module_not_to_convert list as well if this flag is set.
        kwargs (`Dict[str, Any]`, *optional*):
            The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
            `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
            https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques

    Example:

    ```python
    # AOBaseConfig-based configuration
    config = Int4WeightOnlyConfig(group_size=32)
    quantization_config = TorchAoConfig(config)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # String-based configuration
    quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
    # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # autoquant
    # `autoquant` is a convenient way for users to search for the best quantization for each layer
    # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
    # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
    # defaults to None, which means we'll try to get the best performing quantized model without
    # considering accuracy
    quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    # run through example inputs, quantization methods will be selected based on the shape of example input
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
    MAX_NEW_TOKENS = 1000
    model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
    # manually ran finalize_autoquant if needed
    if hasattr(quantized_model, "finalize_autoquant"):
      print("finalizing autoquant")
      quantized_model.finalize_autoquant()

    ```
    """

    def __init__(
        self,
        quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
        modules_to_not_convert: Optional[List] = None,
        include_embedding: bool = False,
        **kwargs,
    ):
        self.quant_method = QuantizationMethod.TORCHAO
        self.quant_type = quant_type
        self.modules_to_not_convert = modules_to_not_convert
        self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
        self.include_embedding = include_embedding
        self.post_init()

    @staticmethod
    def _get_ao_version() -> version.Version:
        """Centralized check for TorchAO availability and version requirements."""
        if not is_torchao_available():
            raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")

        return version.parse(importlib.metadata.version("torchao"))

    def post_init(self):
        """Validate configuration and set defaults."""
        ao_version = self._get_ao_version()

        # Handle quant_type based on type and version
        if isinstance(self.quant_type, str):
            self._validate_string_quant_type()
        elif ao_version > version.parse("0.9.0"):
            from torchao.quantization.quant_api import AOBaseConfig

            if not isinstance(self.quant_type, AOBaseConfig):
                raise ValueError(
                    f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
                )
<|edits_diff|>
--- src/transformers/utils/quantization_config.py
+++ src/transformers/utils/quantization_config.py
@@ -20,6 +20,7 @@
     modules_to_not_convert: Optional[List]
     quant_type_kwargs: Dict[str, Any]
     include_embedding: bool
+    untie_embedding_weights: bool

     """This is a config class for torchao quantization/sparsity techniques.

@@ -34,6 +35,9 @@
         inlcude_embedding (`bool`, default to `False`):
             Whether to include embedding in quantization or not, input embedding will be removed from
             the module_not_to_convert list as well if this flag is set.
+        untie_embedding_weights (`bool`, default to `False`):
+            Whether to untie the weights when we are quantizing input embedding weights that is tied
+            to other weights.
         kwargs (`Dict[str, Any]`, *optional*):
             The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
             `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
@@ -79,6 +83,7 @@
         quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
         modules_to_not_convert: Optional[List] = None,
         include_embedding: bool = False,
+        untie_embedding_weights: bool = False,
         **kwargs,
     ):
         self.quant_method = QuantizationMethod.TORCHAO
<|current_version|>

    def post_init(self):
        r"""
        Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
        """
        if self.bits not in [2, 3, 4]:
            raise ValueError("bits must be 2, 3, or 4")
        if self.p not in [1, 2]:
            raise ValueError("p must be 1 or 2. 2 is always better in practice")
        if self.group_size not in [64, 128, 256]:
            raise ValueError("group_size must be 64, 128, or 256")
        if self.hadamard_size % self.group_size != 0:
            raise ValueError("hadamard_size must be divisible by group_size")

@dataclass
class TorchAoConfig(QuantizationConfigMixin):
    quant_method: QuantizationMethod
    quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
    modules_to_not_convert: Optional[List]
    quant_type_kwargs: Dict[str, Any]
    include_embedding: bool
    untie_embedding_weights: bool

    """This is a config class for torchao quantization/sparsity techniques.

    Args:
        quant_type (`Union[str, AOBaseConfig]`):
            The type of quantization we want to use. Can be either:
            - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
            - An AOBaseConfig instance: for more advanced configuration options.
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision.
        inlcude_embedding (`bool`, default to `False`):
            Whether to include embedding in quantization or not, input embedding will be removed from
            the module_not_to_convert list as well if this flag is set.
        untie_embedding_weights (`bool`, default to `False`):
            Whether to untie the weights when we are quantizing input embedding weights that is tied
            to other weights.
        kwargs (`Dict[str, Any]`, *optional*):
            The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
            `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
            https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques

    Example:

    ```python
    # AOBaseConfig-based configuration
    config = Int4WeightOnlyConfig(group_size=32)
    quantization_config = TorchAoConfig(config)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # String-based configuration
    quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
    # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # autoquant
    # `autoquant` is a convenient way for users to search for the best quantization for each layer
    # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
    # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
    # defaults to None, which means we'll try to get the best performing quantized model without
    # considering accuracy
    quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    # run through example inputs, quantization methods will be selected based on the shape of example input
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
    MAX_NEW_TOKENS = 1000
    model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
    # manually ran finalize_autoquant if needed
    if hasattr(quantized_model, "finalize_autoquant"):
      print("finalizing autoquant")
      quantized_model.finalize_autoquant()

    ```
    """

    def __init__(
        self,
        quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
        modules_to_not_convert: Optional[List] = None,
        include_embedding: bool = False,
        untie_embedding_weights: bool = False,
        **kwargs,
    ):
        self.quant_method = QuantizationMethod.TORCHAO
        self.quant_type = quant_type
        self.modules_to_not_convert = modules_to_not_convert
        self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
        self.include_embedding = include_embedding
        self.post_init()

    @staticmethod
    def _get_ao_version() -> version.Version:
        """Centralized check for TorchAO availability and version requirements."""
        if not is_torchao_available():
            raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")

        return version.parse(importlib.metadata.version("torchao"))

    def post_init(self):
        """Validate configuration and set defaults."""
        ao_version = self._get_ao_version()

        # Handle quant_type based on type and version
        if isinstance(self.quant_type, str):
            self._validate_string_quant_type()
        elif ao_version > version.parse("0.9.0"):
            from torchao.quantization.quant_api import AOBaseConfig

            if not isinstance(self.quant_type, AOBaseConfig):
                raise ValueError(
                    f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
                )
<|next_version|>

    def post_init(self):
        r"""
        Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
        """
        if self.bits not in [2, 3, 4]:
            raise ValueError("bits must be 2, 3, or 4")
        if self.p not in [1, 2]:
            raise ValueError("p must be 1 or 2. 2 is always better in practice")
        if self.group_size not in [64, 128, 256]:
            raise ValueError("group_size must be 64, 128, or 256")
        if self.hadamard_size % self.group_size != 0:
            raise ValueError("hadamard_size must be divisible by group_size")

@dataclass
class TorchAoConfig(QuantizationConfigMixin):
    quant_method: QuantizationMethod
    quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
    modules_to_not_convert: Optional[List]
    quant_type_kwargs: Dict[str, Any]
    include_embedding: bool
    untie_embedding_weights: bool

    """This is a config class for torchao quantization/sparsity techniques.

    Args:
        quant_type (`Union[str, AOBaseConfig]`):
            The type of quantization we want to use. Can be either:
            - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
            - An AOBaseConfig instance: for more advanced configuration options.
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision.
        inlcude_embedding (`bool`, default to `False`):
            Whether to include embedding in quantization or not, input embedding will be removed from
            the module_not_to_convert list as well if this flag is set.
        untie_embedding_weights (`bool`, default to `False`):
            Whether to untie the weights when we are quantizing input embedding weights that is tied
            to other weights.
        kwargs (`Dict[str, Any]`, *optional*):
            The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
            `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
            https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques

    Example:

    ```python
    # AOBaseConfig-based configuration
    config = Int4WeightOnlyConfig(group_size=32)
    quantization_config = TorchAoConfig(config)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # String-based configuration
    quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
    # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # autoquant
    # `autoquant` is a convenient way for users to search for the best quantization for each layer
    # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
    # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
    # defaults to None, which means we'll try to get the best performing quantized model without
    # considering accuracy
    quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    # run through example inputs, quantization methods will be selected based on the shape of example input
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
    MAX_NEW_TOKENS = 1000
    model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
    # manually ran finalize_autoquant if needed
    if hasattr(quantized_model, "finalize_autoquant"):
      print("finalizing autoquant")
      quantized_model.finalize_autoquant()

    ```
    """

    def __init__(
        self,
        quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
        modules_to_not_convert: Optional[List] = None,
        include_embedding: bool = False,
        untie_embedding_weights: bool = False,
        **kwargs,
    ):
        self.quant_method = QuantizationMethod.TORCHAO
        self.quant_type = quant_type
        self.modules_to_not_convert = modules_to_not_convert
        self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
        self.include_embedding = include_embedding
        self.untie_embedding_weights = untie_embedding_weights
        self.post_init()

    @staticmethod
    def _get_ao_version() -> version.Version:
        """Centralized check for TorchAO availability and version requirements."""
        if not is_torchao_available():
            raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")

        return version.parse(importlib.metadata.version("torchao"))

    def post_init(self):
        """Validate configuration and set defaults."""
        ao_version = self._get_ao_version()

        # Handle quant_type based on type and version
        if isinstance(self.quant_type, str):
            self._validate_string_quant_type()
        elif ao_version > version.parse("0.9.0"):
            from torchao.quantization.quant_api import AOBaseConfig

            if not isinstance(self.quant_type, AOBaseConfig):
                raise ValueError(
                    f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
                )


Support `AOPerModuleConfig` and `include_embedding` (#37802)

<|original_code|>
        self.post_init()

    def post_init(self):
        r"""
        Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
        """
        if self.bits not in [2, 3, 4]:
            raise ValueError("bits must be 2, 3, or 4")
        if self.p not in [1, 2]:
            raise ValueError("p must be 1 or 2. 2 is always better in practice")
        if self.group_size not in [64, 128, 256]:
            raise ValueError("group_size must be 64, 128, or 256")
        if self.hadamard_size % self.group_size != 0:
            raise ValueError("hadamard_size must be divisible by group_size")

@dataclass
class TorchAoConfig(QuantizationConfigMixin):
    quant_method: QuantizationMethod
    quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
    modules_to_not_convert: Optional[List]
    quant_type_kwargs: Dict[str, Any]

    """This is a config class for torchao quantization/sparsity techniques.

    Args:
        quant_type (`Union[str, AOBaseConfig]`):
            The type of quantization we want to use. Can be either:
            - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
            - An AOBaseConfig instance: for more advanced configuration options.
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision.
        kwargs (`Dict[str, Any]`, *optional*):
            The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
            `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
            https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques

    Example:

    ```python
    # AOBaseConfig-based configuration
    config = Int4WeightOnlyConfig(group_size=32)
    quantization_config = TorchAoConfig(config)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # String-based configuration
    quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
    # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # autoquant
    # `autoquant` is a convenient way for users to search for the best quantization for each layer
    # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
    # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
    # defaults to None, which means we'll try to get the best performing quantized model without
    # considering accuracy
    quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    # run through example inputs, quantization methods will be selected based on the shape of example input
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
    MAX_NEW_TOKENS = 1000
    model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
    # manually ran finalize_autoquant if needed
    if hasattr(quantized_model, "finalize_autoquant"):
      print("finalizing autoquant")
      quantized_model.finalize_autoquant()

    ```
    """

    def __init__(
        self,
        quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
        modules_to_not_convert: Optional[List] = None,
        **kwargs,
    ):
        self.quant_method = QuantizationMethod.TORCHAO
        self.quant_type = quant_type
        self.modules_to_not_convert = modules_to_not_convert
        self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
        self.post_init()

    @staticmethod
    def _get_ao_version() -> version.Version:
        """Centralized check for TorchAO availability and version requirements."""
        if not is_torchao_available():
            raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")

        return version.parse(importlib.metadata.version("torchao"))

    def post_init(self):
        """Validate configuration and set defaults."""
        ao_version = self._get_ao_version()

        # Handle quant_type based on type and version
        if isinstance(self.quant_type, str):
            self._validate_string_quant_type()
        elif ao_version > version.parse("0.9.0"):
            from torchao.quantization.quant_api import AOBaseConfig

            if not isinstance(self.quant_type, AOBaseConfig):
                raise ValueError(
                    f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
                )
<|edits_diff|>
--- src/transformers/utils/quantization_config.py
+++ src/transformers/utils/quantization_config.py
@@ -20,6 +20,7 @@
     quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
     modules_to_not_convert: Optional[List]
     quant_type_kwargs: Dict[str, Any]
+    include_embedding: bool

     """This is a config class for torchao quantization/sparsity techniques.

@@ -31,6 +32,9 @@
         modules_to_not_convert (`list`, *optional*, default to `None`):
             The list of modules to not quantize, useful for quantizing models that explicitly require to have
             some modules left in their original precision.
+        inlcude_embedding (`bool`, default to `False`):
+            Whether to include embedding in quantization or not, input embedding will be removed from
+            the module_not_to_convert list as well if this flag is set.
         kwargs (`Dict[str, Any]`, *optional*):
             The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
             `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
@@ -75,6 +79,7 @@
         self,
         quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
         modules_to_not_convert: Optional[List] = None,
+        include_embedding: bool = False,
         **kwargs,
     ):
         self.quant_method = QuantizationMethod.TORCHAO
<|current_version|>
        self.post_init()

    def post_init(self):
        r"""
        Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
        """
        if self.bits not in [2, 3, 4]:
            raise ValueError("bits must be 2, 3, or 4")
        if self.p not in [1, 2]:
            raise ValueError("p must be 1 or 2. 2 is always better in practice")
        if self.group_size not in [64, 128, 256]:
            raise ValueError("group_size must be 64, 128, or 256")
        if self.hadamard_size % self.group_size != 0:
            raise ValueError("hadamard_size must be divisible by group_size")

@dataclass
class TorchAoConfig(QuantizationConfigMixin):
    quant_method: QuantizationMethod
    quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
    modules_to_not_convert: Optional[List]
    quant_type_kwargs: Dict[str, Any]
    include_embedding: bool

    """This is a config class for torchao quantization/sparsity techniques.

    Args:
        quant_type (`Union[str, AOBaseConfig]`):
            The type of quantization we want to use. Can be either:
            - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
            - An AOBaseConfig instance: for more advanced configuration options.
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision.
        inlcude_embedding (`bool`, default to `False`):
            Whether to include embedding in quantization or not, input embedding will be removed from
            the module_not_to_convert list as well if this flag is set.
        kwargs (`Dict[str, Any]`, *optional*):
            The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
            `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
            https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques

    Example:

    ```python
    # AOBaseConfig-based configuration
    config = Int4WeightOnlyConfig(group_size=32)
    quantization_config = TorchAoConfig(config)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # String-based configuration
    quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
    # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # autoquant
    # `autoquant` is a convenient way for users to search for the best quantization for each layer
    # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
    # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
    # defaults to None, which means we'll try to get the best performing quantized model without
    # considering accuracy
    quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    # run through example inputs, quantization methods will be selected based on the shape of example input
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
    MAX_NEW_TOKENS = 1000
    model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
    # manually ran finalize_autoquant if needed
    if hasattr(quantized_model, "finalize_autoquant"):
      print("finalizing autoquant")
      quantized_model.finalize_autoquant()

    ```
    """

    def __init__(
        self,
        quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
        modules_to_not_convert: Optional[List] = None,
        include_embedding: bool = False,
        **kwargs,
    ):
        self.quant_method = QuantizationMethod.TORCHAO
        self.quant_type = quant_type
        self.modules_to_not_convert = modules_to_not_convert
        self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
        self.post_init()

    @staticmethod
    def _get_ao_version() -> version.Version:
        """Centralized check for TorchAO availability and version requirements."""
        if not is_torchao_available():
            raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")

        return version.parse(importlib.metadata.version("torchao"))

    def post_init(self):
        """Validate configuration and set defaults."""
        ao_version = self._get_ao_version()

        # Handle quant_type based on type and version
        if isinstance(self.quant_type, str):
            self._validate_string_quant_type()
        elif ao_version > version.parse("0.9.0"):
            from torchao.quantization.quant_api import AOBaseConfig

            if not isinstance(self.quant_type, AOBaseConfig):
                raise ValueError(
                    f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
                )
<|next_version|>
        self.post_init()

    def post_init(self):
        r"""
        Safety checker that arguments are correct - also replaces some NoneType arguments with their default values.
        """
        if self.bits not in [2, 3, 4]:
            raise ValueError("bits must be 2, 3, or 4")
        if self.p not in [1, 2]:
            raise ValueError("p must be 1 or 2. 2 is always better in practice")
        if self.group_size not in [64, 128, 256]:
            raise ValueError("group_size must be 64, 128, or 256")
        if self.hadamard_size % self.group_size != 0:
            raise ValueError("hadamard_size must be divisible by group_size")

@dataclass
class TorchAoConfig(QuantizationConfigMixin):
    quant_method: QuantizationMethod
    quant_type: Union[str, "AOBaseConfig"]  # noqa: F821
    modules_to_not_convert: Optional[List]
    quant_type_kwargs: Dict[str, Any]
    include_embedding: bool

    """This is a config class for torchao quantization/sparsity techniques.

    Args:
        quant_type (`Union[str, AOBaseConfig]`):
            The type of quantization we want to use. Can be either:
            - A string: currently supporting: `int4_weight_only`, `int8_weight_only` and `int8_dynamic_activation_int8_weight`.
            - An AOBaseConfig instance: for more advanced configuration options.
        modules_to_not_convert (`list`, *optional*, default to `None`):
            The list of modules to not quantize, useful for quantizing models that explicitly require to have
            some modules left in their original precision.
        inlcude_embedding (`bool`, default to `False`):
            Whether to include embedding in quantization or not, input embedding will be removed from
            the module_not_to_convert list as well if this flag is set.
        kwargs (`Dict[str, Any]`, *optional*):
            The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
            `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
            https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques

    Example:

    ```python
    # AOBaseConfig-based configuration
    config = Int4WeightOnlyConfig(group_size=32)
    quantization_config = TorchAoConfig(config)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # String-based configuration
    quantization_config = TorchAoConfig("int4_weight_only", group_size=32)
    # int4_weight_only quant is only working with *torch.bfloat16* dtype right now
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)

    # autoquant
    # `autoquant` is a convenient way for users to search for the best quantization for each layer
    # `min_sqnr` is an option to control the accuracy of the model, higher value means the model is more
    # accurate, we can start with 30 and adjust it to larger or smaller (e.g. 40, 20)
    # defaults to None, which means we'll try to get the best performing quantized model without
    # considering accuracy
    quantization_config = TorchAoConfig("autoquant", min_sqnr=30)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    # run through example inputs, quantization methods will be selected based on the shape of example input
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
    MAX_NEW_TOKENS = 1000
    model.generate(**input_ids, max_new_tokens=MAX_NEW_TOKENS, cache_implementation="static")
    # manually ran finalize_autoquant if needed
    if hasattr(quantized_model, "finalize_autoquant"):
      print("finalizing autoquant")
      quantized_model.finalize_autoquant()

    ```
    """

    def __init__(
        self,
        quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
        modules_to_not_convert: Optional[List] = None,
        include_embedding: bool = False,
        **kwargs,
    ):
        self.quant_method = QuantizationMethod.TORCHAO
        self.quant_type = quant_type
        self.modules_to_not_convert = modules_to_not_convert
        self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
        self.include_embedding = include_embedding
        self.post_init()

    @staticmethod
    def _get_ao_version() -> version.Version:
        """Centralized check for TorchAO availability and version requirements."""
        if not is_torchao_available():
            raise ValueError("TorchAoConfig requires torchao to be installed. Install with `pip install torchao`")

        return version.parse(importlib.metadata.version("torchao"))

    def post_init(self):
        """Validate configuration and set defaults."""
        ao_version = self._get_ao_version()

        # Handle quant_type based on type and version
        if isinstance(self.quant_type, str):
            self._validate_string_quant_type()
        elif ao_version > version.parse("0.9.0"):
            from torchao.quantization.quant_api import AOBaseConfig

            if not isinstance(self.quant_type, AOBaseConfig):
                raise ValueError(
                    f"quant_type must be either a string or an AOBaseConfig instance, got {type(self.quant_type)}"
                )


Fix error of HPU TP (#37782)

<|original_code|>
        if tp_plan is not None:
            if not is_torch_greater_or_equal("2.5"):
                raise EnvironmentError("tensor parallel is only supported for `torch>=2.5`.")

            # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.
            device_type = torch._C._get_accelerator().type

            if not torch.distributed.is_initialized():
                try:
                    rank = int(os.environ["RANK"])
                    world_size = int(os.environ["WORLD_SIZE"])
                    if device_type == "cuda":
                        torch.distributed.init_process_group(
                            "nccl", rank=rank, world_size=world_size, init_method="env://"
                        )
                        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
                    elif device_type == "cpu":
                        cpu_backend = "ccl" if int(os.environ.get("CCL_WORKER_COUNT", 0)) else "gloo"
                        torch.distributed.init_process_group(cpu_backend, rank=rank, world_size=world_size)
                    elif device_type == "xpu":
                        torch.distributed.init_process_group("ccl", rank=rank, world_size=world_size)
                        torch.xpu.set_device(int(os.environ["LOCAL_RANK"]))

                except Exception as e:
                    raise EnvironmentError(
                        "We tried to initialize torch.distributed for you, but it failed, make"
                        "sure you init torch distributed in your script to use `tp_plan='auto'`"
                    ) from e

            # Get device with index assuming equal number of devices per host
            if device_type == "xpu":
                index = torch.xpu.current_device()
            else:
                index = None if device_type == "cpu" else torch.cuda.current_device()
            tp_device = torch.device(device_type, index)

            if index is not None and index > 0:
                import sys

                sys.stdout = open(os.devnull, "w")
                sys.stderr = open(os.devnull, "w")
            # This is the easiest way to dispatch to the current process device
            device_map = tp_device

            # Assuming sharding the model onto the world when tp_size not provided
            tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size()
            device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,))

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
<|edits_diff|>
--- src/transformers/modeling_utils.py
+++ src/transformers/modeling_utils.py
@@ -20,6 +20,9 @@
                     elif device_type == "xpu":
                         torch.distributed.init_process_group("ccl", rank=rank, world_size=world_size)
                         torch.xpu.set_device(int(os.environ["LOCAL_RANK"]))
+                    elif device_type == "hpu":
+                        torch.distributed.init_process_group("hccl", rank=rank, world_size=world_size)
+                        torch.hpu.set_device(int(os.environ["LOCAL_RANK"]))

                 except Exception as e:
                     raise EnvironmentError(
<|current_version|>
        if tp_plan is not None:
            if not is_torch_greater_or_equal("2.5"):
                raise EnvironmentError("tensor parallel is only supported for `torch>=2.5`.")

            # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.
            device_type = torch._C._get_accelerator().type

            if not torch.distributed.is_initialized():
                try:
                    rank = int(os.environ["RANK"])
                    world_size = int(os.environ["WORLD_SIZE"])
                    if device_type == "cuda":
                        torch.distributed.init_process_group(
                            "nccl", rank=rank, world_size=world_size, init_method="env://"
                        )
                        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
                    elif device_type == "cpu":
                        cpu_backend = "ccl" if int(os.environ.get("CCL_WORKER_COUNT", 0)) else "gloo"
                        torch.distributed.init_process_group(cpu_backend, rank=rank, world_size=world_size)
                    elif device_type == "xpu":
                        torch.distributed.init_process_group("ccl", rank=rank, world_size=world_size)
                        torch.xpu.set_device(int(os.environ["LOCAL_RANK"]))
                    elif device_type == "hpu":
                        torch.distributed.init_process_group("hccl", rank=rank, world_size=world_size)
                        torch.hpu.set_device(int(os.environ["LOCAL_RANK"]))

                except Exception as e:
                    raise EnvironmentError(
                        "We tried to initialize torch.distributed for you, but it failed, make"
                        "sure you init torch distributed in your script to use `tp_plan='auto'`"
                    ) from e

            # Get device with index assuming equal number of devices per host
            if device_type == "xpu":
                index = torch.xpu.current_device()
            else:
                index = None if device_type == "cpu" else torch.cuda.current_device()
            tp_device = torch.device(device_type, index)

            if index is not None and index > 0:
                import sys

                sys.stdout = open(os.devnull, "w")
                sys.stderr = open(os.devnull, "w")
            # This is the easiest way to dispatch to the current process device
            device_map = tp_device

            # Assuming sharding the model onto the world when tp_size not provided
            tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size()
            device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,))

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
<|next_version|>
        if tp_plan is not None:
            if not is_torch_greater_or_equal("2.5"):
                raise EnvironmentError("tensor parallel is only supported for `torch>=2.5`.")

            # Detect the accelerator on the machine. If no accelerator is available, it returns CPU.
            device_type = torch._C._get_accelerator().type

            if not torch.distributed.is_initialized():
                try:
                    rank = int(os.environ["RANK"])
                    world_size = int(os.environ["WORLD_SIZE"])
                    if device_type == "cuda":
                        torch.distributed.init_process_group(
                            "nccl", rank=rank, world_size=world_size, init_method="env://"
                        )
                        torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
                    elif device_type == "cpu":
                        cpu_backend = "ccl" if int(os.environ.get("CCL_WORKER_COUNT", 0)) else "gloo"
                        torch.distributed.init_process_group(cpu_backend, rank=rank, world_size=world_size)
                    elif device_type == "xpu":
                        torch.distributed.init_process_group("ccl", rank=rank, world_size=world_size)
                        torch.xpu.set_device(int(os.environ["LOCAL_RANK"]))
                    elif device_type == "hpu":
                        torch.distributed.init_process_group("hccl", rank=rank, world_size=world_size)
                        torch.hpu.set_device(int(os.environ["LOCAL_RANK"]))

                except Exception as e:
                    raise EnvironmentError(
                        "We tried to initialize torch.distributed for you, but it failed, make"
                        "sure you init torch distributed in your script to use `tp_plan='auto'`"
                    ) from e

            # Get device with index assuming equal number of devices per host
            if device_type == "xpu":
                index = torch.xpu.current_device()
            elif device_type == "hpu":
                index = torch.hpu.current_device()
            else:
                index = None if device_type == "cpu" else torch.cuda.current_device()
            tp_device = torch.device(device_type, index)

            if index is not None and index > 0:
                import sys

                sys.stdout = open(os.devnull, "w")
                sys.stderr = open(os.devnull, "w")
            # This is the easiest way to dispatch to the current process device
            device_map = tp_device

            # Assuming sharding the model onto the world when tp_size not provided
            tp_size = tp_size if tp_size is not None else torch.distributed.get_world_size()
            device_mesh = torch.distributed.init_device_mesh(tp_device.type, (tp_size,))

        if use_auth_token is not None:
            warnings.warn(
                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers. Please use `token` instead.",
                FutureWarning,
            )
            if token is not None:
                raise ValueError(
                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."


[generate] Fix encoder decoder models attention mask (#36018)

<|original_code|>
                inputs_embeds is not None  # Exception 1
                or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1])  # Exception 3
            ):
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        # 3. Prepare base model inputs
        input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step for every prompt.
        if not self.config.is_encoder_decoder:
            if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:
                model_inputs[input_ids_key] = None
                model_inputs["inputs_embeds"] = inputs_embeds
            else:
                # `clone` calls in this function ensure a consistent stride. See #32227
                model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)
                model_inputs["inputs_embeds"] = None
        else:
            model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)

        # 4. Create missing `position_ids` on the fly
        attention_mask = (
            kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask
        )
        attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
        position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids"
        if (
            attention_mask is not None
            and kwargs.get(position_ids_key) is None
            and position_ids_key in set(inspect.signature(self.forward).parameters.keys())
        ):
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            kwargs[position_ids_key] = position_ids  # placed in kwargs for further processing (see below)

        # 5. Slice model inputs if it's an input that should have the same length as `input_ids`
        for model_input_name in ["position_ids", "token_type_ids", "decoder_position_ids"]:
            model_input = kwargs.get(model_input_name)
            if model_input is not None:
                if past_key_values is not None:
                    current_input_length = (
                        model_inputs["inputs_embeds"].shape[1]
                        if model_inputs.get("inputs_embeds") is not None
                        else model_inputs[input_ids_key].shape[1]
                    )
                    model_input = model_input[:, -current_input_length:]
                    model_input = model_input.clone(memory_format=torch.contiguous_format)
                model_inputs[model_input_name] = model_input

        # 6. Create 4D attention mask is we are using a `StaticCache` (important for performant compiled forward pass)
        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
            if model_inputs["inputs_embeds"] is not None:
                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
                device = model_inputs["inputs_embeds"].device
            else:
                batch_size, sequence_length = model_inputs[input_ids_key].shape
                device = model_inputs[input_ids_key].device

            # Create the causal mask with fixed shape in advance, to reduce recompilations. If the function to create
            # the 4D causal mask exists, it should be present in the base model (XXXModel class).
            base_model = getattr(self, self.base_model_prefix, None)
            if base_model is None:
                causal_mask_creation_function = getattr(
                    self, "_prepare_4d_causal_attention_mask_with_cache_position", None
                )
            else:
                causal_mask_creation_function = getattr(
                    base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None
                )
            if causal_mask_creation_function is None:
                logger.warning_once(
                    f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method "
                    "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're "
                    "writing code, see Llama for an example implementation. If you're a user, please report this "
                    "issue on GitHub."
                )
            else:
                attention_mask = causal_mask_creation_function(
                    attention_mask,
                    sequence_length=sequence_length,
                    target_length=past_key_values.get_max_cache_shape(),
                    dtype=self.dtype,
                    device=device,
                    cache_position=cache_position,
                    batch_size=batch_size,
                    config=self.config,
                    past_key_values=past_key_values,
                )
        if attention_mask is not None:
            model_inputs[attention_mask_key] = attention_mask

        # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
        for key, value in kwargs.items():
            if key not in model_inputs:
                model_inputs[key] = value

        # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
        model_inputs.pop("labels", None)
        return model_inputs

    def _prepare_model_inputs(
        self,
        inputs: Optional[torch.Tensor] = None,
        bos_token_id: Optional[torch.Tensor] = None,
        model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
    ) -> Tuple[torch.Tensor, Optional[str], Dict[str, torch.Tensor]]:
        """
        This function extracts the model-specific `inputs` for generation.
        """
        # 1. retrieve all kwargs that are non-None or non-model input related.
        # some encoder-decoder models have different names for model and encoder
        if (
            self.config.is_encoder_decoder
            and hasattr(self, "encoder")
<|edits_diff|>
--- src/transformers/generation/utils.py
+++ src/transformers/generation/utils.py
@@ -20,6 +20,7 @@
             model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)

         # 4. Create missing `position_ids` on the fly
+        encoder_attention_mask = attention_mask if self.config.is_encoder_decoder else None
         attention_mask = (
             kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask
         )
<|current_version|>
                inputs_embeds is not None  # Exception 1
                or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1])  # Exception 3
            ):
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        # 3. Prepare base model inputs
        input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step for every prompt.
        if not self.config.is_encoder_decoder:
            if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:
                model_inputs[input_ids_key] = None
                model_inputs["inputs_embeds"] = inputs_embeds
            else:
                # `clone` calls in this function ensure a consistent stride. See #32227
                model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)
                model_inputs["inputs_embeds"] = None
        else:
            model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)

        # 4. Create missing `position_ids` on the fly
        encoder_attention_mask = attention_mask if self.config.is_encoder_decoder else None
        attention_mask = (
            kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask
        )
        attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
        position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids"
        if (
            attention_mask is not None
            and kwargs.get(position_ids_key) is None
            and position_ids_key in set(inspect.signature(self.forward).parameters.keys())
        ):
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            kwargs[position_ids_key] = position_ids  # placed in kwargs for further processing (see below)

        # 5. Slice model inputs if it's an input that should have the same length as `input_ids`
        for model_input_name in ["position_ids", "token_type_ids", "decoder_position_ids"]:
            model_input = kwargs.get(model_input_name)
            if model_input is not None:
                if past_key_values is not None:
                    current_input_length = (
                        model_inputs["inputs_embeds"].shape[1]
                        if model_inputs.get("inputs_embeds") is not None
                        else model_inputs[input_ids_key].shape[1]
                    )
                    model_input = model_input[:, -current_input_length:]
                    model_input = model_input.clone(memory_format=torch.contiguous_format)
                model_inputs[model_input_name] = model_input

        # 6. Create 4D attention mask is we are using a `StaticCache` (important for performant compiled forward pass)
        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
            if model_inputs["inputs_embeds"] is not None:
                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
                device = model_inputs["inputs_embeds"].device
            else:
                batch_size, sequence_length = model_inputs[input_ids_key].shape
                device = model_inputs[input_ids_key].device

            # Create the causal mask with fixed shape in advance, to reduce recompilations. If the function to create
            # the 4D causal mask exists, it should be present in the base model (XXXModel class).
            base_model = getattr(self, self.base_model_prefix, None)
            if base_model is None:
                causal_mask_creation_function = getattr(
                    self, "_prepare_4d_causal_attention_mask_with_cache_position", None
                )
            else:
                causal_mask_creation_function = getattr(
                    base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None
                )
            if causal_mask_creation_function is None:
                logger.warning_once(
                    f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method "
                    "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're "
                    "writing code, see Llama for an example implementation. If you're a user, please report this "
                    "issue on GitHub."
                )
            else:
                attention_mask = causal_mask_creation_function(
                    attention_mask,
                    sequence_length=sequence_length,
                    target_length=past_key_values.get_max_cache_shape(),
                    dtype=self.dtype,
                    device=device,
                    cache_position=cache_position,
                    batch_size=batch_size,
                    config=self.config,
                    past_key_values=past_key_values,
                )
        if attention_mask is not None:
            model_inputs[attention_mask_key] = attention_mask

        # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
        for key, value in kwargs.items():
            if key not in model_inputs:
                model_inputs[key] = value

        # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
        model_inputs.pop("labels", None)
        return model_inputs

    def _prepare_model_inputs(
        self,
        inputs: Optional[torch.Tensor] = None,
        bos_token_id: Optional[torch.Tensor] = None,
        model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
    ) -> Tuple[torch.Tensor, Optional[str], Dict[str, torch.Tensor]]:
        """
        This function extracts the model-specific `inputs` for generation.
        """
        # 1. retrieve all kwargs that are non-None or non-model input related.
        # some encoder-decoder models have different names for model and encoder
        if (
            self.config.is_encoder_decoder
            and hasattr(self, "encoder")
<|next_version|>
                inputs_embeds is not None  # Exception 1
                or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1])  # Exception 3
            ):
                input_ids = input_ids[:, -cache_position.shape[0] :]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        # 3. Prepare base model inputs
        input_ids_key = "decoder_input_ids" if self.config.is_encoder_decoder else "input_ids"
        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step for every prompt.
        if not self.config.is_encoder_decoder:
            if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:
                model_inputs[input_ids_key] = None
                model_inputs["inputs_embeds"] = inputs_embeds
            else:
                # `clone` calls in this function ensure a consistent stride. See #32227
                model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)
                model_inputs["inputs_embeds"] = None
        else:
            model_inputs[input_ids_key] = input_ids.clone(memory_format=torch.contiguous_format)

        # 4. Create missing `position_ids` on the fly
        encoder_attention_mask = attention_mask if self.config.is_encoder_decoder else None
        attention_mask = (
            kwargs.pop("decoder_attention_mask", None) if self.config.is_encoder_decoder else attention_mask
        )
        attention_mask_key = "decoder_attention_mask" if self.config.is_encoder_decoder else "attention_mask"
        position_ids_key = "decoder_position_ids" if self.config.is_encoder_decoder else "position_ids"
        if (
            attention_mask is not None
            and kwargs.get(position_ids_key) is None
            and position_ids_key in set(inspect.signature(self.forward).parameters.keys())
        ):
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            kwargs[position_ids_key] = position_ids  # placed in kwargs for further processing (see below)

        # 5. Slice model inputs if it's an input that should have the same length as `input_ids`
        for model_input_name in ["position_ids", "token_type_ids", "decoder_position_ids"]:
            model_input = kwargs.get(model_input_name)
            if model_input is not None:
                if past_key_values is not None:
                    current_input_length = (
                        model_inputs["inputs_embeds"].shape[1]
                        if model_inputs.get("inputs_embeds") is not None
                        else model_inputs[input_ids_key].shape[1]
                    )
                    model_input = model_input[:, -current_input_length:]
                    model_input = model_input.clone(memory_format=torch.contiguous_format)
                model_inputs[model_input_name] = model_input

        # 6. Create 4D attention mask is we are using a `StaticCache` (important for performant compiled forward pass)
        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
            if model_inputs["inputs_embeds"] is not None:
                batch_size, sequence_length, _ = model_inputs["inputs_embeds"].shape
                device = model_inputs["inputs_embeds"].device
            else:
                batch_size, sequence_length = model_inputs[input_ids_key].shape
                device = model_inputs[input_ids_key].device

            # Create the causal mask with fixed shape in advance, to reduce recompilations. If the function to create
            # the 4D causal mask exists, it should be present in the base model (XXXModel class).
            base_model = getattr(self, self.base_model_prefix, None)
            if base_model is None:
                causal_mask_creation_function = getattr(
                    self, "_prepare_4d_causal_attention_mask_with_cache_position", None
                )
            else:
                causal_mask_creation_function = getattr(
                    base_model, "_prepare_4d_causal_attention_mask_with_cache_position", None
                )
            if causal_mask_creation_function is None:
                logger.warning_once(
                    f"{self.__class__.__name__} has no `_prepare_4d_causal_attention_mask_with_cache_position` method "
                    "defined in its base modeling class. Compiled forward passes will be sub-optimal. If you're "
                    "writing code, see Llama for an example implementation. If you're a user, please report this "
                    "issue on GitHub."
                )
            else:
                attention_mask = causal_mask_creation_function(
                    attention_mask,
                    sequence_length=sequence_length,
                    target_length=past_key_values.get_max_cache_shape(),
                    dtype=self.dtype,
                    device=device,
                    cache_position=cache_position,
                    batch_size=batch_size,
                    config=self.config,
                    past_key_values=past_key_values,
                )
        if attention_mask is not None:
            model_inputs[attention_mask_key] = attention_mask

        if encoder_attention_mask is not None:
            model_inputs["attention_mask"] = encoder_attention_mask

        # 7. Forward ALL kwargs that are uninitialized (e.g. `use_cache`).
        for key, value in kwargs.items():
            if key not in model_inputs:
                model_inputs[key] = value

        # 8. Remove unexpected `generate` inputs (TODO @joao: fix trainer and examples)
        model_inputs.pop("labels", None)
        return model_inputs

    def _prepare_model_inputs(
        self,
        inputs: Optional[torch.Tensor] = None,
        bos_token_id: Optional[torch.Tensor] = None,
        model_kwargs: Optional[Dict[str, torch.Tensor]] = None,
    ) -> Tuple[torch.Tensor, Optional[str], Dict[str, torch.Tensor]]:
        """
        This function extracts the model-specific `inputs` for generation.
        """
        # 1. retrieve all kwargs that are non-None or non-model input related.
        # some encoder-decoder models have different names for model and encoder
        if (
            self.config.is_encoder_decoder
            and hasattr(self, "encoder")


Add require_read_token to fp8 tests (#36189)

<|original_code|>
# coding=utf-8
# Copyright 2025 The HuggingFace Team. 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 gc
import tempfile
import unittest

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, FineGrainedFP8Config, OPTForCausalLM
from transformers.testing_utils import (
    require_accelerate,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.utils import is_accelerate_available, is_torch_available

if is_torch_available():
    import torch

if is_accelerate_available():
    from accelerate import init_empty_weights

@require_torch_gpu
class FineGrainedFP8ConfigTest(unittest.TestCase):
    def test_to_dict(self):
        """
        Simple test that checks if one uses a config and converts it to a dict, the dict is the same as the config object
        """
        quantization_config = FineGrainedFP8Config()
        config_to_dict = quantization_config.to_dict()

        for key in config_to_dict:
            self.assertEqual(getattr(quantization_config, key), config_to_dict[key])

    def test_from_dict(self):
        """
        Simple test that checks if one uses a dict and converts it to a config object, the config object is the same as the dict
        """
        dict = {"modules_to_not_convert": ["lm_head.weight"], "quant_method": "fp8"}
        quantization_config = FineGrainedFP8Config.from_dict(dict)

        self.assertEqual(dict["modules_to_not_convert"], quantization_config.modules_to_not_convert)
        self.assertEqual(dict["quant_method"], quantization_config.quant_method)

@slow
@require_accelerate
@require_torch_gpu
class FP8QuantizerTest(unittest.TestCase):
    model_name = "meta-llama/Llama-3.2-1B"
    input_text = "Once upon a time"
    max_new_tokens = 10
    EXPECTED_OUTPUT = "Once upon a time, there was a man who was very rich."
    device_map = "cuda"
    offload_device_map = {
        "model.embed_tokens": 0,
        "model.layers.0": 0,
        "model.layers.1": 0,
        "model.layers.2": 0,
        "model.layers.3": 0,
        "model.layers.4": 0,
        "model.layers.5": 0,
        "model.layers.6": 0,
        "model.layers.7": "cpu",
        "model.layers.8": "cpu",
        "model.layers.9": "cpu",
        "model.layers.10": "cpu",
        "model.layers.11": "cpu",
        "model.layers.12": "cpu",
        "model.layers.13": "cpu",
        "model.layers.14": "cpu",
<|edits_diff|>
--- tests/quantization/finegrained_fp8/test_fp8.py
+++ tests/quantization/finegrained_fp8/test_fp8.py
@@ -20,6 +20,7 @@
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, FineGrainedFP8Config, OPTForCausalLM
 from transformers.testing_utils import (
     require_accelerate,
+    require_read_token,
     require_torch_gpu,
     require_torch_multi_gpu,
     slow,
<|current_version|>
# coding=utf-8
# Copyright 2025 The HuggingFace Team. 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 gc
import tempfile
import unittest

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, FineGrainedFP8Config, OPTForCausalLM
from transformers.testing_utils import (
    require_accelerate,
    require_read_token,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.utils import is_accelerate_available, is_torch_available

if is_torch_available():
    import torch

if is_accelerate_available():
    from accelerate import init_empty_weights

@require_torch_gpu
class FineGrainedFP8ConfigTest(unittest.TestCase):
    def test_to_dict(self):
        """
        Simple test that checks if one uses a config and converts it to a dict, the dict is the same as the config object
        """
        quantization_config = FineGrainedFP8Config()
        config_to_dict = quantization_config.to_dict()

        for key in config_to_dict:
            self.assertEqual(getattr(quantization_config, key), config_to_dict[key])

    def test_from_dict(self):
        """
        Simple test that checks if one uses a dict and converts it to a config object, the config object is the same as the dict
        """
        dict = {"modules_to_not_convert": ["lm_head.weight"], "quant_method": "fp8"}
        quantization_config = FineGrainedFP8Config.from_dict(dict)

        self.assertEqual(dict["modules_to_not_convert"], quantization_config.modules_to_not_convert)
        self.assertEqual(dict["quant_method"], quantization_config.quant_method)

@slow
@require_accelerate
@require_torch_gpu
class FP8QuantizerTest(unittest.TestCase):
    model_name = "meta-llama/Llama-3.2-1B"
    input_text = "Once upon a time"
    max_new_tokens = 10
    EXPECTED_OUTPUT = "Once upon a time, there was a man who was very rich."
    device_map = "cuda"
    offload_device_map = {
        "model.embed_tokens": 0,
        "model.layers.0": 0,
        "model.layers.1": 0,
        "model.layers.2": 0,
        "model.layers.3": 0,
        "model.layers.4": 0,
        "model.layers.5": 0,
        "model.layers.6": 0,
        "model.layers.7": "cpu",
        "model.layers.8": "cpu",
        "model.layers.9": "cpu",
        "model.layers.10": "cpu",
        "model.layers.11": "cpu",
        "model.layers.12": "cpu",
        "model.layers.13": "cpu",
        "model.layers.14": "cpu",
<|next_version|>
# coding=utf-8
# Copyright 2025 The HuggingFace Team. 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 gc
import tempfile
import unittest

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, FineGrainedFP8Config, OPTForCausalLM
from transformers.testing_utils import (
    require_accelerate,
    require_read_token,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.utils import is_accelerate_available, is_torch_available

if is_torch_available():
    import torch

if is_accelerate_available():
    from accelerate import init_empty_weights

@require_torch_gpu
class FineGrainedFP8ConfigTest(unittest.TestCase):
    def test_to_dict(self):
        """
        Simple test that checks if one uses a config and converts it to a dict, the dict is the same as the config object
        """
        quantization_config = FineGrainedFP8Config()
        config_to_dict = quantization_config.to_dict()

        for key in config_to_dict:
            self.assertEqual(getattr(quantization_config, key), config_to_dict[key])

    def test_from_dict(self):
        """
        Simple test that checks if one uses a dict and converts it to a config object, the config object is the same as the dict
        """
        dict = {"modules_to_not_convert": ["lm_head.weight"], "quant_method": "fp8"}
        quantization_config = FineGrainedFP8Config.from_dict(dict)

        self.assertEqual(dict["modules_to_not_convert"], quantization_config.modules_to_not_convert)
        self.assertEqual(dict["quant_method"], quantization_config.quant_method)

@slow
@require_accelerate
@require_read_token
@require_torch_gpu
class FP8QuantizerTest(unittest.TestCase):
    model_name = "meta-llama/Llama-3.2-1B"
    input_text = "Once upon a time"
    max_new_tokens = 10
    EXPECTED_OUTPUT = "Once upon a time, there was a man who was very rich."
    device_map = "cuda"
    offload_device_map = {
        "model.embed_tokens": 0,
        "model.layers.0": 0,
        "model.layers.1": 0,
        "model.layers.2": 0,
        "model.layers.3": 0,
        "model.layers.4": 0,
        "model.layers.5": 0,
        "model.layers.6": 0,
        "model.layers.7": "cpu",
        "model.layers.8": "cpu",
        "model.layers.9": "cpu",
        "model.layers.10": "cpu",
        "model.layers.11": "cpu",
        "model.layers.12": "cpu",
        "model.layers.13": "cpu",
        "model.layers.14": "cpu",


Efficient Inference Kernel for SpQR  (#34976)

<|original_code|>
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from typing import Dict, Optional, Union

from ..models.auto.configuration_auto import AutoConfig
from ..utils import logging
from ..utils.quantization_config import (
    AqlmConfig,
    AwqConfig,
    BitNetConfig,
    BitsAndBytesConfig,
    CompressedTensorsConfig,
    EetqConfig,
    FbgemmFp8Config,
    FineGrainedFP8Config,
    GPTQConfig,
    HiggsConfig,
    HqqConfig,
    QuantizationConfigMixin,
    QuantizationMethod,
    QuantoConfig,
    TorchAoConfig,
    VptqConfig,
)
from .quantizer_aqlm import AqlmHfQuantizer
from .quantizer_awq import AwqQuantizer
from .quantizer_bitnet import BitNetHfQuantizer
from .quantizer_bnb_4bit import Bnb4BitHfQuantizer
from .quantizer_bnb_8bit import Bnb8BitHfQuantizer
from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer
from .quantizer_eetq import EetqHfQuantizer
from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer
from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer
from .quantizer_gptq import GptqHfQuantizer
from .quantizer_higgs import HiggsHfQuantizer
from .quantizer_hqq import HqqHfQuantizer
from .quantizer_quanto import QuantoHfQuantizer
from .quantizer_torchao import TorchAoHfQuantizer
from .quantizer_vptq import VptqHfQuantizer

AUTO_QUANTIZER_MAPPING = {
    "awq": AwqQuantizer,
    "bitsandbytes_4bit": Bnb4BitHfQuantizer,
    "bitsandbytes_8bit": Bnb8BitHfQuantizer,
    "gptq": GptqHfQuantizer,
    "aqlm": AqlmHfQuantizer,
    "quanto": QuantoHfQuantizer,
    "eetq": EetqHfQuantizer,
    "higgs": HiggsHfQuantizer,
    "hqq": HqqHfQuantizer,
    "compressed-tensors": CompressedTensorsHfQuantizer,
    "fbgemm_fp8": FbgemmFp8HfQuantizer,
    "torchao": TorchAoHfQuantizer,
    "bitnet": BitNetHfQuantizer,
    "vptq": VptqHfQuantizer,
    "fp8": FineGrainedFP8HfQuantizer,
}

AUTO_QUANTIZATION_CONFIG_MAPPING = {
    "awq": AwqConfig,
    "bitsandbytes_4bit": BitsAndBytesConfig,
    "bitsandbytes_8bit": BitsAndBytesConfig,
    "eetq": EetqConfig,
    "gptq": GPTQConfig,
    "aqlm": AqlmConfig,
    "quanto": QuantoConfig,
    "hqq": HqqConfig,
    "compressed-tensors": CompressedTensorsConfig,
    "fbgemm_fp8": FbgemmFp8Config,
    "higgs": HiggsConfig,
    "torchao": TorchAoConfig,
    "bitnet": BitNetConfig,
    "vptq": VptqConfig,
    "fp8": FineGrainedFP8Config,
}

logger = logging.get_logger(__name__)

class AutoQuantizationConfig:
    """
    The Auto-HF quantization config class that takes care of automatically dispatching to the correct
    quantization config given a quantization config stored in a dictionary.
    """

    @classmethod
    def from_dict(cls, quantization_config_dict: Dict):
        quant_method = quantization_config_dict.get("quant_method", None)
        # We need a special care for bnb models to make sure everything is BC ..
        if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False):
            suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit"
            quant_method = QuantizationMethod.BITS_AND_BYTES + suffix
        elif quant_method is None:
            raise ValueError(
                "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"
            )

<|edits_diff|>
--- src/transformers/quantizers/auto.py
+++ src/transformers/quantizers/auto.py
@@ -20,6 +20,7 @@
     QuantizationConfigMixin,
     QuantizationMethod,
     QuantoConfig,
+    SpQRConfig,
     TorchAoConfig,
     VptqConfig,
 )
@@ -36,6 +37,7 @@
 from .quantizer_higgs import HiggsHfQuantizer
 from .quantizer_hqq import HqqHfQuantizer
 from .quantizer_quanto import QuantoHfQuantizer
+from .quantizer_spqr import SpQRHfQuantizer
 from .quantizer_torchao import TorchAoHfQuantizer
 from .quantizer_vptq import VptqHfQuantizer

@@ -55,6 +57,7 @@
     "torchao": TorchAoHfQuantizer,
     "bitnet": BitNetHfQuantizer,
     "vptq": VptqHfQuantizer,
+    "spqr": SpQRHfQuantizer,
     "fp8": FineGrainedFP8HfQuantizer,
 }

<|current_version|>
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from typing import Dict, Optional, Union

from ..models.auto.configuration_auto import AutoConfig
from ..utils import logging
from ..utils.quantization_config import (
    AqlmConfig,
    AwqConfig,
    BitNetConfig,
    BitsAndBytesConfig,
    CompressedTensorsConfig,
    EetqConfig,
    FbgemmFp8Config,
    FineGrainedFP8Config,
    GPTQConfig,
    HiggsConfig,
    HqqConfig,
    QuantizationConfigMixin,
    QuantizationMethod,
    QuantoConfig,
    SpQRConfig,
    TorchAoConfig,
    VptqConfig,
)
from .quantizer_aqlm import AqlmHfQuantizer
from .quantizer_awq import AwqQuantizer
from .quantizer_bitnet import BitNetHfQuantizer
from .quantizer_bnb_4bit import Bnb4BitHfQuantizer
from .quantizer_bnb_8bit import Bnb8BitHfQuantizer
from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer
from .quantizer_eetq import EetqHfQuantizer
from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer
from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer
from .quantizer_gptq import GptqHfQuantizer
from .quantizer_higgs import HiggsHfQuantizer
from .quantizer_hqq import HqqHfQuantizer
from .quantizer_quanto import QuantoHfQuantizer
from .quantizer_spqr import SpQRHfQuantizer
from .quantizer_torchao import TorchAoHfQuantizer
from .quantizer_vptq import VptqHfQuantizer

AUTO_QUANTIZER_MAPPING = {
    "awq": AwqQuantizer,
    "bitsandbytes_4bit": Bnb4BitHfQuantizer,
    "bitsandbytes_8bit": Bnb8BitHfQuantizer,
    "gptq": GptqHfQuantizer,
    "aqlm": AqlmHfQuantizer,
    "quanto": QuantoHfQuantizer,
    "eetq": EetqHfQuantizer,
    "higgs": HiggsHfQuantizer,
    "hqq": HqqHfQuantizer,
    "compressed-tensors": CompressedTensorsHfQuantizer,
    "fbgemm_fp8": FbgemmFp8HfQuantizer,
    "torchao": TorchAoHfQuantizer,
    "bitnet": BitNetHfQuantizer,
    "vptq": VptqHfQuantizer,
    "spqr": SpQRHfQuantizer,
    "fp8": FineGrainedFP8HfQuantizer,
}

AUTO_QUANTIZATION_CONFIG_MAPPING = {
    "awq": AwqConfig,
    "bitsandbytes_4bit": BitsAndBytesConfig,
    "bitsandbytes_8bit": BitsAndBytesConfig,
    "eetq": EetqConfig,
    "gptq": GPTQConfig,
    "aqlm": AqlmConfig,
    "quanto": QuantoConfig,
    "hqq": HqqConfig,
    "compressed-tensors": CompressedTensorsConfig,
    "fbgemm_fp8": FbgemmFp8Config,
    "higgs": HiggsConfig,
    "torchao": TorchAoConfig,
    "bitnet": BitNetConfig,
    "vptq": VptqConfig,
    "fp8": FineGrainedFP8Config,
}

logger = logging.get_logger(__name__)

class AutoQuantizationConfig:
    """
    The Auto-HF quantization config class that takes care of automatically dispatching to the correct
    quantization config given a quantization config stored in a dictionary.
    """

    @classmethod
    def from_dict(cls, quantization_config_dict: Dict):
        quant_method = quantization_config_dict.get("quant_method", None)
        # We need a special care for bnb models to make sure everything is BC ..
        if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False):
            suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit"
            quant_method = QuantizationMethod.BITS_AND_BYTES + suffix
        elif quant_method is None:
            raise ValueError(
                "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"
            )

<|next_version|>
# See the License for the specific language governing permissions and
# limitations under the License.
import warnings
from typing import Dict, Optional, Union

from ..models.auto.configuration_auto import AutoConfig
from ..utils import logging
from ..utils.quantization_config import (
    AqlmConfig,
    AwqConfig,
    BitNetConfig,
    BitsAndBytesConfig,
    CompressedTensorsConfig,
    EetqConfig,
    FbgemmFp8Config,
    FineGrainedFP8Config,
    GPTQConfig,
    HiggsConfig,
    HqqConfig,
    QuantizationConfigMixin,
    QuantizationMethod,
    QuantoConfig,
    SpQRConfig,
    TorchAoConfig,
    VptqConfig,
)
from .quantizer_aqlm import AqlmHfQuantizer
from .quantizer_awq import AwqQuantizer
from .quantizer_bitnet import BitNetHfQuantizer
from .quantizer_bnb_4bit import Bnb4BitHfQuantizer
from .quantizer_bnb_8bit import Bnb8BitHfQuantizer
from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer
from .quantizer_eetq import EetqHfQuantizer
from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer
from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer
from .quantizer_gptq import GptqHfQuantizer
from .quantizer_higgs import HiggsHfQuantizer
from .quantizer_hqq import HqqHfQuantizer
from .quantizer_quanto import QuantoHfQuantizer
from .quantizer_spqr import SpQRHfQuantizer
from .quantizer_torchao import TorchAoHfQuantizer
from .quantizer_vptq import VptqHfQuantizer

AUTO_QUANTIZER_MAPPING = {
    "awq": AwqQuantizer,
    "bitsandbytes_4bit": Bnb4BitHfQuantizer,
    "bitsandbytes_8bit": Bnb8BitHfQuantizer,
    "gptq": GptqHfQuantizer,
    "aqlm": AqlmHfQuantizer,
    "quanto": QuantoHfQuantizer,
    "eetq": EetqHfQuantizer,
    "higgs": HiggsHfQuantizer,
    "hqq": HqqHfQuantizer,
    "compressed-tensors": CompressedTensorsHfQuantizer,
    "fbgemm_fp8": FbgemmFp8HfQuantizer,
    "torchao": TorchAoHfQuantizer,
    "bitnet": BitNetHfQuantizer,
    "vptq": VptqHfQuantizer,
    "spqr": SpQRHfQuantizer,
    "fp8": FineGrainedFP8HfQuantizer,
}

AUTO_QUANTIZATION_CONFIG_MAPPING = {
    "awq": AwqConfig,
    "bitsandbytes_4bit": BitsAndBytesConfig,
    "bitsandbytes_8bit": BitsAndBytesConfig,
    "eetq": EetqConfig,
    "gptq": GPTQConfig,
    "aqlm": AqlmConfig,
    "quanto": QuantoConfig,
    "hqq": HqqConfig,
    "compressed-tensors": CompressedTensorsConfig,
    "fbgemm_fp8": FbgemmFp8Config,
    "higgs": HiggsConfig,
    "torchao": TorchAoConfig,
    "bitnet": BitNetConfig,
    "vptq": VptqConfig,
    "spqr": SpQRConfig,
    "fp8": FineGrainedFP8Config,
}

logger = logging.get_logger(__name__)

class AutoQuantizationConfig:
    """
    The Auto-HF quantization config class that takes care of automatically dispatching to the correct
    quantization config given a quantization config stored in a dictionary.
    """

    @classmethod
    def from_dict(cls, quantization_config_dict: Dict):
        quant_method = quantization_config_dict.get("quant_method", None)
        # We need a special care for bnb models to make sure everything is BC ..
        if quantization_config_dict.get("load_in_8bit", False) or quantization_config_dict.get("load_in_4bit", False):
            suffix = "_4bit" if quantization_config_dict.get("load_in_4bit", False) else "_8bit"
            quant_method = QuantizationMethod.BITS_AND_BYTES + suffix
        elif quant_method is None:
            raise ValueError(
                "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized"
            )


Refactoring of ImageProcessorFast (#35069)

<|original_code|>
#!/usr/bin/env python
# Copyright 2020 The HuggingFace Team. 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.

from transformers import HfArgumentParser

from .add_new_model_like import AddNewModelLikeCommand
from .chat import ChatCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands

def main():
    parser = HfArgumentParser(prog="Transformers CLI tool", usage="transformers-cli <command> [<args>]")
    commands_parser = parser.add_subparsers(help="transformers-cli command helpers")

    # Register commands
    ChatCommand.register_subcommand(commands_parser)
    ConvertCommand.register_subcommand(commands_parser)
    DownloadCommand.register_subcommand(commands_parser)
    EnvironmentCommand.register_subcommand(commands_parser)
    RunCommand.register_subcommand(commands_parser)
    ServeCommand.register_subcommand(commands_parser)
    UserCommands.register_subcommand(commands_parser)
    AddNewModelLikeCommand.register_subcommand(commands_parser)
    LfsCommands.register_subcommand(commands_parser)

    # Let's go
    args = parser.parse_args()

    if not hasattr(args, "func"):
        parser.print_help()
        exit(1)

    # Run
    service = args.func(args)
    service.run()

if __name__ == "__main__":
    main()

<|edits_diff|>
--- src/transformers/commands/transformers_cli.py
+++ src/transformers/commands/transformers_cli.py
@@ -15,6 +15,7 @@

 from transformers import HfArgumentParser

+from .add_fast_image_processor import AddFastImageProcessorCommand
 from .add_new_model_like import AddNewModelLikeCommand
 from .chat import ChatCommand
 from .convert import ConvertCommand
<|current_version|>
#!/usr/bin/env python
# Copyright 2020 The HuggingFace Team. 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.

from transformers import HfArgumentParser

from .add_fast_image_processor import AddFastImageProcessorCommand
from .add_new_model_like import AddNewModelLikeCommand
from .chat import ChatCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands

def main():
    parser = HfArgumentParser(prog="Transformers CLI tool", usage="transformers-cli <command> [<args>]")
    commands_parser = parser.add_subparsers(help="transformers-cli command helpers")

    # Register commands
    ChatCommand.register_subcommand(commands_parser)
    ConvertCommand.register_subcommand(commands_parser)
    DownloadCommand.register_subcommand(commands_parser)
    EnvironmentCommand.register_subcommand(commands_parser)
    RunCommand.register_subcommand(commands_parser)
    ServeCommand.register_subcommand(commands_parser)
    UserCommands.register_subcommand(commands_parser)
    AddNewModelLikeCommand.register_subcommand(commands_parser)
    LfsCommands.register_subcommand(commands_parser)

    # Let's go
    args = parser.parse_args()

    if not hasattr(args, "func"):
        parser.print_help()
        exit(1)

    # Run
    service = args.func(args)
    service.run()

if __name__ == "__main__":
    main()

<|next_version|>
#!/usr/bin/env python
# Copyright 2020 The HuggingFace Team. 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.

from transformers import HfArgumentParser

from .add_fast_image_processor import AddFastImageProcessorCommand
from .add_new_model_like import AddNewModelLikeCommand
from .chat import ChatCommand
from .convert import ConvertCommand
from .download import DownloadCommand
from .env import EnvironmentCommand
from .lfs import LfsCommands
from .run import RunCommand
from .serving import ServeCommand
from .user import UserCommands

def main():
    parser = HfArgumentParser(prog="Transformers CLI tool", usage="transformers-cli <command> [<args>]")
    commands_parser = parser.add_subparsers(help="transformers-cli command helpers")

    # Register commands
    ChatCommand.register_subcommand(commands_parser)
    ConvertCommand.register_subcommand(commands_parser)
    DownloadCommand.register_subcommand(commands_parser)
    EnvironmentCommand.register_subcommand(commands_parser)
    RunCommand.register_subcommand(commands_parser)
    ServeCommand.register_subcommand(commands_parser)
    UserCommands.register_subcommand(commands_parser)
    AddNewModelLikeCommand.register_subcommand(commands_parser)
    LfsCommands.register_subcommand(commands_parser)
    AddFastImageProcessorCommand.register_subcommand(commands_parser)

    # Let's go
    args = parser.parse_args()

    if not hasattr(args, "func"):
        parser.print_help()
        exit(1)

    # Run
    service = args.func(args)
    service.run()

if __name__ == "__main__":
    main()


Refactor (and fix) gpt_neox (#35610)

<|original_code|>
from typing import Optional, Tuple

import torch

from ..utils import is_torch_flex_attn_available

if is_torch_flex_attn_available():
    from torch.nn.attention.flex_attention import flex_attention

def flex_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: Optional[float] = None,
    softcap: Optional[float] = None,
    **kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
    causal_mask = attention_mask
    if causal_mask is not None:
        causal_mask = causal_mask[:, :, :, : key.shape[-2]]

    def causal_mod(score, b, h, q_idx, kv_idx):
        if softcap is not None:
            score = softcap * torch.tanh(score / softcap)
        if causal_mask is not None:
            score = score + causal_mask[b][0][q_idx][kv_idx]
        return score

    attn_output, attention_weights = flex_attention(
        query,
        key,
        value,
        score_mod=causal_mod,
        enable_gqa=True,
        scale=scaling,
        # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.
        # For simplification, we thus always return it as no additional computations are introduced.
        return_lse=True,
    )
    # lse is returned in float32
    attention_weights = attention_weights.to(value.dtype)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attention_weights

<|edits_diff|>
--- src/transformers/integrations/flex_attention.py
+++ src/transformers/integrations/flex_attention.py
@@ -17,6 +17,7 @@
     attention_mask: Optional[torch.Tensor],
     scaling: Optional[float] = None,
     softcap: Optional[float] = None,
+    head_mask: Optional[torch.Tensor] = None,
     **kwargs,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     causal_mask = attention_mask
<|current_version|>
from typing import Optional, Tuple

import torch

from ..utils import is_torch_flex_attn_available

if is_torch_flex_attn_available():
    from torch.nn.attention.flex_attention import flex_attention

def flex_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: Optional[float] = None,
    softcap: Optional[float] = None,
    head_mask: Optional[torch.Tensor] = None,
    **kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
    causal_mask = attention_mask
    if causal_mask is not None:
        causal_mask = causal_mask[:, :, :, : key.shape[-2]]

    def causal_mod(score, b, h, q_idx, kv_idx):
        if softcap is not None:
            score = softcap * torch.tanh(score / softcap)
        if causal_mask is not None:
            score = score + causal_mask[b][0][q_idx][kv_idx]
        return score

    attn_output, attention_weights = flex_attention(
        query,
        key,
        value,
        score_mod=causal_mod,
        enable_gqa=True,
        scale=scaling,
        # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.
        # For simplification, we thus always return it as no additional computations are introduced.
        return_lse=True,
    )
    # lse is returned in float32
    attention_weights = attention_weights.to(value.dtype)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attention_weights

<|next_version|>
from typing import Optional, Tuple

import torch

from ..utils import is_torch_flex_attn_available

if is_torch_flex_attn_available():
    from torch.nn.attention.flex_attention import flex_attention

def flex_attention_forward(
    module: torch.nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: Optional[float] = None,
    softcap: Optional[float] = None,
    head_mask: Optional[torch.Tensor] = None,
    **kwargs,
) -> Tuple[torch.Tensor, torch.Tensor]:
    causal_mask = attention_mask
    if causal_mask is not None:
        causal_mask = causal_mask[:, :, :, : key.shape[-2]]

    def causal_mod(score, b, h, q_idx, kv_idx):
        if softcap is not None:
            score = softcap * torch.tanh(score / softcap)
        if causal_mask is not None:
            score = score + causal_mask[b][0][q_idx][kv_idx]
        if head_mask is not None:
            score = score + head_mask[b][h][0][0]
        return score

    attn_output, attention_weights = flex_attention(
        query,
        key,
        value,
        score_mod=causal_mod,
        enable_gqa=True,
        scale=scaling,
        # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.
        # For simplification, we thus always return it as no additional computations are introduced.
        return_lse=True,
    )
    # lse is returned in float32
    attention_weights = attention_weights.to(value.dtype)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attention_weights