mini_config.py

#!/usr/bin/env python3
# coding=utf-8
"""
V2-7way Mini Configuration (Phase 4 PoC).

목적: 1B params 작은 모델로 49 layer 7×7 Latin Square 구조 검증.
- 단일 B200 1대 (192GB)에서 학습 가능
- Mini PoC 1B → 본 학습 30B 검증 후 진행

사용:
  python -c "from mini_config import mini_cfg; print(mini_cfg)"
  또는 train_v2_mini.py에서 import
"""
from aether_pkg.configuration_aether_v2_7way import AETHERV27wayConfig


# =============================================================================
# Mini Scale (1B params)
# =============================================================================
mini_cfg = AETHERV27wayConfig(
    # Hidden / FFN (Mini scale: hidden 4096 → 2048)
    hidden_size=2048,
    intermediate_size=6144,
    expert_intermediate_size=640,

    # 49 layers (구조 그대로 검증)
    num_hidden_layers=49,

    # Attention (mini)
    num_attention_heads=16,
    num_key_value_heads=4,
    head_dim=128,
    sliding_window_size=512,
    compress_block_size=16,

    # MoE (구조 그대로)
    num_experts=25,
    num_experts_per_tok=7,
    use_shared_expert=True,

    # Position
    max_position_embeddings=4096,
    rope_theta=10000.0,

    # Norm + activation
    rms_norm_eps=1e-6,
    hidden_act="silu",
    attention_dropout=0.0,

    # Vocab
    vocab_size=151936,
    pad_token_id=151643,

    # Training
    initializer_range=0.02,
    use_cache=False,
    output_router_logits=True,
    router_aux_loss_coef=0.001,
    tie_word_embeddings=False,
)


# =============================================================================
# Nano Scale (~100M params, 단일 GPU 빠른 검증)
# =============================================================================
nano_cfg = AETHERV27wayConfig(
    hidden_size=512,
    intermediate_size=1536,
    expert_intermediate_size=192,
    num_hidden_layers=49,
    num_attention_heads=8,
    num_key_value_heads=2,
    head_dim=64,
    sliding_window_size=256,
    compress_block_size=8,
    num_experts=25,
    num_experts_per_tok=7,
    use_shared_expert=True,
    max_position_embeddings=2048,
    rope_theta=10000.0,
    rms_norm_eps=1e-6,
    hidden_act="silu",
    attention_dropout=0.0,
    vocab_size=151936,
    pad_token_id=151643,
    initializer_range=0.02,
    use_cache=False,
    output_router_logits=True,
    router_aux_loss_coef=0.001,
    tie_word_embeddings=False,
)


# =============================================================================
# Full Scale (30B params, 본 학습용)
# =============================================================================
full_cfg = AETHERV27wayConfig(
    hidden_size=4096,
    intermediate_size=12288,
    expert_intermediate_size=1280,
    num_hidden_layers=49,
    num_attention_heads=32,
    num_key_value_heads=8,
    head_dim=128,
    sliding_window_size=2048,
    compress_block_size=64,
    num_experts=25,
    num_experts_per_tok=7,
    use_shared_expert=True,
    max_position_embeddings=8192,
    rope_theta=1000000.0,
    rms_norm_eps=1e-6,
    hidden_act="silu",
    attention_dropout=0.0,
    vocab_size=151936,
    pad_token_id=151643,
    initializer_range=0.02,
    use_cache=False,
    output_router_logits=True,
    router_aux_loss_coef=0.001,
    tie_word_embeddings=False,
)


# =============================================================================
# Param 카운트 추정 (참고용)
# =============================================================================
def estimate_params(cfg):
    """Rough parameter count for an AETHER-V2-7way config."""
    h = cfg.hidden_size
    L = cfg.num_hidden_layers
    n_e = cfg.num_experts
    e_int = cfg.expert_intermediate_size
    n_h = cfg.num_attention_heads
    n_kv = getattr(cfg, "num_key_value_heads", n_h)
    head_dim = cfg.head_dim
    vocab = cfg.vocab_size

    # Embed + LM head (tied or not)
    embed = vocab * h
    lm_head = vocab * h

    # Per-layer attention: q,k,v,o
    attn_per_layer = h * (n_h * head_dim) + h * (n_kv * head_dim) * 2 + (n_h * head_dim) * h

    # Per-layer MoE: 25 experts × (gate + up + down)
    expert_per_layer = n_e * (h * e_int * 3)
    # Shared expert
    shared = h * e_int * 3 if cfg.use_shared_expert else 0
    # Router gate
    router = h * n_e

    # Per-layer norms (RMSNorm: 2 weights)
    norms = h * 2

    per_layer = attn_per_layer + expert_per_layer + shared + router + norms
    total = embed + lm_head + per_layer * L + h  # final norm

    return {
        "total": total,
        "total_M": total / 1e6,
        "total_B": total / 1e9,
        "embed": embed,
        "per_layer_attn": attn_per_layer,
        "per_layer_moe": expert_per_layer + shared + router,
        "L": L,
    }


if __name__ == "__main__":
    print("=" * 70)
    print("V2-7way Config Sizes")
    print("=" * 70)
    for name, cfg in [("nano", nano_cfg), ("mini", mini_cfg), ("full", full_cfg)]:
        info = estimate_params(cfg)
        print(f"\n[{name}]")
        print(f"  hidden_size:   {cfg.hidden_size}")
        print(f"  layers:        {cfg.num_hidden_layers}")
        print(f"  experts:       {cfg.num_experts} (top-{cfg.num_experts_per_tok})")
        print(f"  expert dim:    {cfg.expert_intermediate_size}")
        print(f"  attn heads:    {cfg.num_attention_heads} ({cfg.num_key_value_heads} kv)")
        print(f"  vocab:         {cfg.vocab_size}")
        print(f"  estimated params: {info['total_B']:.2f}B")


# =============================================================================
# Onebee Scale (~1.03B params, Phase 2 Chinchilla-optimal - 20B 20:1)
# =============================================================================
onebee_cfg = AETHERV27wayConfig(
    hidden_size=768,
    intermediate_size=2304,
    expert_intermediate_size=256,
    num_hidden_layers=49,
    num_attention_heads=8,
    num_key_value_heads=2,
    head_dim=64,
    sliding_window_size=256,
    compress_block_size=8,
    num_experts=25,
    num_experts_per_tok=7,
    use_shared_expert=True,
    max_position_embeddings=2048,
    rope_theta=10000.0,
    rms_norm_eps=1e-6,
    hidden_act="silu",
    attention_dropout=0.0,
    vocab_size=151936,
    pad_token_id=151643,
    initializer_range=0.02,
    use_cache=False,
    output_router_logits=True,
    router_aux_loss_coef=0.001,
    tie_word_embeddings=False,
)