# TaoNet T2 Configuration for Pretraining
# DeepSeek MLA + RoPE with Hybrid Muon+AdamW Optimizer
# Full BF16 precision (no quantization)

# ============================================================================
# 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 Configuration - Local JSONL
# ============================================================================
dataset:
  local: true
  jsonl_path: /home/student/Data/TaoData/pretrain.jsonl
  text_field: text
  max_samples: 6700000
  samples_per_chunk: 1000
  
  # Tokenizer configuration
  tokenizer_type: sentencepiece
  tokenizer_path: tokenizer/tokenizer.model
  tokenizer_threads: 4

# ============================================================================
# Training Hyperparameters
# ============================================================================
batch_size: 8
num_epochs: 1  # Set to 10 for full training
gradient_accumulation_steps: 32  # Effective batch: 8 × 32 = 256

# Maximum gradient norm for clipping (prevents ternary instability)
max_grad_norm: 1.0

# ============================================================================
# Optimizer - Hybrid Muon + AdamW
# ============================================================================
# Strategy:
# - Muon: For 2D Linear weight matrices (orthogonal/SVD-based optimization)
#   - 2D weights: learning_rate (3e-3)
# - AdamW: For 1D parameters (biases, norms, embeddings)
#   - 1D params: adamw_lr (3e-4) = 1/10 × learning_rate

optimizer:
  optimizer_type: hybrid_muon_adamw
  
  # Learning rate for Muon (2D weight matrices)
  learning_rate: 3e-3
  
  # Learning rate for AdamW (1D parameters)
  # Typically 1/10 of learning_rate to prevent over-updating 1D params
  adamw_lr: 3e-4
  
  # L2 regularization (weight decay)
  weight_decay: 0.01
  
  # Adam betas
  betas: [0.9, 0.999]
  
  # Epsilon for numerical stability
  eps: 1e-8

# ============================================================================
# Learning Rate Scheduler - 3-Phase Cosine with Warmup
# ============================================================================
# Phases:
# 1. Warmup: 0 → 1.0 (300 steps, ~1.4% of training)
# 2. Steady: 1.0 (constant for 5% of training)
# 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 large models
device: cuda     # Use GPU for training

# ============================================================================
# Checkpointing and Validation
# ============================================================================
checkpoint_dir: checkpoints/test
save_every_steps: 81920
save_best_model: true
keep_last_n_checkpoints: 3

# Validation
eval_every_steps: 8192
eval_samples: 8000

# ============================================================================
# Logging
# ============================================================================
log_every_steps: 50
aim_repo: .aim

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