| from dataclasses import dataclass, field | |
| from typing import List | |
| class ModelConfig: | |
| # ── Architecture ────────────────────────────────────────────────────── | |
| vocab_size: int = 5010 # 5000 Words + 4 Specials + 6 Agentic (CoT & Tools) | |
| block_size: int = 2048 # Context window (increased for deeper reasoning) | |
| n_embd: int = 512 # Embedding dimension (384→512 for higher capacity) | |
| n_head: int = 16 # Query heads (12→16) | |
| n_kv_heads: int = 4 # Key/Value heads (GQA 4:1 for KV cache savings) | |
| n_layer: int = 16 # Transformer layers (12→16 for deeper representations) | |
| # ── MoE ─────────────────────────────────────────────────────────────── | |
| n_experts: int = 8 # Total number of routed experts | |
| n_shared_experts: int = 1 # Always-active shared experts (DeepSeek pattern) | |
| top_k: int = 2 # Experts activated per token | |
| moe_intermediate_size: int = 768 # Expert MLP hidden size (512→768) | |
| expert_capacity_factor: float = 1.25 # Capacity buffer to prevent dropping | |
| # ── Training ────────────────────────────────────────────────────────── | |
| batch_size: int = 8 # Micro batch (reduced for 8GB VRAM) | |
| learning_rate: float = 2e-4 # Slightly lower for stability | |
| weight_decay: float = 0.01 | |
| dropout: float = 0.05 # Lower dropout for better generalization | |
| label_smoothing: float = 0.1 | |
| max_grad_norm: float = 1.0 | |
| warmup_steps: int = 100 | |
| # ── Gradient Accumulation ───────────────────────────────────────────── | |
| gradient_accumulation_steps: int = 16 # Effective batch = 8*16 = 128 | |
| # ── RTX 4060 Optimizations ──────────────────────────────────────────── | |
| device: str = 'cuda' | |
| use_checkpointing: bool = True # Gradient checkpointing (cuts VRAM ~2x) | |
| pin_memory: bool = True # DMA pinned memory for async transfer | |
| expert_offloading: bool = True # Offload idle MoE experts to CPU | |
| use_compile: bool = True # torch.compile for CUDA graphs | |
| use_amp: bool = True # Automatic mixed precision (bf16/fp16) | |
| # ── MLA (Multi-head Latent Attention) ───────────────────────────────── | |
| use_mla: bool = True # Compress KV through low-rank bottleneck | |
| mla_latent_dim: int = 128 # Bottleneck dimension for KV compression | |
| # ── EMA (Exponential Moving Average) ────────────────────────────────── | |
| use_ema: bool = True | |
| ema_decay: float = 0.999 | |
| # ── Early Stopping ──────────────────────────────────────────────────── | |
| patience: int = 5 | |
| min_delta: float = 1e-4 | |