# Example configuration for supervised fine-tuning
# Uses TaoNet (MLA+RoPE) architecture loaded from pretrained checkpoint

# ============================================================================
# Model Architecture - TaoNet (DeepSeek MLA + RoPE)
# ============================================================================
model:
  architecture_type: taonet
  vocab_size: 8192
  hidden_dim: 1024
  num_layers: 16
  num_heads: 8
  max_seq_length: 1024
  
  # TaoNet-specific: Multi-head Latent Attention (MLA) compression
  d_latent_kv: 768
  
  # RoPE (Rotary Position Embedding) dimension per head
  # With hidden_dim=1024 and num_heads=8, head_dim = 128
  d_rope: 128
  
  # Feed-forward intermediate dimension
  hidden_dim_ff: 3072
  
  # Dropout rate (low for stability with large models)
  dropout: 0.02
  
  # Grouped Query Attention (1 = standard MLA, >1 = GQA)
  gqa_groups: 1
  
  # Optional: Use factorized embedding for parameter efficiency
  # vocab (8192) → rank (96) → hidden (1024)
  use_factorized_embedding: true
  d_embed_rank: 96
  
  # Weight initialization standard deviation
  init_std: 0.02

dataset:
  split: train
  instruction_column: input
  response_column: output
  
  local: true
  jsonl_path: /home/student/Data/TaoData/sft.jsonl
  samples_per_chunk: 1000
  max_samples: 160000
  cache_dir: .cache/datasets
  instruction_template: "{instruction}\n{response}"

  # Tokenizer configuration
  tokenizer_type: sentencepiece
  tokenizer_path: tokenizer/tokenizer.model
  tokenizer_threads: 4

# SFT-specific configuration (these fields are in SFTConfig)
checkpoint_path: "checkpoints/yarn8k/best_model.pt"
user_token: "<user>"
assistant_token: "<assistant>"
response_loss_only: true

batch_size: 8
num_epochs: 1
gradient_accumulation_steps: 4
max_grad_norm: 1.0

optimizer:
  optimizer_type: adamw
  learning_rate: 5e-5  # Lower LR for fine-tuning (vs 5e-4 pretrain base, 5e-3 Muon)
  weight_decay: 0.01

scheduler:
  scheduler_type: linearWarmup
  warmup_steps: 500  # Less aggressive warmup for fine-tuning

dtype: bfloat16
device: cuda

checkpoint_dir: checkpoints/sft
save_every_steps: 81920
save_best_model: true
keep_last_n_checkpoints: 2

eval_every_steps: 8192
eval_samples: 200

log_every_steps: 10
aim_repo: .aim

seed: 42
num_workers: 0
pin_memory: true