Skip to main content

​
YAML Configuration Files

Configuration files make complex training setups reproducible and shareable.

​
Basic Usage

aitraining --config training.yaml

​
Config File Structure

The YAML config has a specific nested structure: 
# training.yaml
task: llm-sft                        # Task with trainer suffix
backend: local                       # Required: local, spaces-*, etc.

# Model settings
base_model: google/gemma-3-270m      # Note: "base_model", not "model"
project_name: my-gemma-model

# Data settings (nested under "data:")
data:
  path: ./data/conversations.jsonl   # Note: nested under data
  train_split: train
  valid_split: null
  chat_template: tokenizer           # For LLM: tokenizer, chatml, zephyr, none
  column_mapping:                    # Column names
    text_column: text

# Logging
log: wandb

# Hub settings (optional)
hub:
  username: ${HF_USERNAME}
  token: ${HF_TOKEN}
  push_to_hub: false

# All other training parameters go under "params:"
params:
  epochs: 3
  batch_size: 4
  lr: 3e-5
  mixed_precision: bf16
  peft: true
  lora_r: 16
  lora_alpha: 32
  lora_dropout: 0.05
Important Structure Notes:
  • Use base_model, not model
  • Data path is data.path, not data_path
  • Column mappings go under data.column_mapping
  • Training parameters go under params:
  • The backend field is required

​
Task Types

The task field includes the trainer type: 
# LLM tasks (trainer in task name)
task: llm-sft                # SFT training
task: llm-dpo                # DPO training
task: llm-orpo               # ORPO training
task: llm-reward             # Reward model training
task: llm-generic            # Default/pretraining

# Other tasks
task: text-classification    # Text classification
task: image-classification   # Image classification
task: token-classification   # NER
task: seq2seq                # Sequence to sequence
task: tabular                # Tabular data
task: vlm:vqa                # Vision-language (VQA)
task: vlm:captioning         # Vision-language (captioning)
task: sentence-transformers:pair_score  # Sentence transformers

​
LLM Training Configs

​
SFT Training

task: llm-sft
backend: local
base_model: meta-llama/Llama-3.2-1B
project_name: llama-sft

data:
  path: ./conversations.jsonl
  train_split: train
  valid_split: null
  chat_template: tokenizer
  column_mapping:
    text_column: text

log: wandb

hub:
  push_to_hub: false

params:
  epochs: 3
  batch_size: 2
  lr: 3e-5
  peft: true
  lora_r: 16
  lora_alpha: 32

​
DPO Training

task: llm-dpo
backend: local
base_model: meta-llama/Llama-3.2-1B
project_name: llama-dpo

data:
  path: ./preferences.jsonl
  train_split: train
  valid_split: null
  chat_template: tokenizer
  column_mapping:
    prompt_text_column: prompt
    text_column: chosen
    rejected_text_column: rejected

log: wandb

params:
  dpo_beta: 0.1
  max_prompt_length: 128
  max_completion_length: null
  epochs: 1
  batch_size: 2
  peft: true
  lora_r: 16

​
Knowledge Distillation

task: llm-sft
backend: local
base_model: google/gemma-3-270m
project_name: distilled-model

data:
  path: ./prompts.jsonl
  train_split: train
  valid_split: null
  chat_template: tokenizer
  column_mapping:
    text_column: text

log: wandb

params:
  use_distillation: true
  teacher_model: google/gemma-2-2b
  distill_temperature: 3.0
  distill_alpha: 0.7
  epochs: 3

​
Text Classification Config

task: text-classification
backend: local
base_model: bert-base-uncased
project_name: sentiment-classifier

data:
  path: ./reviews.csv
  train_split: train
  valid_split: null
  column_mapping:
    text_column: text
    target_column: target

log: wandb

params:
  epochs: 5
  batch_size: 16
  lr: 5e-5

​
Environment Variables in Configs

Reference environment variables with ${VAR_NAME}: 
hub:
  token: ${HF_TOKEN}
  username: ${HF_USERNAME}
Set them before running: 
export HF_TOKEN="hf_..."
export HF_USERNAME="my-username"
aitraining --config training.yaml

task: llm-sft
backend: local
base_model: meta-llama/Llama-3.2-1B
project_name: production-model

data:
  path: ./conversations.jsonl
  train_split: train
  valid_split: validation
  chat_template: tokenizer
  column_mapping:
    text_column: text

log: wandb

hub:
  push_to_hub: true
  username: ${HF_USERNAME}
  token: ${HF_TOKEN}

params:
  # Training
  epochs: 3
  batch_size: 4
  gradient_accumulation: 4
  lr: 3e-5
  warmup_ratio: 0.1
  mixed_precision: bf16

  # LoRA
  peft: true
  lora_r: 32
  lora_alpha: 64
  lora_dropout: 0.05
  target_modules: all-linear

  # Distribution (for multi-GPU)
  distributed_backend: null        # null for auto (DDP), or "deepspeed"

  # Optimization
  use_flash_attention_2: true
  packing: true
  auto_find_batch_size: true

  # Checkpointing
  logging_steps: 10
  save_strategy: steps
  save_steps: 100
  save_total_limit: 1

​
Minimal Config

The minimal required fields: 
task: llm-sft
backend: local
base_model: google/gemma-3-270m
project_name: my-model

data:
  path: ./data.jsonl
  train_split: train
  valid_split: null
  chat_template: tokenizer
  column_mapping:
    text_column: text

log: wandb

​
Next Steps

LLM Training

LLM parameter reference

Config Templates

Ready-to-use templates