aitraining --config training.yaml

# 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

# 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

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

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

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

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

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

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

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