inherit: configs/base.yaml
experiment_name: chinchilla_baseline

model:
  dim: 2048
  n_layers: 24

attention:
  type: gqa
  num_heads: 16
  num_kv_heads: 4
  head_dim: 128
  qk_norm: true

ffn:
  type: swiglu
  intermediate_mult: 2.667

residual:
  type: prenorm

optim:
  type: adamw
  lr: 3.0e-4
  scheduler: cosine
  warmup_steps: 500

training:
  tokens: 20_000_000_000
  batch_size: 8              # per-GPU; 111GB peak at seq_len=4096
  seq_len: 4096
  grad_accum_steps: 6        # effective batch = 8*8*6 = 384 seqs = 1.57M tok/step
  liger: true
  compile: true
  checkpoint_every_tokens: 2_000_000_000
  hf_repo: GoedelMachines/chinchilla-1b-baseline

data:
  shard_dir: data/fineweb_edu

logging:
  wandb_enabled: true