Trainer

Purpose

Trainer is the all-in-one PyTorch training loop. It supports mixed precision, gradient accumulation, gradient checkpointing, distributed training (DDP, FSDP, DeepSpeed), torch.compile, callbacks, hyperparameter search, evaluation, prediction, and checkpoint resumption. It is one of the largest classes in the library.

Key abstractions

Class / function	File	Role
Trainer	src/transformers/trainer.py (4,418 LOC)	Main training loop
Seq2SeqTrainer	src/transformers/trainer_seq2seq.py	Adds predict_with_generate
TrainingArguments	src/transformers/training_args.py (2,868 LOC)	Every knob
Seq2SeqTrainingArguments	src/transformers/training_args_seq2seq.py	Adds generation flags
TrainerCallback	src/transformers/trainer_callback.py	Hook into the loop
TrainerState, TrainerControl	same	State and control flow
TrainerOptimizerOps	src/transformers/trainer_optimizer.py	Optimizer / scheduler creation
nested_*, LabelSmoother, EvalPrediction, EvalLoopOutput	src/transformers/trainer_utils.py (50K LOC)	Helpers
Distributed helpers	src/transformers/trainer_pt_utils.py (57K LOC)	Sampler, sharded loading, accelerator detect
JIT checkpoint helpers	src/transformers/trainer_jit_checkpoint.py	torch.distributed checkpointing

High-level shape

graph TD
    Init[Trainer.__init__] --> Setup[Resolve TrainingArguments + accelerate.Accelerator]
    Setup --> Build[Wrap model, optimizer, scheduler]
    Build --> Loop[Trainer.train]
    Loop --> Step[for each step]
    Step --> FwdBwd[forward + backward + clip + step]
    Step --> Eval{Evaluate?}
    Eval -->|yes| EvalLoop[Trainer.evaluate]
    EvalLoop --> Save{Save checkpoint?}
    Save -->|yes| Ckpt[trainer._save_checkpoint]
    Step --> Stop{Stop?}
    Stop -->|no| Step
    Stop -->|yes| End[Final eval + save]

TrainingArguments

TrainingArguments is a dataclass with ~250 fields. The major groups:

• Output and logging: output_dir, logging_dir, logging_steps, report_to (tensorboard, wandb, mlflow, comet_ml, clearml, dvclive, swanlab).
• Optimization: learning_rate, weight_decay, adam_*, optim (defaults, fused, paged, 8bit, lion, adafactor, schedule*free, …), lr_scheduler_type, warmup*\*.
• Mixed precision: fp16, bf16, tf32, amp_backend.
• Distributed: ddp_backend, ddp_find_unused_parameters, ddp_bucket_cap_mb, fsdp, fsdp_config, accelerator_config, deepspeed, tp_size, pp_size.
• Eval and saving: eval_strategy, save_strategy, save_steps, save_total_limit, load_best_model_at_end, metric_for_best_model, greater_is_better.
• Memory: gradient_checkpointing, gradient_accumulation_steps, optim, auto_find_batch_size.
• Generation (Seq2Seq): predict_with_generate, generation_max_length, generation_num_beams, generation_config.
• Hyperparameter search: hp_search_*, tune_*.
• Hub integration: push_to_hub, hub_model_id, hub_strategy.

Distributed training

Trainer uses accelerate.Accelerator (configured via AcceleratorConfig from src/transformers/integrations/accelerate.py, 41K LOC) for single-GPU, multi-GPU DDP, multi-node, FSDP, and DeepSpeed. The Trainer does not manage distributed state directly; it delegates to accelerate.

• DDP is the default for multi-GPU.
• FSDP activates via --fsdp and the fsdp_config knobs. The integration uses accelerate's FSDP wrapper.
• DeepSpeed activates via --deepspeed <config.json>. Implementation in src/transformers/integrations/deepspeed.py (33K LOC).
• Tensor parallel activates via tp_size > 1 and model = ... tp_plan="auto". See Tensor parallelism.

Callbacks

TrainerCallback is the extension point. Common built-ins (in src/transformers/trainer_callback.py):

• DefaultFlowCallback — orchestrates step / eval / save flow.
• ProgressCallback — tqdm progress.
• PrinterCallback — log to stdout.
• EarlyStoppingCallback — stop when metric plateaus.

Third-party integrations (W&B, MLFlow, Tensorboard, Comet, ClearML, DVCLive, SwanLab, AzureML, Neptune) live in src/transformers/integrations/integration_utils.py (118K LOC). They register themselves automatically when their package is installed.

Evaluation

trainer.evaluate(eval_dataset) runs the eval loop. For seq2seq tasks, Seq2SeqTrainer.predict_with_generate=True swaps the forward call for generate, useful for BLEU/ROUGE evaluation.

trainer.predict(test_dataset) returns predictions, label_ids, metrics.

Hyperparameter search

trainer.hyperparameter_search(...) integrates with Optuna, Ray Tune, SigOpt, W&B Sweeps. Backed by src/transformers/hyperparameter_search.py and src/transformers/utils/hp_naming.py.

Hub integration

push_to_hub=True uploads the final model + tokenizer + training args + metrics card to the Hub on trainer.train() completion. Behaviour configured by hub_strategy (end, every_save, checkpoint, all_checkpoints).

Resumption

trainer.train(resume_from_checkpoint=True)            # latest in output_dir
trainer.train(resume_from_checkpoint="path/to/ckpt")  # specific

The state stored alongside weights includes optimizer, scheduler, RNG, scaler, callback state, and trainer_state.json (carried by TrainerState).

SageMaker

tests/sagemaker/ and src/transformers/training_args.py declare SageMaker-specific knobs (SageMakerTrainingArguments). Useful when training on AWS SageMaker.

Integration points

• Modeling — Trainer.model is a PreTrainedModel.
• Data — collators turn lists of examples into batches.
• Quantization — quantized models are loaded via from_pretrained and trained with PEFT adapters (LoRA).
• Integrations — accelerate / deepspeed / peft / FSDP / TP plug in here.

Entry points for modification

• New optimizer → register in Trainer.create_optimizer and OptimizerNames in src/transformers/training_args.py.
• New callback → subclass TrainerCallback and pass via Trainer(callbacks=[...]).
• Custom training step → subclass Trainer and override training_step, compute_loss, or prediction_step.
• Custom evaluation → subclass Trainer and override evaluation_loop.