Repository: huggingface/transformers
Pull request: Update TrainingArguments

# What does this PR do?

This PR updates TrainingArguments in several places:
- More tests
- Simplify the post_init  
- update and reorder __init__

For the review, just check that the post_init is correctly modified

Changed files:
filename: src/transformers/training_args.py
patch: @@ -107,26 +107,6 @@
     smp.init()
 
 
-def get_int_from_env(env_keys, default):
-    """Returns the first positive env value found in the `env_keys` list or the default."""
-    for e in env_keys:
-        val = int(os.environ.get(e, "-1"))
-        if val >= 0:
-            return val
-    return default
-
-
-def get_xla_device_type(device: "torch.device") -> str | None:
-    """
-    Returns the xla device type (CPU|GPU|TPU) or None if the device is a non-xla device.
-    """
-    if is_torch_xla_available():
-        if device.type == "cpu":
-            return "CPU"
-        return xm.xla_real_devices([device])[0].split(":")[0]
-    return None
-
-
 class OptimizerNames(ExplicitEnum):
     """
     Stores the acceptable string identifiers for optimizers.
@@ -761,10 +741,8 @@ class TrainingArguments:
             scripts](https://github.com/huggingface/transformers/tree/main/examples) for more details.
     """
 
-    # Sometimes users will pass in a `str` repr of a dict in the CLI
-    # We need to track what fields those can be. Each time a new arg
-    # has a dict type, it must be added to this list.
-    # Important: These should be typed with Optional[Union[dict,str,...]]
+    # Fields that accept dict values via CLI as JSON strings (e.g., '{"key": "value"}').
+    # Any new dict-typed arg must be added here and typed as `dict | str | None`.
     _VALID_DICT_FIELDS = [
         "accelerator_config",
         "fsdp_config",
@@ -773,489 +751,422 @@ class TrainingArguments:
         "lr_scheduler_kwargs",
     ]
 
+    # --- Output ---
     output_dir: str | None = field(
         default=None,
-        metadata={
-            "help": "The output directory where the model predictions and checkpoints will be written. Defaults to 'trainer_output' if not provided."
-        },
-    )
-
-    do_train: bool = field(default=False, metadata={"help": "Whether to run training."})
-    do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."})
-    do_predict: bool = field(default=False, metadata={"help": "Whether to run predictions on the test set."})
-    eval_strategy: IntervalStrategy | str = field(
-        default="no",
-        metadata={"help": "The evaluation strategy to use."},
-    )
-    prediction_loss_only: bool = field(
-        default=False,
-        metadata={"help": "When performing evaluation and predictions, only returns the loss."},
-    )
-
-    per_device_train_batch_size: int = field(
-        default=8, metadata={"help": "Batch size per device accelerator core/CPU for training."}
-    )
-    per_device_eval_batch_size: int = field(
-        default=8, metadata={"help": "Batch size per device accelerator core/CPU for evaluation."}
-    )
-
-    gradient_accumulation_steps: int = field(
-        default=1,
-        metadata={"help": "Number of updates steps to accumulate before performing a backward/update pass."},
-    )
-    eval_accumulation_steps: int | None = field(
-        default=None,
-        metadata={"help": "Number of predictions steps to accumulate before moving the tensors to the CPU."},
-    )
-
-    eval_delay: float = field(
-        default=0,
-        metadata={
-            "help": (
-                "Number of epochs or steps to wait for before the first evaluation can be performed, depending on the"
-                " eval_strategy."
-            )
-        },
-    )
-
-    torch_empty_cache_steps: int | None = field(
-        default=None,
-        metadata={
-            "help": "Number of steps to wait before calling `torch.<device>.empty_cache()`."
-            "This can help avoid CUDA out-of-memory errors by lowering peak VRAM usage at a cost of about [10% slower performance](https://github.com/huggingface/transformers/issues/31372)."
-            "If left unset or set to None, cache will not be emptied."
-        },
+        metadata={"help": "The output directory where the model predictions and checkpoints will be written."},
     )
 
-    learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for AdamW."})
-    weight_decay: float = field(default=0.0, metadata={"help": "Weight decay for AdamW if we apply some."})
-    adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for AdamW optimizer"})
-    adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for AdamW optimizer"})
-    adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon for AdamW optimizer."})
-    max_grad_norm: float = field(default=1.0, metadata={"help": "Max gradient norm."})
-
+    # --- Training Duration and Batch Size ---
+    per_device_train_batch_size: int = field(default=8, metadata={"help": "The batch size per device for training."})
     num_train_epochs: float = field(default=3.0, metadata={"help": "Total number of training epochs to perform."})
     max_steps: int = field(
         default=-1,
-        metadata={"help": "If > 0: set total number of training steps to perform. Override num_train_epochs."},
+        metadata={
+            "help": "Overrides `num_train_epochs`. If set to a positive number, the total number of training steps to perform."
+        },
     )
+
+    # --- Learning Rate & Scheduler ---
+    learning_rate: float = field(default=5e-5, metadata={"help": "The initial learning rate for the optimizer."})
     lr_scheduler_type: SchedulerType | str = field(
         default="linear",
-        metadata={"help": "The scheduler type to use."},
+        metadata={"help": "The learning rate scheduler type to use. See `SchedulerType` for all possible values."},
     )
     lr_scheduler_kwargs: dict | str | None = field(
         default=None,
         metadata={
-            "help": (
-                "Extra parameters for the lr_scheduler such as {'num_cycles': 1} for the cosine with hard restarts."
-            )
+            "help": "The extra arguments for the lr_scheduler. See the documentation of each scheduler for possible values."
         },
     )
-    warmup_ratio: float | None = field(
-        default=None,
+    warmup_steps: float = field(
+        default=0,
         metadata={
-            "help": "This argument is deprecated and will be removed in v5.2. Use `warmup_steps` instead as it also works with float values."
+            "help": "Number of steps for a linear warmup from 0 to `learning_rate`. Can be an integer (exact steps) or a float in [0, 1) (ratio of total steps)."
         },
     )
 
-    warmup_steps: float = field(default=0, metadata={"help": "Linear warmup over warmup_steps."})
+    # --- Optimizer ---
+    default_optim = "adamw_torch"
+    if is_torch_available():
+        from .pytorch_utils import is_torch_greater_or_equal_than_2_8
 
-    log_level: str = field(
-        default="passive",
-        metadata={
-            "help": (
-                "Logger log level to use on the main node. Possible choices are the log levels as strings: 'debug',"
-                " 'info', 'warning', 'error' and 'critical', plus a 'passive' level which doesn't set anything and"
-                " lets the application set the level. Defaults to 'passive'."
-            ),
-            "choices": trainer_log_levels.keys(),
-        },
+        if is_torch_greater_or_equal_than_2_8:
+            default_optim = "adamw_torch_fused"
+    optim: OptimizerNames | str = field(
+        default=default_optim,
+        metadata={"help": "The optimizer to use. See `OptimizerNames` for the complete list."},
     )
-    log_level_replica: str = field(
-        default="warning",
+    optim_args: str | None = field(
+        default=None,
         metadata={
-            "help": "Logger log level to use on replica nodes. Same choices and defaults as ``log_level``",
-            "choices": trainer_log_levels.keys(),
+            "help": "Optional arguments supplied to optimizers such as AnyPrecisionAdamW, AdEMAMix, and GaLore."
         },
     )
-    log_on_each_node: bool = field(
-        default=True,
+    weight_decay: float = field(
+        default=0.0,
         metadata={
-            "help": (
-                "When doing a multinode distributed training, whether to log once per node or just once on the main"
-                " node."
-            )
+            "help": "Weight decay coefficient applied by the optimizer. Automatically excluded from bias and LayerNorm parameters."
         },
     )
-    logging_dir: str | None = field(
-        default=None,
+    adam_beta1: float = field(
+        default=0.9,
         metadata={
-            "help": "Deprecated and will be removed in v5.2. Set env var `TENSORBOARD_LOGGING_DIR` instead. TensorBoard log directory."
+            "help": "The exponential decay rate for the first moment estimates (momentum) in Adam-based optimizers."
         },
     )
-    logging_strategy: IntervalStrategy | str = field(
-        default="steps",
-        metadata={"help": "The logging strategy to use."},
-    )
-    logging_first_step: bool = field(default=False, metadata={"help": "Log the first global_step"})
-    logging_steps: float = field(
-        default=500,
+    adam_beta2: float = field(
+        default=0.999,
         metadata={
-            "help": (
-                "Log every X updates steps. Should be an integer or a float in range `[0,1)`. "
-                "If smaller than 1, will be interpreted as ratio of total training steps."
-            )
+            "help": "The exponential decay rate for the second moment estimates (variance) in Adam-based optimizers."
         },
     )
-    logging_nan_inf_filter: bool = field(default=True, metadata={"help": "Filter nan and inf losses for logging."})
-    save_strategy: SaveStrategy | str = field(
-        default="steps",
-        metadata={"help": "The checkpoint save strategy to use."},
+    adam_epsilon: float = field(
+        default=1e-8, metadata={"help": "Epsilon value for numerical stability in Adam-based optimizers."}
     )
-    save_steps: float = field(
-        default=500,
+    optim_target_modules: None | str | list[str] = field(
+        default=None,
+        metadata={"help": "The target modules to optimize. Currently used for the GaLore and APOLLO algorithms."},
+    )
+
+    # --- Regularization & Training Stability ---
+    gradient_accumulation_steps: int = field(
+        default=1,
         metadata={
             "help": (
-                "Save checkpoint every X updates steps. Should be an integer or a float in range `[0,1)`. "
-                "If smaller than 1, will be interpreted as ratio of total training steps."
+                "Number of update steps to accumulate gradients before performing a backward/update pass."
+                " Effective batch size = per_device_train_batch_size * num_devices * gradient_accumulation_steps."
             )
         },
     )
-    save_total_limit: int | None = field(
-        default=None,
+    average_tokens_across_devices: bool = field(
+        default=True,
         metadata={
-            "help": (
-                "If a value is passed, will limit the total amount of checkpoints. Deletes the older checkpoints in"
-                " `output_dir`. When `load_best_model_at_end` is enabled, the 'best' checkpoint according to"
-                " `metric_for_best_model` will always be retained in addition to the most recent ones. For example,"
-                " 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."
-            )
+            "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to "
+            "synchronize num_tokens_in_batch for precise loss calculation. Reference: "
+            "https://github.com/huggingface/transformers/issues/34242"
         },
     )
-    enable_jit_checkpoint: bool = field(
+    max_grad_norm: float = field(
+        default=1.0, metadata={"help": "Maximum gradient norm for gradient clipping. Set to 0 to disable."}
+    )
+    label_smoothing_factor: float = field(
+        default=0.0, metadata={"help": "Label smoothing factor to prevent overconfidence. Zero means no smoothing."}
+    )
+
+    # --- Mixed Precision ---
+    bf16: bool = field(
         default=False,
         metadata={
-            "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."
-            )
+            "help": "Enable bfloat16 (BF16) mixed precision training. Generally preferred over FP16 due to better numerical stability."
         },
     )
-    save_on_each_node: bool = field(
+    fp16: bool = field(
         default=False,
         metadata={
-            "help": (
-                "When doing multi-node distributed training, whether to save models and checkpoints on each node, or"
-                " only on the main one"
-            )
+            "help": "Enable float16 (FP16) mixed precision training. Consider using BF16 instead if your hardware supports it."
         },
     )
-    save_only_model: bool = field(
+    bf16_full_eval: bool = field(
         default=False,
         metadata={
-            "help": (
-                "When checkpointing, whether to only save the model, or also the optimizer, scheduler & rng state."
-                "Note that when this is true, you won't be able to resume training from checkpoint."
-                "This enables you to save storage by not storing the optimizer, scheduler & rng state."
-                "You can only load the model using from_pretrained with this option set to True."
-            )
+            "help": "Use full BF16 precision for evaluation (not just mixed precision). Faster and saves memory."
         },
     )
-    restore_callback_states_from_checkpoint: bool = field(
+    fp16_full_eval: bool = field(
         default=False,
         metadata={
-            "help": "Whether to restore the callback states from the checkpoint. If `True`, will override callbacks passed to the `Trainer` if they exist in the checkpoint."
+            "help": "Use full FP16 precision for evaluation (not just mixed precision). Faster and saves memory."
         },
     )
-    use_cpu: bool = field(
-        default=False,
+    tf32: bool | None = field(
+        default=None,
         metadata={
-            "help": "Whether or not to use cpu. If left to False, we will use the available torch device/backend (cuda/mps/xpu/hpu etc.)"
+            "help": "Enable TF32 mode on Ampere and newer GPUs. Provides up to 8x speedup with negligible accuracy loss."
         },
     )
-    seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
-    data_seed: int | None = field(default=None, metadata={"help": "Random seed to be used with data samplers."})
-    bf16: bool = field(
+
+    # --- Gradient Checkpointing ---
+    gradient_checkpointing: bool = field(
         default=False,
         metadata={
-            "help": (
-                "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA"
-                " architecture or using CPU (use_cpu) or Ascend NPU. This is an experimental API and it may change."
-            )
+            "help": "Enable gradient checkpointing to trade compute for memory. Reduces memory at the cost of ~20% slower training."
         },
     )
-    fp16: bool = field(
-        default=False,
-        metadata={"help": "Whether to use fp16 (mixed) precision instead of 32-bit"},
+    gradient_checkpointing_kwargs: dict[str, Any] | str | None = field(
+        default=None,
+        metadata={"help": "Keyword arguments passed to `gradient_checkpointing_enable()`."},
     )
 
-    bf16_full_eval: bool = field(
-        default=False,
+    # --- Compilation ---
+    torch_compile: bool = field(
+        default=False, metadata={"help": "Compile the model using `torch.compile()` for faster training."}
+    )
+    torch_compile_backend: str | None = field(
+        default=None,
         metadata={
-            "help": (
-                "Whether to use full bfloat16 evaluation instead of 32-bit. This is an experimental API and it may"
-                " change."
-            )
+            "help": "Backend for `torch.compile()`. If set, automatically enables `torch_compile`.",
         },
     )
-    fp16_full_eval: bool = field(
+    torch_compile_mode: str | None = field(
+        default=None,
+        metadata={
+            "help": "Compilation mode for `torch.compile()`. If set, automatically enables `torch_compile`.",
+        },
+    )
+
+    # --- Kernels ---
+    use_liger_kernel: bool = field(
         default=False,
-        metadata={"help": "Whether to use full float16 evaluation instead of 32-bit"},
+        metadata={
+            "help": "Enable Liger Kernel optimizations. Increases throughput by ~20% and reduces memory by ~60%."
+        },
     )
-    tf32: bool | None = field(
+    liger_kernel_config: dict[str, bool] | None = field(
         default=None,
         metadata={
-            "help": (
-                "Whether to enable tf32 mode, available in Ampere and newer GPU architectures. This is an experimental"
-                " API and it may change."
-            )
+            "help": "Configuration for Liger Kernel. Passed as kwargs to `_apply_liger_kernel_to_instance()`. If None, uses default configuration."
         },
     )
-    local_rank: int = field(
-        default=-1,
+
+    # --- Additional Optimizations ---
+    use_cache: bool = field(
+        default=False,
         metadata={
-            "help": "When using torch.distributed.launch (Deprecated), it will pass `local_rank` in the script, so we need this for the parser. To get the local rank, prefer using the property `local_process_index`"
+            "help": "Whether or not to use cache for the model For training, this is usually not needed apart from some PEFT methods that uses `past_key_values`."
         },
     )
-    ddp_backend: str | None = field(
+    neftune_noise_alpha: float | None = field(
         default=None,
         metadata={
-            "help": "The backend to be used for distributed training",
-            "choices": ["nccl", "gloo", "mpi", "xccl", "hccl", "cncl", "mccl"],
+            "help": "If not None, activates NEFTune noise embeddings. Can drastically improve performance for instruction fine-tuning. Typical range: [5.0, 15.0]."
         },
     )
-    debug: str | list[DebugOption] = field(
-        default="",
+    torch_empty_cache_steps: int | None = field(
+        default=None,
         metadata={
-            "help": (
-                "Whether or not to enable debug mode. Current options: "
-                "`underflow_overflow` (Detect underflow and overflow in activations and weights), "
-                "`tpu_metrics_debug` (print debug metrics on TPU)."
-            )
+            "help": "Number of steps to wait before calling `torch.<device>.empty_cache()`. Helps avoid CUDA OOM at a cost of ~10% slower performance. If None, cache will not be emptied."
+        },
+    )
+    auto_find_batch_size: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to find a batch size that will fit into memory automatically through exponential decay, avoiding CUDA Out-of-Memory errors."
         },
     )
 
-    dataloader_drop_last: bool = field(
-        default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."}
+    # --- Logging & Monitoring ---
+    logging_strategy: IntervalStrategy | str = field(
+        default="steps",
+        metadata={"help": "The logging strategy to adopt during training. Options: 'no', 'epoch', 'steps'."},
     )
-    eval_steps: float | None = field(
-        default=None,
+    logging_steps: float = field(
+        default=500,
         metadata={
             "help": (
-                "Run an evaluation every X steps. Should be an integer or a float in range `[0,1)`. "
+                "Log every X updates steps. Should be an integer or a float in range `[0,1)`. "
                 "If smaller than 1, will be interpreted as ratio of total training steps."
             )
         },
     )
-    dataloader_num_workers: int = field(
-        default=0,
+    logging_first_step: bool = field(
+        default=False, metadata={"help": "Whether to log the first `global_step` or not."}
+    )
+    log_on_each_node: bool = field(
+        default=True,
         metadata={
             "help": (
-                "Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded"
-                " in the main process."
+                "When doing a multinode distributed training, whether to log once per node or just once on the main"
+                " node."
             )
         },
     )
-    dataloader_prefetch_factor: int | None = field(
-        default=None,
+    logging_nan_inf_filter: bool = field(
+        default=True,
         metadata={
-            "help": (
-                "Number of batches loaded in advance by each worker. "
-                "2 means there will be a total of 2 * num_workers batches prefetched across all workers. "
-            )
+            "help": "Filter out NaN and Inf losses when logging. Does not affect gradient computation, only logging."
         },
     )
-
-    run_name: str | None = field(
-        default=None,
+    include_num_input_tokens_seen: str | bool = field(
+        default="no",
         metadata={
             "help": (
-                "An optional descriptor for the run. Notably used for trackio, wandb, mlflow comet and swanlab "
-                "logging."
+                "Whether to track the number of input tokens seen. "
+                "Must be one of [`all`, `non_padding`, `no`] or a boolean value which map to `all` or `no`"
             )
         },
     )
-    disable_tqdm: bool | None = field(
-        default=None, metadata={"help": "Whether or not to disable the tqdm progress bars."}
-    )
 
-    remove_unused_columns: bool = field(
-        default=True, metadata={"help": "Remove columns not required by the model when using an nlp.Dataset."}
-    )
-    label_names: list[str] | None = field(
-        default=None, metadata={"help": "The list of keys in your dictionary of inputs that correspond to the labels."}
-    )
-    load_best_model_at_end: bool = field(
-        default=False,
+    # --- Log Levels ---
+    log_level: str = field(
+        default="passive",
         metadata={
-            "help": (
-                "Whether or not to load the best model found during training at the end of training. When this option"
-                " is enabled, the best checkpoint will always be saved. See `save_total_limit` for more."
-            )
+            "help": "Logging level for the main process. Options: 'debug', 'info', 'warning', 'error', 'critical', 'passive'.",
+            "choices": trainer_log_levels.keys(),
         },
     )
-    metric_for_best_model: str | None = field(
-        default=None, metadata={"help": "The metric to use to compare two different models."}
-    )
-    greater_is_better: bool | None = field(
-        default=None, metadata={"help": "Whether the `metric_for_best_model` should be maximized or not."}
-    )
-    ignore_data_skip: bool = field(
-        default=False,
+    log_level_replica: str = field(
+        default="warning",
         metadata={
-            "help": (
-                "When resuming training, whether or not to skip the first epochs and batches to get to the same"
-                " training data."
-            )
+            "help": "Logging level for replica processes in distributed training. Same options as `log_level`.",
+            "choices": trainer_log_levels.keys(),
         },
     )
-    fsdp: list[FSDPOption] | str | None = field(
+    disable_tqdm: bool | None = field(
         default=None,
+        metadata={"help": "Disable tqdm progress bars. Defaults to True if log_level is warning or lower."},
+    )
+
+    # --- Experiment Tracking ---
+    report_to: None | str | list[str] = field(
+        default="none",
         metadata={
-            "help": (
-                "Whether or not to use PyTorch Fully Sharded Data Parallel (FSDP) training (in distributed training"
-                " only). The base option should be `full_shard`, `shard_grad_op` or `no_shard` and you can add"
-                " CPU-offload to `full_shard` or `shard_grad_op` like this: full_shard offload` or `shard_grad_op"
-                " offload`. You can add auto-wrap to `full_shard` or `shard_grad_op` with the same syntax: full_shard"
-                " auto_wrap` or `shard_grad_op auto_wrap`."
-            ),
+            "help": "The list of integrations to report the results and logs to. Use 'all' for all installed integrations, 'none' for no integrations."
         },
     )
-    fsdp_config: dict[str, Any] | str | None = field(
+    run_name: str | None = field(
         default=None,
         metadata={
             "help": (
-                "Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a "
-                "fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`."
-            )
+                "An optional descriptor for the run. Notably used for trackio, wandb, mlflow comet and swanlab "
+                "logging."
+            )
         },
     )
-    accelerator_config: dict | str | None = field(
+    project: str = field(
+        default="huggingface",
+        metadata={"help": "The name of the project to use for logging. Currently, only used by Trackio."},
+    )
+    trackio_space_id: str | None = field(
+        default="trackio",
+        metadata={
+            "help": "The Hugging Face Space ID to deploy to when using Trackio. Should be a complete Space name like "
+            "'username/reponame' or 'orgname/reponame', or just 'reponame' in which case the Space will be created in "
+            "the currently-logged-in Hugging Face user's namespace. If `None`, will log to a local directory. Note "
+            "that this Space will be public unless you set `hub_private_repo=True` or your organization's "
+            "default is to create private Spaces."
+        },
+    )
+
+    # --- Evaluation ---
+    eval_strategy: IntervalStrategy | str = field(
+        default="no",
+        metadata={"help": "When to run evaluation. Options: 'no', 'steps', 'epoch'."},
+    )
+    eval_steps: float | None = field(
         default=None,
         metadata={
             "help": (
-                "Config to be used with the internal Accelerator object initialization. The value is either a "
-                "accelerator json config file (e.g., `accelerator_config.json`) or an already loaded json file as `dict`."
+                "Number of update steps between evaluations if `eval_strategy='steps'`. Defaults to `logging_steps` if not set."
+                " Should be an integer or a float in range `[0,1)`. If smaller than 1, will be interpreted as ratio of total training steps."
             )
         },
     )
-    parallelism_config: ParallelismConfig | None = field(
-        default=None,
-        metadata={"help": ("Parallelism configuration for the training run. Requires Accelerate `1.12.0`")},
-    )
-    deepspeed: dict | str | None = field(
-        default=None,
+    eval_delay: float = field(
+        default=0,
         metadata={
             "help": (
-                "Enable deepspeed and pass the path to deepspeed json config file (e.g. `ds_config.json`) or an already"
-                " loaded json file as a dict"
+                "Number of epochs or steps to wait for before the first evaluation can be performed, depending on the"
+                " eval_strategy."
             )
         },
     )
-    label_smoothing_factor: float = field(
-        default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."}
+    per_device_eval_batch_size: int = field(
+        default=8, metadata={"help": "The batch size per device (GPU/TPU core/CPU) for evaluation."}
     )
-
-    default_optim = "adamw_torch"
-    if is_torch_available():
-        from .pytorch_utils import is_torch_greater_or_equal_than_2_8
-
-        if is_torch_greater_or_equal_than_2_8:
-            default_optim = "adamw_torch_fused"
-    optim: OptimizerNames | str = field(
-        default=default_optim,
-        metadata={"help": "The optimizer to use."},
+    prediction_loss_only: bool = field(
+        default=False,
+        metadata={"help": "When performing evaluation and generating predictions, only returns the loss."},
     )
-    optim_args: str | None = field(default=None, metadata={"help": "Optional arguments to supply to optimizer."})
-    group_by_length: bool = field(
+    eval_on_start: bool = field(
         default=False,
-        metadata={"help": "Whether or not to group samples of roughly the same length together when batching."},
+        metadata={
+            "help": "Whether to run through the entire `evaluation` step at the very beginning of training as a sanity check."
+        },
     )
-    length_column_name: str = field(
-        default="length",
-        metadata={"help": "Column name with precomputed lengths to use when grouping by length."},
+    eval_do_concat_batches: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, will instead store them as lists, with each batch kept separate."
+        },
     )
-    report_to: None | str | list[str] = field(
-        default="none", metadata={"help": "The list of integrations to report the results and logs to."}
+    eval_use_gather_object: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices."
+        },
     )
-    project: str = field(
-        default="huggingface",
-        metadata={"help": "The name of the project to use for logging. Currenly, only used by Trackio."},
+    eval_accumulation_steps: int | None = field(
+        default=None,
+        metadata={
+            "help": "Number of predictions steps to accumulate the output tensors for, before moving the results to the CPU. If unset, predictions are accumulated on the accelerator before being moved to the CPU."
+        },
     )
-    trackio_space_id: str | None = field(
-        default="trackio",
+
+    # --- Metrics ---
+    include_for_metrics: list[str] = field(
+        default_factory=list,
+        metadata={"help": "Include additional data in the `compute_metrics` function. Options: 'inputs', 'loss'."},
+    )
+    batch_eval_metrics: bool = field(
+        default=False,
+        metadata={"help": "Break eval metrics calculation into batches to save memory."},
+    )
+
+    # --- Checkpointing & Saving ---
+    save_only_model: bool = field(
+        default=False,
         metadata={
-            "help": "The Hugging Face Space ID to deploy to when using Trackio. Should be a complete Space name like "
-            "'username/reponame' or 'orgname/reponame', or just 'reponame' in which case the Space will be created in "
-            "the currently-logged-in Hugging Face user's namespace. If `None`, will log to a local directory. Note "
-            "that this Space will be public unless you set `hub_private_repo=True` or your organization's "
-            "default is to create private Spaces."
+            "help": "Save only model weights, not optimizer/scheduler/RNG state. Prevents resuming training from checkpoint."
         },
     )
-    ddp_find_unused_parameters: bool | None = field(
-        default=None,
+    save_strategy: SaveStrategy | str = field(
+        default="steps",
         metadata={
-            "help": (
-                "When using distributed training, the value of the flag `find_unused_parameters` passed to "
-                "`DistributedDataParallel`."
-            )
+            "help": "The checkpoint save strategy to adopt during training. Options: 'no', 'epoch', 'steps', 'best'."
         },
     )
-    ddp_bucket_cap_mb: int | None = field(
-        default=None,
+    save_steps: float = field(
+        default=500,
         metadata={
             "help": (
-                "When using distributed training, the value of the flag `bucket_cap_mb` passed to "
-                "`DistributedDataParallel`."
+                "Save checkpoint every X updates steps. Should be an integer or a float in range `[0,1)`. "
+                "If smaller than 1, will be interpreted as ratio of total training steps."
             )
         },
     )
-    ddp_broadcast_buffers: bool | None = field(
-        default=None,
+    save_on_each_node: bool = field(
+        default=False,
         metadata={
             "help": (
-                "When using distributed training, the value of the flag `broadcast_buffers` passed to "
-                "`DistributedDataParallel`."
+                "When doing multi-node distributed training, whether to save models and checkpoints on each node, or"
+                " only on the main one"
             )
         },
     )
-    dataloader_pin_memory: bool = field(
-        default=True, metadata={"help": "Whether or not to pin memory for DataLoader."}
+    save_total_limit: int | None = field(
+        default=None,
+        metadata={
+            "help": "Maximum number of checkpoints to keep. Deletes older checkpoints in `output_dir`. The best checkpoint is always retained when `load_best_model_at_end=True`."
+        },
     )
-    dataloader_persistent_workers: bool = field(
+    enable_jit_checkpoint: bool = field(
         default=False,
         metadata={
-            "help": "If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage."
+            "help": "Enable JIT checkpointing on SIGTERM signal for graceful termination on preemptible workloads. Configure your orchestrator's graceful shutdown period accordingly."
         },
     )
-    skip_memory_metrics: bool = field(
-        default=True, metadata={"help": "Whether or not to skip adding of memory profiler reports to metrics."}
-    )
+
+    # --- Hub Integration ---
     push_to_hub: bool = field(
-        default=False, metadata={"help": "Whether or not to upload the trained model to the model hub after training."}
+        default=False, metadata={"help": "Whether or not to push the model to the Hub every time the model is saved."}
     )
-    resume_from_checkpoint: str | None = field(
+    hub_token: str | None = field(
         default=None,
-        metadata={"help": "The path to a folder with a valid checkpoint for your model."},
-    )
-    hub_model_id: str | None = field(
-        default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
-    )
-    hub_strategy: HubStrategy | str = field(
-        default="every_save",
-        metadata={"help": "The hub strategy to use when `--push_to_hub` is activated."},
+        metadata={
+            "help": "The token to use to push the model to the Hub. Defaults to the token from `hf auth login`."
+        },
     )
-    hub_token: str | None = field(default=None, metadata={"help": "The token to use to push to the Model Hub."})
     hub_private_repo: bool | None = field(
         default=None,
         metadata={
@@ -1265,6 +1176,15 @@ class TrainingArguments:
             "private."
         },
     )
+    hub_model_id: str | None = field(
+        default=None, metadata={"help": "The name of the repository to keep in sync with the local `output_dir`."}
+    )
+    hub_strategy: HubStrategy | str = field(
+        default="every_save",
+        metadata={
+            "help": "Defines what and when to push to Hub. Options: 'end', 'every_save', 'checkpoint', 'all_checkpoints'."
+        },
+    )
     hub_always_push: bool = field(
         default=False,
         metadata={"help": "Unless `True`, the Trainer will skip pushes if the previous one wasn't finished yet."},
@@ -1275,192 +1195,300 @@ class TrainingArguments:
             "help": "The revision to use when pushing to the Hub. Can be a branch name, a tag, or a commit hash."
         },
     )
-    gradient_checkpointing: bool = field(
+
+    # --- Best Model Tracking ---
+    load_best_model_at_end: bool = field(
         default=False,
+        metadata={"help": "Load the best checkpoint at the end of training. Requires `eval_strategy` to be set."},
+    )
+    metric_for_best_model: str | None = field(
+        default=None,
         metadata={
-            "help": "If True, use gradient checkpointing to save memory at the expense of slower backward pass."
+            "help": "Metric to use for comparing models when `load_best_model_at_end=True`. Defaults to 'loss'."
         },
     )
-    gradient_checkpointing_kwargs: dict[str, Any] | str | None = field(
+    greater_is_better: bool | None = field(
         default=None,
+        metadata={"help": "Whether higher metric values are better. Defaults based on `metric_for_best_model`."},
+    )
+
+    # --- Resuming Training ---
+    ignore_data_skip: bool = field(
+        default=False,
         metadata={
-            "help": "Gradient checkpointing key word arguments such as `use_reentrant`. Will be passed to `torch.utils.checkpoint.checkpoint` through `model.gradient_checkpointing_enable`."
+            "help": "When resuming training, skip fast-forwarding through the dataset to reach the previous state. If True, training starts from the beginning of the dataset."
         },
     )
-    include_for_metrics: list[str] = field(
-        default_factory=list,
+    restore_callback_states_from_checkpoint: bool = field(
+        default=False,
         metadata={
-            "help": "List of strings to specify additional data to include in the `compute_metrics` function."
-            "Options: 'inputs', 'loss'."
+            "help": "Whether to restore the callback states from the checkpoint. If `True`, will override callbacks passed to the `Trainer` if they exist in the checkpoint."
         },
     )
-    eval_do_concat_batches: bool = field(
-        default=True,
+
+    # --- Reproducibility ---
+    full_determinism: bool = field(
+        default=False,
         metadata={
-            "help": "Whether to recursively concat inputs/losses/labels/predictions across batches. If `False`, will instead store them as lists, with each batch kept separate."
+            "help": (
+                "Whether to call enable_full_determinism instead of set_seed for reproducibility in distributed"
+                " training. Important: this will negatively impact the performance, so only use it for debugging."
+            )
         },
     )
-    auto_find_batch_size: bool = field(
+    seed: int = field(default=42, metadata={"help": "Random seed that will be set at the beginning of training."})
+    data_seed: int | None = field(
+        default=None,
+        metadata={"help": "Random seed to be used with data samplers. If not set, uses the same seed as `seed`."},
+    )
+
+    # --- Hardware ---
+    use_cpu: bool = field(
         default=False,
+        metadata={
+            "help": "Whether or not to use cpu. If set to False, we will use the available torch device/backend."
+        },
+    )
+
+    # --- Accelerate ---
+    accelerator_config: dict | str | None = field(
+        default=None,
+        metadata={
+            "help": "Configuration for the internal Accelerate integration. Can be a path to a JSON config file or a dict."
+        },
+    )
+    parallelism_config: ParallelismConfig | None = field(
+        default=None,
+        metadata={"help": "Parallelism configuration for the training run. Requires Accelerate `1.10.1`."},
+    )
+
+    # --- Dataloader ---
+    dataloader_drop_last: bool = field(
+        default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."}
+    )
+    dataloader_num_workers: int = field(
+        default=0,
         metadata={
             "help": (
-                "Whether to automatically decrease the batch size in half and rerun the training loop again each time"
-                " a CUDA Out-of-Memory was reached"
+                "Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded"
+                " in the main process."
             )
         },
     )
-    full_determinism: bool = field(
+    dataloader_pin_memory: bool = field(
+        default=True, metadata={"help": "Whether or not to pin memory for DataLoader."}
+    )
+    dataloader_persistent_workers: bool = field(
         default=False,
+        metadata={
+            "help": "If True, the data loader will not shut down the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. Can potentially speed up training, but will increase RAM usage."
+        },
+    )
+    dataloader_prefetch_factor: int | None = field(
+        default=None,
         metadata={
             "help": (
-                "Whether to call enable_full_determinism instead of set_seed for reproducibility in distributed"
-                " training. Important: this will negatively impact the performance, so only use it for debugging."
+                "Number of batches loaded in advance by each worker. "
+                "2 means there will be a total of 2 * num_workers batches prefetched across all workers. "
             )
         },
     )
-    ddp_timeout: int = field(
-        default=1800,
+    remove_unused_columns: bool = field(
+        default=True,
+        metadata={"help": "Whether or not to automatically remove the columns unused by the model forward method."},
+    )
+    label_names: list[str] | None = field(
+        default=None, metadata={"help": "The list of keys in your dictionary of inputs that correspond to the labels."}
+    )
+    group_by_length: bool = field(
+        default=False,
         metadata={
-            "help": "Overrides the default timeout for distributed training (value should be given in seconds)."
+            "help": "Whether or not to group samples of roughly the same length together when batching. Only useful if applying dynamic padding."
         },
     )
-    torch_compile: bool = field(
-        default=False, metadata={"help": "If set to `True`, the model will be wrapped in `torch.compile`."}
+    length_column_name: str = field(
+        default="length",
+        metadata={"help": "Column name for precomputed lengths. Ignored unless `group_by_length` is True."},
     )
-    torch_compile_backend: str | None = field(
+
+    # --- DDP ---
+    ddp_find_unused_parameters: bool | None = field(
         default=None,
         metadata={
-            "help": "Which backend to use with `torch.compile`, passing one will trigger a model compilation.",
+            "help": (
+                "When using distributed training, the value of the flag `find_unused_parameters` passed to "
+                "`DistributedDataParallel`."
+            )
         },
     )
-    torch_compile_mode: str | None = field(
+    ddp_bucket_cap_mb: int | None = field(
         default=None,
         metadata={
-            "help": "Which mode to use with `torch.compile`, passing one will trigger a model compilation.",
+            "help": (
+                "When using distributed training, the value of the flag `bucket_cap_mb` passed to "
+                "`DistributedDataParallel`."
+            )
         },
     )
-    include_num_input_tokens_seen: str | bool = field(
-        default="no",
+    ddp_broadcast_buffers: bool | None = field(
+        default=None,
         metadata={
             "help": (
-                "Whether to track the number of input tokens seen. "
-                "Must be one of [`all`, `non_padding`, `no`] or a boolean value which map to `all` or `no`"
+                "When using distributed training, the value of the flag `broadcast_buffers` passed to "
+                "`DistributedDataParallel`."
             )
         },
     )
+    ddp_backend: str | None = field(
+        default=None,
+        metadata={
+            "help": "The backend to use for distributed training. Must be one of 'nccl', 'mpi', 'xccl', 'gloo', 'hccl'.",
+            "choices": ["nccl", "gloo", "mpi", "xccl", "hccl", "cncl", "mccl"],
+        },
+    )
+    ddp_timeout: int = field(
+        default=1800,
+        metadata={"help": "The timeout for `torch.distributed.init_process_group` calls (in seconds)."},
+    )
 
-    neftune_noise_alpha: float | None = field(
+    # --- FSDP ---
+    fsdp: list[FSDPOption] | str | None = field(
+        default=None,
+        metadata={
+            "help": "Enable PyTorch FSDP for distributed training. Options: 'full_shard', 'shard_grad_op', 'hybrid_shard', 'hybrid_shard_zero2', 'offload', 'auto_wrap'.",
+        },
+    )
+    fsdp_config: dict[str, Any] | str | None = field(
         default=None,
         metadata={
-            "help": "Activates neftune noise embeddings into the model. NEFTune has been proven to drastically improve model performances for instruction fine-tuning. Check out the original paper here: https://huggingface.co/papers/2310.05914 and the original code here: https://github.com/neelsjain/NEFTune. Only supported for `PreTrainedModel` and `PeftModel` classes."
+            "help": (
+                "Config to be used with FSDP (Pytorch Fully Sharded Data Parallel). The value is either a "
+                "fsdp json config file (e.g., `fsdp_config.json`) or an already loaded json file as `dict`."
+            )
         },
     )
 
-    optim_target_modules: None | str | list[str] = field(
+    # --- DeepSpeed ---
+    deepspeed: dict | str | None = field(
         default=None,
+        metadata={"help": "Enable DeepSpeed integration. Value is a path to a JSON config file or a dict."},
+    )
+
+    # --- Debugging ---
+    debug: str | list[DebugOption] = field(
+        default="",
+        metadata={
+            "help": "Enable one or more debug features. Options: 'underflow_overflow' (detect overflow in model I/O), 'tpu_metrics_debug' (print TPU metrics)."
+        },
+    )
+    skip_memory_metrics: bool = field(
+        default=True,
         metadata={
-            "help": "Target modules for the optimizer defined in the `optim` argument. Only used for the GaLore optimizer at the moment."
+            "help": "Whether to skip adding memory profiler reports to metrics. Skipped by default because it slows down training."
         },
     )
 
-    batch_eval_metrics: bool = field(
+    # --- External Script Flags ---
+    do_train: bool = field(
         default=False,
-        metadata={"help": "Break eval metrics calculation into batches to save memory."},
+        metadata={
+            "help": "Whether to run training. Not directly used by Trainer; intended for training/evaluation scripts."
+        },
     )
-
-    eval_on_start: bool = field(
+    do_eval: bool = field(
         default=False,
         metadata={
-            "help": "Whether to run through the entire `evaluation` step at the very beginning of training as a sanity check."
+            "help": "Whether to run evaluation. Not directly used by Trainer; intended for training/evaluation scripts."
         },
     )
-
-    use_liger_kernel: bool = field(
+    do_predict: bool = field(
         default=False,
-        metadata={"help": "Whether or not to enable the Liger Kernel for model training."},
+        metadata={
+            "help": "Whether to run predictions on the test set. Not directly used by Trainer; intended for training/evaluation scripts."
+        },
     )
-
-    liger_kernel_config: dict[str, bool] | None = field(
+    resume_from_checkpoint: str | None = field(
         default=None,
         metadata={
-            "help": (
-                "Configuration to be used for Liger Kernel. When use_liger_kernel=True, "
-                "this dict is passed as keyword arguments to the `_apply_liger_kernel_to_instance` function, "
-                "which specifies which kernels to apply. Available options vary by model "
-                "but typically include: 'rope', 'swiglu', 'cross_entropy', 'fused_linear_cross_entropy', "
-                "'rms_norm', etc. If None, use the default kernel configurations."
-            )
+            "help": "Path to a folder with a valid checkpoint for your model. Not directly used by Trainer; intended for training/evaluation scripts."
         },
     )
 
-    eval_use_gather_object: bool = field(
-        default=False,
+    # --- Deprecated / Internal ---
+    warmup_ratio: float | None = field(
+        default=None,
         metadata={
-            "help": "Whether to run recursively gather object in a nested list/tuple/dictionary of objects from all devices."
+            "help": "This argument is deprecated and will be removed in v5.2. Use `warmup_steps` instead as it also works with float values."
         },
     )
-
-    average_tokens_across_devices: bool = field(
-        default=True,
+    logging_dir: str | None = field(
+        default=None,
         metadata={
-            "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to "
-            "synchronize num_tokens_in_batch for precise loss calculation. Reference: "
-            "https://github.com/huggingface/transformers/issues/34242"
+            "help": "Deprecated and will be removed in v5.2. Set env var `TENSORBOARD_LOGGING_DIR` instead. TensorBoard log directory."
         },
     )
-
-    use_cache: bool = field(
-        default=False,
+    local_rank: int = field(
+        default=-1,
         metadata={
-            "help": "Whether or not to use cache for the model For training, this is usually not needed apart from some PEFT methods that uses `past_key_values`."
+            "help": "When using torch.distributed.launch (Deprecated), it will pass `local_rank` in the script, so we need this for the parser. To get the local rank, prefer using the property `local_process_index`"
         },
     )
 
     def __post_init__(self):
-        # Set default output_dir if not provided
+        # ── 1. Defaults & Normalization ──
         if self.output_dir is None:
             self.output_dir = "trainer_output"
             logger.info(
                 "No output directory specified, defaulting to 'trainer_output'. "
                 "To change this behavior, specify --output_dir when creating TrainingArguments."
             )
 
-        # Parse in args that could be `dict` sent in from the CLI as a string
+        # Parse JSON string dict args from CLI (e.g., '{"key": "value"}').
+        # Only parses strings starting with '{'; other strings are treated as file paths.
         for valid_field in self._VALID_DICT_FIELDS:
             passed_value = getattr(self, valid_field)
-            # We only want to do this if the str starts with a bracket to indicate a `dict`
-            # else its likely a filename if supported
             if isinstance(passed_value, str) and passed_value.startswith("{"):
                 loaded_dict = json.loads(passed_value)
-                # Convert str values to types if applicable
                 loaded_dict = _convert_str_dict(loaded_dict)
                 setattr(self, valid_field, loaded_dict)
 
-        # expand paths, if not os.makedirs("~/bar") will make directory
-        # in the current directory instead of the actual home
-        # see https://github.com/huggingface/transformers/issues/10628
+        # Expand ~ in paths so os.makedirs works correctly (#10628)
         if self.output_dir is not None:
             self.output_dir = os.path.expanduser(self.output_dir)
 
         if self.disable_tqdm is None:
             self.disable_tqdm = logger.getEffectiveLevel() > logging.WARN
 
+        if self.warmup_ratio is not None:
+            logger.warning("warmup_ratio is deprecated and will be removed in v5.2. Use `warmup_steps` instead.")
+            self.warmup_steps = self.warmup_ratio
+
+        if self.logging_dir is not None:
+            logger.warning(
+                "`logging_dir` is deprecated and will be removed in v5.2. Please set `TENSORBOARD_LOGGING_DIR` instead."
+            )
+
+        if isinstance(self.include_num_input_tokens_seen, bool):
+            self.include_num_input_tokens_seen = "all" if self.include_num_input_tokens_seen else "no"
+
+        # ── 2. Enum / Type Conversions ──
         self.eval_strategy = IntervalStrategy(self.eval_strategy)
         self.logging_strategy = IntervalStrategy(self.logging_strategy)
         self.save_strategy = SaveStrategy(self.save_strategy)
         self.hub_strategy = HubStrategy(self.hub_strategy)
-
         self.lr_scheduler_type = SchedulerType(self.lr_scheduler_type)
+        self.optim = OptimizerNames(self.optim)
+
+        if isinstance(self.debug, str):
+            self.debug = [DebugOption(s) for s in self.debug.split()]
+        elif self.debug is None:
+            self.debug = []
+
+        # ── 3. Auto-derived Values ──
         if self.do_eval is False and self.eval_strategy != IntervalStrategy.NO:
             self.do_eval = True
 
-        if self.torch_empty_cache_steps is not None:
-            if not (isinstance(self.torch_empty_cache_steps, int) and self.torch_empty_cache_steps > 0):
-                raise ValueError(
-                    f"`torch_empty_cache_steps` must be an integer bigger than 0, got {self.torch_empty_cache_steps}."
-                )
-
-        # eval_steps has to be defined and non-zero, falls back to logging_steps if the latter is non-zero
+        # Fall back to logging_steps if eval_steps is unset
         if self.eval_strategy == IntervalStrategy.STEPS and (self.eval_steps is None or self.eval_steps == 0):
             if self.logging_steps > 0:
                 logger.info(f"using `logging_steps` to initialize `eval_steps` to {self.logging_steps}")
@@ -1471,89 +1499,30 @@ def __post_init__(self):
                     " --logging_steps"
                 )
 
-        # logging_steps must be non-zero for logging_strategy that is other than 'no'
-        if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps == 0:
-            raise ValueError(f"logging strategy {self.logging_strategy} requires non-zero --logging_steps")
-
-        if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps > 1:
-            if self.logging_steps != int(self.logging_steps):
-                raise ValueError(f"--logging_steps must be an integer if bigger than 1: {self.logging_steps}")
-            self.logging_steps = int(self.logging_steps)
-        if self.eval_strategy == IntervalStrategy.STEPS and self.eval_steps > 1:
-            if self.eval_steps != int(self.eval_steps):
-                raise ValueError(f"--eval_steps must be an integer if bigger than 1: {self.eval_steps}")
-            self.eval_steps = int(self.eval_steps)
-        if self.save_strategy == SaveStrategy.STEPS and self.save_steps > 1:
-            if self.save_steps != int(self.save_steps):
-                raise ValueError(f"--save_steps must be an integer if bigger than 1: {self.save_steps}")
-            self.save_steps = int(self.save_steps)
-
-        # Sanity checks for load_best_model_at_end: we require save and eval strategies to be compatible.
-        if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST:
-            if self.eval_strategy != self.save_strategy:
-                raise ValueError(
-                    "--load_best_model_at_end requires the save and eval strategy to match, but found\n- Evaluation "
-                    f"strategy: {self.eval_strategy}\n- Save strategy: {self.save_strategy}"
-                )
-            if self.eval_strategy == IntervalStrategy.STEPS and self.save_steps % self.eval_steps != 0:
-                if self.eval_steps < 1 or self.save_steps < 1:
-                    if not (self.eval_steps < 1 and self.save_steps < 1):
-                        raise ValueError(
-                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
-                            "steps, which cannot get guaranteed when mixing ratio and absolute steps for save_steps "
-                            f"{self.save_steps} and eval_steps {self.eval_steps}."
-                        )
-                    # Work around floating point precision issues
-                    LARGE_MULTIPLIER = 1_000_000
-                    if (self.save_steps * LARGE_MULTIPLIER) % (self.eval_steps * LARGE_MULTIPLIER) != 0:
-                        raise ValueError(
-                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
-                            f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}."
-                        )
-                else:
-                    raise ValueError(
-                        "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation "
-                        f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}."
-                    )
-
         if (
             self.load_best_model_at_end or self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU
         ) and self.metric_for_best_model is None:
             self.metric_for_best_model = "loss"
         if self.greater_is_better is None and self.metric_for_best_model is not None:
             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 (
-                    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")
-
-        if self.fp16_full_eval and self.bf16_full_eval:
-            raise ValueError("At most one of fp16 and bf16 can be True for full eval, but not both")
 
-        if self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU:
-            if self.eval_strategy == IntervalStrategy.NO:
-                raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires an eval strategy")
-            if not is_torch_available():
-                raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires torch>=0.2.0")
+        if self.report_to == "none" or self.report_to == ["none"]:
+            self.report_to = []
+        elif not isinstance(self.report_to, list):
+            self.report_to = [self.report_to]
 
-        self.optim = OptimizerNames(self.optim)
+        # ── 4. Validation ──
+        self._validate_args()
 
-        # We need to get from the env as we are setting deepspeed mixed precison afterwards
+        # ── 5. Mixed Precision ──
+        # Read from env first; DeepSpeed may override this later
         self.mixed_precision = os.environ.get("ACCELERATE_MIXED_PRECISION", "no")
         if self.fp16:
             self.mixed_precision = "fp16"
         elif self.bf16:
             self.mixed_precision = "bf16"
 
+        # ── 6. Torch Compile ──
         if (self.torch_compile_mode is not None or self.torch_compile_backend is not None) and not self.torch_compile:
             self.torch_compile = True
         if self.torch_compile and self.torch_compile_backend is None:
@@ -1563,21 +1532,20 @@ def __post_init__(self):
                 self.torch_compile_backend = "inductor"
 
         if self.torch_compile:
-            # TODO: remove this once we've bumped the minimum accelerate version
+            # TODO: remove env var fallback once minimum accelerate >= 1.2.0
             if not is_accelerate_available("1.2.0"):
                 os.environ["ACCELERATE_DYNAMO_BACKEND"] = self.torch_compile_backend
                 if self.torch_compile_mode is not None:
                     os.environ["ACCELERATE_DYNAMO_MODE"] = self.torch_compile_mode
 
-        # We need to setup the accelerator config here *before* the first call to `self.device`
+        # ── 7. Accelerator Config (must come before self.device) ──
         if is_accelerate_available():
             if not isinstance(self.accelerator_config, AcceleratorConfig):
                 if self.accelerator_config is None:
                     self.accelerator_config = AcceleratorConfig()
                 elif isinstance(self.accelerator_config, dict):
                     self.accelerator_config = AcceleratorConfig(**self.accelerator_config)
-                # Check that a user didn't pass in the class instantiator
-                # such as `accelerator_config = AcceleratorConfig`
+                # Reject uninstantiated class (e.g. AcceleratorConfig instead of AcceleratorConfig())
                 elif isinstance(self.accelerator_config, type):
                     raise NotImplementedError(
                         "Tried passing in a callable to `accelerator_config`, but this is not supported. "
@@ -1591,10 +1559,11 @@ def __post_init__(self):
                     "Batches will be split across all processes equally when using `split_batches=True`."
                 )
 
-        # Initialize device before we proceed
+        # ── 8. Device Init ──
         if is_torch_available():
             self.device
 
+        # ── 9. TF32 ──
         if is_torch_available() and self.torch_compile:
             if is_torch_tf32_available():
                 if self.tf32 is None and not self.fp16 or self.bf16:
@@ -1617,83 +1586,115 @@ def __post_init__(self):
             else:
                 if is_torch_tf32_available():
                     enable_tf32(False)
-                # no need to assert on else
-
-        if self.report_to == "all" or self.report_to == ["all"]:
-            # Import at runtime to avoid a circular import.
-            from .integrations import get_available_reporting_integrations
-
-            self.report_to = get_available_reporting_integrations()
-
-            if "codecarbon" in self.report_to and torch.version.hip:
-                logger.warning(
-                    "When using the Trainer, CodeCarbonCallback requires the `codecarbon` package, which is not compatible with AMD ROCm (https://github.com/mlco2/codecarbon/pull/490). Automatically disabling the codecarbon callback. Reference: https://huggingface.co/docs/transformers/v4.39.3/en/main_classes/trainer#transformers.TrainingArguments.report_to."
-                )
-                self.report_to.remove("codecarbon")
-
-        elif self.report_to == "none" or self.report_to == ["none"]:
-            self.report_to = []
-        elif not isinstance(self.report_to, list):
-            self.report_to = [self.report_to]
+                # TF32 not available, nothing to disable
 
-        if self.warmup_ratio is not None:
-            logger.warning("warmup_ratio is deprecated and will be removed in v5.2. Use `warmup_steps` instead.")
-            self.warmup_steps = self.warmup_ratio
-
-        if self.warmup_steps < 0:
-            raise ValueError("warmup_steps must be an integer or a float")
-
-        if isinstance(self.debug, str):
-            self.debug = [DebugOption(s) for s in self.debug.split()]
-        elif self.debug is None:
-            self.debug = []
+        # ── 10. Hardware Overrides ──
+        if self.use_cpu:
+            self.dataloader_pin_memory = False
 
-        self.fsdp_plugin_args = None
-        # we can't save the plugin due to a pickle issue, so we do not initialize the plugin here
+        # ── 11. FSDP ──
+        # Store args only (not the plugin itself) to avoid pickle issues
         self.fsdp_plugin_args = self._process_fsdp_args()
 
+        # ── 12. DeepSpeed (must be last) ──
         self.deepspeed_plugin = None
         if self.deepspeed:
-            # - must be run very last in arg parsing, since it will use a lot of these settings.
-            # - must be run before the model is created.
-            if not is_accelerate_available():
-                raise ValueError(
-                    f"--deepspeed requires Accelerate to be installed: `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`."
-                )
             from transformers.integrations.deepspeed import HfTrainerDeepSpeedConfig
 
-            # will be used later by the Trainer
-            # note: leave self.deepspeed unmodified in case a user relies on it not to be modified)
+            # Leave self.deepspeed unmodified; users may rely on the original value
             self.hf_deepspeed_config = HfTrainerDeepSpeedConfig(self.deepspeed)
             self.hf_deepspeed_config.trainer_config_process(self)
 
-            # Accelerate DeepSpeed Plugin
             from accelerate.utils import DeepSpeedPlugin
 
             self.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.hf_deepspeed_config)
         elif strtobool(os.environ.get("ACCELERATE_USE_DEEPSPEED", "false")):
-            # Accelerate DeepSpeed Plugin
             from accelerate.utils import DeepSpeedPlugin
 
             self.deepspeed_plugin = DeepSpeedPlugin()
             self.deepspeed_plugin.set_mixed_precision(self.mixed_precision)
             self.deepspeed_plugin.set_deepspeed_weakref()
 
-        if self.use_cpu:
-            self.dataloader_pin_memory = False
+    def _validate_args(self):
+        """Validate argument combinations and value constraints."""
+        if self.torch_empty_cache_steps is not None:
+            if not (isinstance(self.torch_empty_cache_steps, int) and self.torch_empty_cache_steps > 0):
+                raise ValueError(
+                    f"`torch_empty_cache_steps` must be an integer bigger than 0, got {self.torch_empty_cache_steps}."
+                )
+
+        # logging_steps must be non-zero when logging_strategy="steps"
+        if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps == 0:
+            raise ValueError(f"logging strategy {self.logging_strategy} requires non-zero --logging_steps")
+
+        if self.logging_strategy == IntervalStrategy.STEPS and self.logging_steps > 1:
+            if self.logging_steps != int(self.logging_steps):
+                raise ValueError(f"--logging_steps must be an integer if bigger than 1: {self.logging_steps}")
+            self.logging_steps = int(self.logging_steps)
+        if self.eval_strategy == IntervalStrategy.STEPS and self.eval_steps > 1:
+            if self.eval_steps != int(self.eval_steps):
+                raise ValueError(f"--eval_steps must be an integer if bigger than 1: {self.eval_steps}")
+            self.eval_steps = int(self.eval_steps)
+        if self.save_strategy == SaveStrategy.STEPS and self.save_steps > 1:
+            if self.save_steps != int(self.save_steps):
+                raise ValueError(f"--save_steps must be an integer if bigger than 1: {self.save_steps}")
+            self.save_steps = int(self.save_steps)
+
+        # load_best_model_at_end requires compatible save and eval strategies
+        if self.load_best_model_at_end and self.save_strategy != SaveStrategy.BEST:
+            if self.eval_strategy != self.save_strategy:
+                raise ValueError(
+                    "--load_best_model_at_end requires the save and eval strategy to match, but found\n- Evaluation "
+                    f"strategy: {self.eval_strategy}\n- Save strategy: {self.save_strategy}"
+                )
+            if self.eval_strategy == IntervalStrategy.STEPS and self.save_steps % self.eval_steps != 0:
+                if self.eval_steps < 1 or self.save_steps < 1:
+                    if not (self.eval_steps < 1 and self.save_steps < 1):
+                        raise ValueError(
+                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
+                            "steps, which cannot get guaranteed when mixing ratio and absolute steps for save_steps "
+                            f"{self.save_steps} and eval_steps {self.eval_steps}."
+                        )
+                    # Use integer arithmetic to avoid floating point precision issues
+                    LARGE_MULTIPLIER = 1_000_000
+                    if (self.save_steps * LARGE_MULTIPLIER) % (self.eval_steps * LARGE_MULTIPLIER) != 0:
+                        raise ValueError(
+                            "--load_best_model_at_end requires the saving steps to be a multiple of the evaluation "
+                            f"steps, but found {self.save_steps}, which is not a multiple of {self.eval_steps}."
+                        )
+                else:
+                    raise ValueError(
+                        "--load_best_model_at_end requires the saving steps to be a round multiple of the evaluation "
+                        f"steps, but found {self.save_steps}, which is not a round multiple of {self.eval_steps}."
+                    )
+
+        if is_torch_available():
+            if self.bf16 or self.bf16_full_eval:
+                if not self.use_cpu and not is_torch_bf16_gpu_available() and not is_torch_xla_available():
+                    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."
+                    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")
+
+        if self.fp16_full_eval and self.bf16_full_eval:
+            raise ValueError("At most one of fp16 and bf16 can be True for full eval, but not both")
+
+        if self.lr_scheduler_type == SchedulerType.REDUCE_ON_PLATEAU:
+            if self.eval_strategy == IntervalStrategy.NO:
+                raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires an eval strategy")
+            if not is_torch_available():
+                raise ValueError("lr_scheduler_type reduce_lr_on_plateau requires torch>=0.2.0")
+
+        if self.warmup_steps < 0:
+            raise ValueError("warmup_steps must be an integer or a float")
 
         if self.dataloader_num_workers == 0 and self.dataloader_prefetch_factor is not None:
             raise ValueError(
                 "--dataloader_prefetch_factor can only be set when data is loaded in a different process, i.e."
-                " when --dataloader_num_workers > 1."
-            )
-
-        if isinstance(self.include_num_input_tokens_seen, bool):
-            self.include_num_input_tokens_seen = "all" if self.include_num_input_tokens_seen else "no"
-
-        if self.logging_dir is not None:
-            logger.warning(
-                "`logging_dir` is deprecated and will be removed in v5.2. Please set `TENSORBOARD_LOGGING_DIR` instead."
+                " when --dataloader_num_workers > 0."
             )
 
     def __str__(self):
@@ -1739,7 +1740,7 @@ def _setup_devices(self) -> "torch.device":
                     f"Using the `Trainer` with `PyTorch` requires `accelerate>={ACCELERATE_MIN_VERSION}`: "
                     f"Please run `pip install transformers[torch]` or `pip install 'accelerate>={ACCELERATE_MIN_VERSION}'`"
                 )
-        # We delay the init of `PartialState` to the end for clarity
+        # Build kwargs for PartialState; actual init happens below
         accelerator_state_kwargs: dict[str, Any] = {"enabled": True, "use_configured_state": False}
         if isinstance(self.accelerator_config, AcceleratorConfig):
             accelerator_state_kwargs["use_configured_state"] = self.accelerator_config.pop(
@@ -1751,7 +1752,6 @@ def _setup_devices(self) -> "torch.device":
                     "Passing `'use_configured_state':True` to the AcceleratorConfig requires a pre-configured "
                     "`AcceleratorState` or `PartialState` to be defined before calling `TrainingArguments`. "
                 )
-            # We rely on `PartialState` to yell if there's issues here (which it will)
             self.distributed_state = PartialState(cpu=self.use_cpu)
             if self.deepspeed and self.distributed_state.distributed_type != DistributedType.DEEPSPEED:
                 raise RuntimeError(
@@ -1781,11 +1781,11 @@ def _setup_devices(self) -> "torch.device":
             accelerator_state_kwargs["backend"] = self.ddp_backend
             accelerator_state_kwargs["timeout"] = timedelta(seconds=self.ddp_timeout)
 
-        # Now we pop everything
+        # Initialize PartialState with the accumulated kwargs
         if accelerator_state_kwargs.pop("enabled", False) and not accelerator_state_kwargs.pop(
             "use_configured_state", False
         ):
-            # We need to patch this env var when enabling to detect deepspeed
+            # Temporarily set env var so Accelerate detects DeepSpeed
             use_deepspeed = accelerator_state_kwargs.pop("use_deepspeed", False)
             if use_deepspeed:
                 os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
@@ -1803,8 +1803,7 @@ def _setup_devices(self) -> "torch.device":
             device = self.distributed_state.device
             self._n_gpu = 0
         elif is_sagemaker_dp_enabled() or is_sagemaker_mp_enabled():
-            # Already set _n_gpu
-            pass
+            pass  # _n_gpu already set above
         elif self.distributed_state.distributed_type == DistributedType.NO:
             if self.use_cpu:
                 device = torch.device("cpu")
@@ -1826,17 +1825,11 @@ def _setup_devices(self) -> "torch.device":
                 device = torch.device("hpu:0")
                 torch.hpu.set_device(device)
             else:
-                # if n_gpu is > 1 we'll use nn.DataParallel.
-                # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
-                # Explicitly set CUDA to the first (index 0) CUDA device, otherwise `set_device` will
-                # trigger an error that a device index is missing. Index 0 takes into account the
-                # GPUs available in the environment, so `CUDA_VISIBLE_DEVICES=1,2` with `cuda:0`
-                # will use the first GPU in that env, i.e. GPU#1
+                # Default to cuda:0 (respects CUDA_VISIBLE_DEVICES); nn.DataParallel handles n_gpu > 1
                 device = torch.device(
                     "cuda:0" if torch.cuda.is_available() else os.environ.get("ACCELERATE_TORCH_DEVICE", "cpu")
                 )
-                # Sometimes the line in the postinit has not been run before we end up here, so just checking we're not at
-                # the default value.
+                # _n_gpu may not have been set yet if _setup_devices is called early
                 self._n_gpu = torch.cuda.device_count()
                 if device.type == "cuda":
                     torch.cuda.set_device(device)
@@ -1860,7 +1853,7 @@ def n_gpu(self):
             training. For distributed training, it will always be 1.
         """
         requires_backends(self, ["torch"])
-        # Make sure `self._n_gpu` is properly setup.
+        # Ensure _setup_devices has been called
         if not hasattr(self, "_n_gpu"):
             _ = self._setup_devices
         return self._n_gpu
@@ -2068,7 +2061,7 @@ def to_dict(self):
         Serializes this instance while replace `Enum` by their values (for JSON serialization support). It obfuscates
         the token values by removing their value.
         """
-        # filter out fields that are defined as field(init=False)
+        # Exclude non-init fields (they aren't user-facing config)
         d = {field.name: getattr(self, field.name) for field in fields(self) if field.init}
 
         for k, v in d.items():
@@ -2078,10 +2071,10 @@ def to_dict(self):
                 d[k] = [x.value for x in v]
             if k.endswith("_token"):
                 d[k] = f"<{k.upper()}>"
-            # Handle the accelerator_config if passed
+            # Serialize AcceleratorConfig to dict
             if is_accelerate_available() and isinstance(v, AcceleratorConfig):
                 d[k] = v.to_dict()
-            # Handle the quantization_config if passed
+            # Serialize quantization_config if nested inside model_init_kwargs
             if k == "model_init_kwargs" and isinstance(v, dict) and "quantization_config" in v:
                 quantization_config = v.get("quantization_config")
                 if quantization_config and not isinstance(quantization_config, dict):
@@ -2112,7 +2105,7 @@ def to_sanitized_dict(self) -> dict[str, Any]:
 
         return {k: v if type(v) in valid_types else str(v) for k, v in d.items()}
 
-    # The following methods are there to simplify the instantiation of `TrainingArguments`
+    # Convenience setters for grouped configuration
     def set_training(
         self,
         learning_rate: float = 5e-5,
@@ -2686,7 +2679,7 @@ def _process_fsdp_args(self):
 
         self.fsdp_config["min_num_params"] = self.fsdp_config.get("min_num_params", 0)
 
-        # if fsdp_config["transformer_layer_cls_to_wrap"] is specified as a string, convert it to a list with a single object
+        # Normalize transformer_layer_cls_to_wrap from string to list
         if isinstance(self.fsdp_config.get("transformer_layer_cls_to_wrap", None), str):
             self.fsdp_config["transformer_layer_cls_to_wrap"] = [self.fsdp_config["transformer_layer_cls_to_wrap"]]
 
@@ -2707,10 +2700,9 @@ def _process_fsdp_args(self):
         self.fsdp_config["xla_fsdp_grad_ckpt"] = self.fsdp_config.get("xla_fsdp_grad_ckpt", False)
         if self.fsdp_config["xla"]:
             if len(self.fsdp) > 0:
-                # store XLA fsdp configuration parameters into a dictionary
-                # Copy the config to avoid modifying the original config (which may be used for JSON serialization)
+                # Copy to avoid mutating the original (needed for JSON serialization)
                 self.xla_fsdp_config = self.fsdp_config.get("xla_fsdp_settings", {}).copy()
-                # apply appropriate string to torch.dtype conversions for parameters
+                # Convert string dtype names to torch.dtype
                 if "compute_dtype" in self.xla_fsdp_config:
                     self.xla_fsdp_config["compute_dtype"] = getattr(torch, self.xla_fsdp_config["compute_dtype"])
                 if "buffer_dtype" in self.xla_fsdp_config:
@@ -2721,7 +2713,7 @@ def _process_fsdp_args(self):
             if self.fsdp_config["xla_fsdp_grad_ckpt"]:
                 warnings.warn("`--xla_fsdp_grad_ckpt` is useful only when `--xla` is set to true.")
 
-        # accelerate integration for FSDP
+        # Build kwargs for Accelerate's FSDPPlugin
         fsdp_plugin_args = None
         if len(self.fsdp) > 0 and not self.fsdp_config["xla"]:
             from accelerate.utils.constants import (
@@ -2766,11 +2758,10 @@ def _process_fsdp_args(self):
             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()
             if sync_module_states == "false" and cpu_ram_efficient_loading == "true":
-                # In this case, all the processes except the main process would have random weights leading
-                # to unexpected behaviour during training, thus throwing error here to prevent it.
+                # Without sync, non-main processes would have random weights
                 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
+            # Set env var to suppress Accelerate warning and for transformers to read
             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
 

filename: tests/test_training_args.py
patch: @@ -1,67 +0,0 @@
-import os
-import tempfile
-import unittest
-
-from transformers import TrainingArguments
-
-
-class TestTrainingArguments(unittest.TestCase):
-    def test_default_output_dir(self):
-        """Test that output_dir defaults to 'trainer_output' when not specified."""
-        args = TrainingArguments(output_dir=None)
-        self.assertEqual(args.output_dir, "trainer_output")
-
-    def test_custom_output_dir(self):
-        """Test that output_dir is respected when specified."""
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            args = TrainingArguments(output_dir=tmp_dir)
-            self.assertEqual(args.output_dir, tmp_dir)
-
-    def test_output_dir_creation(self):
-        """Test that output_dir is created only when needed."""
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            output_dir = os.path.join(tmp_dir, "test_output")
-
-            # Directory should not exist before creating args
-            self.assertFalse(os.path.exists(output_dir))
-
-            # Create args with save_strategy="no" - should not create directory
-            args = TrainingArguments(
-                output_dir=output_dir,
-                do_train=True,
-                save_strategy="no",
-                report_to=None,
-            )
-            self.assertFalse(os.path.exists(output_dir))
-
-            # Now set save_strategy="steps" - should create directory when needed
-            args.save_strategy = "steps"
-            args.save_steps = 1
-            self.assertFalse(os.path.exists(output_dir))  # Still shouldn't exist
-
-            # Directory should be created when actually needed (e.g. in Trainer)
-
-    def test_torch_empty_cache_steps_requirements(self):
-        """Test that torch_empty_cache_steps is a positive integer or None."""
-
-        # None is acceptable (feature is disabled):
-        args = TrainingArguments(torch_empty_cache_steps=None)
-        self.assertIsNone(args.torch_empty_cache_steps)
-
-        # non-int is unacceptable:
-        with self.assertRaises(ValueError):
-            TrainingArguments(torch_empty_cache_steps=1.0)
-        with self.assertRaises(ValueError):
-            TrainingArguments(torch_empty_cache_steps="none")
-
-        # negative int is unacceptable:
-        with self.assertRaises(ValueError):
-            TrainingArguments(torch_empty_cache_steps=-1)
-
-        # zero is unacceptable:
-        with self.assertRaises(ValueError):
-            TrainingArguments(torch_empty_cache_steps=0)
-
-        # positive int is acceptable:
-        args = TrainingArguments(torch_empty_cache_steps=1)
-        self.assertEqual(args.torch_empty_cache_steps, 1)

filename: tests/trainer/test_training_args.py
patch: @@ -0,0 +1,352 @@
+import os
+import tempfile
+import unittest
+
+from transformers import TrainingArguments
+from transformers.debug_utils import DebugOption
+from transformers.trainer_utils import HubStrategy, IntervalStrategy, SaveStrategy, SchedulerType
+from transformers.training_args import OptimizerNames
+
+
+class TestTrainingArguments(unittest.TestCase):
+    def test_default_output_dir(self):
+        """Test that output_dir defaults to 'trainer_output' when not specified."""
+        args = TrainingArguments(output_dir=None)
+        self.assertEqual(args.output_dir, "trainer_output")
+
+    def test_custom_output_dir(self):
+        """Test that output_dir is respected when specified."""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            args = TrainingArguments(output_dir=tmp_dir)
+            self.assertEqual(args.output_dir, tmp_dir)
+
+    def test_output_dir_creation(self):
+        """Test that output_dir is created only when needed."""
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            output_dir = os.path.join(tmp_dir, "test_output")
+
+            # Directory should not exist before creating args
+            self.assertFalse(os.path.exists(output_dir))
+
+            # Create args with save_strategy="no" - should not create directory
+            args = TrainingArguments(
+                output_dir=output_dir,
+                do_train=True,
+                save_strategy="no",
+                report_to=None,
+            )
+            self.assertFalse(os.path.exists(output_dir))
+
+            # Now set save_strategy="steps" - should create directory when needed
+            args.save_strategy = "steps"
+            args.save_steps = 1
+            self.assertFalse(os.path.exists(output_dir))  # Still shouldn't exist
+
+            # Directory should be created when actually needed (e.g. in Trainer)
+
+    def test_torch_empty_cache_steps_requirements(self):
+        """Test that torch_empty_cache_steps is a positive integer or None."""
+
+        # None is acceptable (feature is disabled):
+        args = TrainingArguments(torch_empty_cache_steps=None)
+        self.assertIsNone(args.torch_empty_cache_steps)
+
+        # non-int is unacceptable:
+        with self.assertRaises(ValueError):
+            TrainingArguments(torch_empty_cache_steps=1.0)
+        with self.assertRaises(ValueError):
+            TrainingArguments(torch_empty_cache_steps="none")
+
+        # negative int is unacceptable:
+        with self.assertRaises(ValueError):
+            TrainingArguments(torch_empty_cache_steps=-1)
+
+        # zero is unacceptable:
+        with self.assertRaises(ValueError):
+            TrainingArguments(torch_empty_cache_steps=0)
+
+        # positive int is acceptable:
+        args = TrainingArguments(torch_empty_cache_steps=1)
+        self.assertEqual(args.torch_empty_cache_steps, 1)
+
+    def test_output_dir_expands_user(self):
+        """Test that ~ in output_dir is expanded to the user's home directory."""
+        args = TrainingArguments(output_dir="~/foo", report_to=None)
+        self.assertEqual(args.output_dir, os.path.expanduser("~/foo"))
+
+    def test_enum_coercions(self):
+        """Test that string values are correctly converted to their enum types."""
+        args = TrainingArguments(
+            output_dir="tmp",
+            eval_strategy="steps",
+            eval_steps=10,
+            logging_strategy="steps",
+            save_strategy="epoch",
+            hub_strategy="end",
+            lr_scheduler_type="linear",
+            optim="adamw_torch",
+            report_to=None,
+        )
+        self.assertEqual(args.eval_strategy, IntervalStrategy.STEPS)
+        self.assertEqual(args.logging_strategy, IntervalStrategy.STEPS)
+        self.assertEqual(args.save_strategy, SaveStrategy.EPOCH)
+        self.assertEqual(args.hub_strategy, HubStrategy.END)
+        self.assertEqual(args.lr_scheduler_type, SchedulerType.LINEAR)
+        self.assertEqual(args.optim, OptimizerNames.ADAMW_TORCH)
+
+        # Invalid string should raise ValueError
+        with self.assertRaises(ValueError):
+            TrainingArguments(output_dir="tmp", eval_strategy="invalid_strategy", report_to=None)
+
+    def test_do_eval_auto_enabled(self):
+        """Test that do_eval is automatically set to True when eval_strategy is not 'no'."""
+        args = TrainingArguments(
+            output_dir="tmp",
+            do_eval=False,
+            eval_strategy="steps",
+            eval_steps=10,
+            report_to=None,
+        )
+        self.assertTrue(args.do_eval)
+
+    def test_eval_steps_fallback_to_logging_steps(self):
+        """Test that eval_steps falls back to logging_steps when not specified."""
+        args = TrainingArguments(
+            output_dir="tmp",
+            eval_strategy="steps",
+            logging_steps=10,
+            report_to=None,
+        )
+        self.assertEqual(args.eval_steps, 10)
+
+    def test_eval_steps_required_when_strategy_steps(self):
+        """Test that eval_strategy='steps' with logging_steps=0 raises ValueError."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                eval_strategy="steps",
+                logging_steps=0,
+                report_to=None,
+            )
+
+    def test_logging_steps_required_nonzero(self):
+        """Test that logging_strategy='steps' with logging_steps=0 raises ValueError."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                logging_strategy="steps",
+                logging_steps=0,
+                report_to=None,
+            )
+
+    def test_steps_must_be_integer_when_greater_than_one(self):
+        """Test that fractional steps >1 raise ValueError, but <=1 are allowed."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                logging_strategy="steps",
+                logging_steps=10.5,
+                report_to=None,
+            )
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                eval_strategy="steps",
+                eval_steps=10.5,
+                report_to=None,
+            )
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                save_strategy="steps",
+                save_steps=10.5,
+                report_to=None,
+            )
+        # Fractional values <=1 (ratios) are allowed
+        args = TrainingArguments(
+            output_dir="tmp",
+            logging_strategy="steps",
+            logging_steps=0.5,
+            report_to=None,
+        )
+        self.assertEqual(args.logging_steps, 0.5)
+
+    def test_load_best_model_requires_matching_strategies(self):
+        """Test load_best_model_at_end validation for strategy and step compatibility."""
+        # Mismatched eval/save strategy should raise
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                load_best_model_at_end=True,
+                eval_strategy="steps",
+                eval_steps=10,
+                save_strategy="epoch",
+                report_to=None,
+            )
+
+        # save_steps not a multiple of eval_steps should raise
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                load_best_model_at_end=True,
+                eval_strategy="steps",
+                eval_steps=10,
+                save_strategy="steps",
+                save_steps=15,
+                report_to=None,
+            )
+
+        # Valid: matching strategies with compatible steps should not raise
+        args = TrainingArguments(
+            output_dir="tmp",
+            load_best_model_at_end=True,
+            eval_strategy="steps",
+            eval_steps=10,
+            save_strategy="steps",
+            save_steps=20,
+            report_to=None,
+        )
+        self.assertTrue(args.load_best_model_at_end)
+
+    def test_metric_for_best_model_defaults(self):
+        """Test default metric_for_best_model and greater_is_better behavior."""
+        # load_best_model_at_end with no metric → defaults to "loss"
+        args = TrainingArguments(
+            output_dir="tmp",
+            load_best_model_at_end=True,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            report_to=None,
+        )
+        self.assertEqual(args.metric_for_best_model, "loss")
+        self.assertFalse(args.greater_is_better)
+
+        # metric ending in "loss" → greater_is_better is False
+        args = TrainingArguments(
+            output_dir="tmp",
+            load_best_model_at_end=True,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            metric_for_best_model="eval_loss",
+            report_to=None,
+        )
+        self.assertFalse(args.greater_is_better)
+
+        # metric not ending in "loss" → greater_is_better is True
+        args = TrainingArguments(
+            output_dir="tmp",
+            load_best_model_at_end=True,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            metric_for_best_model="accuracy",
+            report_to=None,
+        )
+        self.assertTrue(args.greater_is_better)
+
+    def test_fp16_bf16_mutual_exclusivity(self):
+        """Test that fp16 and bf16 cannot both be True."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(output_dir="tmp", fp16=True, bf16=True, report_to=None)
+        with self.assertRaises(ValueError):
+            TrainingArguments(output_dir="tmp", fp16_full_eval=True, bf16_full_eval=True, report_to=None)
+
+    def test_reduce_on_plateau_requires_eval(self):
+        """Test that reduce_lr_on_plateau scheduler requires an eval strategy."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                lr_scheduler_type="reduce_lr_on_plateau",
+                eval_strategy="no",
+                report_to=None,
+            )
+
+    def test_torch_compile_auto_enable(self):
+        """Test that torch_compile is auto-enabled when mode or backend is set."""
+        args = TrainingArguments(
+            output_dir="tmp",
+            torch_compile_mode="reduce-overhead",
+            report_to=None,
+        )
+        self.assertTrue(args.torch_compile)
+
+        args = TrainingArguments(
+            output_dir="tmp",
+            torch_compile_backend="inductor",
+            report_to=None,
+        )
+        self.assertTrue(args.torch_compile)
+
+        # Default backend when torch_compile=True
+        args = TrainingArguments(
+            output_dir="tmp",
+            torch_compile=True,
+            report_to=None,
+        )
+        self.assertEqual(args.torch_compile_backend, "inductor")
+
+    def test_report_to_none_handling(self):
+        """Test report_to normalization for 'none' and string values."""
+        args = TrainingArguments(output_dir="tmp", report_to="none")
+        self.assertEqual(args.report_to, [])
+
+        args = TrainingArguments(output_dir="tmp", report_to=["none"])
+        self.assertEqual(args.report_to, [])
+
+        args = TrainingArguments(output_dir="tmp", report_to="tensorboard")
+        self.assertEqual(args.report_to, ["tensorboard"])
+
+    def test_warmup_steps_validation(self):
+        """Test warmup_steps validation for negative values."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(output_dir="tmp", warmup_steps=-1, report_to=None)
+
+        # Zero and fractional values are valid
+        args = TrainingArguments(output_dir="tmp", warmup_steps=0, report_to=None)
+        self.assertEqual(args.warmup_steps, 0)
+
+        args = TrainingArguments(output_dir="tmp", warmup_steps=0.5, report_to=None)
+        self.assertEqual(args.warmup_steps, 0.5)
+
+    def test_debug_option_parsing(self):
+        """Test debug string parsing into DebugOption enum list."""
+        args = TrainingArguments(output_dir="tmp", debug="underflow_overflow", report_to=None)
+        self.assertEqual(args.debug, [DebugOption.UNDERFLOW_OVERFLOW])
+
+        args = TrainingArguments(output_dir="tmp", debug=None, report_to=None)
+        self.assertEqual(args.debug, [])
+
+    def test_dataloader_prefetch_requires_workers(self):
+        """Test that dataloader_prefetch_factor requires num_workers > 0."""
+        with self.assertRaises(ValueError):
+            TrainingArguments(
+                output_dir="tmp",
+                dataloader_prefetch_factor=2,
+                dataloader_num_workers=0,
+                report_to=None,
+            )
+        # Valid: prefetch with workers > 0
+        args = TrainingArguments(
+            output_dir="tmp",
+            dataloader_prefetch_factor=2,
+            dataloader_num_workers=2,
+            report_to=None,
+        )
+        self.assertEqual(args.dataloader_prefetch_factor, 2)
+
+    def test_use_cpu_disables_pin_memory(self):
+        """Test that use_cpu=True disables dataloader_pin_memory."""
+        args = TrainingArguments(output_dir="tmp", use_cpu=True, report_to=None)
+        self.assertFalse(args.dataloader_pin_memory)
+
+    def test_include_num_input_tokens_seen_coercion(self):
+        """Test bool-to-string coercion for include_num_input_tokens_seen."""
+        args = TrainingArguments(output_dir="tmp", include_num_input_tokens_seen=True, report_to=None)
+        self.assertEqual(args.include_num_input_tokens_seen, "all")
+
+        args = TrainingArguments(output_dir="tmp", include_num_input_tokens_seen=False, report_to=None)
+        self.assertEqual(args.include_num_input_tokens_seen, "no")
+
+    def test_dict_field_parsing(self):
+        """Test that JSON string dict fields are parsed into dicts."""
+        args = TrainingArguments(output_dir="tmp", lr_scheduler_kwargs='{"factor": 0.5}', report_to=None)
+        self.assertEqual(args.lr_scheduler_kwargs, {"factor": 0.5})