# train – Axolotl

# train

`train`

Prepare and train a model on a dataset. Can also infer from a model or merge lora

## Functions

Name | Description
--- | ---
create_model_card | Create a model card for the trained model if needed.
determine_resume_checkpoint | Determine the checkpoint to resume from based on configuration.
execute_training | Execute the training process with appropriate SDP kernel configurations.
handle_untrained_tokens_fix | Apply fixes for untrained tokens if configured.
save_initial_configs | Save initial configurations before training.
save_trained_model | Save the trained model according to configuration and training setup.
setup_model_and_tokenizer | Load the tokenizer, processor (for multimodal models), and model based on
setup_model_and_trainer | Load model, tokenizer, trainer, etc. Helper function to encapsulate the full
setup_model_card | Set up the Axolotl badge and add the Axolotl config to the model card if available.
setup_reference_model | Set up the reference model for RL training if needed.
setup_signal_handler | Set up signal handler for graceful termination.
train | Train a model on the given dataset.

### create_model_card

    train.create_model_card(cfg, trainer)

Create a model card for the trained model if needed.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
trainer | Trainer | The trainer object with model card creation capabilities. | required

### determine_resume_checkpoint

    train.determine_resume_checkpoint(cfg)

Determine the checkpoint to resume from based on configuration.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required

#### Returns

Name | Type | Description
--- | --- | ---
| str | None | Path to the checkpoint to resume from, or `None` if not resuming.

### execute_training

    train.execute_training(cfg, trainer, resume_from_checkpoint)

Execute the training process with appropriate SDP kernel configurations.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
trainer | Any | The configured trainer object. | required
resume_from_checkpoint | str | None | Path to checkpoint to resume from, if applicable. | required

### handle_untrained_tokens_fix

    train.handle_untrained_tokens_fix(
        cfg,
        model,
        tokenizer,
        train_dataset,
        safe_serialization,
    )

Apply fixes for untrained tokens if configured.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
model | PreTrainedModel | The model to apply fixes to. | required
tokenizer | PreTrainedTokenizer | The tokenizer for token identification. | required
train_dataset | Dataset | The training dataset to use. | required
safe_serialization | bool | Whether to use safe serialization when saving. | required

### save_initial_configs

    train.save_initial_configs(cfg, tokenizer, model, peft_config, processor)

Save initial configurations before training.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
tokenizer | PreTrainedTokenizer | The tokenizer to save. | required
model | PreTrainedModel | The model to save configuration for. | required
peft_config | PeftConfig | None | The PEFT configuration to save if applicable. | required

### save_trained_model

    train.save_trained_model(cfg, trainer, model, safe_serialization)

Save the trained model according to configuration and training setup.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
trainer | Any | The trainer object. | required
model | PreTrainedModel | The trained model to save. | required
safe_serialization | bool | Whether to use safe serialization. | required

### setup_model_and_tokenizer

    train.setup_model_and_tokenizer(cfg)

Load the tokenizer, processor (for multimodal models), and model based on configuration.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required

#### Returns

Name | Type | Description
--- | --- | ---
| tuple[PreTrainedModel, PreTrainedTokenizer, PeftConfig | None, ProcessorMixin | None] | Tuple containing model, tokenizer, `peft_config` (if LoRA / QLoRA, else `None`), and processor (if multimodal, else `None`).

### setup_model_and_trainer

    train.setup_model_and_trainer(cfg, dataset_meta)

Load model, tokenizer, trainer, etc. Helper function to encapsulate the full trainer setup.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | The configuration dictionary with training parameters. | required
dataset_meta | TrainDatasetMeta | Object with training, validation datasets and metadata. | required

#### Returns

Name | Type | Description
--- | --- | ---
| tuple['HFRLTrainerBuilder' | 'HFCausalTrainerBuilder', PeftModel | PreTrainedModel, PreTrainedTokenizer, PeftConfig | None, ProcessorMixin | None] | Tuple of: - Trainer (Causal or RLHF) - Model - Tokenizer - PEFT config - Processor

### setup_model_card

    train.setup_model_card(cfg)

Set up the Axolotl badge and add the Axolotl config to the model card if available.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required

### setup_reference_model

    train.setup_reference_model(cfg, tokenizer)

Set up the reference model for RL training if needed.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
tokenizer | PreTrainedTokenizer | The tokenizer to use for the reference model. | required

#### Returns

Name | Type | Description
--- | --- | ---
| PreTrainedModel | None | Reference model if needed for RL training, `None` otherwise.

### setup_signal_handler

    train.setup_signal_handler(cfg, model, safe_serialization)

Set up signal handler for graceful termination.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | Dictionary mapping `axolotl` config keys to values. | required
model | PreTrainedModel | The model to save on termination | required
safe_serialization | bool | Whether to use safe serialization when saving | required

### train

    train.train(cfg, dataset_meta)

Train a model on the given dataset.

#### Parameters

Name | Type | Description | Default
--- | --- | --- | ---
cfg | DictDefault | The configuration dictionary with training parameters | required
dataset_meta | TrainDatasetMeta | Object with training, validation datasets and metadata | required

#### Returns

Name | Type | Description
--- | --- | ---
| tuple[PeftModel | PreTrainedModel, PreTrainedTokenizer, Trainer] | Tuple of (model, tokenizer) after training