code/config.py

from dataclasses import dataclass

@dataclass
class ViuAIConfig:
    vocab_size: int = 64009
    d_model: int = 1280
    n_layers: int = 24
    n_heads: int = 20
    n_kv_heads: int = 4
    ffn_hidden: int = 3456
    context_length: int = 2048
    rope_theta: float = 10000.0
    norm_eps: float = 1e-5
    z_loss_weight: float = 1e-4    # set to 0.0 during SFT
    use_checkpoint: bool = True
    attn_dropout: float = 0.0      # 0.0 for pretraining, 0.1 for SFT
    resid_dropout: float = 0.0     # 0.0 for pretraining, 0.1 for SFT