Files
axolotl/examples/llama-3/diffusion-3.2-1b-pretrain.yaml
2025-08-16 02:44:44 +00:00

62 lines
1.1 KiB
YAML

base_model: meta-llama/Llama-3.2-1B
# Automatically upload checkpoint and final model to HF
# hub_model_id: username/custom_model_name
# Dataset configuration for pretraining
datasets:
- path: wikitext
name: wikitext-103-raw-v1
type: completion
field: text
val_set_size: 0.001
plugins:
- diffusion.DiffusionPlugin
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
output_dir: ./outputs/model-out
sequence_len: 512
sample_packing: false
eval_sample_packing: false
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: 500
save_strategy: steps
eval_strategy: steps
save_steps: 1000
eval_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