# Block-wise AR Diffusion Config for 2 ancestor levels (128 -> 32).

model:
  vocab_size: 50257
  hidden_size: 768
  n_blocks: 12
  n_heads: 12
  cond_dim: 128
  max_seq_len: 512
  block_size: 32
  dropout: 0.0
  # level 0 = leaf token
  # level 1 = 128 ancestors
  # level 2 = 32 ancestors
  num_levels: 3
  level_sizes: [50257, 128, 32]

ancestor:
  # Every ancestor level is selected directly from the leaf token:
  # level 1 top-3, level 2 top-2.
  lut_path: data/ancestor_lut_50257-128-32_leaf_to_each_level_topk_3-2_t1.0.pt
  proto_path: data/hierarchy_prototypes_50257-128-32-8.pt

hierarchy:
  prototype_cache: data/hierarchy_prototypes_50257-128-32-8.pt

loss:
  lambda_ancestor: 0.0
  mask_only: true

training:
  seed: 0
  batch_size: 64
  num_steps: 1_000_000
  lr: 3.0e-4
  lr_min: 3.0e-5
  warmup_steps: 2000
  weight_decay: 0.01
  grad_clip: 1.0
  dtype: bf16
  compile: default

  log_interval: 100
  eval_interval: 5000
  save_interval: 10000

data:
  dataset: openwebtext
  seq_len: 512
  cache_dir: data/owt_cache
  num_workers: 4
  max_train_samples: null
  max_val_samples: 100000
  mode: subsample

logging:
  use_wandb: true
  project: sad_b32_h2_mixed
  save_dir: outputs/sad_b32_h2_mixed