22 - 大语言模型原理¶

# 对比两种Scaling Law的最优分配策略
import numpy as np

def kaplan_optimal(C):
    """Kaplan: 优先扩大模型"""
    # Kaplan建议 N ∝ C^0.73, D ∝ C^0.27
    N = 0.1 * (C ** 0.73)
    D = C / (6 * N)
    return N, D

def chinchilla_optimal(C):
    """Chinchilla: 模型和数据平衡扩展"""
    # Chinchilla定律: C ≈ 6ND, 且N与D等比例增长
    # 即 N ∝ C^0.5, D ∝ C^0.5, 满足 6*N*D = C
    N = (C / 6) ** 0.5  # N = sqrt(C/6)
    D = (C / 6) ** 0.5  # D = sqrt(C/6), 使得 6*N*D = C
    return N, D

# 以GPT-3的计算预算为例 (约3.14 × 10^23 FLOPs)
C = 3.14e23

N_kaplan, D_kaplan = kaplan_optimal(C)
N_chinchilla, D_chinchilla = chinchilla_optimal(C)

print(f"计算预算: {C:.2e} FLOPs")
print(f"\nKaplan最优: N={N_kaplan/1e9:.1f}B, D={D_kaplan/1e9:.1f}B tokens")
print(f"Chinchilla最优: N={N_chinchilla/1e9:.1f}B, D={D_chinchilla/1e9:.1f}B tokens")
print(f"\nGPT-3实际: N=175B, D=300B tokens")
print(f"→ GPT-3更接近Kaplan策略（模型偏大，数据偏少）")

# 涌现能力的量化分析示例
import matplotlib.pyplot as plt
import numpy as np

# 模型规模（参数量，单位：B）
model_sizes = [0.1, 0.3, 1, 3, 10, 30, 100, 280]

# 模拟不同任务的性能曲线
def emergence_curve(N, N_crit, max_perf=100):
    """涌现能力的Sigmoid模型"""
    k = 10 / N_crit  # 陡峭度
    return max_perf / (1 + np.exp(-k * (N - N_crit)))

# 不同任务的临界规模
tasks = {
    "算术运算": 10,
    "思维链推理": 30,
    "多语言理解": 5,
    "代码生成": 10,
    "数学问题": 100,
}

plt.figure(figsize=(12, 6))
for task, N_crit in tasks.items():
    perf = [emergence_curve(N, N_crit) for N in model_sizes]
    plt.plot(model_sizes, perf, 'o-', label=f'{task} (临界点: {N_crit}B)')

plt.xscale('log')
plt.xlabel('模型参数量 (B)')
plt.ylabel('任务性能')
plt.title('涌现能力与模型规模的关系')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

import torch
import torch.nn as nn

class StableTransformerBlock(nn.Module):
    """稳定的Transformer块，包含多种梯度稳定技术"""

    def __init__(self, d_model, n_head, dropout=0.1):
        super().__init__()
        # 1. Pre-LayerNorm（比Post-LN更稳定）
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

        self.attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),  # GELU比ReLU更平滑
            nn.Linear(4 * d_model, d_model),
        )
        self.dropout = nn.Dropout(dropout)

        # 2. 可学习的缩放因子（用于稳定训练初期）
        self.attn_scale = nn.Parameter(torch.ones(1) * 0.1)
        self.ffn_scale = nn.Parameter(torch.ones(1) * 0.1)

    def forward(self, x):
        # Pre-LN结构
        # 3. 残差连接 + 缩放
        attn_out, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
        x = x + self.attn_scale * self.dropout(attn_out)

        ffn_out = self.ffn(self.norm2(x))
        x = x + self.ffn_scale * self.dropout(ffn_out)

        return x


def gradient_clipping(model, max_norm=1.0):
    """梯度裁剪：防止梯度爆炸"""
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
    return model


def adaptive_gradient_clipping(model, clip_factor=0.01):
    """自适应梯度裁剪（AGC）：根据参数范数动态调整"""
    for param in model.parameters():
        if param.grad is not None:
            param_norm = torch.norm(param)
            grad_norm = torch.norm(param.grad)
            max_norm = param_norm * clip_factor
            if grad_norm > max_norm and grad_norm > 0:
                param.grad.data.mul_(max_norm / grad_norm)
    return model

def warmup_lr(step, warmup_steps, base_lr):
    """线性预热学习率"""
    return base_lr * min(1.0, step / warmup_steps)

import math

def cosine_decay_lr(step, total_steps, max_lr, min_lr=0):
    """余弦衰减学习率"""
    ratio = step / total_steps
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * ratio))


def cosine_decay_with_warmup(step, warmup_steps, total_steps, max_lr, min_lr=0):
    """带预热的余弦衰减（LLM训练标准配置）"""
    if step < warmup_steps:
        # 线性预热
        return max_lr * step / warmup_steps
    else:
        # 余弦衰减
        ratio = (step - warmup_steps) / (total_steps - warmup_steps)
        return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * ratio))


# 可视化学习率调度
def plot_lr_schedule():
    import matplotlib.pyplot as plt

    total_steps = 100000
    warmup_steps = 2000
    max_lr = 6e-4
    min_lr = 6e-5

    steps = list(range(total_steps))
    lrs = [cosine_decay_with_warmup(s, warmup_steps, total_steps, max_lr, min_lr)
           for s in steps]

    plt.figure(figsize=(12, 4))
    plt.plot(steps, lrs)
    plt.axvline(x=warmup_steps, color='r', linestyle='--', label=f'Warmup End ({warmup_steps})')
    plt.xlabel('Training Steps')
    plt.ylabel('Learning Rate')
    plt.title('Cosine Decay with Warmup Learning Rate Schedule')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()

plot_lr_schedule()

def wsd_schedule(step, warmup_steps, stable_steps, decay_steps, max_lr, min_lr):
    """
    WSD学习率调度：Warmup -> Stable -> Decay

    优点：
    1. 稳定阶段便于checkpoint选择
    2. 衰减阶段快速收敛
    3. 适合继续预训练
    """
    if step < warmup_steps:
        return max_lr * step / warmup_steps
    elif step < warmup_steps + stable_steps:
        return max_lr
    else:
        decay_progress = (step - warmup_steps - stable_steps) / decay_steps
        return min_lr + (max_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * decay_progress))

from dataclasses import dataclass

@dataclass
class LLMTrainingConfig:
    """LLM训练稳定性配置"""

    # 模型配置
    hidden_size: int = 4096
    num_layers: int = 32
    num_heads: int = 32

    # 优化配置
    learning_rate: float = 3e-4
    weight_decay: float = 0.1
    beta1: float = 0.9
    beta2: float = 0.95  # AdamW的beta2，比默认0.999更小

    # 学习率调度
    warmup_steps: int = 2000
    lr_scheduler: str = "cosine"
    min_lr_ratio: float = 0.1  # 最终LR = max_lr * 0.1

    # 梯度控制
    max_grad_norm: float = 1.0
    grad_clip_type: str = "global"  # global或adaptive

    # 批次配置
    batch_size: int = 4_000_000  # 总token数
    gradient_accumulation_steps: int = 1

    # 精度配置
    dtype: str = "bf16"  # bf16比fp16更稳定
    grad_dtype: str = "fp32"  # 梯度用fp32保持精度