09-大语言模型原理与架构

# GPT-4 Vision示例（伪代码）
response = openai.ChatCompletion.create(
    model="gpt-4-vision-preview",
    messages=[
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "这张图片里有什么？"},
                {"type": "image_url", "image_url": "https://..."}
            ]
        }
    ]
)

# BERT输入嵌入
def get_bert_input(token_ids, segment_ids, position_ids):
    token_embeddings = TokenEmbedding(token_ids)
    segment_embeddings = SegmentEmbedding(segment_ids)
    position_embeddings = PositionEmbedding(position_ids)

    return token_embeddings + segment_embeddings + position_embeddings

# 示例
输入：[CLS] 我 喜欢 AI [SEP] AI 很 有趣 [SEP]
Token IDs: [101, 2769, 1599, 2335, 102, 2335, 2523, 3300, 102]
Segment IDs: [0, 0, 0, 0, 0, 1, 1, 1, 1]
Position IDs: [0, 1, 2, 3, 4, 5, 6, 7, 8]

class MoELayer(nn.Module):
    def __init__(self, num_experts=8, d_model=4096):
        self.experts = nn.ModuleList([
            FeedForward(d_model) for _ in range(num_experts)
        ])
        self.router = nn.Linear(d_model, num_experts)

    def forward(self, x):
        # 计算路由权重
        router_logits = self.router(x)  # [batch, seq, num_experts]
        router_probs = F.softmax(router_logits, dim=-1)

        # 选择Top-2专家
        top2_probs, top2_indices = torch.topk(router_probs, 2, dim=-1)

        # 归一化权重
        top2_probs = top2_probs / top2_probs.sum(dim=-1, keepdim=True)

        # 专家计算
        output = torch.zeros_like(x)
        for i in range(2):
            expert_idx = top2_indices[:, :, i]
            expert_prob = top2_probs[:, :, i:i+1]
            # 批量调用专家
            expert_output = self.experts[expert_idx](x)
            output += expert_prob * expert_output

        return output

# GPT-3 Tokenizer
text = "我喜欢编程"
tokens = ["我", "喜", "欢", "编", "程"]  # 5个tokens

# Qwen Tokenizer
text = "我喜欢编程"
tokens = ["我", "喜欢", "编程"]  # 3个tokens（更高效）

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

class GPTConfig:
    """GPT模型配置"""
    def __init__(
        self,
        vocab_size=50257,      # 词表大小
        n_layer=12,            # 层数
        n_head=12,             # 注意力头数
        n_embd=768,            # 嵌入维度
        block_size=1024,       # 最大序列长度
        dropout=0.1,           # Dropout率
        bias=True,             # 是否使用bias
    ):
        self.vocab_size = vocab_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_embd = n_embd
        self.block_size = block_size
        self.dropout = dropout
        self.bias = bias

class CausalSelfAttention(nn.Module):
    """因果自注意力机制"""

    def __init__(self, config):
        super().__init__()
        assert config.n_embd % config.n_head == 0

        # Q、K、V的投影（合并为一个矩阵提高效率）
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)

        # 输出投影
        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)

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

        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.dropout = config.dropout

        # 因果掩码（下三角矩阵）
        self.register_buffer(
            "bias",
            torch.tril(torch.ones(config.block_size, config.block_size))
            .view(1, 1, config.block_size, config.block_size)
        )

    def forward(self, x):
        B, T, C = x.size()  # batch_size, sequence_length, embedding_dim

        # 计算Q、K、V
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)

        # 分离多头
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)  # (B, nh, T, hs)

        # 计算注意力分数
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))

        # 应用因果掩码
        att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))

        # Softmax
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)

        # 加权求和
        y = att @ v  # (B, nh, T, hs)

        # 合并多头
        y = y.transpose(1, 2).contiguous().view(B, T, C)

        # 输出投影
        y = self.resid_dropout(self.c_proj(y))

        return y

class MLP(nn.Module):
    """前馈网络"""

    def __init__(self, config):
        super().__init__()
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
        self.gelu = nn.GELU()
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.c_fc(x)
        x = self.gelu(x)
        x = self.c_proj(x)
        x = self.dropout(x)
        return x

class Block(nn.Module):
    """Transformer块"""

    def __init__(self, config):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        # Pre-norm架构
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

class GPT(nn.Module):
    """完整的GPT模型"""

    def __init__(self, config):
        super().__init__()
        self.config = config

        self.transformer = nn.ModuleDict(dict(
            wte=nn.Embedding(config.vocab_size, config.n_embd),  # token embedding
            wpe=nn.Embedding(config.block_size, config.n_embd),  # position embedding
            drop=nn.Dropout(config.dropout),
            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
            ln_f=nn.LayerNorm(config.n_embd),
        ))
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # 权重共享
        self.transformer.wte.weight = self.lm_head.weight

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

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, idx, targets=None):
        device = idx.device
        b, t = idx.size()
        assert t <= self.config.block_size, f"序列长度{t}超过最大长度{self.config.block_size}"

        # 位置索引
        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)  # (1, t)

        # 前向传播
        tok_emb = self.transformer.wte(idx)  # token embeddings (b, t, n_embd)
        pos_emb = self.transformer.wpe(pos)  # position embeddings (1, t, n_embd)
        x = self.transformer.drop(tok_emb + pos_emb)

        for block in self.transformer.h:
            x = block(x)

        x = self.transformer.ln_f(x)

        if targets is not None:
            # 训练模式：计算loss
            logits = self.lm_head(x)
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                targets.view(-1),
                ignore_index=-1
            )
        else:
            # 推理模式：只计算最后一个token的logits
            logits = self.lm_head(x[:, [-1], :])
            loss = None

        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        生成文本

        Args:
            idx: (B, T) 输入token序列
            max_new_tokens: 生成的最大token数
            temperature: 温度参数，控制随机性
            top_k: Top-k采样
        """
        for _ in range(max_new_tokens):
            # 截断到最大长度
            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]

            # 前向传播
            logits, _ = self(idx_cond)

            # 取最后一个token的logits
            logits = logits[:, -1, :] / temperature

            # Top-k采样
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')

            # Softmax转换为概率
            probs = F.softmax(logits, dim=-1)

            # 采样
            idx_next = torch.multinomial(probs, num_samples=1)

            # 拼接
            idx = torch.cat((idx, idx_next), dim=1)

        return idx

# 使用示例
if __name__ == "__main__":
    # 创建GPT-2 Small配置
    config = GPTConfig(
        vocab_size=50257,
        n_layer=12,
        n_head=12,
        n_embd=768,
        block_size=1024,
    )

    model = GPT(config)

    # 统计参数量
    n_params = sum(p.numel() for p in model.parameters())
    print(f"参数量: {n_params/1e6:.2f}M")

    # 前向传播测试
    idx = torch.randint(0, config.vocab_size, (2, 128))  # batch=2, seq_len=128
    logits, loss = model(idx, targets=idx)
    print(f"Logits shape: {logits.shape}")

    # 生成测试
    generated = model.generate(idx[:, :10], max_new_tokens=50, temperature=0.8, top_k=50)
    print(f"Generated shape: {generated.shape}")

# 创建因果掩码
def create_causal_mask(seq_len):
    """
    创建因果掩码

    对于序列长度为4的例子：
    [[1, 0, 0, 0],
     [1, 1, 0, 0],
     [1, 1, 1, 0],
     [1, 1, 1, 1]]

    1表示可以关注，0表示不能关注
    """
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask

# 可视化
seq_len = 4
mask = create_causal_mask(seq_len)
print("Causal Mask:")
print(mask)

# 在注意力计算中应用
def apply_causal_mask(attention_scores, mask):
    """
    将掩码应用到注意力分数上

    Args:
        attention_scores: (B, H, T, T) 注意力分数
        mask: (T, T) 因果掩码
    """
    # 将0的位置设为负无穷，softmax后会变成0
    attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))
    return attention_scores

import matplotlib.pyplot as plt
import seaborn as sns

def visualize_attention(attention_weights, tokens):
    """
    可视化注意力权重

    Args:
        attention_weights: (seq_len, seq_len) 注意力权重矩阵
        tokens: 序列的token列表
    """
    plt.figure(figsize=(10, 8))
    sns.heatmap(
        attention_weights,
        xticklabels=tokens,
        yticklabels=tokens,
        cmap='Blues',
        annot=True,
        fmt='.2f'
    )
    plt.xlabel('Key')
    plt.ylabel('Query')
    plt.title('Causal Attention Weights')
    plt.tight_layout()
    plt.show()

# 示例
tokens = ["The", "cat", "sat", "on"]
seq_len = len(tokens)

# 模拟注意力权重（经过softmax后，被掩码的位置为0）
attention = torch.tril(torch.rand(seq_len, seq_len))
# 按行归一化
attention = attention / attention.sum(dim=-1, keepdim=True)

visualize_attention(attention.numpy(), tokens)

import re
from collections import defaultdict, Counter

class BPETokenizer:
    """Byte Pair Encoding分词器"""

    def __init__(self, vocab_size=1000):
        self.vocab_size = vocab_size
        self.merges = {}  # 合并规则
        self.vocab = {}   # 词表

    def get_stats(self, words):
        """统计相邻字符对的频率"""
        pairs = defaultdict(int)
        for word, freq in words.items():
            symbols = word.split()
            for i in range(len(symbols) - 1):
                pairs[symbols[i], symbols[i + 1]] += freq
        return pairs

    def merge_vocab(self, pair, words):
        """合并词表中的字符对"""
        new_words = {}
        bigram = ' '.join(pair)
        replacement = ''.join(pair)

        for word in words:
            new_word = word.replace(bigram, replacement)
            new_words[new_word] = words[word]

        return new_words

    def train(self, texts):
        """训练BPE分词器"""
        # 1. 初始化：统计词频
        word_freqs = Counter()
        for text in texts:
            words = text.split()
            word_freqs.update(words)

        # 2. 将每个词拆分为字符，添加结束符
        vocab = {}
        for word, freq in word_freqs.items():
            vocab[' '.join(list(word)) + ' </w>'] = freq

        # 3. 初始词表：所有字符
        self.vocab = set()
        for word in vocab.keys():
            self.vocab.update(word.split())

        # 4. 迭代合并
        for i in range(self.vocab_size - len(self.vocab)):
            pairs = self.get_stats(vocab)
            if not pairs:
                break

            # 找频率最高的pair
            best_pair = max(pairs, key=pairs.get)

            # 合并
            vocab = self.merge_vocab(best_pair, vocab)

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

            if (i + 1) % 100 == 0:
                print(f"合并 {i + 1}: {best_pair} -> {''.join(best_pair)}")

        # 5. 构建最终词表
        self.vocab = {token: i for i, token in enumerate(sorted(self.vocab))}

        print(f"训练完成！词表大小: {len(self.vocab)}")

    def tokenize(self, text):
        """对文本进行分词"""
        tokens = []
        words = text.split()

        for word in words:
            # 将词拆分为字符
            word_tokens = list(word) + ['</w>']

            # 应用合并规则
            while len(word_tokens) > 1:
                pairs = [(word_tokens[i], word_tokens[i + 1])
                        for i in range(len(word_tokens) - 1)]

                # 找出可以合并的pair（在merges中且优先级最高）
                mergeable_pairs = [p for p in pairs if p in self.merges]
                if not mergeable_pairs:
                    break

                # 选择优先级最高的（最早学到的）
                best_pair = min(mergeable_pairs, key=lambda x: self.merges[x])

                # 执行合并
                new_tokens = []
                i = 0
                while i < len(word_tokens):
                    if i < len(word_tokens) - 1 and \
                       (word_tokens[i], word_tokens[i + 1]) == best_pair:
                        new_tokens.append(word_tokens[i] + word_tokens[i + 1])
                        i += 2
                    else:
                        new_tokens.append(word_tokens[i])
                        i += 1

                word_tokens = new_tokens

            tokens.extend(word_tokens)

        return tokens

    def encode(self, text):
        """将文本编码为token IDs"""
        tokens = self.tokenize(text)
        return [self.vocab.get(token, self.vocab.get('<unk>', 0)) for token in tokens]

    def decode(self, token_ids):
        """将token IDs解码为文本"""
        id_to_token = {i: token for token, i in self.vocab.items()}
        tokens = [id_to_token.get(i, '<unk>') for i in token_ids]
        text = ''.join(tokens).replace('</w>', ' ').strip()
        return text

# 使用示例
if __name__ == "__main__":
    # 训练数据
    texts = [
        "low lower lowest",
        "new newer newest",
        "wide wider widest",
    ] * 100  # 重复以增加频率

    # 训练分词器
    tokenizer = BPETokenizer(vocab_size=100)
    tokenizer.train(texts)

    # 测试分词
    test_text = "lower newest"
    tokens = tokenizer.tokenize(test_text)
    print(f"\n文本: {test_text}")
    print(f"分词: {tokens}")

    # 测试编码解码
    encoded = tokenizer.encode(test_text)
    print(f"编码: {encoded}")
    decoded = tokenizer.decode(encoded)
    print(f"解码: {decoded}")

class WordPieceTokenizer:
    """WordPiece分词器"""

    def __init__(self, vocab_size=1000):
        self.vocab_size = vocab_size
        self.vocab = {}

    def get_pair_scores(self, words):
        """计算字符对的分数（似然增益）"""
        pair_freqs = defaultdict(int)
        token_freqs = defaultdict(int)

        for word, freq in words.items():
            symbols = word.split()
            for i in range(len(symbols) - 1):
                pair_freqs[(symbols[i], symbols[i + 1])] += freq
            for symbol in symbols:
                token_freqs[symbol] += freq

        # 计算分数
        scores = {}
        for pair, freq in pair_freqs.items():
            first, second = pair
            scores[pair] = freq / (token_freqs[first] * token_freqs[second])

        return scores

    def train(self, texts):
        """训练WordPiece分词器"""
        # 初始化词表
        word_freqs = Counter()
        for text in texts:
            words = text.split()
            word_freqs.update(words)

        vocab = {}
        for word, freq in word_freqs.items():
            # 除了第一个字符，其他都加##前缀
            tokens = [word[0]] + [f'##{c}' for c in word[1:]]
            vocab[' '.join(tokens)] = freq

        # 初始词表
        self.vocab = set()
        for word in vocab.keys():
            self.vocab.update(word.split())

        # 迭代合并
        for i in range(self.vocab_size - len(self.vocab)):
            scores = self.get_pair_scores(vocab)
            if not scores:
                break

            # 选择分数最高的pair
            best_pair = max(scores, key=scores.get)

            # 合并
            vocab = self.merge_vocab(best_pair, vocab)
            self.vocab.add(''.join(best_pair))

            if (i + 1) % 100 == 0:
                print(f"合并 {i + 1}: {best_pair}, score={scores[best_pair]:.6f}")

        self.vocab = {token: i for i, token in enumerate(sorted(self.vocab))}
        print(f"训练完成！词表大小: {len(self.vocab)}")

    def merge_vocab(self, pair, words):
        """合并词表中的字符对"""
        new_words = {}
        bigram = ' '.join(pair)
        replacement = ''.join(pair)

        for word in words:
            new_word = word.replace(bigram, replacement)
            new_words[new_word] = words[word]

        return new_words

    def tokenize(self, text):
        """对文本进行分词"""
        tokens = []
        for word in text.split():
            # 贪心最长匹配
            word_tokens = []
            start = 0

            while start < len(word):
                end = len(word)
                found = False

                while start < end:
                    substr = word[start:end]
                    if start > 0:
                        substr = f'##{substr}'

                    if substr in self.vocab:
                        word_tokens.append(substr)
                        start = end
                        found = True
                        break

                    end -= 1

                if not found:
                    # 未知字符
                    word_tokens.append('[UNK]')
                    start += 1

            tokens.extend(word_tokens)

        return tokens

import sentencepiece as spm

# 训练SentencePiece模型
spm.SentencePieceTrainer.train(
    input='corpus.txt',
    model_prefix='tokenizer',
    vocab_size=32000,
    character_coverage=0.9995,
    model_type='bpe',  # 或'unigram'
    pad_id=0,
    unk_id=1,
    bos_id=2,
    eos_id=3,
)

# 加载模型
sp = spm.SentencePieceProcessor()
sp.load('tokenizer.model')

# 编码
text = "Hello, world! 你好世界！"
tokens = sp.encode_as_pieces(text)
ids = sp.encode_as_ids(text)

print(f"Text: {text}")
print(f"Tokens: {tokens}")
print(f"IDs: {ids}")

# 解码
decoded = sp.decode_ids(ids)
print(f"Decoded: {decoded}")
assert text == decoded  # 完美还原

def get_sinusoidal_encoding(seq_len, d_model):
    """
    生成正弦位置编码

    Args:
        seq_len: 序列长度
        d_model: 嵌入维度

    Returns:
        (seq_len, d_model) 的位置编码矩阵
    """
    position = torch.arange(seq_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(
        torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
    )

    pe = torch.zeros(seq_len, d_model)
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)

    return pe

# 可视化
import matplotlib.pyplot as plt

pe = get_sinusoidal_encoding(100, 512)
plt.figure(figsize=(15, 5))
plt.imshow(pe, cmap='RdBu', aspect='auto')
plt.xlabel('Embedding Dimension')
plt.ylabel('Position')
plt.title('Sinusoidal Position Encoding')
plt.colorbar()
plt.show()

class LearnedPositionEmbedding(nn.Module):
    def __init__(self, max_len, d_model):
        super().__init__()
        self.pos_embedding = nn.Embedding(max_len, d_model)

    def forward(self, x):
        seq_len = x.size(1)
        positions = torch.arange(seq_len, device=x.device)
        return self.pos_embedding(positions)

def precompute_freqs_cis(dim, max_seq_len, theta=10000.0):
    """
    预计算RoPE的旋转频率

    Args:
        dim: 头维度
        max_seq_len: 最大序列长度
        theta: 基础频率
    """
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(max_seq_len, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # 复数形式
    return freqs_cis

def reshape_for_broadcast(freqs_cis, x):
    """调整形状以便广播"""
    ndim = x.ndim
    assert 0 <= 1 < ndim
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    return freqs_cis.view(*shape)

def apply_rotary_emb(xq, xk, freqs_cis):
    """
    应用RoPE到Q和K

    Args:
        xq: Query (batch, seq_len, n_heads, head_dim)
        xk: Key (batch, seq_len, n_heads, head_dim)
        freqs_cis: 旋转频率
    """
    # 将实数转换为复数
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # 应用旋转
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)

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

# 使用示例
class RoPEAttention(nn.Module):
    def __init__(self, n_heads, head_dim, max_seq_len):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim

        # 预计算旋转频率
        self.register_buffer(
            "freqs_cis",
            precompute_freqs_cis(head_dim, max_seq_len)
        )

    def forward(self, q, k, v):
        # q, k, v: (batch, seq_len, n_heads, head_dim)
        seq_len = q.shape[1]

        # 应用RoPE
        q, k = apply_rotary_emb(q, k, self.freqs_cis[:seq_len])

        # 计算注意力
        # ... (标准注意力计算)

        return output

class ALiBiAttention(nn.Module):
    def __init__(self, n_heads, max_seq_len):
        super().__init__()
        self.n_heads = n_heads

        # 为每个头计算斜率
        slopes = torch.Tensor(self._get_slopes(n_heads))

        # 创建偏置矩阵
        alibi = slopes.unsqueeze(1).unsqueeze(1) * torch.arange(max_seq_len).unsqueeze(0).unsqueeze(0).expand(n_heads, -1, -1)
        alibi = alibi.view(n_heads, 1, max_seq_len)

        self.register_buffer('alibi', alibi)

    def _get_slopes(self, n_heads):
        """计算每个头的斜率"""
        def get_slopes_power_of_2(n):
            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
            ratio = start
            return [start * ratio ** i for i in range(n)]

        if math.log2(n_heads).is_integer():
            return get_slopes_power_of_2(n_heads)
        else:
            closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
            return (
                get_slopes_power_of_2(closest_power_of_2) +
                self._get_slopes(2 * closest_power_of_2)[0::2][:n_heads - closest_power_of_2]
            )

    def forward(self, q, k, v):
        # q, k: (batch, n_heads, seq_len, head_dim)
        seq_len = q.size(2)

        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1))

        # 添加ALiBi偏置
        scores += self.alibi[:, :, :seq_len, :seq_len]

        # Softmax
        attn = F.softmax(scores, dim=-1)

        # 加权求和
        output = torch.matmul(attn, v)

        return output

# 可视化不同头的斜率
n_heads = 8
slopes = ALiBiAttention(n_heads, 512)._get_slopes(n_heads)
print(f"Slopes for {n_heads} heads: {slopes}")

def compute_loss(model, batch):
    """
    计算语言模型损失

    Args:
        model: GPT模型
        batch: (batch_size, seq_len+1) 的token序列
    """
    # 输入和目标
    inputs = batch[:, :-1]  # 前seq_len个token
    targets = batch[:, 1:]  # 后seq_len个token（向右移一位）

    # 前向传播
    logits, loss = model(inputs, targets=targets)

    return loss

# 训练循环
def train_step(model, optimizer, batch):
    model.train()
    optimizer.zero_grad()

    loss = compute_loss(model, batch)

    loss.backward()

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

    optimizer.step()

    return loss.item()

def clean_common_crawl(text):
    """清洗Common Crawl数据"""
    # 1. 去除HTML标签
    text = remove_html_tags(text)

    # 2. 语言过滤
    if not is_english(text):
        return None

    # 3. 质量过滤
    if len(text.split()) < 50:  # 太短
        return None

    if profanity_score(text) > 0.5:  # 低俗内容
        return None

    # 4. 去重
    if is_duplicate(text):
        return None

    # 5. 个人信息脱敏
    text = remove_pii(text)

    return text

class DataMixer:
    """数据混合器"""

    def __init__(self, datasets, weights):
        """
        Args:
            datasets: 数据集列表
            weights: 采样权重
        """
        self.datasets = datasets
        self.weights = np.array(weights)
        self.weights = self.weights / self.weights.sum()

    def sample_batch(self, batch_size):
        """采样一个批次"""
        # 根据权重选择数据集
        dataset_idx = np.random.choice(len(self.datasets), p=self.weights)
        dataset = self.datasets[dataset_idx]

        # 从选中的数据集采样
        return dataset.sample(batch_size)

# 使用示例
mixer = DataMixer(
    datasets=[common_crawl, books, wikipedia, code, arxiv],
    weights=[0.60, 0.16, 0.03, 0.045, 0.025]
)

batch = mixer.sample_batch(batch_size=32)

def estimate_training_flops(
    n_params,      # 参数量
    n_tokens,      # 训练token数
    n_layers,      # 层数
    d_model,       # 隐藏维度
    seq_len,       # 序列长度
    vocab_size,    # 词表大小
):
    """
    估算训练FLOPs

    Returns:
        总FLOPs数
    """
    # 每个token的前向FLOPs
    flops_per_token = 2 * n_params + 2 * n_layers * seq_len * d_model

    # 训练FLOPs（包含反向传播）
    training_flops = 3 * flops_per_token * n_tokens

    return training_flops

# GPT-3 175B的计算量
n_params = 175e9
n_tokens = 300e9
n_layers = 96
d_model = 12288
seq_len = 2048
vocab_size = 50257

flops = estimate_training_flops(n_params, n_tokens, n_layers, d_model, seq_len, vocab_size)

print(f"GPT-3训练FLOPs: {flops:.2e}")
print(f"GPT-3训练FLOPs: {flops / 1e21:.2f} ZFLOPs")

# 输出: GPT-3训练FLOPs: 3.14e+23 (约314 ZFLOPs)

def estimate_training_time(
    total_flops,        # 总计算量
    gpu_flops,          # 单GPU性能(FLOPS)
    n_gpus,             # GPU数量
    mfu=0.5,            # 模型FLOPs利用率
):
    """
    估算训练时间

    Args:
        total_flops: 总FLOPs
        gpu_flops: 单GPU的FP16 FLOPS
        n_gpus: GPU数量
        mfu: 模型FLOPs利用率（通常30%-50%）

    Returns:
        训练时间（小时）
    """
    effective_flops = gpu_flops * n_gpus * mfu
    training_seconds = total_flops / effective_flops
    training_hours = training_seconds / 3600

    return training_hours

# LLaMA 65B训练时间
total_flops = 6.3e23  # 约630 ZFLOPs
gpu_flops = 312e12    # A100的FP16性能
n_gpus = 2048         # 使用的GPU数量
mfu = 0.4             # 实际利用率

hours = estimate_training_time(total_flops, gpu_flops, n_gpus, mfu)
days = hours / 24

print(f"LLaMA 65B训练时间:")
print(f"  {hours:.0f} 小时")
print(f"  {days:.0f} 天")
print(f"使用 {n_gpus} 张A100 GPU")

# 输出: 约21天，使用2048张A100

def estimate_training_cost(
    training_hours,
    n_gpus,
    gpu_cost_per_hour=3.0,
    power_cost_per_kwh=0.1,
    gpu_power_kw=0.4,
    overhead=0.3,
):
    """
    估算训练成本

    Args:
        training_hours: 训练时间（小时）
        n_gpus: GPU数量
        gpu_cost_per_hour: GPU租用成本（$/小时）
        power_cost_per_kwh: 电费（$/kWh）
        gpu_power_kw: 单GPU功耗（kW）
        overhead: 其他成本占比

    Returns:
        总成本（美元）
    """
    # GPU租用成本
    gpu_cost = training_hours * n_gpus * gpu_cost_per_hour

    # 电费
    power_cost = training_hours * n_gpus * gpu_power_kw * power_cost_per_kwh

    # 总成本（含overhead）
    total_cost = (gpu_cost + power_cost) * (1 + overhead)

    return {
        'gpu_cost': gpu_cost,
        'power_cost': power_cost,
        'total_cost': total_cost,
    }

# LLaMA 65B成本
costs = estimate_training_cost(
    training_hours=21 * 24,
    n_gpus=2048,
    gpu_cost_per_hour=3.0,
)

print(f"LLaMA 65B训练成本:")
print(f"  GPU租用: ${costs['gpu_cost']/1e6:.2f}M")
print(f"  电费: ${costs['power_cost']/1e6:.2f}M")
print(f"  总成本: ${costs['total_cost']/1e6:.2f}M")

# 输出: 约$4M-$5M

def calculate_params(n_layers, d_model, vocab_size):
    """
    计算模型参数量

    Args:
        n_layers: 层数
        d_model: 隐藏维度
        vocab_size: 词表大小

    Returns:
        总参数量
    """
    # Transformer层参数
    transformer_params = n_layers * 12 * d_model ** 2

    # Embedding层（token + position）
    embedding_params = 2 * vocab_size * d_model

    # 输出层（通常与token embedding共享）
    # output_params = vocab_size * d_model

    # 总参数（不重复计算共享的embedding）
    total_params = transformer_params + embedding_params

    return total_params

# 验证GPT-2模型配置
configs = [
    ("GPT-2 Small", 12, 768, 50257),
    ("GPT-2 Medium", 24, 1024, 50257),
    ("GPT-2 Large", 36, 1280, 50257),
    ("GPT-2 XL", 48, 1600, 50257),
]

for name, n_layers, d_model, vocab_size in configs:
    params = calculate_params(n_layers, d_model, vocab_size)
    print(f"{name}: {params/1e6:.1f}M 参数")

# 输出:
# GPT-2 Small: 124.4M 参数
# GPT-2 Medium: 354.8M 参数
# GPT-2 Large: 774.0M 参数
# GPT-2 XL: 1542.0M 参数

model_configs = {
    "7B": {
        "n_layers": 32,
        "d_model": 4096,
        "n_heads": 32,
        "ffn_dim": 11008,  # SwiGLU: ~2.7 * d_model
    },
    "13B": {
        "n_layers": 40,
        "d_model": 5120,
        "n_heads": 40,
        "ffn_dim": 13824,
    },
    "70B": {
        "n_layers": 80,
        "d_model": 8192,
        "n_heads": 64,
        "ffn_dim": 28672,
    },
}

def print_model_stats(name, config):
    """打印模型统计信息"""
    n_layers = config["n_layers"]
    d_model = config["d_model"]
    n_heads = config["n_heads"]
    ffn_dim = config["ffn_dim"]

    # 计算参数量
    attn_params = n_layers * (4 * d_model ** 2)  # QKV + output
    ffn_params = n_layers * (2 * d_model * ffn_dim)  # up + down
    total_params = attn_params + ffn_params

    # 激活值内存（batch=1, seq=2048, dtype=fp16）
    seq_len = 2048
    activation_memory = (
        seq_len * d_model * 2 +  # token + pos embedding
        n_layers * seq_len * d_model * 2 * 4  # 每层4份激活值
    ) * 2 / 1024**3  # 转换为GB (fp16)

    print(f"\n{name}:")
    print(f"  层数: {n_layers}")
    print(f"  隐藏维度: {d_model}")
    print(f"  注意力头数: {n_heads}")
    print(f"  每头维度: {d_model // n_heads}")
    print(f"  FFN维度: {ffn_dim}")
    print(f"  参数量: {total_params/1e9:.1f}B")
    print(f"  激活内存: {activation_memory:.2f}GB")

for name, config in model_configs.items():
    print_model_stats(name, config)

def calculate_weight_memory(n_params, dtype='fp16'):
    """计算权重占用的显存"""
    bytes_per_param = {
        'fp32': 4,
        'fp16': 2,
        'bf16': 2,
        'int8': 1,
        'int4': 0.5,
    }

    memory_bytes = n_params * bytes_per_param[dtype]
    memory_gb = memory_bytes / (1024 ** 3)

    return memory_gb

# LLaMA 7B
params_7b = 7e9
print(f"7B模型权重:")
print(f"  FP32: {calculate_weight_memory(params_7b, 'fp32'):.2f}GB")
print(f"  FP16: {calculate_weight_memory(params_7b, 'fp16'):.2f}GB")
print(f"  INT8: {calculate_weight_memory(params_7b, 'int8'):.2f}GB")
print(f"  INT4: {calculate_weight_memory(params_7b, 'int4'):.2f}GB")

# 输出:
# FP32: 26.01GB
# FP16: 13.01GB
# INT8: 6.50GB
# INT4: 3.25GB

def calculate_training_memory(
    n_params,
    batch_size,
    seq_len,
    n_layers,
    d_model,
    mixed_precision=True,
):
    """
    计算训练显存

    Args:
        n_params: 参数量
        batch_size: 批次大小
        seq_len: 序列长度
        n_layers: 层数
        d_model: 隐藏维度
        mixed_precision: 是否使用混合精度
    """
    # 模型权重 + 梯度 + 优化器状态
    if mixed_precision:
        # 权重FP16: 2P, 梯度FP16: 2P, 优化器FP32: 8P
        model_memory = 12 * n_params
    else:
        # 全FP32: 4P + 4P + 8P = 16P
        model_memory = 16 * n_params

    # 激活值（每层保存）
    # 包括：注意力中间结果、FFN中间结果等
    activation_per_layer = batch_size * seq_len * d_model * (
        2 +  # QK, softmax(QK)
        2 +  # attention output, residual
        4    # FFN (2个中间层 * 2)
    )

    if mixed_precision:
        activation_memory = activation_per_layer * n_layers * 2  # FP16
    else:
        activation_memory = activation_per_layer * n_layers * 4  # FP32

    # 总显存
    total_memory = model_memory + activation_memory
    total_memory_gb = total_memory / (1024 ** 3)

    return {
        'model_gb': model_memory / (1024 ** 3),
        'activation_gb': activation_memory / (1024 ** 3),
        'total_gb': total_memory_gb,
    }

# LLaMA 7B训练显存
memory = calculate_training_memory(
    n_params=7e9,
    batch_size=4,
    seq_len=2048,
    n_layers=32,
    d_model=4096,
    mixed_precision=True,
)

print(f"7B模型训练显存 (batch=4, seq=2048):")
print(f"  模型+梯度+优化器: {memory['model_gb']:.2f}GB")
print(f"  激活值: {memory['activation_gb']:.2f}GB")
print(f"  总计: {memory['total_gb']:.2f}GB")

# 输出示例: 总计约120GB