大语言模型（LLM）核心技术原理

from collections import defaultdict, Counter
import random

class BigramLanguageModel:
    """最简单的二元语言模型：只根据前一个词预测下一个词"""
    def __init__(self):
        # bigram_counts[w1][w2] 表示 w1 后面出现 w2 的次数
        self.bigram_counts = defaultdict(Counter)

    def train(self, sentences):
        for tokens in sentences:
            tokens = ['<s>'] + tokens + ['</s>']  # 加上句子起止标记
            for w1, w2 in zip(tokens[:-1], tokens[1:]):
                self.bigram_counts[w1][w2] += 1

    def predict_next(self, prev_word):
        counter = self.bigram_counts[prev_word]
        total = sum(counter.values())
        if total == 0:
            return None
        # 返回归一化后的概率分布
        return {w: c / total for w, c in counter.items()}

    def generate(self, max_len=20):
        word, result = '<s>', []
        for _ in range(max_len):
            probs = self.predict_next(word)
            if not probs:
                break
            word = random.choices(list(probs), weights=list(probs.values()))[0]
            if word == '</s>':
                break
            result.append(word)
        return result

# 训练与生成示例
corpus = [['我', '喜欢', '机器', '学习'], ['我', '喜欢', '深度', '学习']]
model = BigramLanguageModel()
model.train(corpus)
print(model.predict_next('我'))     # {'喜欢': 1.0}
print(model.generate())              # 随机生成一句话

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

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    计算缩放点积注意力
    Q: 查询矩阵 (batch_size, num_heads, seq_len, d_k)
    K: 键矩阵 (batch_size, num_heads, seq_len, d_k)
    V: 值矩阵 (batch_size, num_heads, seq_len, d_v)
    """
    d_k = Q.size(-1)
    # 计算注意力分数，除以 sqrt(d_k) 防止点积过大导致 softmax 梯度消失
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)

    # 应用掩码（如因果掩码，禁止看到未来的 token）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    # 计算注意力权重（每一行归一化为概率分布）
    attention_weights = F.softmax(scores, dim=-1)

    # 用权重对 Value 加权求和
    output = torch.matmul(attention_weights, V)
    return output, attention_weights

import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.size()

        # 线性映射后拆成多个头：(batch, heads, seq_len, d_k)
        def split_heads(t):
            return t.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        Q = split_heads(self.W_q(x))
        K = split_heads(self.W_k(x))
        V = split_heads(self.W_v(x))

        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = self.dropout(F.softmax(scores, dim=-1))
        out = torch.matmul(attn, V)

        # 合并多头 (batch, seq_len, d_model)
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.W_o(out)

def positional_encoding(seq_len, d_model):
    """经典正弦/余弦位置编码"""
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(0, 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[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe.unsqueeze(0)

def apply_rotary_emb(x, freqs_cos, freqs_sin):
    """
    RoPE 简化实现：把成对的维度看作复数，用位置相关的角度旋转
    x: (batch, seq_len, num_heads, head_dim)
    """
    x1, x2 = x[..., 0::2], x[..., 1::2]   # 偶数维、奇数维
    # 复数旋转：(x1 + i*x2) * (cos + i*sin)
    rotated_x1 = x1 * freqs_cos - x2 * freqs_sin
    rotated_x2 = x1 * freqs_sin + x2 * freqs_cos
    return torch.stack([rotated_x1, rotated_x2], dim=-1).flatten(-2)

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = nn.GELU()   # 现代模型多用 GELU/SwiGLU 而非 ReLU

    def forward(self, x):
        return self.linear2(self.dropout(self.activation(self.linear1(x))))

class TransformerBlock(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = PositionWiseFeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x, mask=None):
        # Pre-LN：残差 + 子层
        x = x + self.attn(self.norm1(x), mask)
        x = x + self.ffn(self.norm2(x))
        return x

class MiniGPT(nn.Module):
    """一个可运行的迷你 GPT，展示完整数据流"""
    def __init__(self, vocab_size, d_model=256, num_layers=6, num_heads=8, max_len=512):
        super().__init__()
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Parameter(torch.zeros(1, max_len, d_model))
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, num_heads, d_model * 4) for _ in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, idx):
        b, t = idx.size()
        x = self.token_emb(idx) + self.pos_emb[:, :t, :]
        # 因果掩码：下三角为 1，保证只能看到自己和之前的 token
        mask = torch.tril(torch.ones(t, t)).view(1, 1, t, t)
        for block in self.blocks:
            x = block(x, mask)
        return self.head(self.norm(x))   # 输出每个位置对整个词表的 logits

class LoRALayer(nn.Module):
    def __init__(self, in_features, out_features, rank=8, alpha=16):
        super().__init__()
        self.lora_A = nn.Parameter(torch.zeros(rank, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        self.scaling = alpha / rank
        # A 用 Kaiming 初始化，B 初始化为 0，保证训练开始时 LoRA 不改变原输出
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        nn.init.zeros_(self.lora_B)

    def forward(self, x):
        # 只学习"增量"，原始权重保持冻结
        return self.scaling * (x @ self.lora_A.T @ self.lora_B.T)

class LinearWithLoRA(nn.Module):
    """把已有的线性层包一层 LoRA，原权重冻结"""
    def __init__(self, linear: nn.Linear, rank=8, alpha=16):
        super().__init__()
        self.linear = linear
        for p in self.linear.parameters():
            p.requires_grad = False
        self.lora = LoRALayer(linear.in_features, linear.out_features, rank, alpha)

    def forward(self, x):
        return self.linear(x) + self.lora(x)

class KVCacheAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.cache_k, self.cache_v = None, None

    def forward(self, x, use_cache=True):
        b, t, _ = x.size()
        q = self.W_q(x).view(b, t, self.num_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(b, t, self.num_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(b, t, self.num_heads, self.d_k).transpose(1, 2)

        if use_cache and self.cache_k is not None:
            # 把历史 K、V 拼接到当前，避免重复计算
            k = torch.cat([self.cache_k, k], dim=2)
            v = torch.cat([self.cache_v, v], dim=2)
        self.cache_k, self.cache_v = k, v

        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v).transpose(1, 2).contiguous().view(b, t, -1)
        return self.W_o(out)

import math

def perplexity(model_log_probs):
    """给定模型对每个真实 token 的对数概率，计算困惑度"""
    n = len(model_log_probs)
    avg_neg_log_prob = -sum(model_log_probs) / n
    return math.exp(avg_neg_log_prob)

# 假设模型对 5 个 token 的对数概率如下（越接近 0 说明越自信且正确）
log_probs = [-0.1, -0.5, -0.2, -1.0, -0.3]
print(f"困惑度 PPL = {perplexity(log_probs):.3f}")

# 奖励模型训练：让"更好的回答"得分高于"更差的回答"
def reward_model_loss(reward_model, prompt, chosen, rejected):
    r_chosen = reward_model(prompt, chosen)      # 人类更喜欢的回答
    r_rejected = reward_model(prompt, rejected)  # 人类不喜欢的回答
    # 排序损失：让 chosen 的分数尽量高于 rejected
    return -F.logsigmoid(r_chosen - r_rejected).mean()

# PPO 阶段的奖励设计：奖励分 - KL 惩罚，防止跑偏
def ppo_reward(reward_model, policy_logprob, ref_logprob, prompt, response, kl_coef=0.1):
    score = reward_model(prompt, response)
    # KL 惩罚：不能离原始 SFT 模型太远，避免 reward hacking
    kl_penalty = kl_coef * (policy_logprob - ref_logprob)
    return score - kl_penalty

def dpo_loss(policy_chosen_logp, policy_rejected_logp,
             ref_chosen_logp, ref_rejected_logp, beta=0.1):
    """
    直接偏好优化损失：拉大 chosen 与 rejected 的相对对数概率差
    logp 为模型对整段回答的对数概率之和
    """
    pi_logratio = policy_chosen_logp - policy_rejected_logp
    ref_logratio = ref_chosen_logp - ref_rejected_logp
    return -F.logsigmoid(beta * (pi_logratio - ref_logratio)).mean()

class TopKMoE(nn.Module):
    def __init__(self, d_model, num_experts=8, top_k=2, d_ff=None):
        super().__init__()
        d_ff = d_ff or d_model * 4
        self.top_k = top_k
        self.gate = nn.Linear(d_model, num_experts)
        self.experts = nn.ModuleList([
            PositionWiseFeedForward(d_model, d_ff) for _ in range(num_experts)
        ])

    def forward(self, x):
        # 路由：为每个 token 选出得分最高的 top_k 个专家
        logits = self.gate(x)                       # (b, t, num_experts)
        weights, idx = torch.topk(logits.softmax(-1), self.top_k, dim=-1)
        out = torch.zeros_like(x)
        for i in range(self.top_k):
            for e, expert in enumerate(self.experts):
                mask = (idx[..., i] == e)
                if mask.any():
                    out[mask] += weights[..., i][mask].unsqueeze(-1) * expert(x[mask])
        return out

from collections import Counter, defaultdict

def get_pair_stats(vocab):
    """统计所有相邻符号对的频率，vocab 形如 {'l o w </w>': 5}"""
    pairs = defaultdict(int)
    for word, freq in vocab.items():
        symbols = word.split()
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return pairs

def merge_pair(pair, vocab):
    """把词表里所有出现的 pair 合并成一个符号"""
    new_vocab = {}
    bigram = ' '.join(pair)
    replacement = ''.join(pair)
    for word, freq in vocab.items():
        new_word = word.replace(bigram, replacement)
        new_vocab[new_word] = freq
    return new_vocab

# 初始语料：词后加 </w> 标记词尾，字符间用空格分隔
corpus = {'l o w </w>': 5, 'l o w e r </w>': 2,
          'n e w e s t </w>': 6, 'w i d e s t </w>': 3}

num_merges = 10
merges = []
for i in range(num_merges):
    stats = get_pair_stats(corpus)
    if not stats:
        break
    best = max(stats, key=stats.get)   # 选出最高频的相邻对
    corpus = merge_pair(best, corpus)
    merges.append(best)
    print(f"第 {i+1} 次合并: {best} (频率 {stats[best]})")

# 前几次合并通常是 ('e','s') -> 'es'，('es','t') -> 'est'，('l','o') -> 'lo' ...
print("学到的合并规则:", merges)

# byte-level 的核心：先转字节，天然无 OOV
text = "你好, AI!"
byte_seq = list(text.encode('utf-8'))
print(byte_seq)          # [228, 189, 160, 230, 156, 189, 44, 32, 65, 73, 33]
print(len(byte_seq))     # 11 个字节：中文每字 3 字节，ASCII 每字符 1 字节
# BBPE 会在这些字节上继续合并高频对，把常见字/词组压成单个 token

# 用 tiktoken 直观感受 token 数（GPT 系列官方分词器）
import tiktoken

enc = tiktoken.get_encoding("cl100k_base")   # GPT-4 使用的编码
for text in ["Hello world", "你好世界", "1234567890"]:
    ids = enc.encode(text)
    print(f"{text!r:20} -> {len(ids)} tokens: {ids}")
# 英文短语常 2~3 token，等长中文往往 token 更多，纯数字容易被切成多段

L(N) ≈ (Nc / N)^αN      # 只受参数量限制时
L(D) ≈ (Dc / D)^αD      # 只受数据量限制时
其中经验指数 αN ≈ 0.076，αD ≈ 0.095（Kaplan et al. 2020）

import numpy as np

def predicted_loss(N, Nc=8.8e13, alpha_N=0.076):
    """根据参数量 N 预测损失（Kaplan 幂律，仅示意）"""
    return (Nc / N) ** alpha_N

for N in [1e8, 1e9, 1e10, 1e11, 1.75e11]:   # 0.1B 到 175B
    print(f"参数量 {N:.0e} -> 预测损失 {predicted_loss(N):.3f}")
# 每放大 10 倍参数，损失以固定比例下降，呈平滑幂律

def chinchilla_optimal(compute_flops):
    """给定算力预算（FLOPs），估算计算最优的参数量与数据量"""
    # 经验近似：C ≈ 6 * N * D，且最优时 D ≈ 20 * N
    # => C ≈ 6 * N * 20N = 120 * N^2 => N = sqrt(C / 120)
    N_opt = (compute_flops / 120) ** 0.5
    D_opt = 20 * N_opt
    return N_opt, D_opt

# 假设有 1e24 FLOPs 的预算
N, D = chinchilla_optimal(1e24)
print(f"最优参数量 ≈ {N:.2e}，最优训练 token ≈ {D:.2e}")

import numpy as np

def softmax_with_temperature(logits, T=1.0):
    logits = np.array(logits, dtype=np.float64) / T
    logits -= logits.max()             # 数值稳定：减去最大值防溢出
    exp = np.exp(logits)
    return exp / exp.sum()

logits = [2.0, 1.0, 0.5, 0.1]
for T in [0.5, 1.0, 2.0]:
    p = softmax_with_temperature(logits, T)
    print(f"T={T}: {np.round(p, 3)}")
# T=0.5 分布更尖锐（第一个 token 概率更高）
# T=2.0 分布更平坦（各 token 概率更接近）

def top_k_top_p_filter(probs, top_k=0, top_p=0.0):
    """对概率分布做 top-k / top-p 截断，返回重新归一化后的分布"""
    probs = np.array(probs, dtype=np.float64)
    idx_sorted = np.argsort(probs)[::-1]          # 概率从高到低的索引
    sorted_probs = probs[idx_sorted]

    keep = np.ones(len(probs), dtype=bool)
    if top_k > 0:
        keep_sorted = np.zeros(len(probs), dtype=bool)
        keep_sorted[:top_k] = True
    else:
        keep_sorted = np.ones(len(probs), dtype=bool)

    if top_p > 0.0:
        cumulative = np.cumsum(sorted_probs)
        # 保留累积概率首次超过 top_p 时（含该 token）的前缀
        cutoff = np.searchsorted(cumulative, top_p) + 1
        keep_sorted[cutoff:] = False

    filtered = np.zeros_like(sorted_probs)
    filtered[keep_sorted] = sorted_probs[keep_sorted]
    # 还原回原始顺序并归一化
    result = np.zeros_like(probs)
    result[idx_sorted] = filtered
    return result / result.sum()

probs = softmax_with_temperature([3.0, 2.0, 1.5, 1.0, 0.2, -1.0])
print("top_p=0.9:", np.round(top_k_top_p_filter(probs, top_p=0.9), 3))
print("top_k=3 :", np.round(top_k_top_p_filter(probs, top_k=3), 3))

def apply_repetition_penalty(logits, generated_ids, penalty=1.2):
    logits = np.array(logits, dtype=np.float64)
    for tid in set(generated_ids):
        if logits[tid] > 0:
            logits[tid] /= penalty     # 正 logit 除以惩罚，变小
        else:
            logits[tid] *= penalty     # 负 logit 乘以惩罚，更负
    return logits

import heapq

def beam_search(step_fn, start_token, eos_token, beam_size=3, max_len=20):
    """
    step_fn(seq) -> 返回下一步每个 token 的对数概率数组
    维护 (累计对数概率, 序列) 的 beam
    """
    beams = [(0.0, [start_token])]
    for _ in range(max_len):
        candidates = []
        for score, seq in beams:
            if seq[-1] == eos_token:
                candidates.append((score, seq))    # 已结束，直接保留
                continue
            log_probs = step_fn(seq)
            # 只看每条 beam 的 top 候选以控制分支数
            top_ids = np.argsort(log_probs)[::-1][:beam_size]
            for tid in top_ids:
                candidates.append((score + log_probs[tid], seq + [int(tid)]))
        # 保留总分最高的 beam_size 条
        beams = heapq.nlargest(beam_size, candidates, key=lambda x: x[0])
        if all(seq[-1] == eos_token for _, seq in beams):
            break
    return max(beams, key=lambda x: x[0])[1]

import numpy as np

def absmax_quantize(x, bits=8):
    """对称量化：以绝对值最大值为界，线性映射到 [-(2^(b-1)-1), 2^(b-1)-1]"""
    q_max = 2 ** (bits - 1) - 1           # INT8 -> 127
    scale = np.abs(x).max() / q_max       # 缩放因子
    q = np.round(x / scale).astype(np.int8)
    return q, scale

def dequantize(q, scale):
    return q.astype(np.float32) * scale

w = np.array([0.5, -1.2, 3.4, -0.1, 2.0], dtype=np.float32)
q, scale = absmax_quantize(w)
recovered = dequantize(q, scale)
print("量化后 INT8:", q)
print("缩放因子:", round(float(scale), 4))
print("反量化误差:", np.round(recovered - w, 3))

def estimate_model_memory(num_params_b, bits):
    """估算纯权重显存占用（GB），num_params_b 单位为十亿"""
    bytes_per_param = bits / 8
    gb = num_params_b * 1e9 * bytes_per_param / (1024 ** 3)
    return gb

for bits in [16, 8, 4]:
    print(f"70B 模型 @ {bits}-bit ≈ {estimate_model_memory(70, bits):.1f} GB")
# 16-bit ≈ 130 GB, 8-bit ≈ 65 GB, 4-bit ≈ 33 GB（另需 KV 缓存与激活的额外开销）

import numpy as np

def get_rope_freqs(head_dim, base=10000, scaling=1.0):
    """计算 RoPE 各维度的旋转频率；scaling>1 表示位置插值扩展上下文"""
    # 频率随维度指数衰减；除以 scaling 相当于把位置"压缩"
    freqs = 1.0 / (base ** (np.arange(0, head_dim, 2) / head_dim))
    return freqs / scaling

base_freqs = get_rope_freqs(128)
extended = get_rope_freqs(128, scaling=4.0)   # 把上下文扩到 4 倍
print("原始最低频率:", round(float(base_freqs[-1]), 8))
print("插值后最低频率:", round(float(extended[-1]), 8))

from math import comb

def pass_at_k(n, c, k):
    """
    n: 生成的样本总数
    c: 其中通过测试的样本数
    k: pass@k 的 k
    返回 pass@k 的无偏估计
    """
    if n - c < k:
        return 1.0                      # 通过的太多，k 个里必有通过
    # 1 - (从 n-c 个失败样本里全抽中 k 个的概率)
    return 1.0 - comb(n - c, k) / comb(n, k)

# 例：生成 20 个，8 个通过
for k in [1, 5, 10]:
    print(f"pass@{k} = {pass_at_k(20, 8, k):.3f}")
# k 越大，pass@k 越高：多试几次总更可能碰到对的

# Few-shot 示例：教模型做情感分类
评论：这家餐厅太棒了，还会再来！ -> 正面
评论：等了一个小时菜还没上，气死了。 -> 负面
评论：环境一般，味道还行。 -> 中性
评论：服务员态度非常好，菜品超出预期！ ->
（模型会据前三个示例的格式，输出：正面）

问题：小明有 5 个苹果，给了小红 2 个，又买了 3 个，现在有几个？

# 不加 CoT，容易直接蒙一个数
# 加 CoT：
让我们一步步思考：
1. 小明开始有 5 个。
2. 给了小红 2 个，剩 5 - 2 = 3 个。
3. 又买了 3 个，共 3 + 3 = 6 个。
所以答案是 6 个。

# 用 system + user 分工构造 prompt，是工程实践的标配
messages = [
    {"role": "system", "content": "你是严谨的法律助手，只依据用户提供的合同原文回答，"
                                   "无法确定时明确说明，不得编造条款。"},
    {"role": "user", "content": "根据以下合同，违约金比例是多少？\n<合同原文...>"},
]
# 关键点：把"角色、约束、防幻觉要求"放进 system，把具体任务与数据放进 user

import torch
import torch.nn as nn

# 1. 准备数据：字符级"分词"
text = "hello world " * 200          # 玩具语料
chars = sorted(set(text))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
data = torch.tensor([stoi[c] for c in text], dtype=torch.long)

# 2. 构造训练批次：输入 x 与右移一位的目标 y
block_size = 32
def get_batch(batch_size=16):
    ix = torch.randint(len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i:i + block_size] for i in ix])
    y = torch.stack([data[i + 1:i + 1 + block_size] for i in ix])
    return x, y

# 3. 初始化模型与优化器
model = MiniGPT(vocab_size, d_model=128, num_layers=4, num_heads=4, max_len=block_size)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)  # 3e-4 是 Transformer 常用起始学习率
criterion = nn.CrossEntropyLoss()

# 4. 训练循环
model.train()
for step in range(500):
    x, y = get_batch()
    logits = model(x)                                # (B, T, vocab_size)
    loss = criterion(logits.view(-1, vocab_size), y.view(-1))
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 梯度裁剪防爆炸
    optimizer.step()
    if step % 100 == 0:
        # 初始损失约等于 ln(vocab_size)（随机猜测），随训练稳步下降
        print(f"step {step}: loss = {loss.item():.4f}")

@torch.no_grad()
def generate(model, start_str, max_new_tokens=100, temperature=0.8, top_k=10):
    model.eval()
    idx = torch.tensor([[stoi[c] for c in start_str]], dtype=torch.long)
    for _ in range(max_new_tokens):
        idx_cond = idx[:, -block_size:]              # 只取最近 block_size 个
        logits = model(idx_cond)[:, -1, :] / temperature   # 取最后一个位置
        if top_k is not None:
            v, _ = torch.topk(logits, top_k)
            logits[logits < v[:, [-1]]] = float('-inf')     # top-k 截断
        probs = torch.softmax(logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)   # 按概率采样
        idx = torch.cat([idx, next_id], dim=1)
    return ''.join(itos[i] for i in idx[0].tolist())

# print(generate(model, "hello"))   # 训练充分后能续写出 "hello world hello world ..."

def training_flops(N_params, D_tokens):
    return 6 * N_params * D_tokens

def gpu_days(flops, gpu_tflops=312, mfu=0.4):
    """gpu_tflops: 单卡峰值（A100 BF16 约 312 TFLOPS）；mfu: 实际利用率约 0.3~0.5"""
    effective = gpu_tflops * 1e12 * mfu
    seconds = flops / effective
    return seconds / 86400

flops = training_flops(7e9, 1e12)          # 7B 模型，1T token
print(f"总算力 ≈ {flops:.2e} FLOPs")
print(f"单卡 A100 需 ≈ {gpu_days(flops):.0f} GPU-天")
print(f"1024 卡集群 ≈ {gpu_days(flops) / 1024:.2f} 天")

def full_finetune_memory_gb(N_params_b):
    """全参数 AdamW 混合精度训练的静态显存估算（不含激活）"""
    return N_params_b * 1e9 * 16 / (1024 ** 3)

for n in [0.5, 7, 13, 70]:
    print(f"{n}B 模型全参训练静态显存 ≈ {full_finetune_memory_gb(n):.0f} GB")
# 0.5B≈7GB, 7B≈104GB, 13B≈194GB, 70B≈1043GB(需多卡分片)

import torch.nn as nn

vocab_size, d_model = 50000, 4096
emb = nn.Embedding(vocab_size, d_model)
print(f"嵌入矩阵参数量: {vocab_size * d_model / 1e6:.0f}M")   # 205M
ids = torch.tensor([[10, 250, 7]])
print(emb(ids).shape)     # torch.Size([1, 3, 4096])：每个 token 查出一个 4096 维向量

import torch.nn.functional as F

def cosine_sim(a, b):
    return F.cosine_similarity(a, b, dim=-1)

# 假设从训练好的模型取出词向量（此处用随机张量示意接口）
vec_cat, vec_dog, vec_car = torch.randn(4096), torch.randn(4096), torch.randn(4096)
print("cat vs dog:", cosine_sim(vec_cat, vec_dog))   # 语义相近的词相似度应更高
print("cat vs car:", cosine_sim(vec_cat, vec_car))

import torch
import torch.nn as nn

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        # 只按均方根缩放，不做中心化，比 LayerNorm 少一次求均值
        rms = x.pow(2).mean(dim=-1, keepdim=True).sqrt()
        return self.weight * x / (rms + self.eps)

x = torch.randn(2, 8, 512)
print(RMSNorm(512)(x).shape)   # torch.Size([2, 8, 512])

import torch.nn.functional as F

class SwiGLU(nn.Module):
    """LLaMA 风格的门控 FFN：一路 SiLU 激活，一路作为门控相乘"""
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.w_gate = nn.Linear(dim, hidden_dim, bias=False)
        self.w_up = nn.Linear(dim, hidden_dim, bias=False)
        self.w_down = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x):
        # SiLU(x) = x * sigmoid(x)，门控路乘以变换路
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))

print(SwiGLU(512, 1376)(torch.randn(2, 8, 512)).shape)
# 注意：为保持总参数量相当，SwiGLU 的 hidden_dim 通常取约 2/3 * 4 * dim

def estimate_bubble_ratio(num_stages, num_microbatches):
    """流水线并行的气泡占比：stage 越多、microbatch 越少，闲置越严重"""
    return (num_stages - 1) / (num_stages - 1 + num_microbatches)

print(f"8 段流水、8 microbatch:  气泡 {estimate_bubble_ratio(8, 8):.1%}")
print(f"8 段流水、32 microbatch: 气泡 {estimate_bubble_ratio(8, 32):.1%}")
# 增大 microbatch 数能显著降低流水线气泡、提升利用率

# 一个极简的工具调用循环示意（不依赖具体 SDK）
def calculator(expr):
    return eval(expr)     # 生产中务必用安全的表达式求值，切勿裸 eval

tools = {"calculator": calculator}

def agent_loop(llm_call, user_query, max_steps=5):
    messages = [{"role": "user", "content": user_query}]
    for _ in range(max_steps):
        resp = llm_call(messages)        # 模型可能返回普通文本或工具调用请求
        if resp.get("tool_call"):
            name = resp["tool_call"]["name"]
            args = resp["tool_call"]["arguments"]
            result = tools[name](**args)      # 执行工具
            # 把工具结果作为观察回填，进入下一轮
            messages.append({"role": "tool", "name": name, "content": str(result)})
        else:
            return resp["content"]           # 模型给出最终答案，结束
    return "达到最大步数，未得出结论"

class VisionProjector(nn.Module):
    """把视觉编码器输出映射到 LLM 词嵌入空间，成为可与文本拼接的视觉 token"""
    def __init__(self, vision_dim=1024, llm_dim=4096):
        super().__init__()
        self.proj = nn.Sequential(
            nn.Linear(vision_dim, llm_dim),
            nn.GELU(),
            nn.Linear(llm_dim, llm_dim),
        )

    def forward(self, image_features):
        # image_features: (batch, num_patches, vision_dim)
        return self.proj(image_features)   # -> (batch, num_patches, llm_dim)

# 之后把这些视觉 token 与文本 token 在序列维度拼接，一起送入 LLM