File size: 3,554 Bytes
38e67d2
da7f1bf
4b89540
38e67d2
 
 
 
 
 
 
 
 
4b89540
 
 
38e67d2
 
4b89540
 
38e67d2
4b89540
 
 
38e67d2
 
 
 
 
 
4b89540
38e67d2
 
 
 
4b89540
 
 
 
38e67d2
 
 
 
 
 
 
 
 
 
4b89540
38e67d2
 
76a2013
 
 
 
4b89540
 
 
38e67d2
76a2013
4b89540
38e67d2
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
"""
Verification Suite for VIU-1.5B-MoE Ultra-Deep 48-Layer Architecture.
Audits parameter counts, active compute, MLA latent dimensions, and expert distributions.
"""
import sys
from pathlib import Path
import torch

sys.path.append(str(Path(__file__).parent))
from viu_moe import ViuArgs, Viu1MoE, load_args

def verify_architecture(config_path: str):
    print("=" * 75)
    print(f"VERIFYING FRONTIER ARCHITECTURE: {config_path}")
    print("=" * 75)
    
    cfg = load_args(config_path)
    print(f"[Config] dim={cfg.dim}, layers={cfg.n_layers}, heads={cfg.n_heads}")
    print(f"[MLA]    q_lora_rank={cfg.q_lora_rank}, kv_lora_rank={cfg.kv_lora_rank}, rope_dim={cfg.rope_dim}")
    print(f"[Config] vocab={cfg.vocab_size}, max_seq={cfg.max_seq_len}, rope_theta={cfg.rope_theta}")
    print(f"[MoE]    routed_experts={cfg.moe_num_routed_experts}, shared_experts={cfg.moe_shared_experts}, top_k={cfg.moe_top_k}")
    print(f"[MoE]    moe_every={cfg.moe_every}, expert_div={cfg.moe_expert_div}, multi_token={cfg.multi_token}")
    print(f"[Caps]   attn_softcap={cfg.attn_logit_softcapping}, final_softcap={cfg.final_logit_softcapping}")
    
    # Instantiate on meta device (0 memory allocation)
    with torch.device('meta'):
        model = Viu1MoE(cfg)
        total_p, active_p = model.param_counts()
        
    print("-" * 75)
    print(f"[Audit] Total Parameters:   {total_p:,} ({total_p/1e6:.2f}M / {total_p/1e9:.3f}B)")
    print(f"[Audit] Active Parameters:  {active_p:,} ({active_p/1e6:.2f}M / {active_p/1e9:.3f}B)")
    print(f"[Audit] Active Sparsity:    {active_p / total_p * 100:.2f}% active per token")
    
    total_experts = cfg.n_layers * (cfg.moe_num_routed_experts + cfg.moe_shared_experts)
    print(f"[Audit] Total Experts:      {total_experts:,} experts ({cfg.moe_shared_experts} shared + {cfg.moe_num_routed_experts} routed per layer)")
    print(f"[Audit] Attention Heads:    {cfg.n_heads} Latent Query Heads (MLA with decoupled RoPE)")
    print(f"[Audit] Head Dimension:     {cfg.head_dim} (Standard LLM head size)")
    
    # Layer type distributions
    n_moe = sum(1 for b in model.layers if b.is_moe)
    n_dense = cfg.n_layers - n_moe
    n_sliding = sum(1 for b in model.layers if b.sliding)
    n_full = sum(1 for b in model.layers if not b.sliding)
    print(f"[Layers] MoE Layers:        {n_moe}/{cfg.n_layers}")
    print(f"[Layers] Dense Layers:      {n_dense}/{cfg.n_layers}")
    print(f"[Layers] Sliding Window:    {n_sliding}/{cfg.n_layers} (window={cfg.sliding_window})")
    print(f"[Layers] Full Attention:    {n_full}/{cfg.n_layers}")
    print("-" * 75)
    
    # Assertions
    assert 1_800_000_000 <= total_p <= 1_920_000_000, f"Total params {total_p} outside expected range!"
    assert 280_000_000 <= active_p <= 320_000_000, f"Active params {active_p} outside expected range!"
    assert cfg.n_layers == 48, f"Expected 48 layers, got {cfg.n_layers}!"
    assert total_experts == 1392, f"Expected 1,392 experts, got {total_experts}!"
    assert cfg.moe_shared_experts == 1, f"Expected 1 shared expert, got {cfg.moe_shared_experts}!"
    assert cfg.moe_num_routed_experts == 28, f"Expected 28 routed experts, got {cfg.moe_num_routed_experts}!"
    assert cfg.q_lora_rank == 256 and cfg.kv_lora_rank == 256, "Expected MLA ranks 256!"
    
    print("SUCCESS: VIU-1.5B-MoE 48-Layer Frontier Architecture is 100% Validated!")
    print("=" * 75)

if __name__ == "__main__":
    default_cfg = Path(__file__).parent.parent / "configs" / "model_config.yaml"
    verify_architecture(str(default_cfg))