从零搭建小型 LLM¶

def analyze_llama_params(dim, n_layers, n_heads, n_kv_heads, vocab_size, multiple_of=256):
    """
    逐层分析 LLaMA 模型的参数量分布

    LLaMA 每层 Transformer 包含:
    - Attention: wq, wk, wv, wo (4个线性层)
    - FFN (SwiGLU): w1, w2, w3 (3个线性层)
    - Norm: 2个 RMSNorm
    """
    head_dim = dim // n_heads
    # SwiGLU hidden_dim = 2/3 * 4 * dim，对齐到 multiple_of
    hidden_dim = int(2 * 4 * dim / 3)
    hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

    # === 每层参数 ===
    # Attention: wq(dim→dim), wk(dim→n_kv_heads*head_dim), wv(dim→n_kv_heads*head_dim), wo(dim→dim)
    attn_params = dim * n_heads * head_dim + 2 * dim * n_kv_heads * head_dim + n_heads * head_dim * dim
    # FFN (SwiGLU): w1(dim→hidden), w2(hidden→dim), w3(dim→hidden)
    ffn_params = dim * hidden_dim + hidden_dim * dim + dim * hidden_dim
    # Norm: 2个 RMSNorm，每个 dim 个参数
    norm_params = 2 * dim

    per_layer = attn_params + ffn_params + norm_params

    # === 全局参数 ===
    embedding_params = vocab_size * dim  # Token Embedding
    # LLaMA 使用 weight tying: LM Head 与 Embedding 共享权重
    final_norm_params = dim  # Final RMSNorm

    total = n_layers * per_layer + embedding_params + final_norm_params

    print(f"{'='*60}")
    print(f"LLaMA 参数分析 (dim={dim}, layers={n_layers}, heads={n_heads}, kv_heads={n_kv_heads})")
    print(f"{'='*60}")
    print(f"  每层 Attention:  {attn_params:>12,} ({attn_params/1e6:.2f}M)")
    print(f"  每层 FFN:        {ffn_params:>12,} ({ffn_params/1e6:.2f}M)")
    print(f"  每层 Norm:       {norm_params:>12,}")
    print(f"  每层合计:        {per_layer:>12,} ({per_layer/1e6:.2f}M)")
    print(f"  ─────────────────────────────")
    print(f"  N层合计:         {n_layers * per_layer:>12,} ({n_layers*per_layer/1e6:.2f}M)")
    print(f"  Embedding:       {embedding_params:>12,} ({embedding_params/1e6:.2f}M)")
    print(f"  Final Norm:      {final_norm_params:>12,}")
    print(f"  ═════════════════════════════")
    print(f"  总参数量:        {total:>12,} ({total/1e6:.2f}M)")
    print(f"  Attention 占比:  {n_layers * attn_params / total * 100:.1f}%")
    print(f"  FFN 占比:        {n_layers * ffn_params / total * 100:.1f}%")
    print(f"  Embedding 占比:  {embedding_params / total * 100:.1f}%")
    return total

# 真实 LLaMA-2 系列配置对比
print(">>> LLaMA-2 7B")
analyze_llama_params(dim=4096, n_layers=32, n_heads=32, n_kv_heads=32, vocab_size=32000)
print()
print(">>> LLaMA-2 13B")
analyze_llama_params(dim=5120, n_layers=40, n_heads=40, n_kv_heads=40, vocab_size=32000)
print()
print(">>> LLaMA-2 70B (GQA)")
analyze_llama_params(dim=8192, n_layers=80, n_heads=64, n_kv_heads=8, vocab_size=32000)
print()
print(">>> 我们的 Mini-LLaMA (~15M)")
analyze_llama_params(dim=512, n_layers=8, n_heads=8, n_kv_heads=4, vocab_size=5000)

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization
    比LayerNorm更快：不需要计算均值，只计算均方根

    公式: RMSNorm(x) = x / RMS(x) * γ
    其中: RMS(x) = sqrt(mean(x²) + ε)
    """
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()  # super()调用父类方法
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))  # 可学习的缩放参数 γ

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args: x [batch_size, seq_len, dim]
        Returns: normalized x, same shape
        """
        # 计算均方根
        rms = torch.sqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        # 归一化并缩放
        return x / rms * self.weight

# 验证
norm = RMSNorm(64)
x = torch.randn(2, 10, 64)
out = norm(x)
print(f"RMSNorm: 输入 {x.shape} → 输出 {out.shape}")
# 验证输出的RMS约为1
rms = torch.sqrt(out.pow(2).mean(dim=-1))
print(f"归一化后RMS均值: {rms.mean().item():.4f}")  # 应接近1.0

class RotaryPositionEmbedding:
    """
    RoPE (Rotary Position Embedding)

    核心思想: 用旋转矩阵编码位置信息
    - 对Q和K分别施加旋转，使得内积自然包含相对位置信息
    - q_m · k_n = f(q, m) · f(k, n) → 只依赖相对位置 m-n
    """

    @staticmethod  # @staticmethod无需实例即可调用
    def precompute_freqs(dim: int, max_seq_len: int, theta: float = 10000.0):
        """
        预计算频率参数

        freqs[i] = 1 / theta^(2i/dim), i = 0, 1, ..., dim/2 - 1
        """
        # [dim/2] 频率
        freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        # [max_seq_len] 位置
        t = torch.arange(max_seq_len).float()
        # [max_seq_len, dim/2] 外积
        freqs = torch.outer(t, freqs)
        # 复数形式: e^(i*θ) = cos(θ) + i*sin(θ)
        freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
        return freqs_cis  # [max_seq_len, dim/2] 复数张量

    @staticmethod
    def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor,
                          freqs_cis: torch.Tensor):
        """
        对Q和K施加旋转位置编码

        Args:
            xq: [B, L, n_heads, head_dim]
            xk: [B, L, n_kv_heads, head_dim]
            freqs_cis: [L, head_dim/2] 复数频率
        """
        # 将实数张量转为复数: [B, L, H, D] → [B, L, H, D/2] (复数)
        xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
        xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

        # 调整freqs_cis的形状以匹配
        # freqs_cis: [L, D/2] → [1, L, 1, D/2]
        freqs = freqs_cis.unsqueeze(0).unsqueeze(2)  # unsqueeze增加一个维度

        # 旋转: 复数乘法 = 旋转
        xq_out = torch.view_as_real(xq_complex * freqs).flatten(-2)
        xk_out = torch.view_as_real(xk_complex * freqs).flatten(-2)

        return xq_out.type_as(xq), xk_out.type_as(xk)

# 验证RoPE
freqs = RotaryPositionEmbedding.precompute_freqs(dim=64, max_seq_len=512)
print(f"RoPE频率表: {freqs.shape}")  # [512, 32]

q = torch.randn(2, 10, 4, 64)  # [B, L, n_heads, head_dim]
k = torch.randn(2, 10, 4, 64)
q_rot, k_rot = RotaryPositionEmbedding.apply_rotary_emb(q, k, freqs[:10])
print(f"RoPE编码后: Q {q_rot.shape}, K {k_rot.shape}")

class SwiGLU_FFN(nn.Module):
    """
    SwiGLU前馈网络 (LLaMA使用)

    公式: FFN(x) = (Swish(xW₁) ⊙ xW₃)W₂

    其中:
    - Swish(x) = x * sigmoid(x) (也叫SiLU)
    - ⊙ 表示逐元素乘法（门控机制）
    - hidden_dim = 8/3 * dim（约2.67倍，取最近的multiple_of对齐）
    """
    def __init__(self, dim: int, hidden_dim: int = None, multiple_of: int = 256):
        super().__init__()
        if hidden_dim is None:
            hidden_dim = int(2 * (4 * dim) / 3)
            # 取最近的multiple_of的倍数
            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False)   # Gate projection
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)    # Down projection
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)    # Up projection

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args: x [B, L, D]
        Returns: [B, L, D]
        """
        # SwiGLU: swish(x @ W1) * (x @ W3) @ W2
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

# 验证
ffn = SwiGLU_FFN(dim=128)
x = torch.randn(2, 10, 128)
out = ffn(x)
print(f"SwiGLU FFN: 输入 {x.shape} → 输出 {out.shape}")
print(f"FFN参数量: {sum(p.numel() for p in ffn.parameters()):,}")
print(f"隐藏层维度: {ffn.w1.out_features}")

class GroupedQueryAttention(nn.Module):
    """
    Grouped-Query Attention (GQA)

    MHA: n_heads个Q, n_heads个K, n_heads个V
    MQA: n_heads个Q, 1个K, 1个V
    GQA: n_heads个Q, n_kv_heads个K, n_kv_heads个V

    每n_heads/n_kv_heads个Q头共享同一组KV → 减少KV Cache
    """
    def __init__(self, dim: int, n_heads: int, n_kv_heads: int = None):
        super().__init__()
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads or n_heads  # 默认MHA
        self.head_dim = dim // n_heads
        self.n_rep = n_heads // self.n_kv_heads  # 每组KV重复次数

        assert dim % n_heads == 0  # assert断言：条件False时抛出AssertionError
        assert n_heads % self.n_kv_heads == 0

        self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(n_heads * self.head_dim, dim, bias=False)

    def _repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
        """将KV头重复以匹配Q头数量"""
        if self.n_rep == 1:
            return x
        B, L, n_kv_heads, head_dim = x.shape
        x = x.unsqueeze(3).expand(B, L, n_kv_heads, self.n_rep, head_dim)
        return x.reshape(B, L, self.n_heads, head_dim)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor,
                mask: torch.Tensor = None):
        """
        Args:
            x: [B, L, D]
            freqs_cis: [L, head_dim/2] RoPE频率
            mask: [L, L] 因果mask
        Returns: [B, L, D]
        """
        B, L, _ = x.shape

        # 线性投影
        q = self.wq(x).view(B, L, self.n_heads, self.head_dim)  # view重塑张量形状
        k = self.wk(x).view(B, L, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(B, L, self.n_kv_heads, self.head_dim)

        # 应用RoPE
        q, k = RotaryPositionEmbedding.apply_rotary_emb(q, k, freqs_cis)

        # GQA: 重复KV头
        k = self._repeat_kv(k)  # [B, L, n_heads, head_dim]
        v = self._repeat_kv(v)

        # 转置: [B, n_heads, L, head_dim]
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Scaled Dot-Product Attention
        scale = math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) / scale  # [B, H, L, L]

        if mask is not None:
            scores = scores + mask  # mask中被遮蔽位置为-inf

        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)  # [B, H, L, head_dim]

        # 合并多头
        output = output.transpose(1, 2).contiguous().view(B, L, -1)
        return self.wo(output)

# 验证
gqa = GroupedQueryAttention(dim=128, n_heads=8, n_kv_heads=2)
x = torch.randn(2, 10, 128)
freqs = RotaryPositionEmbedding.precompute_freqs(dim=128//8, max_seq_len=512)
# 因果mask
L = 10
mask = torch.triu(torch.full((L, L), float('-inf')), diagonal=1)
out = gqa(x, freqs[:L], mask)
print(f"GQA: 输入 {x.shape} → 输出 {out.shape}")
print(f"  Q头数: {gqa.n_heads}, KV头数: {gqa.n_kv_heads}, 重复倍数: {gqa.n_rep}")

class TransformerBlock(nn.Module):
    """LLaMA-2 Transformer Block (Pre-Norm + Residual)"""

    def __init__(self, dim: int, n_heads: int, n_kv_heads: int,
                 ffn_dim_multiplier: float = None, multiple_of: int = 256):
        super().__init__()
        self.attention = GroupedQueryAttention(dim, n_heads, n_kv_heads)
        self.feed_forward = SwiGLU_FFN(dim, multiple_of=multiple_of)
        self.attention_norm = RMSNorm(dim)
        self.ffn_norm = RMSNorm(dim)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor,
                mask: torch.Tensor = None):
        """
        Pre-Norm + Residual连接
        x → RMSNorm → Attention → +x → RMSNorm → FFN → +
        """
        # Self-Attention with residual
        h = x + self.attention(self.attention_norm(x), freqs_cis, mask)
        # FFN with residual
        out = h + self.feed_forward(self.ffn_norm(h))
        return out

from dataclasses import dataclass

@dataclass  # @dataclass自动生成__init__等方法
class MiniLLaMAConfig:
    """Mini-LLaMA模型配置"""
    vocab_size: int = 4096       # 词表大小（小型实验用）
    dim: int = 256               # 隐藏维度
    n_layers: int = 6            # Transformer层数
    n_heads: int = 8             # 注意力头数
    n_kv_heads: int = 2          # KV头数（GQA）
    max_seq_len: int = 512       # 最大序列长度
    multiple_of: int = 64        # FFN维度对齐
    rope_theta: float = 10000.0  # RoPE基础频率
    dropout: float = 0.1         # Dropout率

class MiniLLaMA(nn.Module):
    """
    Mini-LLaMA: 完整的LLaMA-2风格语言模型

    约15M参数，可在单GPU上训练
    """
    def __init__(self, config: MiniLLaMAConfig):
        super().__init__()
        self.config = config

        # Token Embedding
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.dim)

        # Transformer Blocks
        self.layers = nn.ModuleList([
            TransformerBlock(
                dim=config.dim,
                n_heads=config.n_heads,
                n_kv_heads=config.n_kv_heads,
                multiple_of=config.multiple_of
            )
            for _ in range(config.n_layers)
        ])

        # Final RMSNorm
        self.norm = RMSNorm(config.dim)

        # LM Head (与embedding共享权重)
        self.output = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.tok_embeddings.weight = self.output.weight  # weight tying

        # 预计算RoPE频率
        head_dim = config.dim // config.n_heads
        self.freqs_cis = RotaryPositionEmbedding.precompute_freqs(
            head_dim, config.max_seq_len, config.rope_theta
        )

        # Dropout
        self.dropout = nn.Dropout(config.dropout)

        # 初始化权重
        self.apply(self._init_weights)

    def _init_weights(self, module):
        """Xavier均匀初始化"""
        if isinstance(module, nn.Linear):  # isinstance检查类型
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0, std=0.02)

    def forward(self, tokens: torch.Tensor, targets: torch.Tensor = None):
        """
        前向传播

        Args:
            tokens: [B, L] 输入token IDs
            targets: [B, L] 目标token IDs（训练时提供）
        Returns:
            如果有targets: loss (标量)
            否则: logits [B, L, V]
        """
        B, L = tokens.shape
        device = tokens.device

        # Token Embedding
        h = self.tok_embeddings(tokens)  # [B, L, D]
        h = self.dropout(h)

        # 获取RoPE频率
        freqs_cis = self.freqs_cis[:L].to(device)  # .to(device)将数据移至GPU/CPU

        # 构造因果mask
        mask = torch.triu(
            torch.full((L, L), float('-inf'), device=device),
            diagonal=1
        )

        # 通过所有Transformer层
        for layer in self.layers:
            h = layer(h, freqs_cis, mask)

        # Final Norm
        h = self.norm(h)

        # LM Head
        logits = self.output(h)  # [B, L, V]

        if targets is not None:
            # 计算交叉熵损失
            # logits: [B, L, V] → [B*L, V]
            # targets: [B, L] → [B*L]
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-100  # 忽略padding（PyTorch默认值）
            )
            return loss

        return logits

    @torch.no_grad()  # 禁用梯度计算，节省内存（推理时使用）
    def generate(self, prompt_tokens: torch.Tensor, max_new_tokens: int = 100,
                 temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
        """
        自回归文本生成

        Args:
            prompt_tokens: [1, prompt_len] 提示token序列
            max_new_tokens: 最大生成token数
            temperature: 温度参数（越高越随机）
            top_k: Top-K采样
            top_p: Nucleus采样
        """
        self.eval()
        tokens = prompt_tokens.clone()

        for _ in range(max_new_tokens):
            # 截断到最大长度
            context = tokens[:, -self.config.max_seq_len:]

            # 前向传播获取logits
            logits = self(context)

            # 只取最后一个位置的logits
            next_logits = logits[:, -1, :] / temperature  # [1, V]

            # Top-K过滤
            if top_k > 0:
                top_k_val = min(top_k, next_logits.size(-1))
                indices_to_remove = next_logits < torch.topk(next_logits, top_k_val)[0][:, -1:]
                next_logits[indices_to_remove] = float('-inf')

            # Top-P (Nucleus)过滤
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                # 移除累积概率超过top_p的token
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone()
                sorted_indices_to_remove[:, 0] = False
                indices_to_remove = sorted_indices_to_remove.scatter(
                    1, sorted_indices, sorted_indices_to_remove
                )
                next_logits[indices_to_remove] = float('-inf')

            # 采样
            probs = F.softmax(next_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)  # [1, 1]

            tokens = torch.cat([tokens, next_token], dim=-1)

        return tokens

# 创建模型
config = MiniLLaMAConfig()
model = MiniLLaMA(config)

# 统计参数量
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"Mini-LLaMA 模型统计:")
print(f"  总参数量: {total_params:,} ({total_params/1e6:.1f}M)")
print(f"  可训练参数: {trainable_params:,}")
print(f"  配置: dim={config.dim}, layers={config.n_layers}, heads={config.n_heads}")
print(f"  KV头数: {config.n_kv_heads} (GQA, {config.n_heads//config.n_kv_heads}x压缩)")
print(f"  最大序列长度: {config.max_seq_len}")

# 测试前向传播
dummy_input = torch.randint(0, config.vocab_size, (2, 32))
dummy_target = torch.randint(0, config.vocab_size, (2, 32))
loss = model(dummy_input, dummy_target)
print(f"\n前向传播测试: loss = {loss.item():.4f}")
print(f"随机初始化loss理论值: {math.log(config.vocab_size):.4f} (ln({config.vocab_size}))")

import re
from collections import Counter, defaultdict

class SimpleBPETokenizer:
    """
    手写BPE分词器
    支持中英文混合文本
    """

    def __init__(self, vocab_size: int = 4096):
        self.target_vocab_size = vocab_size
        self.merges = {}      # 合并规则: (a, b) → ab
        self.vocab = {}       # token → id
        self.inverse_vocab = {}  # id → token

        # 特殊token
        self.special_tokens = {
            "<pad>": 0,
            "<unk>": 1,
            "<bos>": 2,
            "<eos>": 3,
        }

    def _get_stats(self, word_freqs):
        """统计相邻pair的频率"""
        pairs = Counter()  # Counter统计元素出现次数
        for word, freq in word_freqs.items():
            symbols = word.split()
            for i in range(len(symbols) - 1):
                pairs[(symbols[i], symbols[i+1])] += freq
        return pairs

    def _merge_pair(self, pair, word_freqs):
        """合并一个pair"""
        new_word_freqs = {}
        bigram = ' '.join(pair)
        replacement = ''.join(pair)

        for word, freq in word_freqs.items():
            new_word = word.replace(bigram, replacement)
            new_word_freqs[new_word] = freq

        return new_word_freqs

    def train(self, texts: list, verbose: bool = True):
        """
        训练BPE分词器

        Args:
            texts: 训练文本列表
            verbose: 是否打印训练过程
        """
        # Step 1: 分词并统计词频
        word_freqs = Counter()
        for text in texts:
            # 简单分词：按空格和标点拆分
            words = re.findall(r'[\u4e00-\u9fff]|[a-zA-Z]+|[0-9]+|[^\s\w]', text)  # re.findall正则查找所有匹配项
            for word in words:
                # 将每个词拆分为字符（空格分隔）
                char_word = ' '.join(list(word))
                word_freqs[char_word] += 1

        # Step 2: 初始化词表（所有出现的字符）
        self.vocab = dict(self.special_tokens)
        chars = set()
        for word in word_freqs:
            for char in word.split():
                chars.add(char)

        for char in sorted(chars):
            if char not in self.vocab:
                self.vocab[char] = len(self.vocab)

        initial_vocab_size = len(self.vocab)
        if verbose:
            print(f"初始词表大小: {initial_vocab_size}")

        # Step 3: 迭代合并
        num_merges = self.target_vocab_size - initial_vocab_size

        for i in range(num_merges):
            # 统计pair频率
            pairs = self._get_stats(word_freqs)
            if not pairs:
                break

            # 找出最频繁的pair
            best_pair = max(pairs, key=pairs.get)
            best_freq = pairs[best_pair]

            if best_freq < 2:  # 频率太低则停止
                break

            # 合并
            word_freqs = self._merge_pair(best_pair, word_freqs)

            # 记录合并规则
            merged_token = ''.join(best_pair)
            self.merges[best_pair] = merged_token

            if merged_token not in self.vocab:
                self.vocab[merged_token] = len(self.vocab)

            if verbose and (i + 1) % 100 == 0:
                print(f"  合并 {i+1}/{num_merges}: '{best_pair[0]}' + '{best_pair[1]}' → '{merged_token}' (freq={best_freq})")

        # 构建反向词表
        self.inverse_vocab = {v: k for k, v in self.vocab.items()}

        if verbose:
            print(f"最终词表大小: {len(self.vocab)}")
            print(f"合并规则数: {len(self.merges)}")

    def encode(self, text: str) -> list:
        """将文本编码为token ID序列"""
        # 分词
        words = re.findall(r'[\u4e00-\u9fff]|[a-zA-Z]+|[0-9]+|[^\s\w]', text)

        ids = []
        for word in words:
            # 拆分为字符
            symbols = list(word)

            # 应用合并规则（贪婪）
            while len(symbols) > 1:
                # 找出可以合并的pair
                best_pair = None
                best_idx = -1

                for i in range(len(symbols) - 1):
                    pair = (symbols[i], symbols[i+1])
                    if pair in self.merges:
                        best_pair = pair
                        best_idx = i
                        break  # 按训练顺序优先合并

                if best_pair is None:
                    break

                # 执行合并
                symbols = (symbols[:best_idx] +
                          [self.merges[best_pair]] +
                          symbols[best_idx+2:])

            # 转换为ID
            for sym in symbols:
                ids.append(self.vocab.get(sym, self.special_tokens["<unk>"]))

        return ids

    def decode(self, ids: list) -> str:
        """将token ID序列解码为文本"""
        tokens = [self.inverse_vocab.get(id, "<unk>") for id in ids]
        return ''.join(tokens)

# === 训练Tokenizer ===
# 中文训练语料（示例）
training_texts = [
    "人工智能正在改变世界",
    "深度学习是机器学习的一个分支",
    "大语言模型具有强大的文本生成能力",
    "自然语言处理是人工智能的重要方向",
    "Transformer架构改变了深度学习的格局",
    "注意力机制是现代神经网络的核心组件",
    "预训练语言模型通过大规模数据学习语言知识",
    "GPT系列模型展示了规模扩展的威力",
    "BERT模型在自然语言理解任务上取得了突破",
    "强化学习可以用于优化语言模型的输出",
] * 50  # 重复增加频率

tokenizer = SimpleBPETokenizer(vocab_size=500)
tokenizer.train(training_texts)

# 测试
test_text = "人工智能改变世界"
encoded = tokenizer.encode(test_text)
decoded = tokenizer.decode(encoded)
print(f"\n编码测试:")
print(f"  原文: {test_text}")
print(f"  编码: {encoded}")
print(f"  解码: {decoded}")
print(f"  tokens: {[tokenizer.inverse_vocab.get(id, '?') for id in encoded]}")

import torch
from torch.utils.data import Dataset, DataLoader
import os

class TextDataset(Dataset):
    """
    用于CLM预训练的文本数据集
    将文本编码后切分为固定长度的序列
    """
    def __init__(self, texts: list, tokenizer, seq_len: int = 128):
        self.seq_len = seq_len

        # 编码所有文本
        all_ids = []
        for text in texts:
            ids = tokenizer.encode(text)
            all_ids.extend(ids)
            all_ids.append(tokenizer.special_tokens["<eos>"])

        # 切分为固定长度序列
        self.data = []
        for i in range(0, len(all_ids) - seq_len - 1, seq_len):
            input_ids = all_ids[i : i + seq_len]
            target_ids = all_ids[i + 1 : i + seq_len + 1]
            self.data.append((
                torch.tensor(input_ids, dtype=torch.long),
                torch.tensor(target_ids, dtype=torch.long)
            ))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

print("数据集准备完成")

import time

class MiniLLaMATrainer:
    """Mini-LLaMA训练器"""

    def __init__(self, model, config, tokenizer, device='cpu'):
        self.model = model.to(device)
        self.config = config
        self.tokenizer = tokenizer
        self.device = device

        # 优化器: AdamW
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=3e-4,
            betas=(0.9, 0.95),
            weight_decay=0.1
        )

        # 学习率调度: Cosine with warmup
        self.warmup_steps = 100
        self.total_steps = 0

    def _get_lr(self, step):
        """Cosine学习率调度 with warmup"""
        if step < self.warmup_steps:
            return 3e-4 * step / self.warmup_steps
        # Cosine decay
        progress = (step - self.warmup_steps) / max(1, self.total_steps - self.warmup_steps)
        return 3e-4 * 0.5 * (1 + math.cos(math.pi * progress))

    def train(self, train_dataset, num_epochs=10, batch_size=8,
              log_interval=50, eval_interval=200):
        """
        训练循环
        """
        dataloader = DataLoader(
            train_dataset, batch_size=batch_size, shuffle=True, drop_last=True
        )
        self.total_steps = num_epochs * len(dataloader)

        self.model.train()
        global_step = 0
        best_loss = float('inf')

        print(f"开始训练 Mini-LLaMA")
        print(f"  Epochs: {num_epochs}")
        print(f"  Batch size: {batch_size}")
        print(f"  Total steps: {self.total_steps}")
        print(f"  Device: {self.device}")
        print("=" * 60)

        for epoch in range(num_epochs):
            epoch_loss = 0
            epoch_start = time.time()

            for batch_idx, (input_ids, target_ids) in enumerate(dataloader):  # enumerate同时获取索引和元素
                input_ids = input_ids.to(self.device)
                target_ids = target_ids.to(self.device)

                # 更新学习率
                lr = self._get_lr(global_step)
                for param_group in self.optimizer.param_groups:
                    param_group['lr'] = lr

                # 前向传播
                loss = self.model(input_ids, target_ids)

                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()

                # 梯度裁剪
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

                self.optimizer.step()

                epoch_loss += loss.item()
                global_step += 1

                # 日志
                if global_step % log_interval == 0:
                    avg_loss = epoch_loss / (batch_idx + 1)
                    ppl = math.exp(min(avg_loss, 20))  # 限制防止溢出
                    elapsed = time.time() - epoch_start
                    print(f"  Step {global_step}: loss={avg_loss:.4f}, ppl={ppl:.2f}, lr={lr:.6f}, time={elapsed:.1f}s")

                # 生成样本
                if global_step % eval_interval == 0:
                    self._generate_sample()

            # Epoch结束统计
            avg_epoch_loss = epoch_loss / len(dataloader)
            epoch_time = time.time() - epoch_start
            print(f"\nEpoch {epoch+1}/{num_epochs}: avg_loss={avg_epoch_loss:.4f}, time={epoch_time:.1f}s")

            if avg_epoch_loss < best_loss:
                best_loss = avg_epoch_loss
                # 保存最佳模型
                torch.save(self.model.state_dict(), 'mini_llama_best.pt')
                print(f"  ✓ 保存最佳模型 (loss={best_loss:.4f})")
            print()

        print(f"训练完成! 最佳loss: {best_loss:.4f}")

    def _generate_sample(self):
        """训练中生成样本以观察效果"""
        self.model.eval()

        # 用几个起始token生成
        prompt = torch.tensor([[self.tokenizer.special_tokens["<bos>"]]]).to(self.device)

        generated = self.model.generate(
            prompt, max_new_tokens=50, temperature=0.8, top_k=40
        )

        text = self.tokenizer.decode(generated[0].tolist())
        print(f"  [生成样本] {text[:100]}...")

        self.model.train()

# ===  端到端训练示例 ===
print("端到端预训练流程:")
print("1. 准备语料 → 2. 训练Tokenizer → 3. 构建数据集 → 4. 训练模型 → 5. 生成文本")
print()

# 准备训练语料（实际应用中使用大规模语料）
corpus = [
    "人工智能是计算机科学的一个重要分支，它研究如何让计算机模拟人类的智能行为。",
    "深度学习通过多层神经网络自动学习数据的特征表示，是目前最成功的机器学习方法。",
    "自然语言处理让计算机能够理解和生成人类语言，包括文本分类、机器翻译等任务。",
    "大语言模型通过在海量文本上预训练，学习到了丰富的语言知识和世界知识。",
    "Transformer架构使用自注意力机制，能够并行处理序列中的所有位置。",
    "注意力机制允许模型在处理每个位置时关注序列中的其他相关位置。",
    "预训练加微调的范式极大地提高了模型在各种下游任务上的表现。",
    "强化学习通过试错的方式让智能体学习最优策略，在游戏和机器人控制中表现出色。",
] * 100  # 实际需要大量数据

# 1. 训练Tokenizer
tok = SimpleBPETokenizer(vocab_size=500)
tok.train(corpus, verbose=False)

# 2. 更新模型配置
train_config = MiniLLaMAConfig(
    vocab_size=len(tok.vocab),
    dim=128,
    n_layers=4,
    n_heads=4,
    n_kv_heads=2,
    max_seq_len=128,
    multiple_of=32
)

# 3. 创建模型
train_model = MiniLLaMA(train_config)
params = sum(p.numel() for p in train_model.parameters())
print(f"模型参数量: {params:,} ({params/1e6:.2f}M)")

# 4. 创建数据集
train_dataset = TextDataset(corpus, tok, seq_len=64)
print(f"训练样本数: {len(train_dataset)}")

# 5. 训练（实际训练需要更多数据和更长时间）
# trainer = MiniLLaMATrainer(train_model, train_config, tok, device='cpu')
# trainer.train(train_dataset, num_epochs=5, batch_size=8)

print("\n注意: 完整训练需要：")
print("  - 几GB~几十GB的中文语料")
print("  - GPU加速（推荐至少一张RTX 3090）")
print("  - 数小时到数天的训练时间")
print("  - 更大的模型配置（dim=512+, layers=12+）")

import json
import os

class RobustTrainer:
    """
    支持梯度累积 + 断点续训的训练器

    梯度累积原理：
    ┌─────────────────────────────────────────────────────┐
    │ 等效 batch_size = micro_batch × accumulation_steps   │
    │                                                     │
    │ 例: 想要 batch=64, 但 GPU 只能跑 batch=8            │
    │     → accumulation_steps = 64/8 = 8                 │
    │     → 每8个micro-batch才更新一次参数                  │
    └─────────────────────────────────────────────────────┘
    """
    def __init__(self, model, config, tokenizer, device='cpu',
                 accumulation_steps=4, checkpoint_dir='checkpoints'):
        self.model = model.to(device)
        self.config = config
        self.tokenizer = tokenizer
        self.device = device
        self.accumulation_steps = accumulation_steps
        self.checkpoint_dir = checkpoint_dir
        os.makedirs(checkpoint_dir, exist_ok=True)

        self.optimizer = torch.optim.AdamW(
            model.parameters(), lr=3e-4, betas=(0.9, 0.95), weight_decay=0.1
        )
        self.global_step = 0
        self.best_loss = float('inf')

    def train_epoch(self, dataloader, epoch):
        self.model.train()
        epoch_loss = 0
        self.optimizer.zero_grad()  # 每个epoch开头清零梯度

        for step, (input_ids, target_ids) in enumerate(dataloader):
            input_ids = input_ids.to(self.device)
            target_ids = target_ids.to(self.device)

            # 前向传播（损失已除以 accumulation_steps 以保持等效）
            loss = self.model(input_ids, target_ids) / self.accumulation_steps
            loss.backward()  # 梯度累积，不立即清零

            epoch_loss += loss.item() * self.accumulation_steps

            # 每 accumulation_steps 步执行一次参数更新
            if (step + 1) % self.accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                self.optimizer.step()
                self.optimizer.zero_grad()
                self.global_step += 1

        return epoch_loss / len(dataloader)

    def save_checkpoint(self, epoch, loss):
        """保存完整训练状态，支持精确恢复"""
        checkpoint = {
            'epoch': epoch,
            'global_step': self.global_step,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'loss': loss,
            'best_loss': self.best_loss,
            'config': vars(self.config),
        }
        path = os.path.join(self.checkpoint_dir, f'ckpt_epoch_{epoch}.pt')
        torch.save(checkpoint, path)

        # 同时保存最佳模型
        if loss < self.best_loss:
            self.best_loss = loss
            best_path = os.path.join(self.checkpoint_dir, 'best_model.pt')
            torch.save(checkpoint, best_path)
            print(f"  ✓ 新最佳模型 (loss={loss:.4f})")

    def load_checkpoint(self, path):
        """从检查点恢复训练状态"""
        checkpoint = torch.load(path, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.global_step = checkpoint['global_step']
        self.best_loss = checkpoint['best_loss']
        print(f"  ✓ 从 epoch {checkpoint['epoch']} 恢复, "
              f"step={checkpoint['global_step']}, loss={checkpoint['loss']:.4f}")
        return checkpoint['epoch']

# === 使用示例 ===
print("梯度累积 vs 直接大 batch:")
print("┌────────────────────┬──────────────┬──────────────┬────────────┐")
print("│ 方案               │ micro_batch  │ accum_steps  │ 等效batch  │")
print("├────────────────────┼──────────────┼──────────────┼────────────┤")
print("│ 直接训练           │ 64           │ 1            │ 64         │")
print("│ 梯度累积(显存不足) │ 8            │ 8            │ 64         │")
print("│ 梯度累积(极端情况) │ 1            │ 64           │ 64         │")
print("└────────────────────┴──────────────┴──────────────┴────────────┘")
print()
print("断点续训使用方式:")
print("  trainer = RobustTrainer(model, config, tok, accumulation_steps=8)")
print("  # 如需恢复:")
print("  # start_epoch = trainer.load_checkpoint('checkpoints/best_model.pt')")
print("  # for epoch in range(start_epoch + 1, total_epochs):")

@torch.no_grad()
def evaluate_perplexity(model, dataset, batch_size=8, device='cpu'):
    """
    计算模型在数据集上的困惑度(Perplexity)

    PPL = exp(avg_loss)
    PPL越低，模型越好
    """
    model.eval()
    model = model.to(device)

    dataloader = DataLoader(dataset, batch_size=batch_size)
    total_loss = 0
    total_tokens = 0

    for input_ids, target_ids in dataloader:
        input_ids = input_ids.to(device)
        target_ids = target_ids.to(device)

        loss = model(input_ids, target_ids)

        num_tokens = (target_ids != -100).sum().item()  # 与 ignore_index=-100 保持一致
        total_loss += loss.item() * num_tokens
        total_tokens += num_tokens

    avg_loss = total_loss / total_tokens
    perplexity = math.exp(avg_loss)

    return perplexity, avg_loss

# 如果模型已训练好:
# ppl, loss = evaluate_perplexity(train_model, train_dataset)
# print(f"困惑度: {ppl:.2f}, 平均Loss: {loss:.4f}")