utils.schemas.config

utils.schemas.config

Module with Pydantic models for configuration.

Classes

Name Description
AxolotlConfigWCapabilities Wrapper to valdiate GPU capabilities with the configured options
AxolotlInputConfig Wrapper of all config options.
EBFTConfig Configuration for Energy-Based Fine-Tuning (EBFT)

AxolotlConfigWCapabilities

utils.schemas.config.AxolotlConfigWCapabilities()

Wrapper to valdiate GPU capabilities with the configured options

Methods

Name Description
check_fp8_attention_preflight fp8 attention requires SM90+ and torch >= 2.11 (torchao >= 0.17 is pinned).
check_fp8_attention_preflight
utils.schemas.config.AxolotlConfigWCapabilities.check_fp8_attention_preflight()

fp8 attention requires SM90+ and torch >= 2.11 (torchao >= 0.17 is pinned).

AxolotlInputConfig

utils.schemas.config.AxolotlInputConfig()

Wrapper of all config options.

Methods

Name Description
check_cudagraphs_wo_static_shapes Warn on triton.cudagraphs + sample_packing: cudagraphs need static shapes, packed seqs are dynamic.
check_flash_attention_torch_preflight flash_attention_torch is only usable once transformers registers the backend.
check_torch_compile_options_requires_compile Require torch_compile enabled when torch_compile_options is set.
normalize_attn_implementation Map legacy boolean attention flags to canonical attn_implementation, warn, then strip.
validate_attn_implementation Accept canonical names and hub-kernel paths; reject short-form aliases.
validate_large_head_attention Only the three known policies are valid (a typo must not silently opt into Triton).
validate_torch_compile_options Reject torch_compile_options keys not in INDUCTOR_COMPILE_OPTIONS_ALLOWLIST.
check_cudagraphs_wo_static_shapes
utils.schemas.config.AxolotlInputConfig.check_cudagraphs_wo_static_shapes()

Warn on triton.cudagraphs + sample_packing: cudagraphs need static shapes, packed seqs are dynamic.

check_flash_attention_torch_preflight
utils.schemas.config.AxolotlInputConfig.check_flash_attention_torch_preflight()

flash_attention_torch is only usable once transformers registers the backend.

check_torch_compile_options_requires_compile
utils.schemas.config.AxolotlInputConfig.check_torch_compile_options_requires_compile(
)

Require torch_compile enabled when torch_compile_options is set.

normalize_attn_implementation
utils.schemas.config.AxolotlInputConfig.normalize_attn_implementation(data)

Map legacy boolean attention flags to canonical attn_implementation, warn, then strip.

validate_attn_implementation
utils.schemas.config.AxolotlInputConfig.validate_attn_implementation(value)

Accept canonical names and hub-kernel paths; reject short-form aliases.

validate_large_head_attention
utils.schemas.config.AxolotlInputConfig.validate_large_head_attention(value)

Only the three known policies are valid (a typo must not silently opt into Triton).

validate_torch_compile_options
utils.schemas.config.AxolotlInputConfig.validate_torch_compile_options(value)

Reject torch_compile_options keys not in INDUCTOR_COMPILE_OPTIONS_ALLOWLIST.

EBFTConfig

utils.schemas.config.EBFTConfig()

Configuration for Energy-Based Fine-Tuning (EBFT)