Kiến trúc Nexus Coder / Nexus Coder Architecture
Tổng quan / Overview
Nexus Coder v0.1 sử dụng kiến trúc Mixture of Experts (MoE) Transformer tương tự Mixtral 8x7B và DeepSeek-V3.
Nexus Coder v0.1 uses a Mixture of Experts (MoE) Transformer architecture similar to Mixtral 8x7B and DeepSeek-V3.
Các thành phần / Components
1. Token Embedding
- Vocab size: 32,000
- Hidden size: 2,048
- Tokens được nhúng thành vector 2048 chiều
2. Grouped Query Attention (GQA)
- 16 query heads
- 4 KV heads (ratio 4:1)
- Head dimension: 128
- Giảm 4x memory cho KV cache so với MHA truyền thống
3. Rotary Position Embedding (RoPE)
- Base: 10,000
- Hỗ trợ tối đa 50,000 positions
- Cho phép model hiểu vị trí tương đối giữa các tokens
4. RMSNorm
- Thay thế LayerNorm truyền thống
- Không có bias, không trừ mean
- Nhanh hơn ~10-20%
5. SwiGLU Activation
SiLU(gate(x)) * up(x)- Hiệu quả hơn ReLU/GELU
- Có 3 ma trận: gate, up, down (3 * hidden * intermediate params)
6. Mixture of Experts (MoE) - Cốt lõi
- 24 experts tổng cộng (mỗi expert là một SwiGLU FFN)
- 3 active experts mỗi token (top-3 routing)
- Router: linear layer (hidden_size → num_experts)
- Load balancing loss: auxiliary loss để tránh expert collapse
Routing Algorithm
1. Router tính gate_logits = W_router @ x
2. routing_weights = softmax(gate_logits)
3. top_k_weights, top_k_indices = topk(routing_weights, k=3)
4. Normalize top_k_weights
5. Mỗi token đi qua 3 expert được chọn
6. Output = sum(weight_i * expert_i(x))
Tính toán tham số / Parameter Math
Embedding: vocab_size × hidden = 32000 × 2048 = 65.5M
Per layer attn: 2048² + 2×(2048×512) + 2048² = 10.5M (Q, K, V, O with GQA)
Per expert: 3 × 2048 × 5632 = 34.6M (gate + up + down)
Per layer MoE: 24 × 34.6M = 830M (total)
3 × 34.6M = 104M (active)
Per layer total: 10.5M + 830M = 840.5M
12 layers: 10,086M
LM head: 65.5M
────────────────────────────────────
TOTAL: 10,223M ≈ 10.22B ✓
ACTIVE: 65.5 + 12×(10.5 + 104) + 65.5 = 1,503M ≈ 1.50B ✓
Workflow
Training Workflow
- Tokenize input text → token IDs
- Embed tokens → hidden states [B, L, H]
- For each layer:
- Pre-norm → Attention → residual
- Pre-norm → MoE (router + experts) → residual
- Final norm → LM head → logits
- Compute cross-entropy loss + aux loss
- Backpropagation
Inference Workflow
- Tokenize prompt
- Forward pass through all layers
- Get logits for last position
- Apply temperature, top-k, top-p
- Sample next token
- Append to sequence, repeat
Tối ưu / Optimizations
- KV Cache: Cache K, V từ các token trước để tăng tốc generation
- GQA: Giảm memory và computation cho attention
- Pre-norm: Ổn định hơn post-norm trong training
- Mixed Precision: Hỗ trợ fp16/bf16 để tiết kiệm memory
- Gradient Checkpointing: Đánh đổi compute lấy memory (chưa implement trong v0.1)