# TaoNet T2 Configuration for YaRN Continued Pretraining
# Extended Context: 1024 → 8192 tokens with frequency interpolation
# Built on DeepSeek MLA + RoPE with Hybrid Muon+AdamW Optimizer
# Full BF16 precision (no quantization)

# ============================================================================
# Model Architecture - TaoNet (DeepSeek MLA + RoPE) with YaRN Extension
# ============================================================================
model:
  architecture_type: taonet
  vocab_size: 8192
  hidden_dim: 512
  num_layers: 12
  num_heads: 8
  max_seq_length: 8192  # Extended from 1024 → 8192 (8x longer context)
  
  # TaoNet-specific: Multi-head Latent Attention (MLA) compression
  # KV dimension reduced from 512 to 384 (25% memory savings)
  d_latent_kv: 384
  
  # RoPE (Rotary Position Embedding) dimension per head
  # Default would be 512 / 8 = 64
  d_rope: 64
  
  # Feed-forward intermediate dimension
  # Default would be 4 * 512 = 2048
  hidden_dim_ff: 1024
  
  # Dropout rate (low for stability)
  dropout: 0.02
  
  # Grouped Query Attention (1 = standard MLA, >1 = GQA)
  gqa_groups: 1
  
  # Optional: Use factorized embedding for parameter efficiency
  use_factorized_embedding: false
  d_embed_rank: 96
  
  # Weight initialization standard deviation
  init_std: 0.02
  
  # ========================================================================
  # YaRN (Yet another RoPE eXtension) Configuration
  # ========================================================================
  # Enables frequency interpolation to extend context length from 1024 → 8192
  # The model learns to "pack" RoPE frequencies into the new longer context during training.
  
  # RoPE base scale factor (explicit, previously hardcoded to 40)
  rope_scale: 40.0
  
  # Enable YaRN frequency interpolation
  yarn_enabled: true
  
  # Interpolation smoothness parameter
  # - 1.0 (default): Smooth, gradual interpolation—safer for learning extended context
  # - 0.5: Aggressive interpolation—faster context expansion, higher risk
  # - 2.0: Conservative interpolation—safer but slower adaptation
  # Recommendation: Start with 1.0; tune in follow-up runs if convergence issues
  yarn_alpha: 1.0

# ============================================================================
# Dataset Configuration - Local JSONL (Same as Pretrain)
# ============================================================================
dataset:
  local: true
  jsonl_path: /home/student/Data/TaoData/output.jsonl
  text_field: text
  max_samples: 50000  # Reduced from 1M → 50k for quick YaRN adaptation
  samples_per_chunk: 1000
  
  # Tokenizer configuration (unchanged)
  tokenizer_type: sentencepiece
  tokenizer_path: tokenizer/tokenizer.model
  tokenizer_threads: 4

# ============================================================================
# Training Hyperparameters - Conservative for Context Extension
# ============================================================================
# Strategy: Lower learning rates + smaller batch to prevent catastrophic forgetting
# while the model learns to use 8x longer context.

batch_size: 16  # Reduced from 32 (8192 tokens/seq is memory-intensive)
num_epochs: 1   # 50k samples / effective_batch=256 ≈ 200 updates (1 epoch sufficient for warm-start)

# Gradient accumulation to maintain effective batch size of ~256
# Effective batch = batch_size × gradient_accumulation_steps = 16 × 16 = 256
gradient_accumulation_steps: 16

# Maximum gradient norm for clipping
max_grad_norm: 1.0

# ============================================================================
# Optimizer - Hybrid Muon + AdamW (Conservative LR for Stability)
# ============================================================================
# Strategy: Use 1/2 of pretrain learning rates to:
#   1. Avoid catastrophic forgetting of learned features
#   2. Allow smooth adaptation to YaRN-scaled RoPE frequencies
#   3. Give the model time to learn how to use extended context

optimizer:
  optimizer_type: hybrid_muon_adamw
  
  # Learning rate for Muon (2D weight matrices)
  # Reduced: 5e-3 → 2.5e-3 (50% of pretrain)
  learning_rate: 2.5e-3
  
  # Learning rate for AdamW (1D parameters)
  # Reduced: 5e-4 → 1.25e-4 (25% of pretrain, maintains 1/10 ratio)
  adamw_lr: 1.25e-4
  
  # L2 regularization (weight decay)
  weight_decay: 0.01
  
  # Adam betas (unchanged)
  betas: [0.9, 0.999]
  
  # Epsilon for numerical stability
  eps: 1e-8

# ============================================================================
# Learning Rate Scheduler - 3-Phase Cosine with Warmup (Same as Pretrain)
# ============================================================================
# Phases:
# 1. Warmup: 0 → 1.0 (300 steps, ~1.4% of training)
# 2. Steady: 1.0 (constant for 5% of training steps at peak LR)
# 3. Decay: 1.0 → 0.1 (cosine decay for remaining ~95%)

scheduler:
  scheduler_type: cosineWarmup
  warmup_steps: 300
  warmup_ratio: 0.0  # Ignored if warmup_steps > 0
  steady_ratio: 0.05  # 5% of total training steps at peak LR
  min_lr_ratio: 0.1   # Decay to 10% of peak LR
  num_cycles: 0.5  # For compatibility (not used in 3-phase schedule)

# ============================================================================
# Data Type and Device
# ============================================================================
dtype: bfloat16  # Use BF16 for better convergence with extended context
device: cuda     # Use GPU for training

# ============================================================================
# Checkpointing and Validation
# ============================================================================
# Load pretrained checkpoint and continue training
checkpoint_path: checkpoints/pretrain_taonet/best_model.pt
checkpoint_dir: checkpoints/yarn_taonet
save_every_steps: 512    # More frequent saves for 50k samples (200 updates total)
save_best_model: true
keep_last_n_checkpoints: 3

# Validation every 512 steps (10% of 50k samples)
eval_every_steps: 512
eval_samples: 2500  # Reduced from 8000

# ============================================================================
# Logging
# ============================================================================
log_every_steps: 50  # Log every 50 updates
aim_repo: .aim

# ============================================================================
# Miscellaneous
# ============================================================================
seed: 42
num_workers: 0
pin_memory: true

# ============================================================================
# YaRN Performance Notes
# ============================================================================
# Expected memory usage: ~1.5x of pretrain (8x longer seq, half batch)
# Expected training time: ~50-100 steps/min on H100 (depends on setup)
# Expected convergence: Loss should decrease over 50k samples; monitor perplexity on 8192-length sequences
# 
# Tuning recommendations for iterative improvements:
# 1. If loss is unstable: Reduce learning_rate further (1.25e-3)
# 2. If loss plateaus quickly: Increase max_samples (100k-200k)
# 3. If memory OOM: Reduce batch_size to 8 (maintain grad_accum at 16)
# 4. To speed context expansion: Reduce yarn_alpha to 0.5 (more aggressive)
# 5. For safer training: Increase yarn_alpha to 2.0 (more conservative)