# TaoNet 200M Configuration for YaRN Continued Pretraining (4k -> 8k)

model:
  architecture_type: taonet
  vocab_size: 8192
  hidden_dim: 1024
  num_layers: 16
  num_heads: 8
  max_seq_length: 8192

  d_latent_kv: 768
  d_rope: 128
  hidden_dim_ff: 3072
  dropout: 0.02
  gqa_groups: 1
  use_factorized_embedding: true
  d_embed_rank: 96
  init_std: 0.02

  rope_scale: 40.0
  yarn_enabled: true
  yarn_original_max_seq_length: 4096
  yarn_alpha: 1.0

dataset:
  local: true
  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
  text_field: text
  max_samples: 800000
  samples_per_chunk: 1000

  tokenizer_type: sentencepiece
  tokenizer_path: tokenizer/tokenizer.model
  tokenizer_threads: 4

checkpoint_path: checkpoints/yarn4k/best_model.pt

batch_size: 2
num_epochs: 1
gradient_accumulation_steps: 32
max_grad_norm: 1.0

optimizer:
  optimizer_type: hybrid_muon_adamw
  learning_rate: 7.5e-4
  adamw_lr: 7.5e-5
  weight_decay: 0.01
  betas: [0.9, 0.999]
  eps: 1e-8

scheduler:
  scheduler_type: cosineWarmup
  warmup_steps: 300
  warmup_ratio: 0.0
  steady_ratio: 0.05
  min_lr_ratio: 0.1
  num_cycles: 0.5

dtype: bfloat16
device: cuda

checkpoint_dir: checkpoints/yarn8k
save_every_steps: 2048
save_best_model: true
keep_last_n_checkpoints: 3

eval_every_steps: 2048
eval_samples: 2000

log_every_steps: 50
aim_repo: .aim

seed: 42
num_workers: 0
pin_memory: true