kaira.training.Trainer

Inheritance diagram for Trainer

class kaira.training.Trainer(model, args: TrainingArguments | TrainingArguments | DictConfig | dict, **kwargs)

Bases: Trainer

Unified trainer for all communication models.

This trainer automatically adapts to different model types and supports multiple configuration systems for training arguments: - Hugging Face TrainingArguments - Kaira TrainingArguments - Hydra DictConfig - Plain Python dictionaries

Models are responsible for their own configuration, channel simulation, constraints, and domain-specific logic via their config systems.

The trainer focuses on training mechanics. All domain-specific metrics and loss functions should be provided by the user via the compute_metrics and loss function parameters.

Methods

__init__	Initialize trainer.
add_callback	Add a callback to the current list of [~transformers.TrainerCallback].
autocast_smart_context_manager	A helper wrapper that creates an appropriate context manager for autocast while feeding it the desired arguments, depending on the situation.
call_model_init
compare_trainer_and_checkpoint_args
compute_loss	How the loss is computed by Trainer.
compute_loss_context_manager	A helper wrapper to group together context managers.
create_accelerator_and_postprocess
create_model_card	Creates a draft of a model card using the information available to the Trainer.
create_optimizer	Setup the optimizer.
create_optimizer_and_scheduler	Setup the optimizer and the learning rate scheduler.
create_scheduler	Setup the scheduler.
evaluate	Run evaluation and returns metrics.
evaluation_loop	Prediction/evaluation loop, shared by Trainer.evaluate() and Trainer.predict().
floating_point_ops	For models that inherit from [PreTrainedModel], uses that method to compute the number of floating point operations for every backward + forward pass.
from_hydra_config	Create trainer from Hydra configuration.
from_training_args	Create trainer from Hugging Face TrainingArguments.
get_batch_samples	Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.
get_decay_parameter_names	Get all parameter names that weight decay will be applied to.
get_eval_dataloader	Returns the evaluation [~torch.utils.data.DataLoader].
get_learning_rates	Returns the learning rate of each parameter from self.optimizer.
get_num_trainable_parameters	Get the number of trainable parameters.
get_optimizer_cls_and_kwargs	Returns the optimizer class and optimizer parameters based on the training arguments.
get_optimizer_group	Returns optimizer group for a parameter if given, else returns all optimizer groups for params.
get_test_dataloader	Returns the test [~torch.utils.data.DataLoader].
get_total_train_batch_size	Calculates total batch size (micro_batch * grad_accum * dp_world_size).
get_tp_size	Get the tensor parallel size from either the model or DeepSpeed config.
get_train_dataloader	Returns the training [~torch.utils.data.DataLoader].
hyperparameter_search	Launch an hyperparameter search using optuna or Ray Tune or SigOpt.
init_hf_repo	Initializes a git repo in self.args.hub_model_id.
is_local_process_zero	Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several machines) main process.
is_world_process_zero	Whether or not this process is the global main process (when training in a distributed fashion on several machines, this is only going to be True for one process).
log	Log logs on the various objects watching training.
log_metrics	Log metrics in a specially formatted way.
metrics_format	Reformat Trainer metrics values to a human-readable format.
num_examples	Helper to get number of samples in a [~torch.utils.data.DataLoader] by accessing its dataset.
num_tokens	Helper to get number of tokens in a [~torch.utils.data.DataLoader] by enumerating dataloader.
pop_callback	Remove a callback from the current list of [~transformers.TrainerCallback] and returns it.
predict	Run prediction and returns predictions and potential metrics.
prediction_loop	Prediction/evaluation loop, shared by Trainer.evaluate() and Trainer.predict().
prediction_step	Perform an evaluation step on model using inputs.
propagate_args_to_deepspeed	Sets values in the deepspeed plugin based on the Trainer args
remove_callback	Remove a callback from the current list of [~transformers.TrainerCallback].
save_metrics	Save metrics into a json file for that split, e.g. train_results.json.
save_model	Save model and optionally upload to Hub.
save_state	Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model.
set_initial_training_values	Calculates and returns the following values: - num_train_epochs - num_update_steps_per_epoch - num_examples - num_train_samples - epoch_based - len_dataloader - max_steps
store_flos
torch_jit_model_eval
train	Main training entry point.
training_step	Perform a training step on a batch of inputs.

Attributes

tokenizer

Examples using kaira.training.Trainer

Original DeepJSCC Model (Bourtsoulatze 2019) with Training

Original DeepJSCC Model (Bourtsoulatze 2019) with Training

Deep Joint Source-Channel Coding (DeepJSCC) Model - Bourtsoulatze2019 Implementation

Deep Joint Source-Channel Coding (DeepJSCC) Model - Bourtsoulatze2019 Implementation

__init__(model, args: TrainingArguments | TrainingArguments | DictConfig | dict, **kwargs)

Initialize trainer.

Parameters:

• model – BaseModel instance to train (handles domain-specific logic internally)

• args – Training arguments (TrainingArguments, HFTrainingArguments, DictConfig, or dict)

• **kwargs – Additional arguments for base Trainer

save_model(output_dir: str | None = None, _internal_call: bool = False)

Save model and optionally upload to Hub.

Parameters:

• output_dir – Optional output directory

• _internal_call – Internal call flag

classmethod from_hydra_config(hydra_cfg: DictConfig, model, **kwargs)

Create trainer from Hydra configuration.

classmethod from_training_args(training_args: TrainingArguments, model, **kwargs)

Create trainer from Hugging Face TrainingArguments.

add_callback(callback)

Add a callback to the current list of [~transformers.TrainerCallback].

Parameters:

callback (type or [~transformers.TrainerCallback]) – A [~transformers.TrainerCallback] class or an instance of a [~transformers.TrainerCallback]. In the first case, will instantiate a member of that class.

autocast_smart_context_manager(cache_enabled: bool | None = True)

A helper wrapper that creates an appropriate context manager for autocast while feeding it the desired arguments, depending on the situation.

call_model_init(trial=None)

compare_trainer_and_checkpoint_args(training_args, trainer_state)

compute_loss(model: Module, inputs: dict[str, Tensor | Any], return_outputs: bool = False, num_items_in_batch: Tensor | None = None)

How the loss is computed by Trainer. By default, all models return the loss in the first element.

Parameters:

• model (nn.Module) – The model to compute the loss for.

• inputs (dict[str, Union[torch.Tensor, Any]]) – The input data for the model.

• return_outputs (bool, optional, defaults to False) – Whether to return the model outputs along with the loss.

• num_items_in_batch (Optional[torch.Tensor], optional) – The number of items in the batch. If num_items_in_batch is not passed,

Returns:

The loss of the model along with its output if return_outputs was set to True

Subclass and override for custom behavior. If you are not using num_items_in_batch when computing your loss, make sure to overwrite self.model_accepts_loss_kwargs to False. Otherwise, the loss calculationg might be slightly inacurate when performing gradient accumulation.

compute_loss_context_manager()

A helper wrapper to group together context managers.

create_accelerator_and_postprocess()

create_model_card(language: str | None = None, license: str | None = None, tags: str | list[str] | None = None, model_name: str | None = None, finetuned_from: str | None = None, tasks: str | list[str] | None = None, dataset_tags: str | list[str] | None = None, dataset: str | list[str] | None = None, dataset_args: str | list[str] | None = None)

Creates a draft of a model card using the information available to the Trainer.

Parameters:

• language (str, optional) – The language of the model (if applicable)

• license (str, optional) – The license of the model. Will default to the license of the pretrained model used, if the original model given to the Trainer comes from a repo on the Hub.

• tags (str or list[str], optional) – Some tags to be included in the metadata of the model card.

• model_name (str, optional) – The name of the model.

• finetuned_from (str, optional) – The name of the model used to fine-tune this one (if applicable). Will default to the name of the repo of the original model given to the Trainer (if it comes from the Hub).

• tasks (str or list[str], optional) – One or several task identifiers, to be included in the metadata of the model card.

• dataset_tags (str or list[str], optional) – One or several dataset tags, to be included in the metadata of the model card.

• dataset (str or list[str], optional) – One or several dataset identifiers, to be included in the metadata of the model card.

• dataset_args (str or list[str], optional) – One or several dataset arguments, to be included in the metadata of the model card.

create_optimizer()

Setup the optimizer.

We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the Trainer’s init through optimizers, or subclass and override this method in a subclass.

create_optimizer_and_scheduler(num_training_steps: int)

Setup the optimizer and the learning rate scheduler.

We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the Trainer’s init through optimizers, or subclass and override this method (or create_optimizer and/or create_scheduler) in a subclass.

create_scheduler(num_training_steps: int, optimizer: Optimizer = None)

Setup the scheduler. The optimizer of the trainer must have been set up either before this method is called or passed as an argument.

Parameters:

num_training_steps (int) – The number of training steps to do.

evaluate(eval_dataset: Dataset | dict[str, Dataset] | None = None, ignore_keys: list[str] | None = None, metric_key_prefix: str = 'eval') → dict[str, float]

Run evaluation and returns metrics.

The calling script will be responsible for providing a method to compute metrics, as they are task-dependent (pass it to the init compute_metrics argument).

You can also subclass and override this method to inject custom behavior.

Parameters:

• eval_dataset (Union[Dataset, dict[str, Dataset]), optional) –

  Pass a dataset if you wish to override self.eval_dataset. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. If it is a dictionary, it will evaluate on each dataset, prepending the dictionary key to the metric name. Datasets must implement the __len__ method.

  <Tip>

  If you pass a dictionary with names of datasets as keys and datasets as values, evaluate will run separate evaluations on each dataset. This can be useful to monitor how training affects other datasets or simply to get a more fine-grained evaluation. When used with load_best_model_at_end, make sure metric_for_best_model references exactly one of the datasets. If you, for example, pass in {“data1”: data1, “data2”: data2} for two datasets data1 and data2, you could specify metric_for_best_model=”eval_data1_loss” for using the loss on data1 and metric_for_best_model=”eval_data2_loss” for the loss on data2.

  </Tip>

• ignore_keys (list[str], optional) – A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions.

• metric_key_prefix (str, optional, defaults to “eval”) – An optional prefix to be used as the metrics key prefix. For example the metrics “bleu” will be named “eval_bleu” if the prefix is “eval” (default)

Returns:

A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The dictionary also contains the epoch number which comes from the training state.

evaluation_loop(dataloader: DataLoader, description: str, prediction_loss_only: bool | None = None, ignore_keys: list[str] | None = None, metric_key_prefix: str = 'eval') → EvalLoopOutput

Prediction/evaluation loop, shared by Trainer.evaluate() and Trainer.predict().

Works both with or without labels.

floating_point_ops(inputs: dict[str, Tensor | Any])

For models that inherit from [PreTrainedModel], uses that method to compute the number of floating point operations for every backward + forward pass. If using another model, either implement such a method in the model or subclass and override this method.

Parameters:

inputs (dict[str, Union[torch.Tensor, Any]]) – The inputs and targets of the model.

Returns:

The number of floating-point operations.

Return type:

int

get_batch_samples(epoch_iterator: Iterator, num_batches: int, device: device) → tuple[list, Tensor | None]

Collects a specified number of batches from the epoch iterator and optionally counts the number of items in the batches to properly scale the loss.

get_decay_parameter_names(model) → list[str]

Get all parameter names that weight decay will be applied to.

This function filters out parameters in two ways: 1. By layer type (instances of layers specified in ALL_LAYERNORM_LAYERS) 2. By parameter name patterns (containing ‘bias’, or variation of ‘norm’)

get_eval_dataloader(eval_dataset: str | Dataset | None = None) → DataLoader

Returns the evaluation [~torch.utils.data.DataLoader].

Subclass and override this method if you want to inject some custom behavior.

Parameters:

eval_dataset (str or torch.utils.data.Dataset, optional) – If a str, will use self.eval_dataset[eval_dataset] as the evaluation dataset. If a Dataset, will override self.eval_dataset and must implement __len__. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed.

get_learning_rates()

Returns the learning rate of each parameter from self.optimizer.

get_num_trainable_parameters()

Get the number of trainable parameters.

static get_optimizer_cls_and_kwargs(args: TrainingArguments, model: PreTrainedModel | None = None) → tuple[Any, Any]

Returns the optimizer class and optimizer parameters based on the training arguments.

Parameters:

args (transformers.training_args.TrainingArguments) – The training arguments for the training session.

get_optimizer_group(param: str | Parameter | None = None)

Returns optimizer group for a parameter if given, else returns all optimizer groups for params.

Parameters:

param (str or torch.nn.parameter.Parameter, optional) – The parameter for which optimizer group needs to be returned.

get_test_dataloader(test_dataset: Dataset) → DataLoader

Returns the test [~torch.utils.data.DataLoader].

Subclass and override this method if you want to inject some custom behavior.

Parameters:

test_dataset (torch.utils.data.Dataset, optional) – The test dataset to use. If it is a [~datasets.Dataset], columns not accepted by the model.forward() method are automatically removed. It must implement __len__.

get_total_train_batch_size(args) → int

Calculates total batch size (micro_batch * grad_accum * dp_world_size).

Note: Only considers DP and TP (dp_world_size = world_size // tp_size).

get_tp_size() → int

Get the tensor parallel size from either the model or DeepSpeed config.

get_train_dataloader() → DataLoader

Returns the training [~torch.utils.data.DataLoader].

Will use no sampler if train_dataset does not implement __len__, a random sampler (adapted to distributed training if necessary) otherwise.

Subclass and override this method if you want to inject some custom behavior.

hyperparameter_search(hp_space: Callable[[optuna.Trial], dict[str, float]] | None = None, compute_objective: Callable[[dict[str, float]], float] | None = None, n_trials: int = 20, direction: str | list[str] = 'minimize', backend: str | HPSearchBackend | None = None, hp_name: Callable[[optuna.Trial], str] | None = None, **kwargs) → BestRun | list[BestRun]

Launch an hyperparameter search using optuna or Ray Tune or SigOpt. The optimized quantity is determined by compute_objective, which defaults to a function returning the evaluation loss when no metric is provided, the sum of all metrics otherwise.

<Tip warning={true}>

To use this method, you need to have provided a model_init when initializing your [Trainer]: we need to reinitialize the model at each new run. This is incompatible with the optimizers argument, so you need to subclass [Trainer] and override the method [~Trainer.create_optimizer_and_scheduler] for custom optimizer/scheduler.

</Tip>

Parameters:

• hp_space (Callable[[“optuna.Trial”], dict[str, float]], optional) – A function that defines the hyperparameter search space. Will default to [~trainer_utils.default_hp_space_optuna] or [~trainer_utils.default_hp_space_ray] or [~trainer_utils.default_hp_space_sigopt] depending on your backend.

• compute_objective (Callable[[dict[str, float]], float], optional) – A function computing the objective to minimize or maximize from the metrics returned by the evaluate method. Will default to [~trainer_utils.default_compute_objective].

• n_trials (int, optional, defaults to 100) – The number of trial runs to test.

• direction (str or list[str], optional, defaults to “minimize”) – If it’s single objective optimization, direction is str, can be “minimize” or “maximize”, you should pick “minimize” when optimizing the validation loss, “maximize” when optimizing one or several metrics. If it’s multi objectives optimization, direction is list[str], can be List of “minimize” and “maximize”, you should pick “minimize” when optimizing the validation loss, “maximize” when optimizing one or several metrics.

• backend (str or [~training_utils.HPSearchBackend], optional) – The backend to use for hyperparameter search. Will default to optuna or Ray Tune or SigOpt, depending on which one is installed. If all are installed, will default to optuna.

• hp_name (Callable[[“optuna.Trial”], str]], optional) – A function that defines the trial/run name. Will default to None.

• kwargs (dict[str, Any], optional) –

  Additional keyword arguments for each backend:

  • optuna: parameters from [optuna.study.create_study](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.create_study.html) and also the parameters timeout, n_jobs and gc_after_trial from [optuna.study.Study.optimize](https://optuna.readthedocs.io/en/stable/reference/generated/optuna.study.Study.html#optuna.study.Study.optimize)

  • ray: parameters from [tune.run](https://docs.ray.io/en/latest/tune/api_docs/execution.html#tune-run). If resources_per_trial is not set in the kwargs, it defaults to 1 CPU core and 1 GPU (if available). If progress_reporter is not set in the kwargs, [ray.tune.CLIReporter](https://docs.ray.io/en/latest/tune/api/doc/ray.tune.CLIReporter.html) is used.

  • sigopt: the parameter proxies from [sigopt.Connection.set_proxies](https://docs.sigopt.com/support/faq#how-do-i-use-sigopt-with-a-proxy).

Returns:

All the information about the best run or best runs for multi-objective optimization. Experiment summary can be found in run_summary attribute for Ray backend.

Return type:

[trainer_utils.BestRun or list[trainer_utils.BestRun]]

init_hf_repo(token: str | None = None)

Initializes a git repo in self.args.hub_model_id.

is_local_process_zero() → bool

Whether or not this process is the local (e.g., on one machine if training in a distributed fashion on several machines) main process.

is_world_process_zero() → bool

Whether or not this process is the global main process (when training in a distributed fashion on several machines, this is only going to be True for one process).

log(logs: dict[str, float], start_time: float | None = None) → None

Log logs on the various objects watching training.

Subclass and override this method to inject custom behavior.

Parameters:

• logs (dict[str, float]) – The values to log.

• start_time (Optional[float]) – The start of training.

log_metrics(split, metrics)

Log metrics in a specially formatted way.

Under distributed environment this is done only for a process with rank 0.

Parameters:

• split (str) – Mode/split name: one of train, eval, test

• metrics (dict[str, float]) – The metrics returned from train/evaluate/predictmetrics: metrics dict

Notes on memory reports:

In order to get memory usage report you need to install psutil. You can do that with pip install psutil.

Now when this method is run, you will see a report that will include:

` init_mem_cpu_alloc_delta   =     1301MB init_mem_cpu_peaked_delta  =      154MB init_mem_gpu_alloc_delta   =      230MB init_mem_gpu_peaked_delta  =        0MB train_mem_cpu_alloc_delta  =     1345MB train_mem_cpu_peaked_delta =        0MB train_mem_gpu_alloc_delta  =      693MB train_mem_gpu_peaked_delta =        7MB `

Understanding the reports:

• the first segment, e.g., train__, tells you which stage the metrics are for. Reports starting with init_

  will be added to the first stage that gets run. So that if only evaluation is run, the memory usage for the __init__ will be reported along with the eval_ metrics.

• the third segment, is either cpu or gpu, tells you whether it’s the general RAM or the gpu0 memory

  metric.

• *_alloc_delta - is the difference in the used/allocated memory counter between the end and the start of the

  stage - it can be negative if a function released more memory than it allocated.

• *_peaked_delta - is any extra memory that was consumed and then freed - relative to the current allocated

  memory counter - it is never negative. When you look at the metrics of any stage you add up alloc_delta + peaked_delta and you know how much memory was needed to complete that stage.

The reporting happens only for process of rank 0 and gpu 0 (if there is a gpu). Typically this is enough since the main process does the bulk of work, but it could be not quite so if model parallel is used and then other GPUs may use a different amount of gpu memory. This is also not the same under DataParallel where gpu0 may require much more memory than the rest since it stores the gradient and optimizer states for all participating GPUs. Perhaps in the future these reports will evolve to measure those too.

The CPU RAM metric measures RSS (Resident Set Size) includes both the memory which is unique to the process and the memory shared with other processes. It is important to note that it does not include swapped out memory, so the reports could be imprecise.

The CPU peak memory is measured using a sampling thread. Due to python’s GIL it may miss some of the peak memory if that thread didn’t get a chance to run when the highest memory was used. Therefore this report can be less than reality. Using tracemalloc would have reported the exact peak memory, but it doesn’t report memory allocations outside of python. So if some C++ CUDA extension allocated its own memory it won’t be reported. And therefore it was dropped in favor of the memory sampling approach, which reads the current process memory usage.

The GPU allocated and peak memory reporting is done with torch.cuda.memory_allocated() and torch.cuda.max_memory_allocated(). This metric reports only “deltas” for pytorch-specific allocations, as torch.cuda memory management system doesn’t track any memory allocated outside of pytorch. For example, the very first cuda call typically loads CUDA kernels, which may take from 0.5 to 2GB of GPU memory.

Note that this tracker doesn’t account for memory allocations outside of [Trainer]’s __init__, train, evaluate and predict calls.

Because evaluation calls may happen during train, we can’t handle nested invocations because torch.cuda.max_memory_allocated is a single counter, so if it gets reset by a nested eval call, train’s tracker will report incorrect info. If this [pytorch issue](https://github.com/pytorch/pytorch/issues/16266) gets resolved it will be possible to change this class to be re-entrant. Until then we will only track the outer level of train, evaluate and predict methods. Which means that if eval is called during train, it’s the latter that will account for its memory usage and that of the former.

This also means that if any other tool that is used along the [Trainer] calls torch.cuda.reset_peak_memory_stats, the gpu peak memory stats could be invalid. And the [Trainer] will disrupt the normal behavior of any such tools that rely on calling torch.cuda.reset_peak_memory_stats themselves.

For best performance you may want to consider turning the memory profiling off for production runs.

metrics_format(metrics: dict[str, float]) → dict[str, float]

Reformat Trainer metrics values to a human-readable format.

Parameters:

metrics (dict[str, float]) – The metrics returned from train/evaluate/predict

Returns:

The reformatted metrics

Return type:

metrics (dict[str, float])

num_examples(dataloader: DataLoader) → int

Helper to get number of samples in a [~torch.utils.data.DataLoader] by accessing its dataset. When dataloader.dataset does not exist or has no length, estimates as best it can

static num_tokens(train_dl: DataLoader, max_steps: int | None = None) → int

Helper to get number of tokens in a [~torch.utils.data.DataLoader] by enumerating dataloader.

pop_callback(callback)

Remove a callback from the current list of [~transformers.TrainerCallback] and returns it.

If the callback is not found, returns None (and no error is raised).

Parameters:

callback (type or [~transformers.TrainerCallback]) – A [~transformers.TrainerCallback] class or an instance of a [~transformers.TrainerCallback]. In the first case, will pop the first member of that class found in the list of callbacks.

Returns:

The callback removed, if found.

Return type:

[~transformers.TrainerCallback]

predict(test_dataset: Dataset, ignore_keys: list[str] | None = None, metric_key_prefix: str = 'test') → PredictionOutput

Run prediction and returns predictions and potential metrics.

Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method will also return metrics, like in evaluate().

Parameters:

• test_dataset (Dataset) – Dataset to run the predictions on. If it is an datasets.Dataset, columns not accepted by the model.forward() method are automatically removed. Has to implement the method __len__

• ignore_keys (list[str], optional) – A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions.

• metric_key_prefix (str, optional, defaults to “test”) – An optional prefix to be used as the metrics key prefix. For example the metrics “bleu” will be named “test_bleu” if the prefix is “test” (default)

<Tip>

If your predictions or labels have different sequence length (for instance because you’re doing dynamic padding in a token classification task) the predictions will be padded (on the right) to allow for concatenation into one array. The padding index is -100.

</Tip>

Returns: NamedTuple A namedtuple with the following keys:

• predictions (np.ndarray): The predictions on test_dataset.

• label_ids (np.ndarray, optional): The labels (if the dataset contained some).

• metrics (dict[str, float], optional): The potential dictionary of metrics (if the dataset contained labels).

prediction_loop(dataloader: DataLoader, description: str, prediction_loss_only: bool | None = None, ignore_keys: list[str] | None = None, metric_key_prefix: str = 'eval') → EvalLoopOutput

Prediction/evaluation loop, shared by Trainer.evaluate() and Trainer.predict().

Works both with or without labels.

prediction_step(model: Module, inputs: dict[str, Tensor | Any], prediction_loss_only: bool, ignore_keys: list[str] | None = None) → tuple[Tensor | None, Tensor | None, Tensor | None]

Perform an evaluation step on model using inputs.

Subclass and override to inject custom behavior.

Parameters:

• model (nn.Module) – The model to evaluate.

• inputs (dict[str, Union[torch.Tensor, Any]]) –

  The inputs and targets of the model.

  The dictionary will be unpacked before being fed to the model. Most models expect the targets under the argument labels. Check your model’s documentation for all accepted arguments.

• prediction_loss_only (bool) – Whether or not to return the loss only.

• ignore_keys (list[str], optional) – A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions.

Returns:

A tuple with the loss, logits and labels (each being optional).

Return type:

tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]

propagate_args_to_deepspeed(auto_find_batch_size=False)

Sets values in the deepspeed plugin based on the Trainer args

remove_callback(callback)

Remove a callback from the current list of [~transformers.TrainerCallback].

Parameters:

callback (type or [~transformers.TrainerCallback]) – A [~transformers.TrainerCallback] class or an instance of a [~transformers.TrainerCallback]. In the first case, will remove the first member of that class found in the list of callbacks.

save_metrics(split, metrics, combined=True)

Save metrics into a json file for that split, e.g. train_results.json.

Under distributed environment this is done only for a process with rank 0.

Parameters:

• split (str) – Mode/split name: one of train, eval, test, all

• metrics (dict[str, float]) – The metrics returned from train/evaluate/predict

• combined (bool, optional, defaults to True) – Creates combined metrics by updating all_results.json with metrics of this call

To understand the metrics please read the docstring of [~Trainer.log_metrics]. The only difference is that raw unformatted numbers are saved in the current method.

save_state()

Saves the Trainer state, since Trainer.save_model saves only the tokenizer with the model.

Under distributed environment this is done only for a process with rank 0.

set_initial_training_values(args: TrainingArguments, dataloader: DataLoader, total_train_batch_size: int)

Calculates and returns the following values: - num_train_epochs - num_update_steps_per_epoch - num_examples - num_train_samples - epoch_based - len_dataloader - max_steps

store_flos()

property tokenizer: PreTrainedTokenizerBase | None

torch_jit_model_eval(model, dataloader, training=False)

train(resume_from_checkpoint: bool | str | None = None, trial: optuna.Trial | dict[str, Any] | None = None, ignore_keys_for_eval: list[str] | None = None, **kwargs)

Main training entry point.

Parameters:

• resume_from_checkpoint (str or bool, optional) – If a str, local path to a saved checkpoint as saved by a previous instance of [Trainer]. If a bool and equals True, load the last checkpoint in args.output_dir as saved by a previous instance of [Trainer]. If present, training will resume from the model/optimizer/scheduler states loaded here.

• trial (optuna.Trial or dict[str, Any], optional) – The trial run or the hyperparameter dictionary for hyperparameter search.

• ignore_keys_for_eval (list[str], optional) – A list of keys in the output of your model (if it is a dictionary) that should be ignored when gathering predictions for evaluation during the training.

• kwargs (dict[str, Any], optional) – Additional keyword arguments used to hide deprecated arguments

training_step(model: Module, inputs: dict[str, Tensor | Any], num_items_in_batch: Tensor | None = None) → Tensor

Perform a training step on a batch of inputs.

Subclass and override to inject custom behavior.

Parameters:

• model (nn.Module) – The model to train.

• inputs (dict[str, Union[torch.Tensor, Any]]) –

  The inputs and targets of the model.

  The dictionary will be unpacked before being fed to the model. Most models expect the targets under the argument labels. Check your model’s documentation for all accepted arguments.

Returns:

The tensor with training loss on this batch.

Return type:

torch.Tensor