# byrne_100m_ultrax_mc.yaml -- Byrne-100M-Looped, DENSE + looped, on UltraX,
# with Memory Caching (arXiv 2602.24281) enabled: a linear-attention memory
# branch parallel to MLA giving growing cross-segment memory (Gated Residual).
#
# vs byrne_100m (Stack/code): dataset is pure UltraX web (code/math/dclm = 0),
# and use_memory_cache=true. The MC branch is added through a zero-init gate, so
# at init the model is identical to the no-MC baseline and can only help from there.
# vocab_size omitted on purpose (trainer derives 16512 from the tokenizer).

model:
  hidden_size: 768
  num_hidden_layers: 16
  num_attention_heads: 12
  num_key_value_heads: 2
  head_dim: 64
  qk_rope_head_dim: 16
  q_lora_rank: 128
  o_lora_rank: 128

  tie_word_embeddings: true
  use_moe: false
  moe_layers: []
  moe_intermediate_size: 2048
  scoring_func: sqrtsoftplus
  num_hash_layers: 2
  moe_aux_loss_coef: 0.01

  use_hyper_connections: true
  hc_mult: 2
  num_nextn_predict_layers: 1
  mtp_loss_weight: 0.3

  use_derf: false
  use_xsa: true
  use_qk_norm: true
  zloss_coef: 1.0e-4
  use_value_embed: false

  use_engram: true
  engram_compress_dim: 32
  engram_num_heads: 2
  engram_table_size: 4096
  engram_max_ngram: 3
  engram_gate_init_bias: -1.0

  use_elo: false
  use_fractal_rope: true
  fractal_rope_gamma: 1.0

  use_hrm_refine: true
  hrm_refine_steps: 1
  hrm_refine_dim: 128
  hrm_deep_supervision: true
  use_abstain_head: false

  loop_count: 3
  loop_pass_embed: true

  # --- Memory Caching (arXiv 2602.24281) ---
  use_memory_cache: true
  mc_segment_len: 256          # seq 1024 -> 4 segments (growing memory across them)
  mc_num_heads: 4
  mc_head_dim: 32
  mc_gate_dim: 64

# Training -- Chinchilla for ~100M dense: ~2.0B tokens.
steps: 61000
batch_size: 6
grad_accum: 4
lr: 3.0e-4
min_lr_frac: 0.1
warmup_steps: 2000
weight_decay: 0.01
grad_clip: 1.0
seq_len: 1024                  # 4 x mc_segment_len(256); MC cross-seg cost is O(N^2)
muon_lr_mult: 10.0

# Blend: UltraX 50.15% + DCLM 26.13% + Cosmopedia-v2 13.53% + FineMath-4+ 7.88%
# + NPset-2 Python-Edu 2.31%. code/math weights 0 -- every non-UltraX source is
# an explicit extra_source below; UltraX web = 1 - sum(extra_sources) = 0.5015.
# Knowledge-heavy blend: FineWeb-Edu (100B sample) REPLACES UltraX as the base
# web source, + Wikipedia added, to target weak factual recall. web_source gets
# the leftover weight = 1 - sum(extra_sources) = 1 - 0.60 = 0.40.
#   FineWeb-Edu 40% | Wikipedia 15% | DCLM 15% | Cosmopedia 10% | FineMath 10% | Code 10%
dataset: ultrax-ultra-fineweb   # unused label; base source is web_source below
code_weight: 0.0
math_weight: 0.0
web_source: {path: HuggingFaceFW/fineweb-edu, name: sample-100BT, text_key: text}
extra_sources:
  - {path: wikimedia/wikipedia, name: 20231101.en, text_key: text, weight: 0.15}
  - {path: mlfoundations/dclm-baseline-1.0, text_key: text, weight: 0.15}
  - {path: HuggingFaceTB/smollm-corpus, name: cosmopedia-v2, text_key: text, weight: 0.10}
  - {path: HuggingFaceTB/finemath, name: finemath-4plus, text_key: text, weight: 0.10}
  - {path: AxiomicLabs/NPset-2-Python-Edu, text_key: code, weight: 0.10}
tokenizer: ./tokenizer.json

compile: false
gradient_checkpointing: true
num_workers: 4

output_dir: checkpoints_100m_ultrax_mc
save_every: 500
eval_every: 1000
log_every: 50
keep_checkpoints: 3