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:
  - diffusion.DiffusionPlugin
noise_schedule: cosine
min_mask_ratio: 0.15
max_mask_ratio: 0.85
eps: 5e-4
importance_weighting: true
mask_token_id: 128002
generate_samples: true
generation_interval: 10

output_dir: ./outputs/model-out

sequence_len: 512
sample_packing: true

gradient_accumulation_steps: 8
micro_batch_size: 4
max_steps: 10000

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

bf16: auto
tf32: true

gradient_checkpointing: true
resume_from_checkpoint:
logging_steps: 1
sdp_attention: true

warmup_steps: 1000

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