数值计算与稳定性

import numpy as np

# 演示机器精度
eps_32 = np.finfo(np.float32).eps
eps_64 = np.finfo(np.float64).eps
print(f"FP32 机器精度: {eps_32}")  # 约 1.19e-7
print(f"FP64 机器精度: {eps_64}")  # 约 2.22e-16

# 灾难性抵消演示
a = np.float32(1e8)
b = np.float32(1e8 + 1)
print(f"(1e8+1) - 1e8 = {b - a}")  # 可能为 0！

# 浮点数表示范围
info_32 = np.finfo(np.float32)
print(f"FP32 最大值: {info_32.max:.2e}")   # 约 3.4e+38
print(f"FP32 最小正值: {info_32.tiny:.2e}")  # 约 1.18e-38
print(f"FP32 最小正规数: {info_32.smallest_normal:.2e}")

# 大数吃小数问题
x = np.float32(1e8)
y = np.float32(1.0)
print(f"1e8 + 1.0 = {x + y}")  # 结果等于 1e8

import numpy as np

# 上溢演示
x = np.float32(100.0)
print(f"exp(100) = {np.exp(x)}")  # inf!

# 下溢演示
x = np.float32(-100.0)
print(f"exp(-100) = {np.exp(x)}")  # 0!

# 对数下溢
p = np.exp(np.float32(-100.0))
print(f"log(exp(-100)) = {np.log(p)}")  # -inf!

# 安全的指数计算（上溢保护）
def safe_exp(x, dtype=np.float32):
    info = np.finfo(dtype)
    x = np.clip(x, a_min=None, a_max=info.max * 0.9)
    x = np.clip(x, a_min=info.min * 0.9, a_max=None)
    return np.exp(x)

# 安全的对数计算（下溢保护）
def safe_log(x, eps=1e-7):
    return np.log(np.clip(x, eps, None))

# 测试
print(f"safe_exp(100) = {safe_exp(100.0)}")
print(f"safe_log(0) = {safe_log(0.0)}")

import numpy as np

# 不稳定的 Softmax 实现
def unstable_softmax(x):
    e = np.exp(x)
    return e / np.sum(e)

# 数值稳定的 Softmax 实现
def stable_softmax(x):
    shifted = x - np.max(x)  # 关键技巧
    e = np.exp(shifted)
    return e / np.sum(e)

# 测试
x = np.array([1000, 2000, 3000], dtype=np.float64)
# print(unstable_softmax(x))  # 崩溃！
print(stable_softmax(x))  # [2.06e-435, ..., 1.0] 近似 [0, 0, 1]

# 批量矩阵的数值稳定 Softmax（沿最后一维）
def batch_stable_softmax(logits):
    """logits shape: (batch, seq_len, vocab_size)"""
    max_vals = np.max(logits, axis=-1, keepdims=True)
    shifted = logits - max_vals
    exp_vals = np.exp(shifted)
    sum_vals = np.sum(exp_vals, axis=-1, keepdims=True)
    return exp_vals / (sum_vals + 1e-8)

# 模拟 Transformer 注意力中的 Softmax
batch_size, seq_len, vocab_size = 2, 128, 10000
logits = np.random.randn(batch_size, seq_len, vocab_size) * 5
probs = batch_stable_softmax(logits)
print(f"概率和: {np.sum(probs, axis=-1)[0, 0]:.6f}")  # 应为 1.0

import numpy as np
from scipy.special import logsumexp

# 不稳定的 LSE
def unstable_lse(x):
    return np.log(np.sum(np.exp(x)))

# 数值稳定的 LSE
def stable_lse(x):
    c = np.max(x)
    return c + np.log(np.sum(np.exp(x - c)))

# 测试
x = np.array([1000, 200, 300], dtype=np.float64)
# unstable_lse(x)  → 上溢崩溃
print(f"stable LSE: {stable_lse(x):.4f}")
print(f"scipy LSE:  {logsumexp(x):.4f}")

# 用 LSE 计算数值稳定的交叉熵损失
def stable_cross_entropy(logits, target_idx):
    """
    交叉熵 = -log(softmax(x)[target])
           = log(sum(exp(x))) - x[target]
           = LSE(x) - x[target]
    """
    lse = logsumexp(logits)
    return lse - logits[target_idx]

# 测试
logits = np.array([2.0, 1.0, 0.1])
target = 0
loss = stable_cross_entropy(logits, target)
print(f"交叉熵损失: {loss:.4f}")

# 对比：直接用 softmax 再取 log 可能因下溢而不准确
softmax_probs = np.exp(logits - np.max(logits))
softmax_probs /= np.sum(softmax_probs)
naive_loss = -np.log(softmax_probs[target])
print(f"naive 损失: {naive_loss:.4f}")

import numpy as np

# 演示梯度消失：sigmoid 的连乘效应
def sigmoid_grad(x):
    s = 1 / (1 + np.exp(-x))
    return s * (1 - s)

# 模拟 20 层网络的梯度传播
initial_grad = np.float64(1.0)
for layer in range(20):
    initial_grad *= sigmoid_grad(0.0)  # 最大值 0.25
    print(f"Layer {layer+1}: grad = {initial_grad:.2e}")

# 20 层后: ~0.25^20 ≈ 9e-13，接近下溢阈值

# 梯度裁剪（Gradient Clipping）
def clip_gradient_norm(grads, max_norm=1.0):
    """按全局范数裁剪梯度"""
    total_norm = np.sqrt(sum(np.sum(g ** 2) for g in grads))
    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1.0:
        return [g * clip_coef for g in grads]
    return grads

# 模拟梯度爆炸场景
explosive_grads = [np.random.randn(100) * 50 for _ in range(5)]
original_norm = np.sqrt(sum(np.sum(g ** 2) for g in explosive_grads))
clipped = clip_gradient_norm(explosive_grads, max_norm=5.0)
clipped_norm = np.sqrt(sum(np.sum(g ** 2) for g in clipped))
print(f"原始范数: {original_norm:.1f}, 裁剪后: {clipped_norm:.1f}")

import torch

# PyTorch 原生 AMP（自动混合精度）
scaler = torch.cuda.amp.GradScaler()

for inputs, targets in dataloader:
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# BF16 训练（推荐 Ampere+ GPU）
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

# 手动 Loss Scaling 实现
class ManualLossScaler:
    def __init__(self, init_scale=2**16, growth_factor=2.0, backoff_factor=0.5):
        self.scale = init_scale
        self.growth_factor = growth_factor
        self.backoff_factor = backoff_factor
        self.good_steps = 0

    def step(self, found_inf):
        if found_inf:
            self.scale *= self.backoff_factor
            self.good_steps = 0
        else:
            self.good_steps += 1
            if self.good_steps >= 2000:
                self.scale *= self.growth_factor
                self.good_steps = 0

    def unscale_grads(self, grads):
        return [g / self.scale for g in grads]

scaler = ManualLossScaler()
print(f"当前缩放系数: {scaler.scale}")

import torch
import torch.nn as nn

# 数值稳定性调试工具
def train_with_debug(model, loader, optimizer):
    torch.autograd.set_detect_anomaly(True)  # 追踪 NaN 来源

    for batch in loader:
        optimizer.zero_grad()
        out = model(batch)
        loss = nn.functional.cross_entropy(out, targets)

        # 检查损失
        if torch.isnan(loss):
            print("NaN 检测!")
            break

        loss.backward()

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

        # 检查梯度
        for name, param in model.named_parameters():
            if param.grad is not None and torch.isnan(param.grad).any():
                print(f"NaN 梯度: {name}")

# 数值稳定初始化 + 训练模板
def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.LayerNorm):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)

model.apply(init_weights)

# 稳定性 checklist
print("数值稳定性检查:")
print("1. 使用 log_softmax + NLLLoss 代替 softmax + MSELoss")
print("2. 设置合适的学习率 + warmup")
print("3. 启用 AMP 混合精度")
print("4. 添加梯度裁剪")
print("5. 使用 LayerNorm 而不是 BatchNorm（Transformer）")
print("6. 监控 loss 和梯度范数")