Axolotl - Api

Axolotl - Api

Pages: 150

cli.cloud.modal_

URL: https://docs.axolotl.ai/docs/api/cli.cloud.modal_.html

Contents:

• cli.cloud.modal_
• Classes
  • ModalCloud
• Functions
  • run_cmd

Modal Cloud support from CLI

Modal Cloud implementation.

Run a command inside a folder, with Modal Volume reloading before and commit on success.

Examples:

Example 1 (python):

cli.cloud.modal_.ModalCloud(config, app=None)

Example 2 (python):

cli.cloud.modal_.run_cmd(cmd, run_folder, volumes=None)

core.trainers.base

URL: https://docs.axolotl.ai/docs/api/core.trainers.base.html

Contents:

• core.trainers.base
• Classes
  • AxolotlTrainer
    • Methods
      • log
        • Parameters
      • push_to_hub
      • store_metrics
        • Parameters

Module for customized trainers

Extend the base Trainer for axolotl helpers

Log logs on the various objects watching training, including stored metrics.

Overwrite the push_to_hub method in order to force-add the tags when pushing the model on the Hub. Please refer to ~transformers.Trainer.push_to_hub for more details.

Store metrics with specified reduction type.

Examples:

Example 1 (python):

core.trainers.base.AxolotlTrainer(
    *_args,
    bench_data_collator=None,
    eval_data_collator=None,
    dataset_tags=None,
    **kwargs,
)

Example 2 (python):

core.trainers.base.AxolotlTrainer.log(logs, start_time=None)

Example 3 (python):

core.trainers.base.AxolotlTrainer.push_to_hub(*args, **kwargs)

Example 4 (python):

core.trainers.base.AxolotlTrainer.store_metrics(
    metrics,
    train_eval='train',
    reduction='mean',
)

prompt_strategies.input_output

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.input_output.html

Contents:

• prompt_strategies.input_output
• Classes
  • RawInputOutputPrompter
  • RawInputOutputStrategy

prompt_strategies.input_output

Module for plain input/output prompt pairs

prompter for raw i/o data

Prompt Strategy class for input/output pairs

Examples:

Example 1 (python):

prompt_strategies.input_output.RawInputOutputPrompter()

Example 2 (python):

prompt_strategies.input_output.RawInputOutputStrategy(
    *args,
    eos_token=None,
    **kwargs,
)

prompt_strategies.completion

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.completion.html

Contents:

• prompt_strategies.completion
• Classes
  • CompletionPromptTokenizingStrategy
  • CompletionPrompter

prompt_strategies.completion

Basic completion text

Tokenizing strategy for Completion prompts.

Prompter for completion

Examples:

Example 1 (python):

prompt_strategies.completion.CompletionPromptTokenizingStrategy(
    *args,
    max_length=None,
    **kwargs,
)

Example 2 (python):

prompt_strategies.completion.CompletionPrompter()

utils.collators.core

URL: https://docs.axolotl.ai/docs/api/utils.collators.core.html

Contents:

• utils.collators.core

basic shared collator constants

monkeypatch.data.batch_dataset_fetcher

URL: https://docs.axolotl.ai/docs/api/monkeypatch.data.batch_dataset_fetcher.html

Contents:

• monkeypatch.data.batch_dataset_fetcher
• Functions
  • apply_multipack_dataloader_patch
  • patch_fetchers
  • patched_worker_loop
  • remove_multipack_dataloader_patch

monkeypatch.data.batch_dataset_fetcher

Monkey patches for the dataset fetcher to handle batches of packed indexes.

This patch allows DataLoader to correctly process batches that contain multiple bins of packed sequences.

Apply patches to PyTorch’s DataLoader components.

Worker loop that ensures patches are applied in worker processes.

Remove the monkeypatch and restore original PyTorch DataLoader behavior.

Examples:

Example 1 (python):

monkeypatch.data.batch_dataset_fetcher.apply_multipack_dataloader_patch()

Example 2 (python):

monkeypatch.data.batch_dataset_fetcher.patch_fetchers()

Example 3 (python):

monkeypatch.data.batch_dataset_fetcher.patched_worker_loop(*args, **kwargs)

Example 4 (python):

monkeypatch.data.batch_dataset_fetcher.remove_multipack_dataloader_patch()

core.datasets.chat

URL: https://docs.axolotl.ai/docs/api/core.datasets.chat.html

Contents:

• core.datasets.chat
• Classes
  • TokenizedChatDataset

Tokenized chat dataset

Examples:

Example 1 (python):

core.datasets.chat.TokenizedChatDataset(
    data,
    model_transform,
    *args,
    message_transform=None,
    formatter=None,
    process_count=None,
    keep_in_memory=False,
    **kwargs,
)

utils.freeze

URL: https://docs.axolotl.ai/docs/api/utils.freeze.html

Contents:

• utils.freeze
• Classes
  • LayerNamePattern
    • Methods
      • match
• Functions
  • freeze_layers_except

module to freeze/unfreeze parameters by name

Represents a regex pattern for layer names, potentially including a parameter index range.

Checks if the given layer name matches the regex pattern.

Parameters: - name (str): The layer name to check.

Returns: - bool: True if the layer name matches the pattern, False otherwise.

Freezes all layers of the given model except for the layers that match given regex patterns. Periods in the patterns are treated as literal periods, not as wildcard characters.

Parameters: - model (nn.Module): The PyTorch model to be modified. - regex_patterns (list of str): List of regex patterns to match layer names to keep unfrozen. Note that you cannot use a dot as a wildcard character in the patterns since it is reserved for separating layer names. Also, to match the entire layer name, the pattern should start with “^” and end with “(", otherwise it will match any part of the layer name. The range pattern part is optional and it is not compiled as a regex pattern which means you must put ")” before the range pattern if you want to match the entire layer name. E.g., [“^model.embed_tokens.weight([:32000]", "layers.2[0-9]+.block_sparse_moe.gate.[a-z]+)”]

Returns: None; the model is modified in place.

Examples:

Example 1 (python):

utils.freeze.LayerNamePattern(pattern)

Example 2 (python):

utils.freeze.LayerNamePattern.match(name)

Example 3 (python):

utils.freeze.freeze_layers_except(model, regex_patterns)

monkeypatch.unsloth_

URL: https://docs.axolotl.ai/docs/api/monkeypatch.unsloth_.html

Contents:

• monkeypatch.unsloth_

module for patching with unsloth optimizations

utils.schemas.datasets

URL: https://docs.axolotl.ai/docs/api/utils.schemas.datasets.html

Contents:

• utils.schemas.datasets
• Classes
  • DPODataset
  • KTODataset
  • PretrainingDataset
  • SFTDataset
    • Methods
      • handle_legacy_message_fields
  • StepwiseSupervisedDataset
  • UserDefinedDPOType

utils.schemas.datasets

Pydantic models for datasets-related configuration

DPO configuration subset

KTO configuration subset

Pretraining dataset configuration subset

SFT configuration subset

Handle backwards compatibility between legacy message field mapping and new property mapping system.

Stepwise supervised dataset configuration subset

User defined typing for DPO

User defined typing for KTO

Structure for user defined prompt types

Examples:

Example 1 (python):

utils.schemas.datasets.DPODataset()

Example 2 (python):

utils.schemas.datasets.KTODataset()

Example 3 (python):

utils.schemas.datasets.PretrainingDataset()

Example 4 (python):

utils.schemas.datasets.SFTDataset()

core.chat.format.llama3x

URL: https://docs.axolotl.ai/docs/api/core.chat.format.llama3x.html

Contents:

• core.chat.format.llama3x

core.chat.format.llama3x

Llama 3.x chat formatting functions for MessageContents

datasets

URL: https://docs.axolotl.ai/docs/api/datasets.html

Contents:

• datasets
• Classes
  • TokenizedPromptDataset
    • Parameters

Module containing dataset functionality.

We want this to be a wrapper for an existing dataset that we have loaded. Lets use the concept of middlewares to wrap each dataset. We’ll use the collators later on to pad the datasets.

Dataset that returns tokenized prompts from a stream of text files.

Examples:

Example 1 (python):

datasets.TokenizedPromptDataset(
    prompt_tokenizer,
    dataset,
    process_count=None,
    keep_in_memory=False,
    **kwargs,
)

prompt_strategies.bradley_terry.llama3

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.bradley_terry.llama3.html

Contents:

• prompt_strategies.bradley_terry.llama3
• Functions
  • icr

prompt_strategies.bradley_terry.llama3

chatml transforms for datasets with system, input, chosen, rejected to match llama3 chat template

chatml transforms for datasets with system, input, chosen, rejected ex. https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs

Examples:

Example 1 (python):

prompt_strategies.bradley_terry.llama3.icr(cfg, **kwargs)

common.datasets

URL: https://docs.axolotl.ai/docs/api/common.datasets.html

Contents:

• common.datasets
• Classes
  • TrainDatasetMeta
• Functions
  • load_datasets
    • Parameters
    • Returns
  • load_preference_datasets
    • Parameters
    • Returns

Dataset loading utilities.

Dataclass with fields for training and validation datasets and metadata.

Loads one or more training or evaluation datasets, calling axolotl.utils.data.prepare_datasets. Optionally, logs out debug information.

Loads one or more training or evaluation datasets for RL training using paired preference data, calling axolotl.utils.data.rl.prepare_preference_datasets. Optionally, logs out debug information.

Randomly sample num_samples samples with replacement from dataset.

Examples:

Example 1 (python):

common.datasets.TrainDatasetMeta(
    train_dataset,
    eval_dataset=None,
    total_num_steps=None,
)

Example 2 (python):

common.datasets.load_datasets(cfg, cli_args=None, debug=False)

Example 3 (python):

common.datasets.load_preference_datasets(cfg, cli_args=None)

Example 4 (python):

common.datasets.sample_dataset(dataset, num_samples)

cli.train

URL: https://docs.axolotl.ai/docs/api/cli.train.html

Contents:

• cli.train
• Functions
  • do_cli
    • Parameters
  • do_train
    • Parameters

CLI to run training on a model.

Parses axolotl config, CLI args, and calls do_train.

Trains a transformers model by first loading the dataset(s) specified in the axolotl config, and then calling axolotl.train.train. Also runs the plugin manager’s post_train_unload once training completes.

Examples:

Example 1 (python):

cli.train.do_cli(config=Path('examples/'), **kwargs)

Example 2 (python):

cli.train.do_train(cfg, cli_args)

cli.utils.fetch

URL: https://docs.axolotl.ai/docs/api/cli.utils.fetch.html

Contents:

• cli.utils.fetch
• Functions
  • fetch_from_github
    • Parameters

Utilities for axolotl fetch CLI command.

Sync files from a specific directory in the GitHub repository. Only downloads files that don’t exist locally or have changed.

Examples:

Example 1 (python):

cli.utils.fetch.fetch_from_github(dir_prefix, dest_dir=None, max_workers=5)

utils.tokenization

URL: https://docs.axolotl.ai/docs/api/utils.tokenization.html

Contents:

• utils.tokenization
• Functions
  • color_token_for_rl_debug
  • process_tokens_for_rl_debug

Module for tokenization utilities

Helper function to color tokens based on their type.

Helper function to process and color tokens.

Examples:

Example 1 (python):

utils.tokenization.color_token_for_rl_debug(
    decoded_token,
    encoded_token,
    color,
    text_only,
)

Example 2 (python):

utils.tokenization.process_tokens_for_rl_debug(
    tokens,
    color,
    tokenizer,
    text_only,
)

core.trainers.grpo.sampler

URL: https://docs.axolotl.ai/docs/api/core.trainers.grpo.sampler.html

Contents:

• core.trainers.grpo.sampler
• Classes
  • SequenceParallelRepeatRandomSampler
    • Parameters
    • Methods
      • set_epoch
        • Parameters

core.trainers.grpo.sampler

Repeat random sampler (similar to the one implemented in https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py) that adds sequence parallelism functionality; i.e., duplicating data across ranks in the same sequence parallel group.

Sampler for GRPO training with sequence parallelism.

This sampler ensures: - Ranks in the same sequence parallel (SP) group receive identical data. - Each index is repeated multiple times for sampling different completions. - Entire batches are repeated for reuse in multiple updates. - Data is properly distributed across SP groups.

In the table below, the values represent dataset indices. Each SP group has context_parallel_size = 2 GPUs working together on the same data. There are 2 SP groups (SP0 and SP1), with world_size = 4 total GPUs.

grad_accum=2 ▲ ▲ 0 0 [0 0 0 1 1 1] [2 2 2 3 3 3] <- SP groups get different data ▼ | 0 1 [0 0 0 1 1 1] [2 2 2 3 3 3] <- Same data for each SP group GPU | | 1 2 [0 0 0 1 1 1] [2 2 2 3 3 3] <- Repeat same indices for iterations num_iterations=2 ▼ 1 3 [0 0 0 1 1 1] [2 2 2 3 3 3] <- When using gradient accumulation

Sets the epoch for this sampler.

Examples:

Example 1 (python):

core.trainers.grpo.sampler.SequenceParallelRepeatRandomSampler(
    dataset,
    mini_repeat_count,
    world_size,
    rank,
    batch_size=1,
    repeat_count=1,
    context_parallel_size=1,
    shuffle=True,
    seed=0,
    drop_last=False,
)

Example 2 (unknown):

Sequence Parallel Groups
                                |       SP0        |       SP1        |
                                |  GPU 0  |  GPU 1 |  GPU 2  |  GPU 3 |
            global_step  step    <---> mini_repeat_count=3
                                    <----------> batch_size=2 per SP group

Example 3 (unknown):

2       4         [4 4 4  5 5 5]     [6 6 6  7 7 7]   <- New batch of data indices
                 2       5         [4 4 4  5 5 5]     [6 6 6  7 7 7]
                                    ...

Example 4 (python):

core.trainers.grpo.sampler.SequenceParallelRepeatRandomSampler.set_epoch(epoch)

evaluate

URL: https://docs.axolotl.ai/docs/api/evaluate.html

Contents:

• evaluate
• Functions
  • evaluate
    • Parameters
    • Returns
  • evaluate_dataset
    • Parameters
    • Returns

Module for evaluating models.

Evaluate a model on training and validation datasets.

Helper function to evaluate a single dataset.

Examples:

Example 1 (python):

evaluate.evaluate(cfg, dataset_meta)

Example 2 (python):

evaluate.evaluate_dataset(trainer, dataset, dataset_type, flash_optimum=False)

utils.optimizers.adopt

URL: https://docs.axolotl.ai/docs/api/utils.optimizers.adopt.html

Contents:

• utils.optimizers.adopt
• Functions
  • adopt

utils.optimizers.adopt

Copied from https://github.com/iShohei220/adopt

ADOPT: Modified Adam Can Converge with Any β2 with the Optimal Rate (2024) Taniguchi, Shohei and Harada, Keno and Minegishi, Gouki and Oshima, Yuta and Jeong, Seong Cheol and Nagahara, Go and Iiyama, Tomoshi and Suzuki, Masahiro and Iwasawa, Yusuke and Matsuo, Yutaka

Functional API that performs ADOPT algorithm computation.

Examples:

Example 1 (python):

utils.optimizers.adopt.adopt(
    params,
    grads,
    exp_avgs,
    exp_avg_sqs,
    state_steps,
    foreach=None,
    capturable=False,
    differentiable=False,
    fused=None,
    grad_scale=None,
    found_inf=None,
    has_complex=False,
    *,
    beta1,
    beta2,
    lr,
    clip_lambda,
    weight_decay,
    decouple,
    eps,
    maximize,
)

prompt_tokenizers

URL: https://docs.axolotl.ai/docs/api/prompt_tokenizers.html

Contents:

• prompt_tokenizers
• Classes
  • AlpacaMultipleChoicePromptTokenizingStrategy
  • AlpacaPromptTokenizingStrategy
  • AlpacaReflectionPTStrategy
  • DatasetWrappingStrategy
  • GPTeacherPromptTokenizingStrategy
  • InstructionPromptTokenizingStrategy
  • InvalidDataException
  • JeopardyPromptTokenizingStrategy

Module containing PromptTokenizingStrategy and Prompter classes

Tokenizing strategy for Alpaca Multiple Choice prompts.

Tokenizing strategy for Alpaca prompts.

Tokenizing strategy for Alpaca Reflection prompts.

Abstract class for wrapping datasets for Chat Messages

Tokenizing strategy for GPTeacher prompts.

Tokenizing strategy for instruction-based prompts.

Exception raised when the data is invalid

Tokenizing strategy for Jeopardy prompts.

Tokenizing strategy for NomicGPT4All prompts.

Tokenizing strategy for OpenAssistant prompts.

Abstract class for tokenizing strategies

Tokenizing strategy for Reflection prompts.

Tokenizing strategy for SummarizeTLDR prompts.

Parses the tokenized prompt and append the tokenized input_ids, attention_mask and labels to the result

Returns the default values for the tokenize prompt function

Examples:

Example 1 (python):

prompt_tokenizers.AlpacaMultipleChoicePromptTokenizingStrategy(
    prompter,
    tokenizer,
    train_on_inputs=False,
    sequence_len=2048,
)

Example 2 (python):

prompt_tokenizers.AlpacaPromptTokenizingStrategy(
    prompter,
    tokenizer,
    train_on_inputs=False,
    sequence_len=2048,
)

Example 3 (python):

prompt_tokenizers.AlpacaReflectionPTStrategy(
    prompter,
    tokenizer,
    train_on_inputs=False,
    sequence_len=2048,
)

Example 4 (python):

prompt_tokenizers.DatasetWrappingStrategy()

cli.art

URL: https://docs.axolotl.ai/docs/api/cli.art.html

Contents:

• cli.art
• Functions
  • print_axolotl_text_art

Axolotl ASCII logo utils.

Prints axolotl ASCII art.

Examples:

Example 1 (python):

cli.art.print_axolotl_text_art()

utils.callbacks.perplexity

URL: https://docs.axolotl.ai/docs/api/utils.callbacks.perplexity.html

Contents:

• utils.callbacks.perplexity
• Classes
  • Perplexity
    • Methods
      • compute

utils.callbacks.perplexity

callback to calculate perplexity as an evaluation metric.

Calculate perplexity as defined in https://huggingface.co/docs/transformers/en/perplexity. This is a custom variant that doesn’t re-tokenize the input or re-load the model.

Compute perplexity in a fixed length sliding window across the sequence.

Examples:

Example 1 (python):

utils.callbacks.perplexity.Perplexity(tokenizer, max_seq_len, stride=512)

Example 2 (python):

utils.callbacks.perplexity.Perplexity.compute(model, references=None)

cli.utils.train

URL: https://docs.axolotl.ai/docs/api/cli.utils.train.html

Contents:

• cli.utils.train
• Functions
  • build_command
    • Parameters
    • Returns
  • generate_config_files
    • Parameters
  • launch_training

Utilities for axolotl train CLI command.

Build command list from base command and options.

Generate list of configuration files to process. Yields a tuple of the configuration file name and a boolean indicating whether this is a group of configurations (i.e., a sweep).

Execute training with the given configuration.

Examples:

Example 1 (python):

cli.utils.train.build_command(base_cmd, options)

Example 2 (python):

cli.utils.train.generate_config_files(config, sweep)

Example 3 (python):

cli.utils.train.launch_training(
    cfg_file,
    launcher,
    cloud,
    kwargs,
    launcher_args=None,
    use_exec=False,
)

cli.vllm_serve

URL: https://docs.axolotl.ai/docs/api/cli.vllm_serve.html

Contents:

• cli.vllm_serve
• Classes
  • AxolotlScriptArguments
• Functions
  • do_vllm_serve
    • Returns

CLI to start the vllm server for online RL

Additional arguments for the VLLM server

Starts the VLLM server for serving LLM models used for online RL

Args :param cfg: Parsed doct of the YAML config :param cli_args: dict of additional command-line arguments of type VllmServeCliArgs

Examples:

Example 1 (python):

cli.vllm_serve.AxolotlScriptArguments(
    reasoning_parser='',
    enable_reasoning=None,
)

Example 2 (python):

cli.vllm_serve.do_vllm_serve(config, cli_args)

convert

URL: https://docs.axolotl.ai/docs/api/convert.html

Contents:

• convert
• Classes
  • FileReader
  • FileWriter
  • JsonParser
  • JsonToJsonlConverter
  • JsonlSerializer
  • StdoutWriter

Module containing File Reader, File Writer, Json Parser, and Jsonl Serializer classes

Reads a file and returns its contents as a string

Writes a string to a file

Parses a string as JSON and returns the result

Converts a JSON file to JSONL

Serializes a list of JSON objects into a JSONL string

Writes a string to stdout

Examples:

Example 1 (python):

convert.FileReader()

Example 2 (python):

convert.FileWriter(file_path)

Example 3 (python):

convert.JsonParser()

Example 4 (python):

convert.JsonToJsonlConverter(
    file_reader,
    file_writer,
    json_parser,
    jsonl_serializer,
)

monkeypatch.utils

URL: https://docs.axolotl.ai/docs/api/monkeypatch.utils.html

Contents:

• monkeypatch.utils
• Functions
  • get_cu_seqlens
  • get_cu_seqlens_from_pos_ids
  • mask_2d_to_4d

Shared utils for the monkeypatches

generate a cumulative sequence length mask for flash attention using attn mask

generate a cumulative sequence length mask for flash attention using pos ids

Expands attention_mask from [bsz, seq_len] to [bsz, 1, tgt_seq_len, src_seq_len]. This expansion handles packed sequences so that sequences share the same attention mask integer value when they attend to each other within that sequence. This expansion transforms the mask to lower triangular form to prevent future peeking.

Examples:

Example 1 (python):

monkeypatch.utils.get_cu_seqlens(attn_mask)

Example 2 (python):

monkeypatch.utils.get_cu_seqlens_from_pos_ids(position_ids)

Example 3 (python):

monkeypatch.utils.mask_2d_to_4d(mask, dtype, tgt_len=None)

prompt_strategies.pygmalion

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.pygmalion.html

Contents:

• prompt_strategies.pygmalion
• Classes
  • PygmalionPromptTokenizingStrategy
  • PygmalionPrompter

prompt_strategies.pygmalion

Module containing the PygmalionPromptTokenizingStrategy and PygmalionPrompter class

Tokenizing strategy for Pygmalion.

Prompter for Pygmalion.

Examples:

Example 1 (python):

prompt_strategies.pygmalion.PygmalionPromptTokenizingStrategy(
    prompter,
    tokenizer,
    *args,
    **kwargs,
)

Example 2 (python):

prompt_strategies.pygmalion.PygmalionPrompter(*args, **kwargs)

utils.callbacks.mlflow_

URL: https://docs.axolotl.ai/docs/api/utils.callbacks.mlflow_.html

Contents:

• utils.callbacks.mlflow_
• Classes
  • SaveAxolotlConfigtoMlflowCallback

utils.callbacks.mlflow_

MLFlow module for trainer callbacks

Callback to save axolotl config to mlflow

Examples:

Example 1 (python):

utils.callbacks.mlflow_.SaveAxolotlConfigtoMlflowCallback(axolotl_config_path)

loaders.adapter

URL: https://docs.axolotl.ai/docs/api/loaders.adapter.html

Contents:

• loaders.adapter
• Functions
  • setup_quantized_meta_for_peft
  • setup_quantized_peft_meta_for_training

Adapter loading functionality, including LoRA / QLoRA and associated utils

Replaces quant_state.to with a dummy function to prevent PEFT from moving quant_state to meta device

Replaces dummy quant_state.to method with the original function to allow training to continue

Examples:

Example 1 (python):

loaders.adapter.setup_quantized_meta_for_peft(model)

Example 2 (python):

loaders.adapter.setup_quantized_peft_meta_for_training(model)

cli.cloud.base

URL: https://docs.axolotl.ai/docs/api/cli.cloud.base.html

Contents:

• cli.cloud.base
• Classes
  • Cloud

base class for cloud platforms from cli

Abstract base class for cloud platforms.

Examples:

Example 1 (python):

cli.cloud.base.Cloud()

monkeypatch.llama_attn_hijack_flash

URL: https://docs.axolotl.ai/docs/api/monkeypatch.llama_attn_hijack_flash.html

Contents:

• monkeypatch.llama_attn_hijack_flash
• Functions
  • flashattn_forward_with_s2attn

monkeypatch.llama_attn_hijack_flash

Flash attention monkey patch for llama model

Input shape: Batch x Time x Channel

From: https://github.com/dvlab-research/LongLoRA/blob/main/llama_attn_replace.py

attention_mask: [bsz, q_len]

cu_seqlens will be ignored if provided max_seqlen will be ignored if provided

Examples:

Example 1 (python):

monkeypatch.llama_attn_hijack_flash.flashattn_forward_with_s2attn(
    self,
    hidden_states,
    attention_mask=None,
    position_ids=None,
    past_key_value=None,
    output_attentions=False,
    use_cache=False,
    padding_mask=None,
    cu_seqlens=None,
    max_seqlen=None,
)

monkeypatch.llama_patch_multipack

URL: https://docs.axolotl.ai/docs/api/monkeypatch.llama_patch_multipack.html

Contents:

• monkeypatch.llama_patch_multipack

monkeypatch.llama_patch_multipack

Patched LlamaAttention to use torch.nn.functional.scaled_dot_product_attention

cli.inference

URL: https://docs.axolotl.ai/docs/api/cli.inference.html

Contents:

• cli.inference
• Functions
  • do_cli
    • Parameters
  • do_inference
    • Parameters
  • do_inference_gradio
    • Parameters
  • get_multi_line_input
    • Returns

CLI to run inference on a trained model.

Parses axolotl config, CLI args, and calls do_inference or do_inference_gradio.

Runs inference on the command line in a loop. User input is accepted, a chat template is (optionally) applied, and the model specified in the axolotl config is used to generate completions according to a default generation config.

Runs inference in a Gradio interface. User input is accepted, a chat template is (optionally) applied, and the model specified in the axolotl config is used to generate completions according to a default generation config.

Gets multi-line input from terminal.

Examples:

Example 1 (python):

cli.inference.do_cli(config=Path('examples/'), gradio=False, **kwargs)

Example 2 (python):

cli.inference.do_inference(cfg, cli_args)

Example 3 (python):

cli.inference.do_inference_gradio(cfg, cli_args)

Example 4 (python):

cli.inference.get_multi_line_input()

loaders.tokenizer

URL: https://docs.axolotl.ai/docs/api/loaders.tokenizer.html

Contents:

• loaders.tokenizer
• Functions
  • load_tokenizer
  • modify_tokenizer_files
    • Parameters
    • Returns

Tokenizer loading functionality and associated utils

Load and configure the tokenizer based on the provided config.

Modify tokenizer files to replace added_tokens strings, save to output directory, and return the path to the modified tokenizer.

This only works with reserved tokens that were added to the tokenizer, not tokens already part of the vocab.

Ref: https://github.com/huggingface/transformers/issues/27974#issuecomment-1854188941

Examples:

Example 1 (python):

loaders.tokenizer.load_tokenizer(cfg)

Example 2 (python):

loaders.tokenizer.modify_tokenizer_files(
    tokenizer_path,
    token_mappings,
    output_dir,
)

cli.utils.sweeps

URL: https://docs.axolotl.ai/docs/api/cli.utils.sweeps.html

Contents:

• cli.utils.sweeps
• Functions
  • generate_sweep_configs
    • Parameters
    • Returns
    • Example

Utilities for handling sweeps over configs for axolotl train CLI command

Recursively generates all possible configurations by applying sweeps to the base config.

sweeps_config = { ‘learning_rate’: [0.1, 0.01], ’_’: [ {‘load_in_8bit’: True, ‘adapter’: ‘lora’}, {‘load_in_4bit’: True, ‘adapter’: ‘qlora’} ] }

Examples:

Example 1 (python):

cli.utils.sweeps.generate_sweep_configs(base_config, sweeps_config)

prompt_strategies.dpo.chatml

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.dpo.chatml.html

Contents:

• prompt_strategies.dpo.chatml
• Functions
  • argilla_chat
  • icr
  • intel
  • ultra

prompt_strategies.dpo.chatml

DPO strategies for chatml

for argilla/dpo-mix-7k conversations

chatml transforms for datasets with system, input, chosen, rejected ex. https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs

For Intel Orca DPO Pairs

for ultrafeedback binarized conversations

Examples:

Example 1 (python):

prompt_strategies.dpo.chatml.argilla_chat(cfg, **kwargs)

Example 2 (python):

prompt_strategies.dpo.chatml.icr(cfg, **kwargs)

Example 3 (python):

prompt_strategies.dpo.chatml.intel(cfg, **kwargs)

Example 4 (python):

prompt_strategies.dpo.chatml.ultra(cfg, **kwargs)

cli.quantize

URL: https://docs.axolotl.ai/docs/api/cli.quantize.html

Contents:

• cli.quantize
• Functions
  • do_quantize
    • Parameters

CLI to post-training quantize a model using torchao

Quantizes a model’s model’s weights

Examples:

Example 1 (python):

cli.quantize.do_quantize(config, cli_args)

utils.dict

URL: https://docs.axolotl.ai/docs/api/utils.dict.html

Contents:

• utils.dict
• Classes
  • DictDefault
• Functions
  • remove_none_values

Module containing the DictDefault class

A Dict that returns None instead of returning empty Dict for missing keys.

Remove null from a dictionary-like obj or list. These can appear due to Dataset loading causing schema merge. See https://github.com/axolotl-ai-cloud/axolotl/pull/2909

Examples:

Example 1 (python):

utils.dict.DictDefault()

Example 2 (python):

utils.dict.remove_none_values(obj)

API Reference

URL: https://docs.axolotl.ai/docs/api/

Contents:

• API Reference
• Core
• CLI
• Trainers
• Model Loading
• Mixins
• Context Managers
• Prompt Strategies
• Kernels
• Monkey Patches

Core functionality for training

Command-line interface

Training implementations

Functionality for loading and patching models, tokenizers, etc.

Mixin classes for augmenting trainers

Context managers for altering trainer behaviors

Prompt formatting strategies

Low-level performance optimizations

Runtime patches for model optimizations

Pydantic data models for Axolotl config

Third-party integrations and extensions

Common utilities and shared functionality

Custom model implementations

Data processing utilities

monkeypatch.lora_kernels

URL: https://docs.axolotl.ai/docs/api/monkeypatch.lora_kernels.html

Contents:

• monkeypatch.lora_kernels
• Classes
  • FakeMLP
• Functions
  • apply_lora_kernel_patches
    • Parameters
    • Returns
    • Raises
    • Note
  • get_attention_cls_from_config

monkeypatch.lora_kernels

Module for patching custom LoRA Triton kernels and torch.autograd functions.

placeholder MLP for triton patching

Applies optimized Triton kernel patches to a PEFT model.

Patches a PEFT model with optimized implementations for MLP and attention computations. The optimizations include custom Triton kernels for activation functions and specialized autograd functions for LoRA computations.

The optimizations require LoRA adapters with no dropout and no bias terms. The function will skip patching if these conditions aren’t met.

Get the appropriate attention class by inspecting the model config. Uses dynamic import to support any model architecture that follows the standard transformers naming convention.

Get the layers of the model. Handles text-only and multimodal models.

Original implementation of output projection without optimizations.

Original implementation of QKV projection without optimizations.

Given an axolotl config, this method patches the inferred attention class forward pass with optimized LoRA implementations.

It modifies the attention class to use optimized QKV and output projections. The original implementation is preserved and can be restored if needed.

Examples:

Example 1 (python):

monkeypatch.lora_kernels.FakeMLP(gate_proj, up_proj, down_proj)

Example 2 (python):

monkeypatch.lora_kernels.apply_lora_kernel_patches(model, cfg)

Example 3 (python):

monkeypatch.lora_kernels.get_attention_cls_from_config(cfg)

Example 4 (python):

monkeypatch.lora_kernels.get_layers(model)

monkeypatch.stablelm_attn_hijack_flash

URL: https://docs.axolotl.ai/docs/api/monkeypatch.stablelm_attn_hijack_flash.html

Contents:

• monkeypatch.stablelm_attn_hijack_flash
• Functions
  • repeat_kv
  • rotate_half

monkeypatch.stablelm_attn_hijack_flash

PyTorch StableLM Epoch model.

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)

Rotates half the hidden dims of the input.

Examples:

Example 1 (python):

monkeypatch.stablelm_attn_hijack_flash.repeat_kv(hidden_states, n_rep)

Example 2 (python):

monkeypatch.stablelm_attn_hijack_flash.rotate_half(x)

core.trainers.mixins.rng_state_loader

URL: https://docs.axolotl.ai/docs/api/core.trainers.mixins.rng_state_loader.html

Contents:

• core.trainers.mixins.rng_state_loader
• Classes
  • RngLoaderMixin

core.trainers.mixins.rng_state_loader

Temporary fix/override for bug in resume from checkpoint

See https://github.com/huggingface/transformers/pull/37162

TODO: Remove when upstream added PR to release

mixin for method override to load RNG states from a checkpoint

Examples:

Example 1 (python):

core.trainers.mixins.rng_state_loader.RngLoaderMixin()

core.trainers.utils

URL: https://docs.axolotl.ai/docs/api/core.trainers.utils.html

Contents:

• core.trainers.utils

Utils for Axolotl trainers

core.training_args

URL: https://docs.axolotl.ai/docs/api/core.training_args.html

Contents:

• core.training_args
• Classes
  • AxolotlCPOConfig
  • AxolotlKTOConfig
  • AxolotlORPOConfig
  • AxolotlPRMConfig
  • AxolotlRewardConfig
  • AxolotlTrainingArguments

extra axolotl specific training args

CPO config for CPO training

KTO config for KTO training

ORPO config for ORPO training

PRM config for PRM training

Reward config for Reward training

Training arguments for Causal trainer

This code is duplicated due to HF TrainingArguments not setting output_dir with a default value so it can’t be used as a mixin.

Examples:

Example 1 (python):

core.training_args.AxolotlCPOConfig(simpo_gamma=None)

Example 2 (python):

core.training_args.AxolotlKTOConfig()

Example 3 (python):

core.training_args.AxolotlORPOConfig()

Example 4 (python):

core.training_args.AxolotlPRMConfig()

monkeypatch.btlm_attn_hijack_flash

URL: https://docs.axolotl.ai/docs/api/monkeypatch.btlm_attn_hijack_flash.html

Contents:

• monkeypatch.btlm_attn_hijack_flash

monkeypatch.btlm_attn_hijack_flash

Flash attention monkey patch for cerebras btlm model

prompt_strategies.dpo.passthrough

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.dpo.passthrough.html

Contents:

• prompt_strategies.dpo.passthrough

prompt_strategies.dpo.passthrough

DPO prompt strategies passthrough/zero-processing strategy

kernels.swiglu

URL: https://docs.axolotl.ai/docs/api/kernels.swiglu.html

Contents:

• kernels.swiglu
• Functions
  • swiglu_backward
    • Parameters
    • Returns
  • swiglu_forward
    • Parameters
    • Returns

Module for definition of SwiGLU Triton kernels.

See “GLU Variants Improve Transformer” (https://arxiv.org/abs/2002.05202).

Credit to unsloth (https://unsloth.ai/) for inspiration for this implementation.

SwiGLU backward pass using in-place operations.

SwiGLU forward pass. Computes SwiGLU activation: x * sigmoid(x) * up, where x is the gate tensor.

Examples:

Example 1 (python):

kernels.swiglu.swiglu_backward(grad_output, gate, up)

Example 2 (python):

kernels.swiglu.swiglu_forward(gate, up)

core.trainers.grpo.trainer

URL: https://docs.axolotl.ai/docs/api/core.trainers.grpo.trainer.html

Contents:

• core.trainers.grpo.trainer
• Classes
  • AxolotlGRPOSequenceParallelTrainer
    • Methods
      • get_train_dataloader
  • AxolotlGRPOTrainer

core.trainers.grpo.trainer

Axolotl GRPO trainers (with and without sequence parallelism handling)

Extend the base GRPOTrainer for sequence parallelism handling

Get dataloader for training

Extend the base GRPOTrainer for axolotl helpers

Examples:

Example 1 (python):

core.trainers.grpo.trainer.AxolotlGRPOSequenceParallelTrainer(
    model,
    reward_funcs,
    args=None,
    train_dataset=None,
    eval_dataset=None,
    processing_class=None,
    reward_processing_classes=None,
    callbacks=None,
    optimizers=(None, None),
    peft_config=None,
    optimizer_cls_and_kwargs=None,
)

Example 2 (python):

core.trainers.grpo.trainer.AxolotlGRPOSequenceParallelTrainer.get_train_dataloader(
)

Example 3 (python):

core.trainers.grpo.trainer.AxolotlGRPOTrainer(*args, **kwargs)

prompt_strategies.user_defined

URL: https://docs.axolotl.ai/docs/api/prompt_strategies.user_defined.html

Contents:

• prompt_strategies.user_defined
• Classes
  • UserDefinedDatasetConfig
  • UserDefinedPromptTokenizationStrategy

prompt_strategies.user_defined

User Defined prompts with configuration from the YML config

dataclass configuration representing a userdefined dataset type

Prompt Tokenization Strategy for user defined prompts

Examples:

Example 1 (python):

prompt_strategies.user_defined.UserDefinedDatasetConfig(
    system_prompt='',
    field_system='system',
    field_instruction='instruction',
    field_input='input',
    field_output='output',
    format='{instruction} {input} ',
    no_input_format='{instruction} ',
    system_format='{system}',
)

Example 2 (python):

prompt_strategies.user_defined.UserDefinedPromptTokenizationStrategy(
    prompter,
    tokenizer,
    train_on_inputs=False,
    sequence_len=2048,
)

utils.schemas.training

URL: https://docs.axolotl.ai/docs/api/utils.schemas.training.html

Contents:

• utils.schemas.training
• Classes
  • HyperparametersConfig
  • JaggedLRConfig
  • LrGroup

utils.schemas.training

Pydantic models for training hyperparameters

Training hyperparams configuration subset

JaggedLR configuration subset, can be used w/ ReLoRA training

Custom learning rate group configuration

Examples:

Example 1 (python):

utils.schemas.training.HyperparametersConfig()

Example 2 (python):

utils.schemas.training.JaggedLRConfig()

Example 3 (python):

utils.schemas.training.LrGroup()

utils.quantization

URL: https://docs.axolotl.ai/docs/api/utils.quantization.html

Contents:

• utils.quantization
• Functions
  • convert_qat_model
  • get_quantization_config
    • Parameters
    • Returns
    • Raises
  • prepare_model_for_qat
    • Parameters
    • Raises

Utilities for quantization including QAT and PTQ using torchao.

This function converts a QAT model which has fake quantized layers back to the original model.

…(truncated)