From 3aa8e9d205c5ec46ecd0026ec571b5753e03d3a2 Mon Sep 17 00:00:00 2001 From: XiaLi Bot Date: Sun, 21 Jun 2026 04:40:59 +0000 Subject: [PATCH] =?UTF-8?q?XiaLi-Next=E7=BB=9F=E4=B8=80=E6=9E=B6=E6=9E=84:?= =?UTF-8?q?=20RoPE+GQA+SwiGLU+WeightTying?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- xiaoli_next.py | 422 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 422 insertions(+) create mode 100644 xiaoli_next.py diff --git a/xiaoli_next.py b/xiaoli_next.py new file mode 100644 index 0000000..209ec29 --- /dev/null +++ b/xiaoli_next.py @@ -0,0 +1,422 @@ +#!/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}") \ No newline at end of file