浮点数的深渊：深度学习数值稳定性的完整解析

import numpy as np

# FP32正常工作
print(np.float32([1000, 1000, 1000]))  # [1000. 1000. 1000.]

# FP16溢出
print(np.float16([1000, 1000, 1000]))  # [inf inf inf]

# FP16下溢
print(np.float16([1e-10, 1e-10, 1e-10]))  # [0. 0. 0.]

def unstable_softmax(x):
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x)

# 正常情况
unstable_softmax([1, 2, 3])  # [0.090, 0.245, 0.665]

# 数值爆炸
unstable_softmax([1000, 2000, 3000])  # [nan, nan, nan]
# 因为 exp(1000) = inf, exp(2000) = inf, exp(3000) = inf
# inf / inf = nan

# FP16下的不稳定Softmax
x = np.array([6.0, -3, 15], dtype=np.float16)
unstable_softmax(x)  # [0., 0., nan]

def stable_softmax(x):
    x_max = np.max(x)
    exp_x = np.exp(x - x_max)
    return exp_x / np.sum(exp_x)

# 现在可以处理任意大的输入
stable_softmax([1000, 2000, 3000])  # [0., 0., 1.]
stable_softmax(np.array([6.0, -3, 15], dtype=np.float16))  # [1.23e-04, 0., 1.]

def logsumexp(x):
    c = np.max(x)
    return c + np.log(np.sum(np.exp(x - c)))

# 可以处理任意范围的输入
logsumexp([1000, 1000, 1000])  # 1001.099 (不是inf)
logsumexp([-1000, -1000, -1000])  # -998.901 (不是-inf)

def log_softmax(x):
    x_max = np.max(x)
    return x - x_max - np.log(np.sum(np.exp(x - x_max)))

def cross_entropy_loss(logits, target):
    return -log_softmax(logits)[target]

# 即使target对应的softmax概率接近0，也能得到有限值
cross_entropy_loss([-1000, 1000], 0)  # 2000.0 (不是inf)

# 维护FP32主权重
master_weights = model.parameters().float().clone()  # FP32副本

for input, target in data:
    # 前向传播使用FP16
    fp16_weights = master_weights.half()
    output = model(input, fp16_weights)
    loss = criterion(output, target)
    
    # 反向传播
    loss.backward()
    
    # 更新FP32主权重
    for p, master_p in zip(model.parameters(), master_weights):
        master_p += lr * p.grad.float()  # 梯度转FP32后更新
    # 同步回FP16
    p.data = master_p.half()

scale_factor = 8192  # 放大因子

# 前向传播
output = model(input)
loss = criterion(output, target)

# 缩放损失
scaled_loss = loss * scale_factor

# 反向传播（梯度自动放大scale_factor倍）
scaled_loss.backward()

# 反缩放梯度
for p in model.parameters():
    p.grad = p.grad / scale_factor

# 更新权重
optimizer.step()

class DynamicGradScaler:
    def __init__(self, init_scale=2**15, growth_factor=2.0, 
                 backoff_factor=0.5, growth_interval=2000):
        self.scale = init_scale
        self.growth_factor = growth_factor
        self.backoff_factor = backoff_factor
        self.growth_interval = growth_interval
        self._growth_tracker = 0
    
    def scale_loss(self, loss):
        return loss * self.scale
    
    def unscale_grads(self, optimizer):
        for p in optimizer.param_groups[0]['params']:
            if p.grad is not None:
                p.grad = p.grad / self.scale
    
    def update(self, found_inf):
        if found_inf:
            # 检测到溢出，减小缩放因子
            self.scale *= self.backoff_factor
            self._growth_tracker = 0
        else:
            # 正常更新，尝试增大缩放因子
            self._growth_tracker += 1
            if self._growth_tracker >= self.growth_interval:
                self.scale *= self.growth_factor
                self._growth_tracker = 0

# PyTorch中使用BF16
model = model.to(torch.bfloat16)

# BF16训练通常不需要GradScaler
# 但仍需要FP32主权重
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for batch in dataloader:
    optimizer.zero_grad()
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        output = model(batch)
        loss = criterion(output)
    loss.backward()
    optimizer.step()

def online_softmax(x_blocks):
    """在线Softmax算法"""
    global_max = -float('inf')
    global_sum = 0.0
    
    for block in x_blocks:
        block_max = np.max(block)
        
        if global_max == -float('inf'):
            global_max = block_max
            global_sum = np.sum(np.exp(block - block_max))
        else:
            # 关键：需要重新缩放之前的累积和
            if block_max > global_max:
                scale = np.exp(global_max - block_max)
                global_sum = global_sum * scale + np.sum(np.exp(block - block_max))
                global_max = block_max
            else:
                scale = np.exp(block_max - global_max)
                global_sum = global_sum + np.sum(np.exp(block - global_max))
    
    return global_max + np.log(global_sum)

def flash_attention(Q, K, V, block_size=256):
    """简化的Flash Attention实现"""
    seq_len = Q.shape[0]
    output = np.zeros_like(Q)
    
    for i in range(0, seq_len, block_size):
        q_block = Q[i:i+block_size]
        block_max = np.full(block_size, -float('inf'))
        block_sum = np.zeros(block_size)
        block_output = np.zeros_like(q_block)
        
        for j in range(0, seq_len, block_size):
            k_block = K[j:j+block_size]
            v_block = V[j:j+block_size]
            
            # 计算当前块的注意力分数
            scores = np.dot(q_block, k_block.T)
            
            # 在线更新最大值和归一化因子
            new_max = np.maximum(block_max, np.max(scores, axis=1))
            scale = np.exp(block_max - new_max)
            scale_scores = np.exp(scores - new_max[:, None])
            
            block_sum = block_sum * scale + np.sum(scale_scores, axis=1)
            block_output = block_output * scale[:, None] + np.dot(scale_scores, v_block)
            block_max = new_max
        
        output[i:i+block_size] = block_output / block_sum[:, None]
    
    return output

def clip_grad_norm_(parameters, max_norm, norm_type=2):
    """梯度范数裁剪"""
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    
    parameters = [p for p in parameters if p.grad is not None]
    max_norm = float(max_norm)
    norm_type = float(norm_type)
    
    if norm_type == float('inf'):
        total_norm = max(p.grad.data.abs().max() for p in parameters)
    else:
        total_norm = 0
        for p in parameters:
            param_norm = p.grad.data.norm(norm_type)
            total_norm += param_norm.item() ** norm_type
        total_norm = total_norm ** (1. / norm_type)
    
    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1:
        for p in parameters:
            p.grad.data.mul_(clip_coef)
    
    return total_norm

class LossSpikeHandler:
    def __init__(self, threshold=5.0, patience=3):
        self.threshold = threshold  # 相对于滑动平均的倍数
        self.patience = patience
        self.loss_history = []
        self.spike_count = 0
    
    def check_spike(self, loss):
        self.loss_history.append(loss)
        
        if len(self.loss_history) < 10:
            return False
        
        # 计算滑动平均和标准差
        recent = self.loss_history[-10:]
        mean = np.mean(recent)
        std = np.std(recent)
        
        # 检测异常
        is_spike = loss > mean + self.threshold * std
        
        if is_spike:
            self.spike_count += 1
            if self.spike_count >= self.patience:
                return True  # 需要干预
        else:
            self.spike_count = 0
        
        return False
    
    def intervene(self, optimizer, original_lr):
        """干预策略"""
        # 降低学习率
        for param_group in optimizer.param_groups:
            param_group['lr'] = original_lr * 0.1

class ZClip:
    """自适应梯度裁剪，专门针对Loss Spike"""
    def __init__(self, z_threshold=3.0, min_grad_norm=1e-6):
        self.z_threshold = z_threshold
        self.grad_norm_history = []
        self.min_grad_norm = min_grad_norm
    
    def clip(self, parameters):
        # 计算当前梯度范数
        total_norm = 0
        for p in parameters:
            if p.grad is not None:
                total_norm += p.grad.data.norm() ** 2
        total_norm = total_norm ** 0.5
        
        self.grad_norm_history.append(total_norm)
        
        if len(self.grad_norm_history) < 10:
            return total_norm
        
        # 计算Z分数
        recent = self.grad_norm_history[-10:]
        mean = np.mean(recent)
        std = np.std(recent)
        
        z_score = (total_norm - mean) / (std + self.min_grad_norm)
        
        # 如果Z分数超过阈值，裁剪
        if z_score > self.z_threshold:
            clip_coef = mean / (total_norm + self.min_grad_norm)
            for p in parameters:
                if p.grad is not None:
                    p.grad.data.mul_(clip_coef)
            return mean
        
        return total_norm

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler(init_scale=2**15, growth_factor=2.0, 
                    backoff_factor=0.5, growth_interval=2000)

for epoch in range(num_epochs):
    for batch, target in dataloader:
        optimizer.zero_grad()
        
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(batch)
            loss = criterion(output, target)
        
        # 缩放损失并反向传播
        scaler.scale(loss).backward()
        
        # 反缩放梯度并裁剪
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        # 更新权重
        scaler.step(optimizer)
        scaler.update()

import torch.nn.functional as F

# 不要分开使用
logits = model(x)
probs = F.softmax(logits, dim=-1)
log_probs = torch.log(probs)  # 可能出现log(0)

# 使用组合版本
log_probs = F.log_softmax(logits, dim=-1)  # 数值稳定

# 交叉熵损失也内置了稳定的组合
loss = F.cross_entropy(logits, targets)  # 等价于 log_softmax + nll_loss

# LayerNorm实现
def layer_norm(x, normalized_shape, weight, bias, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    x_norm = (x - mean) / torch.sqrt(var + eps)  # eps防止除零
    return weight * x_norm + bias

# FP16训练可能需要更大的epsilon
layer_norm = nn.LayerNorm(hidden_size, eps=1e-6)  # 或更大

def check_tensor_health(tensor, name="tensor"):
    """检查张量的数值健康状况"""
    has_nan = torch.isnan(tensor).any()
    has_inf = torch.isinf(tensor).any()
    
    if has_nan or has_inf:
        print(f"[ALERT] {name}: NaN={has_nan.item()}, Inf={has_inf.item()}")
        print(f"  Shape: {tensor.shape}")
        print(f"  Min: {tensor[~torch.isnan(tensor) & ~torch.isinf(tensor)].min() if (~torch.isnan(tensor) & ~torch.isinf(tensor)).any() else 'N/A'}")
        print(f"  Max: {tensor[~torch.isnan(tensor) & ~torch.isinf(tensor)].max() if (~torch.isnan(tensor) & ~torch.isinf(tensor)).any() else 'N/A'}")
        return False
    return True

# 使用钩子监控每层输出
def register_health_checks(model):
    for name, module in model.named_modules():
        module.register_forward_hook(
            lambda m, i, o, n=name: check_tensor_health(o, n) if o is not None else None
        )

class FP8DynamicScaler:
    """FP8动态缩放"""
    def __init__(self, margin=3):
        self.margin = margin  # 偏移量，避免溢出
        self.scale = 1.0
    
    def compute_scale(self, tensor, fp8_max=448.0):
        """计算最优缩放因子"""
        tensor_max = tensor.abs().max()
        if tensor_max == 0:
            return self.scale
        
        # 计算缩放因子使最大值接近FP8最大值
        target_max = fp8_max / (2 ** self.margin)
        new_scale = target_max / tensor_max
        
        # 平滑更新
        self.scale = 0.9 * self.scale + 0.1 * new_scale
        return self.scale
    
    def quantize(self, tensor):
        scale = self.compute_scale(tensor)
        return (tensor * scale).to(torch.int8).to(torch.float32) / scale