大模型训练稳定性：梯度爆炸、Loss Spike 与恢复策略

训练不稳定性事件案例:
┌──────────────────┬─────────────────────────────────────────┐
│ 模型             │ 不稳定事件                              │
├──────────────────┼─────────────────────────────────────────┤
│ GPT-3 175B       │ 多次 Loss Spike, 部分无法恢复           │
│ OPT 175B         │ 训练崩溃 10+ 次, 需要回退检查点         │
│ BLOOM 176B       │ 梯度爆炸导致 NaN, 引入更多稳定化手段    │
│ LLaMA-2 70B      │ 7次显著 Loss Spike, 4次需手动干预       │
│ GLM-130B         │ Embedding 层梯度爆炸, 引入 DeepNorm     │
└──────────────────┴─────────────────────────────────────────┘

# 梯度爆炸/消失的理论分析
# 在深度网络中, 梯度通过链式法则传播:
# ∂L/∂x_0 = ∂L/∂x_n · ∏(i=1 to n) ∂x_i/∂x_{i-1}

# 如果每层的雅可比矩阵的谱半径 > 1: 梯度爆炸
# 如果每层的雅可比矩阵的谱半径 < 1: 梯度消失

def analyze_gradient_flow(model, input_data):
    """分析各层梯度流"""
    layer_grad_norms = {}
    
    for name, param in model.named_parameters():
        if param.grad is not None:
            layer_grad_norms[name] = param.grad.norm().item()
    
    # 按层分组统计
    return layer_grad_norms

# 典型的梯度分布 (100层Transformer):
# Layer 0  (输入): grad_norm = 0.001  ← 梯度消失
# Layer 50 (中间): grad_norm = 0.5
# Layer 99 (输出): grad_norm = 12.3   ← 梯度较大

class ResidualBlock(nn.Module):
    """残差连接: 梯度可以通过捷径直接传播"""
    def __init__(self, d_model, dropout=0.1):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.attention = MultiHeadAttention(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model),
        )
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        # 注意力子层 + 残差
        x = x + self.dropout(self.attention(self.norm(x)))
        # FFN子层 + 残差
        x = x + self.dropout(self.ffn(self.norm2(x)))
        return x

# 梯度传播: dL/dx = dL/dx + dL/d(sublayer(x)) * d(sublayer)/dx
# 即使 sublayer 梯度小, 主路径梯度仍然可以通过

class GradientClipping:
    """多种梯度裁剪策略"""
    
    @staticmethod
    def clip_by_global_norm(parameters, max_norm=1.0):
        """
        全局梯度裁剪 (最常用)
        如果所有梯度L2范数之和超过max_norm, 则等比例缩放
        """
        total_norm = torch.norm(
            torch.stack([p.grad.norm() for p in parameters if p.grad is not None])
        )
        clip_coef = min(1.0, max_norm / (total_norm + 1e-6))
        
        for p in parameters:
            if p.grad is not None:
                p.grad.mul_(clip_coef)
        
        return total_norm
    
    @staticmethod
    def clip_by_value(parameters, min_val=-1.0, max_val=1.0):
        """逐元素裁剪"""
        for p in parameters:
            if p.grad is not None:
                p.grad.clamp_(min_val, max_val)
    
    @staticmethod
    def adaptive_clip(parameters, target_ratio=0.5):
        """
        自适应梯度裁剪 (AGC, Brock et al. 2021)
        基于参数范数动态确定裁剪阈值
        """
        for p in parameters:
            if p.grad is None:
                continue
            
            param_norm = p.norm()
            grad_norm = p.grad.norm()
            
            # 裁剪阈值 = target_ratio * max(param_norm, min_norm)
            max_grad = max(target_ratio * param_norm, 1e-3)
            
            if grad_norm > max_grad:
                p.grad.mul_(max_grad / (grad_norm + 1e-6))

典型 Loss Spike 模式:

Loss
  │
  │     ╱╲          ╱╲
  │    ╱  ╲        ╱  ╲
  │   ╱    ╲      ╱    ╲
  │──╱──────╲────╱──────╲────── 正常轨迹
  │          ╲  ╱
  │           ╲╱
  │
  └────────────────────────────── Steps
  
类型1: V型恢复 (自然恢复, ~100步)
类型2: 缓慢恢复 (需要数百步)
类型3: 永久发散 (无法恢复, 需回退检查点)

class LossSpikeAnalyzer:
    """Loss Spike 分析器"""
    
    def __init__(self):
        self.spike_events = []
    
    def analyze_spike(self, log_data, spike_step):
        """
        分析 Loss Spike 的可能原因
        
        log_data: 训练日志, 包含各指标随时间变化
        spike_step: Spike 发生的步数
        """
        # 获取 spike 前后的指标
        pre_window = log_data[max(0, spike_step-50):spike_step]
        post_window = log_data[spike_step:spike_step+100]
        
        analysis = {
            'spike_magnitude': post_window[0]['loss'] / pre_window[-1]['loss'],
            'recovery_steps': self._find_recovery(pre_window[-1]['loss'], post_window),
            'possible_causes': [],
        }
        
        # 检查各指标异常
        checks = [
            ('grad_norm', self._check_grad_spike),
            ('learning_rate', self._check_lr_schedule),
            ('data_quality', self._check_data_anomaly),
            ('attention_entropy', self._check_attention_collapse),
            ('logits_scale', self._check_logits_explosion),
        ]
        
        for metric, check_fn in checks:
            cause = check_fn(pre_window, post_window)
            if cause:
                analysis['possible_causes'].append(cause)
        
        return analysis
    
    def _check_grad_spike(self, pre, post):
        """检查梯度范数异常"""
        pre_grad = np.mean([x['grad_norm'] for x in pre])
        post_grad = post[0]['grad_norm']
        if post_grad > 5 * pre_grad:
            return f"梯度爆炸: {pre_grad:.2f} → {post_grad:.2f}"
        return None
    
    def _check_logits_explosion(self, pre, post):
        """检查 logits 尺度爆炸"""
        pre_scale = np.mean([x['logits_norm'] for x in pre])
        post_scale = post[0]['logits_norm']
        if post_scale > 3 * pre_scale:
            return f"Logits 爆炸: {pre_scale:.2f} → {post_scale:.2f}"
        return None

# FP16 和 BF16 的精度对比
precision_comparison = {
    "FP16": {
        "exponent_bits": 5,
        "mantissa_bits": 10,
        "max_value": 65504,
        "min_normal": 6.1e-5,  # 容易下溢!
        "issues": ["梯度下溢", "舍入误差累积"],
    },
    "BF16": {
        "exponent_bits": 8,    # 与 FP32 相同
        "mantissa_bits": 7,
        "max_value": 3.4e38,   # 几乎不会溢出
        "min_normal": 1.2e-38, # 不会下溢
        "issues": ["精度低"],
    },
    "FP32": {
        "exponent_bits": 8,
        "mantissa_bits": 23,
        "max_value": 3.4e38,
        "min_normal": 1.2e-38,
        "issues": ["显存占用大"],
    }
}

# FP16 训练的典型问题:
# 1. 梯度下溢: 小于 6.1e-5 的梯度变成 0
# 2. 激活值溢出: 大于 65504 的值变成 inf
# 3. 舍入误差: 10位尾数精度不足

# 解决方案: 混合精度训练
class MixedPrecisionTrainer:
    def __init__(self, model, optimizer, use_bf16=True):
        self.model = model
        self.optimizer = optimizer
        
        if use_bf16 and torch.cuda.is_bf16_supported():
            # BF16: 不需要 GradScaler, 更稳定
            self.dtype = torch.bfloat16
            self.scaler = None
        else:
            # FP16: 需要 GradScaler 防止梯度下溢
            self.dtype = torch.float16
            self.scaler = torch.cuda.amp.GradScaler()
    
    def training_step(self, batch):
        self.model.train()
        
        with torch.autocast(device_type='cuda', dtype=self.dtype):
            outputs = self.model(**batch)
            loss = outputs.loss
        
        if self.scaler:
            # FP16: 缩放损失 → 反向传播 → 裁剪 → 取消缩放 → 更新
            self.scaler.scale(loss).backward()
            self.scaler.unscale_(self.optimizer)
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            # BF16: 直接反向传播
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.optimizer.step()
        
        self.optimizer.zero_grad()
        return loss.item()

class GradScalerManual:
    """手动实现的梯度缩放器"""
    def __init__(self, init_scale=2**16, 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(self, loss):
        """缩放损失, 防止梯度下溢"""
        return loss * self.scale
    
    def unscale_(self, optimizer):
        """取消梯度缩放"""
        for group in optimizer.param_groups:
            for p in group['params']:
                if p.grad is not None:
                    p.grad.div_(self.scale)
    
    def step(self, optimizer):
        """执行优化器步骤, 检查溢出"""
        # 检查梯度是否有 inf/nan
        found_inf = False
        for group in optimizer.param_groups:
            for p in group['params']:
                if p.grad is not None:
                    if torch.isinf(p.grad).any() or torch.isnan(p.grad).any():
                        found_inf = True
                        break
        
        if found_inf:
            # 跳过此步, 减小缩放因子
            self.scale *= self.backoff_factor
            return
        else:
            # 正常更新, 可能增大缩放因子
            optimizer.step()
            self._growth_tracker += 1
            if self._growth_tracker >= self.growth_interval:
                self.scale *= self.growth_factor
                self._growth_tracker = 0

# 不当初始化导致的典型问题
init_problems = {
    "标准初始化 (Xavier/Glorot)": {
        "问题": "在深层 Transformer 中, 注意力输出的方差逐层增长",
        "症状": "前几层训练正常, 后续层 Loss 波动增大",
    },
    "过大方差初始化": {
        "问题": "初始 logits 过大, softmax 趋近 one-hot",
        "症状": "初始 Loss 极高, 训练初期梯度爆炸",
    },
    "过小方差初始化": {
        "问题": "所有 token 表示趋近相同",
        "症状": "注意力塌缩, 训练停滞",
    },
}

def init_weights_normal(module, std=0.02):
    """GPT 风格初始化: 截断正态分布"""
    if isinstance(module, nn.Linear):
        torch.nn.init.normal_(module.weight, mean=0.0, std=std)
        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=std)

def init_weights_scaled(module, n_layers, std=0.02):
    """带深度缩放的初始化: 深层网络更小方差"""
    if isinstance(module, nn.Linear):
        # 关键: 残差路径的权重缩放
        if hasattr(module, '_is_residual'):
            torch.nn.init.normal_(module.weight, mean=0.0, 
                                  std=std / math.sqrt(2 * n_layers))
        else:
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)

class DeepNormInit:
    """
    DeepNorm (Microsoft, 2022): 支持上千层 Transformer
    关键: 对残差路径进行特殊缩放
    """
    def __init__(self, n_layers, hidden_size, alpha=None, beta=None):
        self.n_layers = n_layers
        # alpha 和 beta 的公式推导来自理论分析
        if alpha is None:
            # alpha = (2N)^(1/4) for pre-norm
            self.alpha = (2 * n_layers) ** 0.25
        if beta is None:
            self.beta = (8 * n_layers) ** -0.25
    
    def apply(self, model):
        for name, param in model.named_parameters():
            if 'attention' in name and 'weight' in name:
                # 注意力投影: 用 beta 缩放
                torch.nn.init.xavier_normal_(param, gain=self.beta)
            elif 'ffn' in name and 'weight' in name:
                # FFN 投影: 用 beta 缩放
                torch.nn.init.xavier_normal_(param, gain=self.beta)
        
        # 修改前向传播: x = LayerNorm(x * alpha + sublayer(x))
        # alpha > 1 放大残差路径, 稳定深层训练

class QKNormAttention(nn.Module):
    """QK-Norm: 在计算注意力分数前归一化 Q 和 K"""
    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)
        
        # 关键: 为 Q 和 K 添加 LayerNorm
        self.q_norm = nn.LayerNorm(self.d_k)
        self.k_norm = nn.LayerNorm(self.d_k)
    
    def forward(self, x):
        B, N, D = x.shape
        q = self.W_q(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, N, self.num_heads, self.d_k).transpose(1, 2)
        
        # QK-Norm: 归一化后再计算注意力
        q = self.q_norm(q)
        k = self.k_norm(k)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn = F.softmax(scores, dim=-1)
        output = torch.matmul(attn, v)
        
        return output.transpose(1, 2).reshape(B, N, D)

# AdamW 的关键超参数对稳定性的影响
optimizer_config = {
    "beta1": 0.9,       # 一阶动量衰减
    "beta2": 0.95,      # 二阶动量衰减 (LLM常用0.95, 而非默认0.999)
    "eps": 1e-8,        # 数值稳定项
    "weight_decay": 0.1, # 权重衰减
    "lr": 3e-4,         # 峰值学习率
}

# beta2 的选择对训练稳定性影响显著:
# beta2=0.999 (默认): 动量累积过慢, 大batch下不稳定
# beta2=0.95:         动量更新更快, 适应大batch训练
# beta2=0.9:          动量非常短视, 可能导致震荡

class StableAdamW:
    """带稳定性增强的 AdamW"""
    def __init__(self, params, lr=3e-4, betas=(0.9, 0.95), 
                 eps=1e-8, weight_decay=0.1, max_grad_norm=1.0):
        self.params = list(params)
        self.lr = lr
        self.beta1, self.beta2 = betas
        self.eps = eps
        self.weight_decay = weight_decay
        self.max_grad_norm = max_grad_norm
        
        # 状态
        self.m = [torch.zeros_like(p) for p in self.params]  # 一阶动量
        self.v = [torch.zeros_like(p) for p in self.params]  # 二阶动量
        self.t = 0  # 时间步
    
    def step(self):
        self.t += 1
        
        # 梯度裁剪
        global_grad_norm = torch.norm(
            torch.stack([p.grad.norm() for p in self.params if p.grad is not None])
        )
        clip_coef = min(1.0, self.max_grad_norm / (global_grad_norm + 1e-6))
        
        for i, p in enumerate(self.params):
            if p.grad is None:
                continue
            
            grad = p.grad * clip_coef
            
            # 更新动量
            self.m[i] = self.beta1 * self.m[i] + (1 - self.beta1) * grad
            self.v[i] = self.beta2 * self.v[i] + (1 - self.beta2) * grad**2
            
            # 偏差校正
            m_hat = self.m[i] / (1 - self.beta1**self.t)
            v_hat = self.v[i] / (1 - self.beta2**self.t)
            
            # 参数更新 + 权重衰减 (decoupled)
            p.data -= self.lr * (m_hat / (v_hat.sqrt() + self.eps) + 
                                  self.weight_decay * p.data)

class LRSchedule:
    """大模型训练的标准学习率调度"""
    
    def __init__(self, optimizer, max_lr=3e-4, min_lr=3e-5, 
                 warmup_steps=2000, total_steps=100000, decay_type='cosine'):
        self.optimizer = optimizer
        self.max_lr = max_lr
        self.min_lr = min_lr
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.decay_type = decay_type
    
    def get_lr(self, step):
        if step < self.warmup_steps:
            # 线性 Warmup
            return self.max_lr * step / self.warmup_steps
        elif self.decay_type == 'cosine':
            # Cosine 衰减
            progress = (step - self.warmup_steps) / (self.total_steps - self.warmup_steps)
            return self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + math.cos(math.pi * progress))
        elif self.decay_type == 'wsd':
            # WSD: Warmup-Stable-Decay
            stable_end = int(self.total_steps * 0.8)
            if step < stable_end:
                return self.max_lr  # 稳定阶段
            else:
                # 衰减阶段
                decay_progress = (step - stable_end) / (self.total_steps - stable_end)
                return self.min_lr + (self.max_lr - self.min_lr) * 0.5 * (1 + math.cos(math.pi * decay_progress))
    
    def step(self):
        lr = self.get_lr(self.current_step)
        for group in self.optimizer.param_groups:
            group['lr'] = lr
        self.current_step += 1

class SpikeRecoveryManager:
    """训练 Spike 自动恢复管理器"""
    
    def __init__(self, model, optimizer, patience=20, 
                 recovery_lr_factor=0.5, save_every=500):
        self.model = model
        self.optimizer = optimizer
        self.patience = patience
        self.recovery_lr_factor = recovery_lr_factor
        self.save_every = save_every
        
        self.loss_history = []
        self.best_state = None
        self.best_loss = float('inf')
        self.spike_count = 0
    
    def check_and_recover(self, current_loss, step):
        """检查是否需要触发恢复"""
        self.loss_history.append(current_loss)
        
        # 更新最佳状态
        if current_loss < self.best_loss:
            self.best_loss = current_loss
            self.best_state = self._save_state()
        
        # 检测 Spike
        if len(self.loss_history) > 10:
            recent_avg = np.mean(self.loss_history[-10:-1])
            if current_loss > 3 * recent_avg:  # Loss 超过3倍
                self.spike_count += 1
                return self._handle_spike(step)
        
        return False
    
    def _handle_spike(self, step):
        """处理 Spike"""
        print(f"⚠ Loss Spike detected at step {step}")
        
        # 策略1: 回退到最佳检查点
        if self.spike_count > 3:
            print("  → 多次 Spike, 回退到最佳检查点")
            self._restore_state(self.best_state)
            return True
        
        # 策略2: 降低学习率
        for group in self.optimizer.param_groups:
            group['lr'] *= self.recovery_lr_factor
        print(f"  → 降低学习率到 {self.optimizer.param_groups[0]['lr']:.2e}")
        
        return False
    
    def _save_state(self):
        """保存模型和优化器状态"""
        return {
            'model': {k: v.clone() for k, v in self.model.state_dict().items()},
            'optimizer': self.optimizer.state_dict(),
            'loss': self.best_loss,
        }
    
    def _restore_state(self, state):
        """恢复状态"""
        self.model.load_state_dict(state['model'])
        self.optimizer.load_state_dict(state['optimizer'])
        self.best_loss = state['loss']

def reset_optimizer_state(optimizer, reset_ratio=0.5):
    """
    在 Spike 后部分重置优化器状态
    Adam 的二阶动量 v 可能因异常梯度而污染
    """
    for group in optimizer.param_groups:
        for p in group['params']:
            if p in optimizer.state:
                state = optimizer.state[p]
                if 'exp_avg' in state:
                    # 部分衰减一阶动量
                    state['exp_avg'].mul_(reset_ratio)
                if 'exp_avg_sq' in state:
                    # 完全重置二阶动量
                    state['exp_avg_sq'].zero_()
                if 'step' in state:
                    # 减小步数, 恢复偏差校正强度
                    state['step'] = max(1, int(state['step'] * reset_ratio))

训练前稳定性检查:
□ 初始化策略: 使用缩放初始化 (std/sqrt(2N))
□ 学习率: Warmup ≥ 2000 步, 峰值 LR 在 1e-4 ~ 6e-4
□ 梯度裁剪: max_norm = 1.0
□ 精度: 优先使用 BF16, FP16 需 GradScaler
□ 优化器: AdamW, beta2=0.95, weight_decay=0.1
□ Embedding: 考虑 DeepNorm 或 Embedding 裁剪
□ QK-Norm: 深层模型 (>50层) 建议启用

训练中监控:
□ Loss 曲线: 是否有突然跳变
□ 梯度范数: 是否稳定在 0.1-10 之间
□ 注意力熵: 是否塌缩到少数 token
□ Logits 范数: 是否持续增长
□ 各层激活值方差: 是否有层消失/爆炸
□ 优化器二阶动量: 是否异常累积