transformers 5.0.0rc0__py3-none-any.whl → 5.0.0rc2__py3-none-any.whl

transformers/trainer_jit_checkpoint.py

@@ -0,0 +1,126 @@
+import os
+import signal
+import threading
+from typing import Optional
+
+from .trainer_callback import TrainerCallback
+from .trainer_utils import PREFIX_CHECKPOINT_DIR
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CheckpointManager:
+    def __init__(self, trainer, kill_wait: int = 3):
+        """
+        Initialize the CheckpointManager for Just-In-Time checkpoint handling.
+
+        Args:
+            trainer: The Trainer instance that will be used to save checkpoints when SIGTERM is received.
+            kill_wait (`int`, *optional*, defaults to 3): Grace period to distinguish between SIGTERM and SIGKILL.
+        """
+        self.trainer = trainer
+        self.is_checkpoint_requested = False
+        self._original_sigterm_handler = None
+        self.kill_wait = kill_wait
+
+    def setup_signal_handler(self):
+        self._original_sigterm_handler = signal.signal(signal.SIGTERM, self._sigterm_handler)
+        logger.info("JIT checkpoint signal handler registered for SIGTERM")
+
+    def _sigterm_handler(self, signum, frame):
+        if self.is_checkpoint_requested:
+            return
+
+        logger.info(f"SIGTERM received, will request JIT checkpoint after {self.kill_wait}s")
+        threading.Timer(self.kill_wait, self._enable_checkpoint).start()
+
+    def _enable_checkpoint(self):
+        logger.info("Kill wait period elapsed, requesting checkpoint")
+        self.is_checkpoint_requested = True
+
+    def execute_jit_checkpoint(self):
+        try:
+            # Set checkpoint flag to False to avoid multiple checkpoints getting triggered by other callbacks
+            self.is_checkpoint_requested = False
+
+            logger.info("Starting JIT checkpointing...")
+            current_step = self.trainer.state.global_step
+            logger.info(f"Saving JIT checkpoint at step {current_step}")
+
+            output_dir = self.trainer._get_output_dir(trial=None)
+            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{current_step}"
+            checkpoint_path = os.path.join(output_dir, checkpoint_folder)
+
+            # Create checkpoint directory
+            os.makedirs(checkpoint_path, exist_ok=True)
+
+            # Create a sentinel file to indicate checkpointing is in progress
+            sentinel_file = os.path.join(output_dir, checkpoint_folder, "checkpoint-is-incomplete.txt")
+            with open(sentinel_file, "w") as f:
+                f.write(f"Checkpoint started at step {current_step} and in progress...")
+            logger.info(f"Created checkpoint progress sentinel marker file: {sentinel_file}")
+
+            # Invoke the trainer's checkpoint method directly
+            self.trainer._save_checkpoint(self.trainer.model, trial=None)
+
+            # Remove sentinel file upon successful checkpointing
+            if os.path.exists(sentinel_file):
+                os.remove(sentinel_file)
+                logger.info("Sentinel marker file removed")
+
+            logger.info("Immediate JIT checkpoint completed successfully")
+
+        except Exception as e:
+            logger.error(f"Failed to save JIT checkpoint: {e}")
+            raise
+
+
+class JITCheckpointCallback(TrainerCallback):
+    """
+    Callback for Just-In-Time checkpointing on SIGTERM signals.
+
+    When SIGTERM is received, the checkpoint manager sets `is_checkpoint_requested=True`.
+    The callbacks detect this flag and set `control.should_training_stop=True`, which signals
+    the Trainer's training loop to exit gracefully after saving the checkpoint.
+    """
+
+    def __init__(self):
+        self.trainer = None
+        self.jit_manager: Optional[CheckpointManager] = None
+
+    def set_trainer(self, trainer):
+        self.trainer = trainer
+        if trainer.args.enable_jit_checkpoint:
+            self.jit_manager = CheckpointManager(trainer=trainer)
+            self.jit_manager.setup_signal_handler()
+            logger.info("JIT checkpointing enabled")
+
+    def on_pre_optimizer_step(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_save = False
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        if self.jit_manager and self.jit_manager.is_checkpoint_requested:
+            control.should_save = False
+            control.should_training_stop = True
+            self.jit_manager.execute_jit_checkpoint()
+
+    def on_train_end(self, args, state, control, **kwargs):
+        #  Restore original SIGTERM handler
+        if self.jit_manager and self.jit_manager._original_sigterm_handler is not None:
+            signal.signal(signal.SIGTERM, self.jit_manager._original_sigterm_handler)
+            logger.info("Restored original SIGTERM handler after training completion")

transformers/trainer_seq2seq.py

@@ -333,7 +333,11 @@ class Seq2SeqTrainer(Trainer):
             self.model.generation_config._from_model_config = False
 
         # Retrieves GenerationConfig from model.generation_config
+        # Update with defaults because earlier the generation config used ot be init
+        # with default values. Now we init it with `None` and keep defaults for BC
         gen_config = self.model.generation_config
+        default_gen_config = gen_config._get_default_generation_params()
+        gen_config.update(**default_gen_config, defaults_only=True)
         # in case the batch is shorter than max length, the output should be padded
         if generated_tokens.shape[-1] < gen_config.max_length:
             generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_config.max_length)

transformers/trainer_utils.py

@@ -924,7 +924,7 @@ def load_sharded_checkpoint(model, folder, strict=True, prefer_safe=True):
     shard_files = list(set(index["weight_map"].values()))
 
     # If strict=True, error before loading any of the state dicts.
-    # TODO: Here, update the weigth map with the config.dynamic_weight_conversion
+    # TODO: Here, update the weight map with the config.dynamic_weight_conversion
     loaded_keys = index["weight_map"].keys()
     model_keys = model.state_dict().keys()
     missing_keys = [key for key in model_keys if key not in loaded_keys]

transformers/training_args.py

@@ -340,9 +340,17 @@ class TrainingArguments:
             `save_total_limit=5` and `load_best_model_at_end`, the four last checkpoints will always be retained
             alongside the best model. When `save_total_limit=1` and `load_best_model_at_end`, it is possible that two
             checkpoints are saved: the last one and the best one (if they are different).
-        save_safetensors (`bool`, *optional*, defaults to `True`):
-            Use [safetensors](https://huggingface.co/docs/safetensors) saving and loading for state dicts instead of
-            default `torch.load` and `torch.save`.
+        enable_jit_checkpoint (`bool`, *optional*, defaults to `False`):
+            Whether to enable Just-In-Time (JIT) checkpointing on SIGTERM signal. When enabled, training will
+            checkpoint upon receiving SIGTERM, allowing for graceful termination without losing
+            progress. This is particularly useful for shared clusters with preemptible workloads (e.g., Kueue).
+            **Important**: You must configure your orchestrator's graceful shutdown period to allow sufficient time
+            for checkpoint completion. For Kubernetes, set `terminationGracePeriodSeconds` in your job definition
+            (method varies by cloud-native trainer: Kubeflow, Ray, etc.). Note: the default is only 30 seconds,
+            which is typically insufficient. For Slurm, use `--signal=USR1@<seconds>` in your sbatch script to send
+            SIGTERM with adequate time before the job time limit. Calculate the required grace period as: longest
+            possible iteration time + checkpoint saving time. For example, if an iteration takes 2 minutes and
+            checkpoint saving takes 2 minutes, set at least 4 minutes (240 seconds) of grace time.
         save_on_each_node (`bool`, *optional*, defaults to `False`):
             When doing multi-node distributed training, whether to save models and checkpoints on each node, or only on
             the main one.
@@ -585,9 +593,9 @@ class TrainingArguments:
             instance of `Dataset`.
         report_to (`str` or `list[str]`, *optional*, defaults to `"none"`):
             The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,
-            `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, `"neptune"`,
-            `"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations
-            installed, `"none"` for no integrations.
+            `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`, `"swanlab"`,
+            `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations installed, `"none"`
+            for no integrations.
         project (`str`, *optional*, defaults to `"huggingface"`):
             The name of the project to use for logging. Currently, only used by Trackio.
         trackio_space_id (`str` or `None`, *optional*, defaults to `"trackio"`):
@@ -852,7 +860,7 @@ class TrainingArguments:
     warmup_ratio: float | None = field(
         default=None,
         metadata={
-            "help": "This argument is deprecated and will be removed in v5. Use `warmup_steps` instead as it also works with float values."
+            "help": "This argument is deprecated and will be removed in v5.2. Use `warmup_steps` instead as it also works with float values."
         },
     )
 
@@ -929,14 +937,24 @@ class TrainingArguments:
                 " for `save_total_limit=5` and `load_best_model_at_end=True`, the four last checkpoints will always be"
                 " retained alongside the best model. When `save_total_limit=1` and `load_best_model_at_end=True`,"
                 " it is possible that two checkpoints are saved: the last one and the best one (if they are different)."
-                " Default is unlimited checkpoints"
+                " Default is unlimited checkpoints."
             )
         },
     )
-    save_safetensors: bool = field(
-        default=True,
+    enable_jit_checkpoint: bool = field(
+        default=False,
         metadata={
-            "help": "Use safetensors saving and loading for state dicts instead of default torch.load and torch.save."
+            "help": (
+                "Whether to enable Just-In-Time (JIT) checkpointing on SIGTERM signal. "
+                "When enabled, training will checkpoint upon receiving SIGTERM, "
+                "allowing for graceful termination without losing progress. "
+                "This is particularly useful for shared clusters with preemptible workloads (Kueue). "
+                "IMPORTANT: You must configure your orchestrator's graceful shutdown period. "
+                "Kubernetes: set terminationGracePeriodSeconds (default 30s is insufficient!) in your job definition. "
+                "Slurm: use --signal=USR1@<seconds> in sbatch to send SIGTERM before time limit. "
+                "Calculate required grace period as: iteration time + checkpoint saving time. "
+                "Example: 2min iteration + 2min checkpoint = 240 seconds minimum."
+            )
         },
     )
     save_on_each_node: bool = field(
@@ -1504,14 +1522,6 @@ class TrainingArguments:
                         f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}."
                     )
 
-        if not self.save_safetensors:
-            logger.info(
-                f"Found safetensors installation, but --save_safetensors={self.save_safetensors}. "
-                f"Safetensors should be a preferred weights saving format due to security and performance reasons. "
-                f"If your model cannot be saved by safetensors please feel free to open an issue at "
-                f"https://github.com/huggingface/safetensors!"
-            )
-
         if (
             self.load_best_model_at_end or self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU
         ) and self.metric_for_best_model is None:
@@ -1520,16 +1530,14 @@ class TrainingArguments:
             self.greater_is_better = not self.metric_for_best_model.endswith("loss")
         if is_torch_available():
             if self.bf16 or self.bf16_full_eval:
-                if self.use_cpu and not is_torch_xla_available():
-                    # cpu
-                    raise ValueError("Your setup doesn't support bf16/(cpu, tpu, neuroncore). You need torch>=1.10")
-                elif not self.use_cpu:
-                    if not is_torch_bf16_gpu_available() and not is_torch_xla_available():  # added for tpu support
-                        error_message = "Your setup doesn't support bf16/gpu."
-                        if is_torch_cuda_available():
-                            error_message += " You need Ampere+ GPU with cuda>=11.0"
-                        # gpu
-                        raise ValueError(error_message)
+                if (
+                    not self.use_cpu and not is_torch_bf16_gpu_available() and not is_torch_xla_available()
+                ):  # added for tpu support
+                    error_message = "Your setup doesn't support bf16/gpu. You need to assign use_cpu if you want to train the model on CPU"
+                    if is_torch_cuda_available():
+                        error_message += " You need Ampere+ GPU with cuda>=11.0"
+                    # gpu
+                    raise ValueError(error_message)
 
         if self.fp16 and self.bf16:
             raise ValueError("At most one of fp16 and bf16 can be True, but not both")
@@ -2359,8 +2367,8 @@ class TrainingArguments:
             report_to (`str` or `list[str]`, *optional*, defaults to `"none"`):
                 The list of integrations to report the results and logs to. Supported platforms are `"azure_ml"`,
                 `"clearml"`, `"codecarbon"`, `"comet_ml"`, `"dagshub"`, `"dvclive"`, `"flyte"`, `"mlflow"`,
-                `"neptune"`, `"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all
-                integrations installed, `"none"` for no integrations.
+                `"swanlab"`, `"tensorboard"`, `"trackio"` and `"wandb"`. Use `"all"` to report to all integrations
+                installed, `"none"` for no integrations.
             first_step (`bool`, *optional*, defaults to `False`):
                 Whether to log and evaluate the first `global_step` or not.
             nan_inf_filter (`bool`, *optional*, defaults to `True`):
@@ -2565,7 +2573,7 @@ class TrainingArguments:
         ```
         """
         if warmup_ratio is not None:
-            logger.warning("warmup_ratio is deprecated and will be removed in v5. Use `warmup_steps` instead.")
+            logger.warning("warmup_ratio is deprecated and will be removed in v5.2 . Use `warmup_steps` instead.")
             warmup_steps = warmup_ratio
 
         self.lr_scheduler_type = SchedulerType(name)
@@ -2742,10 +2750,24 @@ class TrainingArguments:
                         fsdp_plugin_args["transformer_cls_names_to_wrap"] = ",".join(
                             self.fsdp_config["transformer_layer_cls_to_wrap"]
                         )
-            fsdp_plugin_args["fsdp_version"] = self.fsdp_config.get("fsdp_version", 1)
+            fsdp_version = int(self.fsdp_config.get("version", 1))
+            fsdp_plugin_args["fsdp_version"] = fsdp_version
             prefetch_policy = self.fsdp_config.get("backward_prefetch", "NO_PREFETCH")
-            fsdp_plugin_args["backward_prefetch"] = prefetch_policy.upper()
-            fsdp_plugin_args["forward_prefetch"] = str(self.fsdp_config.get("forward_prefetch", "false")).lower()
+            if fsdp_version == 2:
+                fsdp_plugin_args["reshard_after_forward"] = str_to_bool(
+                    str(self.fsdp_config.get("reshard_after_forward", "false")).lower()
+                )
+            else:
+                fsdp_plugin_args["forward_prefetch"] = str_to_bool(
+                    str(self.fsdp_config.get("forward_prefetch", "false")).lower()
+                )
+                fsdp_plugin_args["backward_prefetch"] = prefetch_policy.upper()
+                fsdp_plugin_args["reshard_after_forward"] = str(
+                    self.fsdp_config.get("reshard_after_forward", "FULL_SHARD")
+                ).lower()
+                fsdp_plugin_args["use_orig_params"] = str_to_bool(
+                    str(self.fsdp_config.get("use_orig_params", "true")).lower()
+                )
 
             sync_module_states = str(self.fsdp_config.get("sync_module_states", "true")).lower()
             cpu_ram_efficient_loading = str(self.fsdp_config.get("cpu_ram_efficient_loading", "false")).lower()
@@ -2755,11 +2777,10 @@ class TrainingArguments:
                 raise ValueError('`sync_module_states` must be `"True"` if `cpu_ram_efficient_loading` is `"True"`')
 
             # we need to set the env here as otherwise we get a warning in accelerate + we need to set it for transformers
-            fsdp_plugin_args["cpu_ram_efficient_loading"] = cpu_ram_efficient_loading
+            fsdp_plugin_args["cpu_ram_efficient_loading"] = str_to_bool(cpu_ram_efficient_loading)
             os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = cpu_ram_efficient_loading
 
-            fsdp_plugin_args["sync_module_states"] = sync_module_states
-            fsdp_plugin_args["use_orig_params"] = str(self.fsdp_config.get("use_orig_params", "true")).lower()
+            fsdp_plugin_args["sync_module_states"] = str_to_bool(sync_module_states)
 
         return fsdp_plugin_args
 
@@ -2771,3 +2792,18 @@ class ParallelMode(Enum):
     SAGEMAKER_MODEL_PARALLEL = "sagemaker_model_parallel"
     SAGEMAKER_DATA_PARALLEL = "sagemaker_data_parallel"
     TPU = "tpu"
+
+
+def str_to_bool(value, to_bool: bool = True) -> int | bool:
+    """
+    Converts a string representation of truth to `True` (1) or `False` (0).
+
+    True values are `y`, `yes`, `t`, `true`, `on`, and `1`; False value are `n`, `no`, `f`, `false`, `off`, and `0`;
+    """
+    value = value.lower()
+    if value in ("y", "yes", "t", "true", "on", "1"):
+        return 1 if not to_bool else True
+    elif value in ("n", "no", "f", "false", "off", "0"):
+        return 0 if not to_bool else False
+    else:
+        raise ValueError(f"invalid truth value {value}")

transformers/utils/__init__.py

@@ -49,6 +49,7 @@ from .generic import (
     PaddingStrategy,
     TensorType,
     TransformersKwargs,
+    _is_tensor_or_array_like,
     can_return_loss,
     can_return_tuple,
     expand_dims,
@@ -91,7 +92,6 @@ from .hub import (
     extract_commit_hash,
     has_file,
     http_user_agent,
-    is_offline_mode,
     list_repo_templates,
     try_to_load_from_cache,
 )
@@ -114,8 +114,6 @@ from .import_utils import (
     is_apex_available,
     is_apollo_torch_available,
     is_aqlm_available,
-    is_auto_awq_available,
-    is_auto_gptq_available,
     is_auto_round_available,
     is_av_available,
     is_bitsandbytes_available,
@@ -129,7 +127,8 @@ from .import_utils import (
     is_datasets_available,
     is_decord_available,
     is_detectron2_available,
-    is_eetq_available,
+    is_env_variable_false,
+    is_env_variable_true,
     is_essentia_available,
     is_faiss_available,
     is_fbgemm_gpu_available,
@@ -146,6 +145,7 @@ from .import_utils import (
     is_gguf_available,
     is_gptqmodel_available,
     is_grokadamw_available,
+    is_grouped_mm_available,
     is_habana_gaudi1,
     is_hadamard_available,
     is_hqq_available,
@@ -161,6 +161,7 @@ from .import_utils import (
     is_libcst_available,
     is_librosa_available,
     is_liger_kernel_available,
+    is_llm_awq_available,
     is_lomo_available,
     is_matplotlib_available,
     is_mistral_common_available,
@@ -169,6 +170,7 @@ from .import_utils import (
     is_ninja_available,
     is_nltk_available,
     is_num2words_available,
+    is_numba_available,
     is_onnx_available,
     is_openai_available,
     is_optimum_available,
@@ -183,6 +185,7 @@ from .import_utils import (
     is_pyctcdecode_available,
     is_pytesseract_available,
     is_pytest_available,
+    is_pytest_order_available,
     is_pytorch_quantization_available,
     is_quanto_greater,
     is_quark_available,

transformers/utils/attention_visualizer.py

@@ -21,7 +21,7 @@ from ..models.auto.auto_factory import _get_model_class
 from ..models.auto.configuration_auto import AutoConfig
 from ..models.auto.modeling_auto import MODEL_FOR_PRETRAINING_MAPPING, MODEL_MAPPING
 from ..models.auto.processing_auto import PROCESSOR_MAPPING_NAMES, AutoProcessor
-from ..models.auto.tokenization_auto import TOKENIZER_MAPPING_NAMES, AutoTokenizer
+from ..models.auto.tokenization_auto import AutoTokenizer
 from .import_utils import is_torch_available
 
 
@@ -199,12 +199,12 @@ class AttentionMaskVisualizer:
             if "token_type_ids" in inputs:  # TODO inspect signature of update causal mask
                 kwargs["token_type_ids"] = inputs["token_type_ids"]
             tokens = processor.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
-        elif self.config.model_type in TOKENIZER_MAPPING_NAMES:
+        else:
             tokenizer = AutoTokenizer.from_pretrained(self.repo_id)
             tokens = tokenizer.tokenize(input_sentence)
             attention_mask = tokenizer(input_sentence, return_tensors="pt")["attention_mask"]
-        else:
-            raise ValueError(f"Model type {model.config.model_type} does not support attention visualization")
+            if attention_mask is None:
+                raise ValueError(f"Model type {self.config.model_type} does not support attention visualization")
 
         model.config._attn_implementation = "eager"
         model.train()

transformers/utils/auto_docstring.py

@@ -17,7 +17,8 @@ import inspect
 import os
 import textwrap
 from pathlib import Path
-from typing import get_args
+from types import UnionType
+from typing import Union, get_args, get_origin
 
 import regex as re
 
@@ -67,6 +68,7 @@ HARDCODED_CONFIG_FOR_MODELS = {
     "donut": "DonutSwinConfig",
     "esmfold": "EsmConfig",
     "parakeet": "ParakeetCTCConfig",
+    "lasr": "LasrCTCConfig",
 }
 
 _re_checkpoint = re.compile(r"\[(.+?)\]\((https://huggingface\.co/.+?)\)")
@@ -1279,38 +1281,46 @@ def _get_model_info(func, parent_class):
     return model_name_lowercase, class_name, config_class
 
 
-def _process_parameter_type(param, param_name, func):
+def _process_parameter_type(param):
     """
     Process and format a parameter's type annotation.
 
     Args:
         param (`inspect.Parameter`): The parameter from the function signature
-        param_name (`str`): The name of the parameter
-        func (`function`): The function the parameter belongs to
     """
     optional = False
-    if param.annotation != inspect.Parameter.empty:
-        param_type = param.annotation
-        if "typing" in str(param_type):
-            param_type = "".join(str(param_type).split("typing.")).replace("transformers.", "~")
-        elif hasattr(param_type, "__module__"):
-            param_type = f"{param_type.__module__.replace('transformers.', '~').replace('builtins', '')}.{param.annotation.__name__}"
-            if param_type[0] == ".":
-                param_type = param_type[1:]
-        else:
-            if False:
-                print(
-                    f"[ERROR] {param_type} for {param_name} of {func.__qualname__} in file {func.__code__.co_filename} has an invalid type"
-                )
-        if "ForwardRef" in param_type:
-            param_type = re.sub(r"ForwardRef\('([\w.]+)'\)", r"\1", param_type)
-        if "Optional" in param_type:
-            param_type = re.sub(r"Optional\[(.*?)\]", r"\1", param_type)
+    if param.annotation == inspect.Parameter.empty:
+        return "", False
+    elif param.annotation is None:
+        return "None", True
+    # This is, astonishingly, the right way to do it: https://docs.python.org/3/library/typing.html#typing.Union
+    elif get_origin(param.annotation) is Union or get_origin(param.annotation) is UnionType:
+        subtypes = get_args(param.annotation)
+    else:
+        subtypes = [param.annotation]  # Just pretend it's a single-element union so we don't need two code paths
+    out_str = []
+    for subtype in subtypes:
+        if subtype is type(None):
             optional = True
+            continue
+        if hasattr(subtype, "__module__") and hasattr(subtype, "__name__"):
+            subtype = f"{subtype.__module__.replace('transformers.', '~').replace('builtins', '').replace('typing.', '')}.{subtype.__name__}".removeprefix(
+                "."
+            )
+        else:
+            subtype = str(subtype)  # Just give up
+        if "ForwardRef" in subtype:
+            subtype = re.sub(r"ForwardRef\('([\w.]+)'\)", r"\1", subtype)
+        out_str.append(subtype)
+
+    if param.default is not inspect.Parameter.empty:
+        optional = True
+    if not out_str:
+        return "", optional
+    elif len(out_str) == 1:
+        return out_str[0], optional
     else:
-        param_type = ""
-
-    return param_type, optional
+        return f"Union[{', '.join(out_str)}]", optional
 
 
 def _get_parameter_info(param_name, documented_params, source_args_dict, param_type, optional):
@@ -1391,7 +1401,7 @@ def _process_regular_parameters(
             continue
 
         # Process parameter type and optional status
-        param_type, optional = _process_parameter_type(param, param_name, func)
+        param_type, optional = _process_parameter_type(param)
 
         # Check for default value
         param_default = ""

transformers/utils/generic.py

@@ -21,7 +21,7 @@ import os
 import warnings
 from collections import OrderedDict, UserDict, defaultdict
 from collections.abc import Callable, Iterable, MutableMapping
-from contextlib import AbstractContextManager, ExitStack
+from contextlib import AbstractContextManager, ExitStack, nullcontext
 from dataclasses import dataclass, fields, is_dataclass
 from enum import Enum
 from functools import partial, wraps
@@ -42,6 +42,7 @@ _is_torch_available = False
 if is_torch_available():
     # required for @can_return_tuple decorator to work with torchdynamo
     import torch
+    from torch.types import _dtype
 
     from ..model_debugging_utils import model_addition_debugger_context
 
@@ -154,6 +155,48 @@ def is_torch_dtype(x):
     return isinstance(x, torch.dtype)
 
 
+def _is_tensor_or_array_like(value):
+    """
+    Check if a value is array-like (includes ragged arrays)
+    """
+    if is_numpy_array(value):
+        return True
+    if is_torch_tensor(value):
+        return True
+    if isinstance(value, (int, float, bool, np.number)):
+        return True
+
+    if isinstance(value, (list, tuple)):
+        if len(value) == 0:
+            # consider empty list or nested list as array-like
+            return True
+        return _is_tensor_or_array_like(value[0])
+
+    return False
+
+
+def maybe_autocast(
+    device_type: str,
+    dtype: Optional["_dtype"] = None,
+    enabled: bool = True,
+    cache_enabled: Optional[bool] = None,
+):
+    """
+    Context manager that only autocasts if:
+
+    - `autocast` is already enabled in this context
+    - Or this call to `maybe_autocast` has `enabled=True`
+
+    This prevents `autocast` being added to the graph when it is effectively a no-op.
+    Which makes graph splitting in `torch.compile` more flexible as it removes the
+    requirement that partition IDs be monotonically increasing.
+    """
+    if torch.is_autocast_enabled(device_type) or enabled:
+        return torch.autocast(device_type, dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)
+    else:
+        return nullcontext()
+
+
 def _is_mlx(x):
     import mlx.core as mx
 
@@ -680,6 +723,8 @@ class TransformersKwargs(TypedDict, total=False):
             Maximum sequence length for query state.
         max_length_k (`int`, *optional*):
             Maximum sequence length for key state.
+        position_ids (`torch.LongTensor`, *optional*)
+            Indices of positions of each input sequence tokens.
     """
 
     num_items_in_batch: Optional["torch.Tensor"]
@@ -690,6 +735,7 @@ class TransformersKwargs(TypedDict, total=False):
     cu_seq_lens_k: Optional["torch.LongTensor"]
     max_length_q: int | None
     max_length_k: int | None
+    position_ids: Optional["torch.LongTensor"]
 
 
 def is_timm_config_dict(config_dict: dict[str, Any]) -> bool: