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

pretraining_dataset:
  - path: wikitext
    name: wikitext-103-raw-v1
    type: completion
    field: text

plugins:
  - axolotl.integrations.diffusion.DiffusionPlugin

diffusion:
  noise_schedule: cosine
  min_mask_ratio: 0.15
  max_mask_ratio: 0.85
  num_diffusion_steps: 128
  eps: 5e-4
  importance_weighting: true
  mask_token_id: 128002
  generate_samples: true
  generation_interval: 250

output_dir: ./outputs/model-out

sequence_len: 512
sample_packing: true

gradient_accumulation_steps: 8
micro_batch_size: 4
max_steps: 10000
warmup_ratio: 0.1

optimizer: adamw_8bit
lr_scheduler: cosine
learning_rate: 3e-4
sdp_attention: true

bf16: auto
tf32: true

logging_steps: 1
save_strategy: steps
save_steps: 1000

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