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
- log
- Methods
- AxolotlTrainer
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
- Methods
- LayerNamePattern
- 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
- Methods
- 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
- TokenizedPromptDataset
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
- load_datasets
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
- do_cli
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
- fetch_from_github
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
- set_epoch
- SequenceParallelRepeatRandomSampler
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
- evaluate
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
- Methods
- Perplexity
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
- build_command
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
- do_vllm_serve
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
- do_cli
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
- generate_sweep_configs
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
- do_quantize
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
- apply_lora_kernel_patches
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
- swiglu_backward
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
- Methods
- AxolotlGRPOTrainer
- AxolotlGRPOSequenceParallelTrainer
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)