第 10 章：训练崩溃与 NaN 问题

# 紧急保存当前状态
torch.save({
    'step': current_step,
    'model_state': model.state_dict(),
    'optimizer_state': optimizer.state_dict(),
    'loss_history': loss_history[-100:],  # 最近100个step的loss
    'grad_norm_history': grad_norm_history[-100:],
}, f'debug_checkpoint_step_{current_step}.pt')

模式 1: 突然跳跃
Loss: 2.1 → 2.0 → 1.9 → 8734.5 → NaN
原因: 单个异常样本或数值溢出

模式 2: 指数增长
Loss: 2.1 → 2.3 → 2.8 → 4.5 → 12.3 → 89.7 → NaN
原因: 学习率过大或梯度累积错误

模式 3: 震荡发散
Loss: 2.1 → 1.8 → 2.5 → 1.6 → 3.2 → 1.4 → 5.8 → NaN
原因: 优化器状态损坏或数值不稳定

def check_model_health(model):
    """快速检查模型各层的健康状态"""
    issues = []

    for name, param in model.named_parameters():
        # 检查参数本身
        if torch.isnan(param).any():
            issues.append(f"NaN in parameter: {name}")
        if torch.isinf(param).any():
            issues.append(f"Inf in parameter: {name}")

        # 检查梯度
        if param.grad is not None:
            if torch.isnan(param.grad).any():
                issues.append(f"NaN in gradient: {name}")
            if torch.isinf(param.grad).any():
                issues.append(f"Inf in gradient: {name}")

            # 检查梯度范数
            grad_norm = param.grad.norm().item()
            if grad_norm > 1000:
                issues.append(f"Large gradient norm ({grad_norm:.2f}): {name}")

    return issues

# 检查注意力权重
attention_weights = torch.softmax(scores / math.sqrt(d_k), dim=-1)
if (attention_weights == 0).all(dim=-1).any():
    print("警告：出现全零注意力权重（数值下溢）")
if torch.isnan(attention_weights).any():
    print("警告：注意力权重包含 NaN")

# 添加数值稳定性
logits = model(inputs)
# 错误：可能导致 log(0)
loss = -torch.log(probs[target])
# 正确：添加 epsilon
loss = -torch.log(probs[target] + 1e-8)

# 检查 LayerNorm 是否稳定
def stable_layer_norm(x, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var = x.var(dim=-1, keepdim=True, unbiased=False)
    # 确保方差不为零
    return (x - mean) / torch.sqrt(var + eps)

发现 Loss 爆炸/NaN
│
├─ 是否在前 1000 步内？
│  ├─ 是 → 检查初始化和学习率预热
│  └─ 否 → 继续诊断
│
├─ 是否使用混合精度？
│  ├─ 是 → 检查 loss scaling 和 dtype 转换
│  └─ 否 → 检查数值溢出
│
├─ 是否有异常大的梯度？
│  ├─ 是 → 降低学习率或增强 gradient clipping
│  └─ 否 → 检查数据和损失函数
│
└─ 是否可以从 checkpoint 恢复？
   ├─ 是 → 调整超参数后恢复训练
   └─ 否 → 降级到更保守的配置重新开始

class GradientMonitor:
    """梯度监控器，实时跟踪梯度统计信息"""

    def __init__(self, model, logger=None):
        self.model = model
        self.logger = logger
        self.history = defaultdict(list)
        self.anomaly_threshold = {
            'max_norm': 100.0,
            'min_norm': 1e-8,
            'nan_tolerance': 0,
        }

    def check_gradients(self, step):
        """检查当前步的梯度健康状态"""
        alerts = []

        for name, param in self.model.named_parameters():
            if param.grad is None:
                continue

            grad = param.grad.data

            # 计算统计信息
            grad_norm = grad.norm().item()
            grad_mean = grad.mean().item()
            grad_std = grad.std().item()

            # 记录历史
            self.history[name].append({
                'step': step,
                'norm': grad_norm,
                'mean': grad_mean,
                'std': grad_std
            })

            # 异常检测
            if torch.isnan(grad).any():
                alerts.append(f"Step {step}: NaN gradient in {name}")

            if grad_norm > self.anomaly_threshold['max_norm']:
                alerts.append(f"Step {step}: Large gradient norm {grad_norm:.2f} in {name}")

            if grad_norm < self.anomaly_threshold['min_norm'] and grad_norm > 0:
                alerts.append(f"Step {step}: Vanishing gradient {grad_norm:.2e} in {name}")

        return alerts

def analyze_gradient_flow(model, input_batch, target_batch):
    """分析梯度在模型中的流动情况"""

    model.zero_grad()
    output = model(input_batch)
    loss = compute_loss(output, target_batch)
    loss.backward()

    # 收集每层的梯度信息
    gradient_flow = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_data = param.grad.data
            gradient_flow.append({
                'layer': name,
                'grad_norm': grad_data.norm().item(),
                'grad_mean': grad_data.mean().item(),
                'grad_max': grad_data.max().item(),
                'grad_min': grad_data.min().item(),
                'percent_zeros': (grad_data == 0).float().mean().item() * 100
            })

    # 可视化梯度流
    import matplotlib.pyplot as plt

    layers = [g['layer'].split('.')[-1] for g in gradient_flow]
    grad_norms = [g['grad_norm'] for g in gradient_flow]

    plt.figure(figsize=(12, 6))
    plt.semilogy(grad_norms, label='Gradient Norm')
    plt.xticks(range(len(layers)), layers, rotation=45, ha='right')
    plt.xlabel('Layers')
    plt.ylabel('Gradient Norm (log scale)')
    plt.title('Gradient Flow Through Network')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()

    return gradient_flow

FP16 (半精度浮点)
├─ 优点：硬件支持广泛，速度快
├─ 缺点：数值范围小 (±65,504)，容易溢出
└─ 适用：稳定的模型，充分的 loss scaling

BF16 (Brain Float 16)
├─ 优点：数值范围大 (±3.4×10^38)，与FP32相同
├─ 缺点：精度较低，需要新硬件（A100+）
└─ 适用：大模型训练，数值稳定性要求高

class DynamicLossScaler:
    """自适应的 loss scaling，防止梯度下溢/上溢"""

    def __init__(self, init_scale=2**16, scale_factor=2.0, 
                 scale_window=2000, tolerance=0.05):
        self.scale = init_scale
        self.scale_factor = scale_factor
        self.scale_window = scale_window
        self.tolerance = tolerance
        self.overflow_counter = 0
        self.step_counter = 0

    def scale_loss(self, loss):
        """放大loss防止梯度下溢"""
        return loss * self.scale

    def unscale_gradients(self, optimizer):
        """缩小梯度到正确范围"""
        for group in optimizer.param_groups:
            for param in group['params']:
                if param.grad is not None:
                    param.grad.data.div_(self.scale)

    def update_scale(self, overflow):
        """根据溢出情况动态调整scale"""
        if overflow:
            # 发生溢出，减小scale
            self.scale /= self.scale_factor
            self.overflow_counter += 1
            print(f"Gradient overflow! Reducing scale to {self.scale}")
            return True

        self.step_counter += 1
        if self.step_counter % self.scale_window == 0:
            # 长时间无溢出，尝试增大scale
            self.scale *= self.scale_factor
            print(f"Increasing scale to {self.scale}")

        return False

def configure_mixed_precision(model):
    """配置混合精度训练的层级精度"""

    # 始终保持 FP32 的层
    fp32_layers = [
        'layer_norm',      # LayerNorm 对精度敏感
        'softmax',         # Softmax 需要高精度
        'loss',            # 损失计算
        'positional',      # 位置编码
    ]

    for name, module in model.named_modules():
        # 检查是否需要FP32
        need_fp32 = any(fp_layer in name.lower() 
                        for fp_layer in fp32_layers)

        if need_fp32:
            # 强制使用FP32
            module.float()
            for param in module.parameters():
                param.data = param.data.float()
        else:
            # 可以使用FP16/BF16
            module.half()  # or module.bfloat16()

    return model

def stable_gradient_accumulation(model, optimizer, data_loader, 
                                accumulation_steps=4):
    """稳定的梯度累积实现"""

    scaler = torch.cuda.amp.GradScaler()
    accumulated_loss = 0

    for step, batch in enumerate(data_loader):
        # 判断是否是累积的最后一步
        is_accumulation_boundary = (step + 1) % accumulation_steps == 0

        with torch.cuda.amp.autocast():
            outputs = model(batch['input'])
            loss = compute_loss(outputs, batch['target'])
            # 重要：除以累积步数
            loss = loss / accumulation_steps

        # Scale loss并反向传播
        scaler.scale(loss).backward()
        accumulated_loss += loss.item()

        if is_accumulation_boundary:
            # 梯度裁剪（在unscale之后，step之前）
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            # 优化器步进
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

            # 记录
            print(f"Step {step}: Loss = {accumulated_loss:.4f}")
            accumulated_loss = 0

class CheckpointManager:
    """全面的检查点管理器"""

    def __init__(self, save_dir, keep_last_n=3, save_interval=1000):
        self.save_dir = save_dir
        self.keep_last_n = keep_last_n
        self.save_interval = save_interval
        self.checkpoints = []

    def save_checkpoint(self, model, optimizer, scheduler, 
                       epoch, step, metrics, extra_state=None):
        """保存完整的训练状态"""

        checkpoint = {
            'epoch': epoch,
            'step': step,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
            'metrics': metrics,
            'rng_state': {
                'python': random.getstate(),
                'numpy': np.random.get_state(),
                'torch': torch.get_rng_state(),
                'cuda': torch.cuda.get_rng_state_all(),
            },
            'timestamp': datetime.now().isoformat(),
        }

        if extra_state:
            checkpoint['extra_state'] = extra_state

        # 保存checkpoint
        checkpoint_path = os.path.join(
            self.save_dir, 
            f'checkpoint_step_{step}.pt'
        )
        torch.save(checkpoint, checkpoint_path)
        self.checkpoints.append(checkpoint_path)

        # 清理旧的checkpoints
        if len(self.checkpoints) > self.keep_last_n:
            old_checkpoint = self.checkpoints.pop(0)
            if os.path.exists(old_checkpoint):
                os.remove(old_checkpoint)

        return checkpoint_path

    def load_checkpoint(self, checkpoint_path, model, optimizer=None, 
                       scheduler=None, strict=True):
        """恢复训练状态"""

        checkpoint = torch.load(checkpoint_path, map_location='cpu')

        # 恢复模型
        model.load_state_dict(checkpoint['model_state_dict'], strict=strict)

        # 恢复优化器
        if optimizer and 'optimizer_state_dict' in checkpoint:
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

        # 恢复学习率调度器
        if scheduler and 'scheduler_state_dict' in checkpoint:
            scheduler.load_state_dict(checkpoint['scheduler_state_dict'])

        # 恢复随机数状态
        if 'rng_state' in checkpoint:
            random.setstate(checkpoint['rng_state']['python'])
            np.random.set_state(checkpoint['rng_state']['numpy'])
            torch.set_rng_state(checkpoint['rng_state']['torch'])
            torch.cuda.set_rng_state_all(checkpoint['rng_state']['cuda'])

        return checkpoint

def resume_training(checkpoint_path, model, optimizer, data_loader):
    """安全的断点续训流程"""

    # 1. 加载checkpoint
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

    # 2. 恢复到正确的数据位置
    start_epoch = checkpoint['epoch']
    start_step = checkpoint['step']

    # 3. 验证恢复是否成功
    validation_batch = next(iter(data_loader))
    with torch.no_grad():
        output = model(validation_batch['input'])
        loss = compute_loss(output, validation_batch['target'])

    print(f"Validation loss after resume: {loss.item():.4f}")

    # 4. 检查是否需要降级配置
    if checkpoint.get('crashed', False):
        print("Previous training crashed. Applying conservative settings...")
        # 降低学习率
        for param_group in optimizer.param_groups:
            param_group['lr'] *= 0.5
        # 增强梯度裁剪
        max_grad_norm = 0.5
    else:
        max_grad_norm = 1.0

    return start_epoch, start_step, max_grad_norm

class CrashRecoveryTrainer:
    """具有崩溃恢复能力的训练器"""

    def __init__(self, model, config):
        self.model = model
        self.config = config
        self.crash_counter = 0
        self.max_crashes = 3

    def train_with_recovery(self, data_loader):
        """带自动恢复的训练循环"""

        while self.crash_counter < self.max_crashes:
            try:
                # 正常训练
                self._train_epoch(data_loader)
                self.crash_counter = 0  # 重置计数器

            except (RuntimeError, ValueError) as e:
                self.crash_counter += 1
                print(f"Training crashed (attempt {self.crash_counter}/{self.max_crashes}): {e}")

                # 崩溃恢复策略
                recovery_actions = self._get_recovery_strategy(e)
                for action in recovery_actions:
                    action()

                # 从最近的checkpoint恢复
                if self.last_checkpoint:
                    self.load_checkpoint(self.last_checkpoint)
                else:
                    print("No checkpoint available, restarting training...")
                    self._reinitialize_training()

    def _get_recovery_strategy(self, error):
        """根据错误类型确定恢复策略"""

        strategies = []

        if "CUDA out of memory" in str(error):
            strategies.append(self._reduce_batch_size)
            strategies.append(self._enable_gradient_checkpointing)

        elif "nan" in str(error).lower():
            strategies.append(self._reduce_learning_rate)
            strategies.append(self._reset_optimizer_state)
            strategies.append(self._switch_to_fp32)

        elif "gradient" in str(error).lower():
            strategies.append(self._enhance_gradient_clipping)
            strategies.append(self._reduce_accumulation_steps)

        return strategies

序列A: 1.8, 1.7, 1.6, 1.5, 1.4, 87234.5, NaN
序列B: 2.1, 2.2, 2.5, 3.1, 4.8, 9.2, 23.5, 156.7, NaN
序列C: 2.0, 1.8, 2.2, 1.6, 2.5, 1.4, 3.2, 1.2, 5.8, NaN