base_model: meta-llama/Llama-3.2-1B
# Automatically upload checkpoint and final model to HF
# hub_model_id: username/custom_model_name

datasets:
  - path: teknium/GPT4-LLM-Cleaned
    type: alpaca
val_set_size: 0.05

plugins:
  - axolotl.integrations.diffusion.DiffusionPlugin

diffusion:
  noise_schedule: cosine
  min_mask_ratio: 0.1
  max_mask_ratio: 0.9
  num_diffusion_steps: 128
  eps: 1e-3
  importance_weighting: true
  mask_token_id: 128002
  generate_samples: true
  generation_interval: 250

output_dir: ./outputs/model-out

sequence_len: 512
sample_packing: true
eval_sample_packing: true

gradient_accumulation_steps: 4
micro_batch_size: 4
num_epochs: 1
warmup_ratio: 0.1

optimizer: adamw_8bit
lr_scheduler: cosine
learning_rate: 1e-5

bf16: auto
tf32: true

gradient_checkpointing: true
resume_from_checkpoint:
sdp_attention: true

logging_steps: 1
save_strategy: epoch
eval_strategy: epoch

special_tokens:
  pad_token: "<|end_of_text|>"

wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

# save_first_step: true  # uncomment this to validate checkpoint saving works with your config