transformers 5.0.0rc2__py3-none-any.whl → 5.1.0__py3-none-any.whl

transformers/processing_utils.py

@@ -24,7 +24,7 @@ import sys
 import typing
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Any, Literal, Optional, TypedDict, TypeVar, Union
+from typing import Annotated, Any, Literal, TypedDict, TypeVar, Union
 
 import numpy as np
 import typing_extensions
@@ -199,26 +199,26 @@ class TextKwargs(TypedDict, total=False):
             - `'np'`: Return NumPy `np.ndarray` objects.
     """
 
-    text_pair: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]
-    text_target: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]
-    text_pair_target: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]]
-    add_special_tokens: Optional[bool]
-    padding: Annotated[Optional[Union[bool, str, PaddingStrategy]], padding_validator()]
-    truncation: Annotated[Optional[Union[bool, str, TruncationStrategy]], truncation_validator()]
-    max_length: Annotated[Optional[int], positive_int()]
-    stride: Annotated[Optional[int], positive_int()]
-    is_split_into_words: Optional[bool]
-    pad_to_multiple_of: Annotated[Optional[int], positive_int()]
-    return_token_type_ids: Optional[bool]
-    return_attention_mask: Optional[bool]
-    return_overflowing_tokens: Optional[bool]
-    return_special_tokens_mask: Optional[bool]
-    return_offsets_mapping: Optional[bool]
-    return_length: Optional[bool]
-    verbose: Optional[bool]
-    padding_side: Optional[Literal["left", "right"]]
-    return_mm_token_type_ids: Optional[bool]
-    return_tensors: Annotated[Optional[Union[str, TensorType]], tensor_type_validator()]
+    text_pair: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None
+    text_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None
+    text_pair_target: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None
+    add_special_tokens: bool | None
+    padding: Annotated[bool | str | PaddingStrategy | None, padding_validator()]
+    truncation: Annotated[bool | str | TruncationStrategy | None, truncation_validator()]
+    max_length: Annotated[int | None, positive_int()]
+    stride: Annotated[int | None, positive_int()]
+    is_split_into_words: bool | None
+    pad_to_multiple_of: Annotated[int | None, positive_int()]
+    return_token_type_ids: bool | None
+    return_attention_mask: bool | None
+    return_overflowing_tokens: bool | None
+    return_special_tokens_mask: bool | None
+    return_offsets_mapping: bool | None
+    return_length: bool | None
+    verbose: bool | None
+    padding_side: Literal["left", "right"] | None
+    return_mm_token_type_ids: bool | None
+    return_tensors: Annotated[str | TensorType | None, tensor_type_validator()]
 
 
 class ImagesKwargs(TypedDict, total=False):
@@ -272,25 +272,25 @@ class ImagesKwargs(TypedDict, total=False):
             Added for backward compatibility but this should be set as a processor attribute in future models.
     """
 
-    do_convert_rgb: Optional[bool]
-    do_resize: Optional[bool]
-    size: Annotated[Optional[Union[int, list[int], tuple[int, ...], dict[str, int]]], image_size_validator()]
-    crop_size: Annotated[Optional[Union[int, list[int], tuple[int, ...], dict[str, int]]], image_size_validator()]
-    resample: Annotated[Optional[Union["PILImageResampling", int]], resampling_validator()]
-    do_rescale: Optional[bool]
-    rescale_factor: Optional[float]
-    do_normalize: Optional[bool]
-    image_mean: Optional[Union[float, list[float], tuple[float, ...]]]
-    image_std: Optional[Union[float, list[float], tuple[float, ...]]]
-    do_pad: Optional[bool]
-    pad_size: Annotated[Optional[Union[int, list[int], tuple[int, ...], dict[str, int]]], image_size_validator()]
-    do_center_crop: Optional[bool]
-    data_format: Optional[Union[str, ChannelDimension]]
-    input_data_format: Optional[Union[str, ChannelDimension]]
-    device: Annotated[Optional[Union[str, "torch.device"]], device_validator()]
-    return_tensors: Annotated[Optional[Union[str, TensorType]], tensor_type_validator()]
-    disable_grouping: Optional[bool]
-    image_seq_length: Optional[int]
+    do_convert_rgb: bool | None
+    do_resize: bool | None
+    size: Annotated[int | list[int] | tuple[int, ...] | dict[str, int] | None, image_size_validator()]
+    crop_size: Annotated[int | list[int] | tuple[int, ...] | dict[str, int] | None, image_size_validator()]
+    resample: Annotated[Union["PILImageResampling", int] | None, resampling_validator()]
+    do_rescale: bool | None
+    rescale_factor: float | None
+    do_normalize: bool | None
+    image_mean: float | list[float] | tuple[float, ...] | None
+    image_std: float | list[float] | tuple[float, ...] | None
+    do_pad: bool | None
+    pad_size: Annotated[int | list[int] | tuple[int, ...] | dict[str, int] | None, image_size_validator()]
+    do_center_crop: bool | None
+    data_format: str | ChannelDimension | None
+    input_data_format: str | ChannelDimension | None
+    device: Annotated[Union[str, "torch.device"] | None, device_validator()]
+    return_tensors: Annotated[str | TensorType | None, tensor_type_validator()]
+    disable_grouping: bool | None
+    image_seq_length: int | None
 
 
 class VideosKwargs(TypedDict, total=False):
@@ -346,28 +346,28 @@ class VideosKwargs(TypedDict, total=False):
             - `'np'`: Return NumPy `np.ndarray` objects.
     """
 
-    do_convert_rgb: Optional[bool]
-    do_resize: Optional[bool]
-    size: Annotated[Optional[Union[int, list[int], tuple[int, ...], dict[str, int]]], image_size_validator()]
-    default_to_square: Optional[bool]
-    resample: Annotated[Optional[Union["PILImageResampling", int]], resampling_validator()]
-    do_rescale: Optional[bool]
-    rescale_factor: Optional[float]
-    do_normalize: Optional[bool]
-    image_mean: Optional[Union[float, list[float], tuple[float, ...]]]
-    image_std: Optional[Union[float, list[float], tuple[float, ...]]]
-    do_center_crop: Optional[bool]
-    do_pad: Optional[bool]
-    crop_size: Annotated[Optional[Union[int, list[int], tuple[int, ...], dict[str, int]]], image_size_validator()]
-    data_format: Optional[Union[str, ChannelDimension]]
-    input_data_format: Optional[Union[str, ChannelDimension]]
-    device: Annotated[Optional[Union[str, "torch.device"]], device_validator()]
-    do_sample_frames: Optional[bool]
-    video_metadata: Annotated[Optional[VideoMetadataType], video_metadata_validator()]
-    fps: Annotated[Optional[Union[int, float]], positive_any_number()]
-    num_frames: Annotated[Optional[int], positive_int()]
-    return_metadata: Optional[bool]
-    return_tensors: Annotated[Optional[Union[str, TensorType]], tensor_type_validator()]
+    do_convert_rgb: bool | None
+    do_resize: bool | None
+    size: Annotated[int | list[int] | tuple[int, ...] | dict[str, int] | None, image_size_validator()]
+    default_to_square: bool | None
+    resample: Annotated[Union["PILImageResampling", int] | None, resampling_validator()]
+    do_rescale: bool | None
+    rescale_factor: float | None
+    do_normalize: bool | None
+    image_mean: float | list[float] | tuple[float, ...] | None
+    image_std: float | list[float] | tuple[float, ...] | None
+    do_center_crop: bool | None
+    do_pad: bool | None
+    crop_size: Annotated[int | list[int] | tuple[int, ...] | dict[str, int] | None, image_size_validator()]
+    data_format: str | ChannelDimension | None
+    input_data_format: str | ChannelDimension | None
+    device: Annotated[Union[str, "torch.device"] | None, device_validator()]
+    do_sample_frames: bool | None
+    video_metadata: Annotated[VideoMetadataType | None, video_metadata_validator()]
+    fps: Annotated[int | float | None, positive_any_number()]
+    num_frames: Annotated[int | None, positive_int()]
+    return_metadata: bool | None
+    return_tensors: Annotated[str | TensorType | None, tensor_type_validator()]
 
 
 class AudioKwargs(TypedDict, total=False):
@@ -404,14 +404,14 @@ class AudioKwargs(TypedDict, total=False):
             - `'np'`: Return NumPy `np.ndarray` objects.
     """
 
-    sampling_rate: Annotated[Optional[int], positive_int()]
-    raw_speech: Optional[Union["np.ndarray", list[float], list["np.ndarray"], list[list[float]]]]
-    padding: Annotated[Optional[Union[bool, str, PaddingStrategy]], padding_validator()]
-    max_length: Annotated[Optional[int], positive_int()]
-    truncation: Annotated[Optional[Union[bool, str, TruncationStrategy]], truncation_validator()]
-    pad_to_multiple_of: Annotated[Optional[int], positive_int()]
-    return_attention_mask: Optional[bool]
-    return_tensors: Annotated[Optional[Union[str, TensorType]], tensor_type_validator()]
+    sampling_rate: Annotated[int | None, positive_int()]
+    raw_speech: Union["np.ndarray", list[float], list["np.ndarray"], list[list[float]]] | None
+    padding: Annotated[bool | str | PaddingStrategy | None, padding_validator()]
+    max_length: Annotated[int | None, positive_int()]
+    truncation: Annotated[bool | str | TruncationStrategy | None, truncation_validator()]
+    pad_to_multiple_of: Annotated[int | None, positive_int()]
+    return_attention_mask: bool | None
+    return_tensors: Annotated[str | TensorType | None, tensor_type_validator()]
 
 
 class ProcessingKwargs(TypedDict, total=False):
@@ -501,29 +501,14 @@ class TokenizerChatTemplateKwargs(TypedDict, total=False):
         This functionality is only available for chat templates that support it via the `{% generation %}` keyword.
     """
 
-    tools: Optional[list[dict]] = None
-    documents: Optional[list[dict[str, str]]] = None
-    add_generation_prompt: Optional[bool] = False
-    continue_final_message: Optional[bool] = False
-    return_assistant_tokens_mask: Optional[bool] = False
+    tools: list[dict] | None = None
+    documents: list[dict[str, str]] | None = None
+    add_generation_prompt: bool | None = False
+    continue_final_message: bool | None = False
+    return_assistant_tokens_mask: bool | None = False
 
 
-class ChatTemplateLoadKwargs(TypedDict, total=False):
-    """
-    Keyword arguments used to load multimodal data in processor chat templates.
-
-    num_frames (`int`, *optional*):
-        Number of frames to sample uniformly. If not passed, the whole video is loaded.
-    load_audio_from_video (`bool`, *optional*):
-            Whether to use the audio track of input video. If `True` the audio track will be loaded and passed to the
-            processor. This flag has no effect if the model doesn't support audio modality.
-    """
-
-    sampling_rate: Optional[int] = 16_000
-    load_audio_from_video: Optional[bool] = False
-
-
-class ProcessorChatTemplateKwargs(ChatTemplateLoadKwargs, TokenizerChatTemplateKwargs, total=False):
+class ProcessorChatTemplateKwargs(TokenizerChatTemplateKwargs, total=False):
     """
     Keyword arguments for processor's `apply_chat_template`.
 
@@ -531,15 +516,18 @@ class ProcessorChatTemplateKwargs(ChatTemplateLoadKwargs, TokenizerChatTemplateK
         Whether to tokenize the output or not.
     return_dict (`bool`, defaults to `False`):
         Whether to return a dictionary with named outputs. Has no effect if tokenize is `False`.
+    load_audio_from_video (`bool`, *optional*, defaults to `False`):
+        Whether to use the audio track of input video. If `True` the audio track will be loaded and passed to the
+        processor. This flag has no effect if the model doesn't support audio modality.
     """
 
-    tokenize: Optional[bool] = False
-    return_dict: Optional[bool] = False
+    tokenize: bool | None = False
+    return_dict: bool | None = False
+    load_audio_from_video: bool | None = False
 
 
 class AllKwargsForChatTemplate(TypedDict, total=False):
     processor_kwargs: ProcessingKwargs
-    mm_load_kwargs: ChatTemplateLoadKwargs
     template_kwargs: ProcessorChatTemplateKwargs
 
 
@@ -555,10 +543,10 @@ class MultiModalData:
     and we might change its API in the future.
     """
 
-    num_image_tokens: Optional[list[int]] = None
-    num_video_tokens: Optional[list[int]] = None
-    num_audio_tokens: Optional[list[int]] = None
-    num_image_patches: Optional[list[int]] = None
+    num_image_tokens: list[int] | None = None
+    num_video_tokens: list[int] | None = None
+    num_audio_tokens: list[int] | None = None
+    num_image_patches: list[int] | None = None
 
     def __contains__(self, key):
         return hasattr(self, key) and getattr(self, key) is not None
@@ -616,10 +604,10 @@ class ProcessorMixin(PushToHubMixin):
 
     def __call__(
         self,
-        images: Optional[ImageInput] = None,
-        text: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] = None,
-        videos: Optional[VideoInput] = None,
-        audio: Optional[AudioInput] = None,
+        images: ImageInput | None = None,
+        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None,
+        videos: VideoInput | None = None,
+        audio: AudioInput | None = None,
         **kwargs: Unpack[ProcessingKwargs],
     ):
         """
@@ -773,7 +761,7 @@ class ProcessorMixin(PushToHubMixin):
 
         return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"
 
-    def to_json_file(self, json_file_path: Union[str, os.PathLike]):
+    def to_json_file(self, json_file_path: str | os.PathLike):
         """
         Save this instance to a JSON file.
 
@@ -897,7 +885,7 @@ class ProcessorMixin(PushToHubMixin):
 
     @classmethod
     def get_processor_dict(
-        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+        cls, pretrained_model_name_or_path: str | os.PathLike, **kwargs
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         """
         From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used for instantiating a
@@ -1190,7 +1178,7 @@ class ProcessorMixin(PushToHubMixin):
     def _merge_kwargs(
         self,
         ModelProcessorKwargs: ProcessingKwargs,
-        tokenizer_init_kwargs: Optional[dict] = None,
+        tokenizer_init_kwargs: dict | None = None,
         **kwargs,
     ) -> dict[str, dict]:
         """
@@ -1233,7 +1221,8 @@ class ProcessorMixin(PushToHubMixin):
 
         """
         # holding a copy to avoid mutating user-provided arguments
-        kwargs = kwargs.copy()
+        # Use deepcopy to also copy nested dicts (like videos_kwargs) that will be modified via pop()
+        kwargs = copy.deepcopy(kwargs)
 
         # Initialize dictionaries
         output_kwargs = {
@@ -1364,11 +1353,11 @@ class ProcessorMixin(PushToHubMixin):
     @classmethod
     def from_pretrained(
         cls: type[SpecificProcessorType],
-        pretrained_model_name_or_path: Union[str, os.PathLike],
-        cache_dir: Optional[Union[str, os.PathLike]] = None,
+        pretrained_model_name_or_path: str | os.PathLike,
+        cache_dir: str | os.PathLike | None = None,
         force_download: bool = False,
         local_files_only: bool = False,
-        token: Optional[Union[str, bool]] = None,
+        token: str | bool | None = None,
         revision: str = "main",
         **kwargs,
     ) -> SpecificProcessorType:
@@ -1520,7 +1509,9 @@ class ProcessorMixin(PushToHubMixin):
                 if "PixtralProcessor" in cls.__name__:
                     from .tokenization_utils_tokenizers import TokenizersBackend
 
-                    tokenizer = TokenizersBackend.from_pretrained(pretrained_model_name_or_path, **kwargs)
+                    tokenizer = TokenizersBackend.from_pretrained(
+                        pretrained_model_name_or_path, subfolder=subfolder, **kwargs
+                    )
                 else:
                     tokenizer = cls._load_tokenizer_from_pretrained(
                         sub_processor_type, pretrained_model_name_or_path, subfolder=subfolder, **kwargs
@@ -1630,8 +1621,8 @@ class ProcessorMixin(PushToHubMixin):
 
     def apply_chat_template(
         self,
-        conversation: Union[list[dict[str, str]], list[list[dict[str, str]]]],
-        chat_template: Optional[str] = None,
+        conversation: list[dict[str, str]] | list[list[dict[str, str]]],
+        chat_template: str | None = None,
         **kwargs: Unpack[AllKwargsForChatTemplate],
     ) -> str:
         """
@@ -1708,22 +1699,25 @@ class ProcessorMixin(PushToHubMixin):
             else:
                 kwargs["return_offsets_mapping"] = True  # force offset mapping so we can infer token boundaries
 
-        # Fill sets of kwargs that should be used by different parts of template
-        processed_kwargs = {
-            "mm_load_kwargs": {},
-            "template_kwargs": {},
-        }
-
-        for kwarg_type in processed_kwargs:
-            for key in AllKwargsForChatTemplate.__annotations__[kwarg_type].__annotations__:
-                kwarg_type_defaults = AllKwargsForChatTemplate.__annotations__[kwarg_type]
-                default_value = getattr(kwarg_type_defaults, key, None)
-                value = kwargs.pop(key, default_value)
-                if value is not None and not isinstance(value, dict):
-                    processed_kwargs[kwarg_type][key] = value
+        # Fill sets of kwargs that should be used by jinja template, filtering out kwargs used in `processor.__call__`
+        # NOTE: we don't only filter but also set the default values here. Without default values, we can remove it
+        template_kwargs = {}
+        for key in AllKwargsForChatTemplate.__annotations__["template_kwargs"].__annotations__:
+            kwarg_type_defaults = AllKwargsForChatTemplate.__annotations__["template_kwargs"]
+            default_value = getattr(kwarg_type_defaults, key, None)
+            value = kwargs.pop(key, default_value)
+            if value is not None and not isinstance(value, dict):
+                template_kwargs[key] = value
 
         # Pass unprocessed custom kwargs
-        processed_kwargs["template_kwargs"].update(kwargs)
+        template_kwargs.update(kwargs)
+
+        # Set the sampling rate to load the audio files if user hasn't already passed with `kwargs`
+        if "sampling_rate" not in template_kwargs:
+            if hasattr(self, "feature_extractor") and hasattr(self.feature_extractor, "sampling_rate"):
+                template_kwargs["sampling_rate"] = self.feature_extractor.sampling_rate
+            else:
+                template_kwargs["sampling_rate"] = 16_000
 
         if isinstance(conversation, (list, tuple)) and (
             isinstance(conversation[0], (list, tuple)) or hasattr(conversation[0], "content")
@@ -1734,9 +1728,8 @@ class ProcessorMixin(PushToHubMixin):
             is_batched = False
             conversations = [conversation]
 
-        tokenize = processed_kwargs["template_kwargs"].pop("tokenize", False)
-        return_dict = processed_kwargs["template_kwargs"].pop("return_dict", True)
-        mm_load_kwargs = processed_kwargs["mm_load_kwargs"]
+        tokenize = template_kwargs.pop("tokenize", False)
+        return_dict = template_kwargs.pop("return_dict", True)
 
         if tokenize:
             batch_images, batch_videos = [], []
@@ -1767,12 +1760,12 @@ class ProcessorMixin(PushToHubMixin):
                     videos.extend(video_fnames)
 
                     # Audio models do not accept nested list of audios (yet!) so we construct a flat input audio list
-                    if not mm_load_kwargs["load_audio_from_video"]:
+                    if not template_kwargs["load_audio_from_video"]:
                         for fname in audio_fnames:
-                            batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"]))
+                            batch_audios.append(load_audio(fname, sampling_rate=template_kwargs["sampling_rate"]))
                     else:
                         for fname in video_fnames:
-                            batch_audios.append(load_audio(fname, sampling_rate=mm_load_kwargs["sampling_rate"]))
+                            batch_audios.append(load_audio(fname, sampling_rate=template_kwargs["sampling_rate"]))
 
                 # Currently all processors can accept nested list of batches, but not flat list of visuals
                 # So we'll make a batched list of images and let the processor handle it
@@ -1783,14 +1776,12 @@ class ProcessorMixin(PushToHubMixin):
         if hasattr(self, "tokenizer") and hasattr(self.tokenizer, "special_tokens_map"):
             special_tokens = self.tokenizer.special_tokens_map
             # Filter out tokens that conflict with template kwargs
-            special_tokens_map = {
-                k: v for k, v in special_tokens.items() if k not in processed_kwargs["template_kwargs"]
-            }
+            special_tokens_map = {k: v for k, v in special_tokens.items() if k not in template_kwargs}
 
         prompt, generation_indices = render_jinja_template(
             conversations=conversations,
             chat_template=chat_template,
-            **processed_kwargs["template_kwargs"],  # different flags such as `return_assistant_mask`
+            **template_kwargs,  # different flags such as `return_assistant_mask`
             **special_tokens_map,  # tokenizer special tokens are used by some templates
         )
 
@@ -1826,7 +1817,7 @@ class ProcessorMixin(PushToHubMixin):
             )
 
             if return_dict:
-                if processed_kwargs["template_kwargs"].get("return_assistant_tokens_mask", False):
+                if template_kwargs.get("return_assistant_tokens_mask", False):
                     assistant_masks = []
                     offset_mapping = out.pop("offset_mapping")
                     input_ids = out["input_ids"]

transformers/pytorch_utils.py

@@ -35,10 +35,10 @@ logger = logging.get_logger(__name__)
 
 is_torch_greater_or_equal_than_2_8 = is_torch_greater_or_equal("2.8", accept_dev=True)
 is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True)
-is_torch_greater_or_equal_than_2_4 = is_torch_greater_or_equal("2.4", accept_dev=True)
-is_torch_greater_or_equal_than_2_3 = is_torch_greater_or_equal("2.3", accept_dev=True)
 
 # For backwards compatibility (e.g. some remote codes on Hub using those variables).
+is_torch_greater_or_equal_than_2_4 = is_torch_greater_or_equal("2.4", accept_dev=True)
+is_torch_greater_or_equal_than_2_3 = is_torch_greater_or_equal("2.3", accept_dev=True)
 is_torch_greater_or_equal_than_2_2 = is_torch_greater_or_equal("2.2", accept_dev=True)
 is_torch_greater_or_equal_than_2_1 = is_torch_greater_or_equal("2.1", accept_dev=True)
 is_torch_greater_or_equal_than_2_0 = is_torch_greater_or_equal("2.0", accept_dev=True)
@@ -238,30 +238,6 @@ def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]:
     return tensor.device, unique_id, storage_size(tensor)
 
 
-def isin_mps_friendly(elements: torch.Tensor, test_elements: torch.Tensor | int) -> torch.Tensor:
-    """
-    Same as `torch.isin` without flags, but MPS-friendly. We can remove this function when we stop supporting
-    torch <= 2.3. See https://github.com/pytorch/pytorch/issues/77764#issuecomment-2067838075
-
-    Args:
-        elements (`torch.Tensor`): Input elements
-        test_elements (`torch.Tensor` or `int`): The elements to check against.
-
-    Returns:
-        `torch.Tensor`: A boolean tensor of the same shape as `elements` that is True for `elements` in `test_elements`
-        and False otherwise
-    """
-
-    if elements.device.type == "mps" and not is_torch_greater_or_equal_than_2_4:
-        test_elements = torch.tensor(test_elements)
-        if test_elements.ndim == 0:
-            test_elements = test_elements.unsqueeze(0)
-        return elements.tile(test_elements.shape[0], 1).eq(test_elements.unsqueeze(1)).sum(dim=0).bool().squeeze()
-    else:
-        # Note: don't use named arguments in `torch.isin`, see https://github.com/pytorch/pytorch/issues/126045
-        return torch.isin(elements, test_elements)
-
-
 @wraps(lru_cache)
 def compile_compatible_method_lru_cache(*lru_args, **lru_kwargs):
     """

transformers/quantizers/base.py

@@ -58,6 +58,14 @@ def get_keys_to_not_convert(model) -> list:
     return list(modules_to_not_convert)
 
 
+def _assign_is_quantized(model):
+    from ..modeling_utils import PreTrainedModel
+
+    for module in model.modules():
+        if isinstance(module, PreTrainedModel):
+            module.config._is_quantized = True
+
+
 class HfQuantizer(ABC):
     """
     Abstract class of the HuggingFace quantizer. Supports for now quantizing HF transformers models for inference and/or quantization.
@@ -176,6 +184,8 @@ class HfQuantizer(ABC):
 
         if self.pre_quantized and getattr(self.quantization_config, "dequantize", False):
             self.remove_quantization_config(model)
+        else:
+            _assign_is_quantized(model)
 
         return self._process_model_after_weight_loading(model, **kwargs)
 

transformers/quantizers/quantizer_compressed_tensors.py

@@ -84,13 +84,15 @@ class CompressedTensorsHfQuantizer(HfQuantizer):
         ) or self.quantization_config.is_sparsification_compressed:
             self.compressor.decompress_model(model=model)
 
+    # NOTE: TP plan override for compressed tensors removed - unsupported styles were used.
+    # TODO: Implement proper TP support for compressed tensors quantization
     def update_tp_plan(self, config):
         additional_plan = {
-            "layers.*.feed_forward.experts.*.gate_proj.weight": "local_colwise",
-            "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "local_colwise",
-            "layers.*.feed_forward.experts.*.up_proj.weight": "local_colwise",
-            "layers.*.feed_forward.experts.*.up_proj.weight_scale": "local_colwise",
-            "layers.*.feed_forward.experts.*.down_proj.weight": "local_rowwise",
+            "layers.*.feed_forward.experts.*.gate_proj.weight": "colwise",
+            "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "colwise",
+            "layers.*.feed_forward.experts.*.up_proj.weight": "colwise",
+            "layers.*.feed_forward.experts.*.up_proj.weight_scale": "colwise",
+            "layers.*.feed_forward.experts.*.down_proj.weight": "rowwise",
         }
         if config.get_text_config() is not None and config.get_text_config().base_model_tp_plan is not None:
             config.get_text_config().base_model_tp_plan.update(additional_plan)

transformers/quantizers/quantizer_fbgemm_fp8.py

@@ -133,14 +133,13 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
                 # We are using a different tp plan with local_colwise and local_rowwise for the attention because fbgemm operations cannot be parallelized
                 # With local_colwise and local_rowwise, all the operations are done locally, and we add a gather operation to gather the results instead of
                 # using dtensors
-                "layers.*.self_attn.q_proj.weight": "local_colwise",
-                "layers.*.self_attn.q_proj.weight_scale": "local_colwise",
-                "layers.*.self_attn.k_proj.weight": "local_colwise",
-                "layers.*.self_attn.k_proj.weight_scale": "local_colwise",
-                "layers.*.self_attn.v_proj.weight": "local_colwise",
-                "layers.*.self_attn.v_proj.weight_scale": "local_colwise",
-                "layers.*.self_attn.o_proj.weight": "local_rowwise",
-                "layers.*.self_attn": "gather",
+                "layers.*.self_attn.q_proj.weight": "colwise",
+                "layers.*.self_attn.q_proj.weight_scale": "colwise",
+                "layers.*.self_attn.k_proj.weight": "colwise",
+                "layers.*.self_attn.k_proj.weight_scale": "colwise",
+                "layers.*.self_attn.v_proj.weight": "colwise",
+                "layers.*.self_attn.v_proj.weight_scale": "colwise",
+                "layers.*.self_attn.o_proj.weight": "rowwise",
                 # We keep the same sequence_parallel plan for layernorms
                 "layers.*.input_layernorm.weight": "sequence_parallel",
                 "layers.*.post_attention_layernorm.weight": "sequence_parallel",
@@ -148,23 +147,21 @@ class FbgemmFp8HfQuantizer(HfQuantizer):
                 # We keep the same local_colwise and local_rowwise plan for the feed forward shared expert
                 # We also add scales for the shared expert, for local_colwise the scale is also local_colwise
                 # For local_rowwise the scale is replicated, so we don't need to add it
-                "layers.*.feed_forward.shared_expert.gate_proj.weight": "local_colwise",
-                "layers.*.feed_forward.shared_expert.gate_proj.weight_scale": "local_colwise",
-                "layers.*.feed_forward.shared_expert.up_proj.weight": "local_colwise",
-                "layers.*.feed_forward.shared_expert.up_proj.weight_scale": "local_colwise",
-                "layers.*.feed_forward.shared_expert.down_proj.weight": "local_rowwise",
-                "layers.*.feed_forward.experts": "local",
-                "layers.*.feed_forward": "gather",
-                "layers.*.feed_forward.experts.*.gate_proj.weight": "local_colwise",
-                "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "local_colwise",
-                "layers.*.feed_forward.experts.*.up_proj.weight": "local_colwise",
-                "layers.*.feed_forward.experts.*.up_proj.weight_scale": "local_colwise",
-                "layers.*.feed_forward.experts.*.down_proj.weight": "local_rowwise",
+                "layers.*.feed_forward.shared_expert.gate_proj.weight": "colwise",
+                "layers.*.feed_forward.shared_expert.gate_proj.weight_scale": "colwise",
+                "layers.*.feed_forward.shared_expert.up_proj.weight": "colwise",
+                "layers.*.feed_forward.shared_expert.up_proj.weight_scale": "colwise",
+                "layers.*.feed_forward.shared_expert.down_proj.weight": "rowwise",
+                "layers.*.feed_forward.experts.*.gate_proj.weight": "colwise",
+                "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "colwise",
+                "layers.*.feed_forward.experts.*.up_proj.weight": "colwise",
+                "layers.*.feed_forward.experts.*.up_proj.weight_scale": "colwise",
+                "layers.*.feed_forward.experts.*.down_proj.weight": "rowwise",
                 # For Fused implementation we use local_packed_rowwise for the gate_up_proj, and the same for the packed scales
                 # We use local_colwise for the down_proj, and the scales are replicated so we don't add them
-                "layers.*.feed_forward.experts.gate_up_proj": "local_packed_rowwise",
-                "layers.*.feed_forward.experts.gate_up_proj_scale": "local_packed_rowwise",
-                "layers.*.feed_forward.experts.down_proj": "local_colwise",
+                "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise",
+                "layers.*.feed_forward.experts.gate_up_proj_scale": "packed_rowwise",
+                "layers.*.feed_forward.experts.down_proj": "colwise",
             }
             if config.get_text_config() is not None:
                 config.get_text_config().base_model_tp_plan = text_plan

transformers/quantizers/quantizer_finegrained_fp8.py

@@ -110,29 +110,23 @@ class FineGrainedFP8HfQuantizer(HfQuantizer):
             pre_quantized=self.pre_quantized,
         )
 
-    # NOTE: TP is applied before quantization so this is only to add hooks.
-    # Quantization is incompatible with DTensors, so we have to anyway have
-    # gathers! But it should be model independent -> figure out where to put
-    # the gather and that's it.
     def update_tp_plan(self, config):
         if "Qwen3" in config.__class__.__name__:
             text_plan = {
-                "layers.*.self_attn.q_proj.weight": "local_colwise",
-                "layers.*.self_attn.q_proj.weight_scale_inv": "local_colwise",
-                "layers.*.self_attn.k_proj.weight": "local_colwise",
-                "layers.*.self_attn.k_proj.weight_scale_inv": "local_colwise",
-                "layers.*.self_attn.v_proj.weight": "local_colwise",
-                "layers.*.self_attn.v_proj.weight_scale_inv": "local_colwise",
-                "layers.*.self_attn.o_proj.weight": "local_rowwise",
-                "layers.*.self_attn.o_proj.weight_scale_inv": "local_rowwise",
-                "layers.*.self_attn": "gather",
-                "layers.*.mlp.gate_proj.weight": "local_colwise",
-                "layers.*.mlp.gate_proj.weight_scale_inv": "local_colwise",
-                "layers.*.mlp.up_proj.weight": "local_colwise",
-                "layers.*.mlp.up_proj.weight_scale_inv": "local_colwise",
-                "layers.*.mlp.down_proj.weight": "local_rowwise",
-                "layers.*.mlp.down_proj.weight_scale_inv": "local_rowwise",
-                "layers.*.mlp": "gather",
+                "layers.*.self_attn.q_proj.weight": "colwise",
+                "layers.*.self_attn.q_proj.weight_scale_inv": "colwise",
+                "layers.*.self_attn.k_proj.weight": "colwise",
+                "layers.*.self_attn.k_proj.weight_scale_inv": "colwise",
+                "layers.*.self_attn.v_proj.weight": "colwise",
+                "layers.*.self_attn.v_proj.weight_scale_inv": "colwise",
+                "layers.*.self_attn.o_proj.weight": "rowwise",
+                "layers.*.self_attn.o_proj.weight_scale_inv": "rowwise",
+                "layers.*.mlp.gate_proj.weight": "colwise",
+                "layers.*.mlp.gate_proj.weight_scale_inv": "colwise",
+                "layers.*.mlp.up_proj.weight": "colwise",
+                "layers.*.mlp.up_proj.weight_scale_inv": "colwise",
+                "layers.*.mlp.down_proj.weight": "rowwise",
+                "layers.*.mlp.down_proj.weight_scale_inv": "rowwise",
             }
 
             config.base_model_tp_plan = text_plan

transformers/quantizers/quantizer_mxfp4.py

@@ -55,7 +55,7 @@ class Mxfp4HfQuantizer(HfQuantizer):
             try:
                 from ..integrations.hub_kernels import get_kernel
 
-                self.triton_kernels_hub = get_kernel("kernels-community/triton_kernels")
+                self.triton_kernels_hub = get_kernel("kernels-community/gpt-oss-triton-kernels")
             except ImportError:
                 raise ImportError("kernels package is required for MXFP4 quantization")
         return self.triton_kernels_hub