# Neuron TRL Trainers

[TRL](https://huggingface.co/docs/trl/en/index)-compatible trainers for AWS Trainium accelerators.

## NeuronSFTTrainer

### NeuronSFTConfig[[optimum.neuron.NeuronSFTConfig]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.NeuronSFTConfig</name><anchor>optimum.neuron.NeuronSFTConfig</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1020/optimum/neuron/trainers/sft_config.py#L34</source><parameters>[{"name": "output_dir", "val": ": str | None = None"}, {"name": "overwrite_output_dir", "val": ": bool = False"}, {"name": "do_train", "val": ": bool = False"}, {"name": "do_eval", "val": ": bool = False"}, {"name": "eval_strategy", "val": ": transformers.trainer_utils.IntervalStrategy | str = 'no'"}, {"name": "per_device_train_batch_size", "val": ": int = 1"}, {"name": "per_device_eval_batch_size", "val": ": int = 1"}, {"name": "gradient_accumulation_steps", "val": ": int = 1"}, {"name": "learning_rate", "val": ": float = 5e-05"}, {"name": "weight_decay", "val": ": float = 0.0"}, {"name": "adam_beta1", "val": ": float = 0.9"}, {"name": "adam_beta2", "val": ": float = 0.999"}, {"name": "adam_epsilon", "val": ": float = 1e-08"}, {"name": "max_grad_norm", "val": ": float = 1.0"}, {"name": "num_train_epochs", "val": ": float = 3.0"}, {"name": "max_steps", "val": ": int = -1"}, {"name": "lr_scheduler_type", "val": ": transformers.trainer_utils.SchedulerType | str = 'linear'"}, {"name": "lr_scheduler_kwargs", "val": ": dict[str, typing.Any] | str | None = <factory>"}, {"name": "warmup_ratio", "val": ": float = 0.0"}, {"name": "warmup_steps", "val": ": int = 0"}, {"name": "log_level", "val": ": str = 'info'"}, {"name": "log_level_replica", "val": ": str = 'silent'"}, {"name": "logging_dir", "val": ": str | None = None"}, {"name": "logging_strategy", "val": ": transformers.trainer_utils.IntervalStrategy | str = 'steps'"}, {"name": "logging_first_step", "val": ": bool = False"}, {"name": "logging_steps", "val": ": float = 500"}, {"name": "save_strategy", "val": ": transformers.trainer_utils.SaveStrategy | str = 'steps'"}, {"name": "save_steps", "val": ": float = 500"}, {"name": "save_total_limit", "val": ": int | None = None"}, {"name": "save_only_model", "val": ": bool = False"}, {"name": "restore_callback_states_from_checkpoint", "val": ": bool = False"}, {"name": "seed", "val": ": int = 42"}, {"name": "bf16", "val": ": bool = False"}, {"name": "dataloader_drop_last", "val": ": bool = False"}, {"name": "eval_steps", "val": ": float | None = None"}, {"name": "dataloader_num_workers", "val": ": int = 0"}, {"name": "dataloader_prefetch_factor", "val": ": int | None = None"}, {"name": "run_name", "val": ": str | None = None"}, {"name": "disable_tqdm", "val": ": bool | None = None"}, {"name": "remove_unused_columns", "val": ": bool | None = True"}, {"name": "label_names", "val": ": list[str] | None = None"}, {"name": "accelerator_config", "val": ": dict | str | None = None"}, {"name": "label_smoothing_factor", "val": ": float = 0.0"}, {"name": "optim", "val": ": transformers.training_args.OptimizerNames | str = 'adamw_torch'"}, {"name": "optim_args", "val": ": str | None = None"}, {"name": "report_to", "val": ": None | str | list[str] = None"}, {"name": "resume_from_checkpoint", "val": ": str | None = None"}, {"name": "gradient_checkpointing", "val": ": bool = False"}, {"name": "gradient_checkpointing_kwargs", "val": ": dict[str, typing.Any] | str | None = None"}, {"name": "use_liger_kernel", "val": ": bool | None = False"}, {"name": "average_tokens_across_devices", "val": ": bool | None = False"}, {"name": "dataloader_prefetch_size", "val": ": int = None"}, {"name": "skip_cache_push", "val": ": bool = False"}, {"name": "use_autocast", "val": ": bool = False"}, {"name": "zero_1", "val": ": bool = True"}, {"name": "stochastic_rounding_enabled", "val": ": bool = True"}, {"name": "optimizer_use_master_weights", "val": ": bool = True"}, {"name": "optimizer_use_fp32_grad_acc", "val": ": bool = True"}, {"name": "optimizer_save_master_weights_in_ckpt", "val": ": bool = False"}, {"name": "tensor_parallel_size", "val": ": int = 1"}, {"name": "disable_sequence_parallel", "val": ": bool = False"}, {"name": "pipeline_parallel_size", "val": ": int = 1"}, {"name": "pipeline_parallel_num_microbatches", "val": ": int = -1"}, {"name": "kv_size_multiplier", "val": ": int | None = None"}, {"name": "num_local_ranks_per_step", "val": ": int = 8"}, {"name": "use_xser", "val": ": bool = True"}, {"name": "async_save", "val": ": bool = False"}, {"name": "fuse_qkv", "val": ": bool = False"}, {"name": "recompute_causal_mask", "val": ": bool = True"}, {"name": "enable_throughput_metrics", "val": ": bool = True"}, {"name": "enable_mfu_metrics", "val": ": bool = True"}, {"name": "enable_efficiency_metrics", "val": ": bool = True"}, {"name": "metrics_window_size", "val": ": int = 50"}]</parameters></docstring>

Configuration class for Neuron-optimized SFT training.

Inherits from both NeuronTrainingArguments (for Trainium-specific settings) and
trl's SFTConfig (for SFT-specific settings).

Key Neuron-specific behavior:
- padding_free is always set to False to avoid recompilation on Trainium devices
- All other SFT parameters from trl 0.24.0+ are supported


</div>

### NeuronSFTTrainer[[optimum.neuron.NeuronSFTTrainer]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.NeuronSFTTrainer</name><anchor>optimum.neuron.NeuronSFTTrainer</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1020/optimum/neuron/trainers/sft_trainer.py#L132</source><parameters>[{"name": "model", "val": ": transformers.modeling_utils.PreTrainedModel | torch.nn.modules.module.Module | str"}, {"name": "args", "val": ": optimum.neuron.trainers.sft_trainer.SFTConfig | None = None"}, {"name": "data_collator", "val": ": typing.Optional[typing.Callable[[list[typing.Any]], dict[str, typing.Any]]] = None"}, {"name": "train_dataset", "val": ": Dataset | IterableDataset | datasets.Dataset | None = None"}, {"name": "eval_dataset", "val": ": Dataset | dict[str, Dataset] | datasets.Dataset | None = None"}, {"name": "processing_class", "val": ": transformers.tokenization_utils_base.PreTrainedTokenizerBase | transformers.processing_utils.ProcessorMixin | None = None"}, {"name": "compute_loss_func", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "compute_metrics", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "callbacks", "val": ": list[transformers.trainer_callback.TrainerCallback] | None = None"}, {"name": "optimizers", "val": ": tuple[torch.optim.optimizer.Optimizer | None, torch.optim.lr_scheduler.LambdaLR | None] = (None, None)"}, {"name": "optimizer_cls_and_kwargs", "val": ": tuple[type[torch.optim.optimizer.Optimizer], dict[str, typing.Any]] | None = None"}, {"name": "preprocess_logits_for_metrics", "val": ": typing.Optional[typing.Callable] = None"}, {"name": "peft_config", "val": ": peft.config.PeftConfig | None = None"}, {"name": "formatting_func", "val": ": typing.Optional[typing.Callable] = None"}]</parameters></docstring>

`SFTTrainer` adapted for Neuron (Trainium) devices.



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>compute_loss</name><anchor>optimum.neuron.NeuronSFTTrainer.compute_loss</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1020/optimum/neuron/trainers/sft_trainer.py#L403</source><parameters>[{"name": "model", "val": ""}, {"name": "inputs", "val": ""}, {"name": "return_outputs", "val": " = False"}, {"name": "num_items_in_batch", "val": " = None"}]</parameters></docstring>

Compute training loss for Neuron-optimized training.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>log</name><anchor>optimum.neuron.NeuronSFTTrainer.log</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1020/optimum/neuron/trainers/sft_trainer.py#L387</source><parameters>[{"name": "logs", "val": ": dict[str, float]"}]</parameters></docstring>

Override SFTTrainer's log method to use NeuronTrainer's implementation.

SFTTrainer has custom metrics tracking that we don't use for Neuron training.


</div>
<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>training_step</name><anchor>optimum.neuron.NeuronSFTTrainer.training_step</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1020/optimum/neuron/trainers/sft_trainer.py#L413</source><parameters>[{"name": "model", "val": ": Module"}, {"name": "inputs", "val": ": dict[str, typing.Any]"}, {"name": "num_items_in_batch", "val": ": int | None = None"}]</parameters></docstring>

Perform a training step for Neuron-optimized training.


</div></div>