422 lines
17 KiB
Python
422 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
XiaLi-Next 统一架构 v1.0
|
||
=========================
|
||
同一套架构，从 20M 到 7B+ 通用
|
||
核心设计：RoPE + SwiGLU + Pre-RMSNorm + GQA + Weight Tying
|
||
|
||
用法:
|
||
from xiaoli_next import XiaLiConfig, XiaLiModel
|
||
|
||
# 20M 超小模型
|
||
cfg = XiaLiConfig.preset('20m')
|
||
model = XiaLiModel(cfg)
|
||
|
||
# 1B 中等模型
|
||
cfg = XiaLiConfig.preset('1b')
|
||
model = XiaLiModel(cfg)
|
||
|
||
# 7B 大模型
|
||
cfg = XiaLiConfig.preset('7b')
|
||
model = XiaLiModel(cfg)
|
||
"""
|
||
|
||
import math
|
||
import json
|
||
from typing import Optional, Tuple
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 1. 配置系统
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
class XiaLiConfig:
|
||
"""统一配置 - 通过 preset 快速切换规模"""
|
||
|
||
def __init__(self):
|
||
# 架构
|
||
self.d_model: int = 320
|
||
self.n_layers: int = 8
|
||
self.n_heads: int = 8
|
||
self.n_kv_heads: int = 4 # GQA: KV头数（=n_heads 就是MHA，=1 就是MQA）
|
||
self.d_ff: int = 0 # 0=自动 = int(d_model * 8/3 * 2) SwiGLU需要额外gate
|
||
self.max_seq_len: int = 2048 # 最大序列长度
|
||
self.dropout: float = 0.0 # 大模型dropout=0
|
||
self.rope_theta: float = 10000.0 # RoPE base frequency
|
||
self.tie_weights: bool = True # Weight Tying
|
||
self.vocab_size: int = 20007
|
||
|
||
# 训练
|
||
self.norm_eps: float = 1e-5
|
||
|
||
@classmethod
|
||
def preset(cls, name: str) -> 'XiaLiConfig':
|
||
"""按规模取配置"""
|
||
cfgs = {
|
||
# ── 超小模型 ──
|
||
'20m': dict(d_model=288, n_layers=6, n_heads=6, n_kv_heads=3, max_seq_len=1024),
|
||
'26m': dict(d_model=320, n_layers=8, n_heads=8, n_kv_heads=4, max_seq_len=2048),
|
||
# ── 百M级 ──
|
||
'100m': dict(d_model=512, n_layers=12, n_heads=8, n_kv_heads=4, max_seq_len=2048),
|
||
'300m': dict(d_model=768, n_layers=16, n_heads=12, n_kv_heads=4, max_seq_len=4096),
|
||
# ── 十亿级 ──
|
||
'500m': dict(d_model=1024, n_layers=24, n_heads=16, n_kv_heads=8, max_seq_len=8192),
|
||
'1b': dict(d_model=1536, n_layers=24, n_heads=16, n_kv_heads=8, max_seq_len=8192),
|
||
'3b': dict(d_model=2560, n_layers=32, n_heads=24, n_kv_heads=8, max_seq_len=8192),
|
||
# ── 十亿级+ ──
|
||
'7b': dict(d_model=4096, n_layers=32, n_heads=32, n_kv_heads=8, max_seq_len=8192, rope_theta=500000.0),
|
||
'13b': dict(d_model=5120, n_layers=40, n_heads=40, n_kv_heads=10, max_seq_len=8192, rope_theta=500000.0),
|
||
'70b': dict(d_model=8192, n_layers=80, n_heads=64, n_kv_heads=8, max_seq_len=8192, rope_theta=500000.0),
|
||
}
|
||
cfg = cls()
|
||
if name in cfgs:
|
||
for k, v in cfgs[name].items():
|
||
setattr(cfg, k, v)
|
||
if cfg.d_ff == 0:
|
||
# SwiGLU 标准: 8/3 * d_model（参数量等价于 4x 标准 FFN）
|
||
cfg.d_ff = int(cfg.d_model * 8 / 3)
|
||
return cfg
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 2. RMSNorm
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
class XiaLiRMSNorm(nn.Module):
|
||
"""RMS Layer Normalization"""
|
||
def __init__(self, dim: int, eps: float = 1e-5):
|
||
super().__init__()
|
||
self.eps = eps
|
||
self.weight = nn.Parameter(torch.ones(dim))
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||
return x / rms * self.weight
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 3. RoPE（旋转位置编码）
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
def precompute_rope_freqs(dim: int, max_len: int, theta: float = 10000.0) -> Tuple[torch.Tensor, torch.Tensor]:
|
||
"""预计算 RoPE 频率"""
|
||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
|
||
t = torch.arange(max_len, dtype=torch.float32)
|
||
angles = torch.outer(t, freqs) # [max_len, dim//2]
|
||
cos = angles.cos().unsqueeze(0).unsqueeze(2) # [1, max_len, 1, dim//2]
|
||
sin = angles.sin().unsqueeze(0).unsqueeze(2)
|
||
return cos, sin
|
||
|
||
def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
|
||
"""对 x 应用 RoPE。x: [batch, seq, n_heads, head_dim]"""
|
||
head_dim = x.size(-1)
|
||
# cos/sin 是每对维度一个值 (head_dim//2)，需要扩张到 head_dim
|
||
cos = torch.repeat_interleave(cos, 2, dim=-1) # [..., d//2] -> [..., d]
|
||
sin = torch.repeat_interleave(sin, 2, dim=-1)
|
||
x1 = x[..., :head_dim // 2] # [..., d//2]
|
||
x2 = x[..., head_dim // 2:] # [..., d//2]
|
||
rotated = torch.cat([-x2, x1], dim=-1) # [..., d]
|
||
return x * cos + rotated * sin
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 4. GQA 注意力（分组查询注意力）
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
class GroupedQueryAttention(nn.Module):
|
||
"""分组查询注意力 - 小模型=MHA，大模型节省KV Cache"""
|
||
|
||
def __init__(self, config: XiaLiConfig):
|
||
super().__init__()
|
||
self.n_heads = config.n_heads
|
||
self.n_kv_heads = config.n_kv_heads # GQA: KV头数
|
||
self.head_dim = config.d_model // config.n_heads
|
||
self.n_rep = self.n_heads // self.n_kv_heads # 每个KV头被几个Q头共享
|
||
|
||
# 投影（无bias）
|
||
self.wq = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=False)
|
||
self.wk = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False)
|
||
self.wv = nn.Linear(config.d_model, self.n_kv_heads * self.head_dim, bias=False)
|
||
self.wo = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=False)
|
||
self.attn_dropout = nn.Dropout(config.dropout)
|
||
|
||
def forward(
|
||
self,
|
||
x: torch.Tensor,
|
||
cos: torch.Tensor,
|
||
sin: torch.Tensor,
|
||
mask: Optional[torch.Tensor] = None,
|
||
) -> torch.Tensor:
|
||
batch, seq_len, _ = x.shape
|
||
|
||
# QKV 投影
|
||
q = self.wq(x).view(batch, seq_len, self.n_heads, self.head_dim)
|
||
k = self.wk(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
|
||
v = self.wv(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
|
||
|
||
# RoPE
|
||
cos_ = cos[:, :seq_len]
|
||
sin_ = sin[:, :seq_len]
|
||
q = apply_rope(q, cos_, sin_)
|
||
k = apply_rope(k, cos_, sin_)
|
||
|
||
# GQA: KV 头扩展到 Q 头数
|
||
if self.n_rep > 1:
|
||
k = k.repeat_interleave(self.n_rep, dim=2)
|
||
v = v.repeat_interleave(self.n_rep, dim=2)
|
||
|
||
# 转置为 [batch, n_heads, seq, head_dim]
|
||
q = q.transpose(1, 2)
|
||
k = k.transpose(1, 2)
|
||
v = v.transpose(1, 2)
|
||
|
||
# 注意力计算
|
||
scale = self.head_dim ** -0.5
|
||
attn = torch.matmul(q, k.transpose(-2, -1)) * scale
|
||
|
||
if mask is not None:
|
||
attn = attn + mask[:, :, :seq_len, :seq_len]
|
||
|
||
attn = F.softmax(attn, dim=-1, dtype=torch.float32).to(q.dtype)
|
||
attn = self.attn_dropout(attn)
|
||
|
||
out = torch.matmul(attn, v) # [batch, n_heads, seq, head_dim]
|
||
out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
|
||
|
||
return self.wo(out)
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 5. SwiGLU 前馈网络
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
class SwiGLU(nn.Module):
|
||
"""SwiGLU 前馈网络 - 比 ReLU FFN 质量更高"""
|
||
|
||
def __init__(self, config: XiaLiConfig):
|
||
super().__init__()
|
||
# d_ff 已按 8/3*d_model 计算。SwiGLU: w1和w3同维度，w2降回
|
||
hidden = config.d_ff
|
||
self.w1 = nn.Linear(config.d_model, hidden, bias=False)
|
||
self.w2 = nn.Linear(hidden, config.d_model, bias=False)
|
||
self.w3 = nn.Linear(config.d_model, hidden, bias=False)
|
||
self.dropout = nn.Dropout(config.dropout)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
# SwiGLU(x) = (SiLU(x @ W1) * (x @ W3)) @ W2
|
||
return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 6. Transformer 块
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
class XiaLiBlock(nn.Module):
|
||
"""Transformer 块 - Pre-RMSNorm + GQA + SwiGLU"""
|
||
|
||
def __init__(self, config: XiaLiConfig, layer_idx: int):
|
||
super().__init__()
|
||
self.layer_idx = layer_idx
|
||
|
||
self.attention = GroupedQueryAttention(config)
|
||
self.feed_forward = SwiGLU(config)
|
||
|
||
self.attn_norm = XiaLiRMSNorm(config.d_model, config.norm_eps)
|
||
self.ff_norm = XiaLiRMSNorm(config.d_model, config.norm_eps)
|
||
|
||
def forward(
|
||
self,
|
||
x: torch.Tensor,
|
||
cos: torch.Tensor,
|
||
sin: torch.Tensor,
|
||
mask: Optional[torch.Tensor] = None,
|
||
) -> torch.Tensor:
|
||
h = x + self.attention(self.attn_norm(x), cos, sin, mask)
|
||
out = h + self.feed_forward(self.ff_norm(h))
|
||
return out
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 7. 主模型
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
class XiaLiModel(nn.Module):
|
||
"""
|
||
XiaLi-Next 统一架构
|
||
|
||
核心设计：
|
||
- RoPE 位置编码（无学习的 position embedding）
|
||
- Pre-RMSNorm（训练稳定）
|
||
- SwiGLU FFN（同等参数量更高质量）
|
||
- GQA 注意力（大模型省 KV Cache）
|
||
- Weight Tying（embedding 与输出层共享权重）
|
||
- 所有 Linear 无 bias（更干净）
|
||
- 从 20M 到 7B+ 统一架构
|
||
"""
|
||
|
||
def __init__(self, config: XiaLiConfig):
|
||
super().__init__()
|
||
self.config = config
|
||
|
||
# 词嵌入
|
||
self.embedding = nn.Embedding(config.vocab_size, config.d_model)
|
||
|
||
# 预计算 RoPE
|
||
cos, sin = precompute_rope_freqs(
|
||
config.d_model // config.n_heads,
|
||
config.max_seq_len * 2,
|
||
config.rope_theta,
|
||
)
|
||
self.register_buffer('rope_cos', cos, persistent=False)
|
||
self.register_buffer('rope_sin', sin, persistent=False)
|
||
|
||
# Transformer 块
|
||
self.blocks = nn.ModuleList([
|
||
XiaLiBlock(config, i) for i in range(config.n_layers)
|
||
])
|
||
|
||
# 最终归一化
|
||
self.norm = XiaLiRMSNorm(config.d_model, config.norm_eps)
|
||
|
||
# 输出层（与 embedding 共享权重）
|
||
self.fc_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
|
||
|
||
if config.tie_weights:
|
||
self.fc_out.weight = self.embedding.weight # Weight Tying
|
||
|
||
self._init_weights()
|
||
|
||
def _init_weights(self):
|
||
"""初始化权重"""
|
||
for name, param in self.named_parameters():
|
||
if param.dim() >= 2:
|
||
nn.init.normal_(param, mean=0.0, std=0.02)
|
||
else:
|
||
nn.init.zeros_(param)
|
||
|
||
def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
|
||
batch, seq_len = x.shape
|
||
|
||
# 词嵌入
|
||
h = self.embedding(x)
|
||
|
||
# RoPE
|
||
cos = self.rope_cos[:, :seq_len]
|
||
sin = self.rope_sin[:, :seq_len]
|
||
|
||
# Transformer 层
|
||
for block in self.blocks:
|
||
h = block(h, cos, sin, mask)
|
||
|
||
# 最终归一化
|
||
h = self.norm(h)
|
||
|
||
# 输出
|
||
return self.fc_out(h)
|
||
|
||
def generate_mask(self, seq_len: int) -> torch.Tensor:
|
||
"""生成因果掩码 [1, 1, seq, seq]"""
|
||
mask = torch.full((seq_len, seq_len), float('-inf'))
|
||
mask = torch.triu(mask, diagonal=1)
|
||
return mask.unsqueeze(0).unsqueeze(0)
|
||
|
||
@torch.no_grad()
|
||
def generate(
|
||
self,
|
||
input_ids: torch.Tensor,
|
||
max_new_tokens: int = 100,
|
||
temperature: float = 0.7,
|
||
top_k: int = 40,
|
||
eos_id: int = 3,
|
||
) -> torch.Tensor:
|
||
"""自回归生成"""
|
||
self.eval()
|
||
for _ in range(max_new_tokens):
|
||
# 截断到 max_seq_len
|
||
seq = input_ids if input_ids.size(1) <= self.config.max_seq_len else input_ids[:, -self.config.max_seq_len:]
|
||
|
||
mask = self.generate_mask(seq.size(1)).to(seq.device)
|
||
logits = self.forward(seq, mask)
|
||
logits = logits[:, -1, :] / temperature
|
||
|
||
# Top-k 采样
|
||
if top_k > 0:
|
||
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
|
||
logits[logits < v[:, -1:]] = float('-inf')
|
||
|
||
probs = F.softmax(logits, dim=-1)
|
||
next_token = torch.multinomial(probs, 1)
|
||
input_ids = torch.cat([input_ids, next_token], dim=-1)
|
||
|
||
if next_token.item() == eos_id:
|
||
break
|
||
|
||
return input_ids
|
||
|
||
def count_params(self) -> dict:
|
||
"""统计参数量"""
|
||
total = sum(p.numel() for p in self.parameters())
|
||
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
|
||
# 各模块
|
||
emb = self.embedding.weight.numel()
|
||
attn = sum(p.numel() for n, p in self.named_parameters() if 'attention' in n)
|
||
ffn = sum(p.numel() for n, p in self.named_parameters() if 'feed_forward' in n)
|
||
norm = sum(p.numel() for n, p in self.named_parameters() if 'norm' in n)
|
||
return {
|
||
'total': total,
|
||
'trainable': trainable,
|
||
'embedding': emb,
|
||
'attention': attn,
|
||
'ffn': ffn,
|
||
'norm': norm,
|
||
'total_m': total / 1e6,
|
||
}
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 8. 与原有系统兼容的训练适配
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
# 导出兼容名（方便从原系统切换）
|
||
XiaLiNextModel = XiaLiModel
|
||
XiaLiNextConfig = XiaLiConfig
|
||
|
||
|
||
# ════════════════════════════════════════════════════════════
|
||
# 9. 快速测试
|
||
# ════════════════════════════════════════════════════════════
|
||
|
||
if __name__ == "__main__":
|
||
print("=" * 65)
|
||
print(" XiaLi-Next 架构测试")
|
||
print("=" * 65)
|
||
|
||
for name in ['20m', '26m', '100m']:
|
||
cfg = XiaLiConfig.preset(name)
|
||
try:
|
||
model = XiaLiModel(cfg)
|
||
stats = model.count_params()
|
||
print(f"\n [{name}] {cfg.d_model}d × {cfg.n_layers}L × {cfg.n_heads}H (KV={cfg.n_kv_heads})")
|
||
print(f" 参数量: {stats['total_m']:.2f}M")
|
||
print(f" 其中: emb={stats['embedding']/1e6:.2f}M, attn={stats['attention']/1e6:.2f}M, "
|
||
f"ffn={stats['ffn']/1e6:.2f}M")
|
||
|
||
# 前向测试
|
||
x = torch.randint(0, cfg.vocab_size - 1, (1, 64))
|
||
try:
|
||
out = model(x)
|
||
print(f" 前向: ✅ {out.shape}")
|
||
except Exception as e:
|
||
print(f" 前向: ❌ {e}")
|
||
except Exception as e:
|
||
print(f"\n [{name}] ❌ {e}")
|
||
|
||
print(f"\n{'='*65}")
|
||
print(f" ✅ XiaLi-Next 就绪！")
|
||
print(f" 从 20M 到 7B+ 统一架构")
|
||
print(f"{'='*65}") |