config.json

{
  "arch": "stack",
  "model": {
    "vocab_size": 16384,
    "d_model": 384,
    "n_pre": 1,
    "n_loop": 1,
    "n_post": 1,
    "n_heads": 6,
    "n_kv_heads": 2,
    "head_dim": 64,
    "d_mlp": 720,
    "max_seq_len": 512,
    "rope_theta": 10000.0,
    "dropout": 0.0,
    "norm_eps": 1e-06,
    "init_std": 0.02,
    "scale_resid_init": true
  },
  "train": {
    "K": 1,
    "k_schedule": "const",
    "k_end": 1,
    "k_decay_start_frac": 0.5,
    "k_decay_end_frac": 0.9,
    "k_min": 1,
    "drop_eps": 0.01,
    "eval_k_max": 8,
    "batch_size": 16,
    "grad_accum": 2,
    "seq_len": 512,
    "max_tokens": 60000000,
    "lr": 0.001,
    "min_lr_ratio": 0.1,
    "warmup_steps": 100,
    "weight_decay": 0.1,
    "beta1": 0.9,
    "beta2": 0.95,
    "grad_clip": 1.0,
    "ce_weight": 1.0,
    "teacher_run": "",
    "teacher_k": 0,
    "init_from_run": "",
    "distill_weight": 0.0,
    "distill_space": "hidden",
    "distill_depths": "first",
    "distill_start_frac": 0.0,
    "distill_tau": 1.0,
    "anytime_ce_weight": 0.0,
    "grad_checkpoint": false,
    "budget_locked": true,
    "log_every": 20,
    "eval_every": 150,
    "val_batches": 32,
    "ckpt_every": 1000,
    "best_metric": "val_ce",
    "seed": 1337
  },
  "K": 1,
  "params": 9962496,
  "training_tokens": 60014592
}