Cosmos Predict2 FP16 支持调试经验总结

背景信息

• 模型: Anima (基于 NVIDIA Cosmos-Predict2-2B-Text2Image)
• 架构: Cosmos Predict2 (DiT Transformer with AdaLN)
• 问题: FP16 精度下生成雪花图或黑图
• 参考方案: Lumina FP16 Overflow Fix (PR #11187)

---

调试过程记录

第一阶段：直接套用 Lumina 方案

配置：
- Attention 输出: 除以 4
- FeedForward 输出: 除以 32

结果：
- 4步: 雪花图
- 30步: 雪花图

结论: Lumina 的除数对 Cosmos 太激进

---

第二阶段：逐步降低除数

发现: 3.48 是关键的转折点，4步能看到内容说明方向正确

---

第三阶段：添加 Clamp 保护 + 调整除数

配置：
- Attention 输出: 除以 1.5
- FeedForward 输出: 除以 3.48
- 添加 clamp_fp16() 函数 (NaN→0, Inf→±65504)
- Block 输出: 添加 clamp 保护

结果：
- 4步: 斑纹图（块状噪点）
- 30步: 雪花图

---

第四阶段：缩放-还原策略

配置：
- Attention 输出: 除以 1.5 → clamp → 还原 × 1.5
- FeedForward 输出: 除以 3.48 → clamp → 还原 × 3.48

结果：
- 4步和30步: 雪花图

结论: 还原策略没有改善，因为问题在于中间累积而非最终输出

---

第五阶段：添加调试日志

添加了 clamp_fp16(x, debug_name) 函数来记录数值统计：

def clamp_fp16(x, debug_name=None):
    if x.dtype == torch.float16:
        has_nan = torch.isnan(x).any().item()
        has_inf = torch.isinf(x).any().item()
        max_val = x[~torch.isnan(x) & ~torch.isinf(x)].abs().max().item()
        if has_nan or has_inf or max_val > 60000:
            print(f"[FP16 Debug] {debug_name}: max_abs={max_val}, ...")
        return torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
    return x

关键发现：

• 数值都在安全范围内（< 60000），没有触发日志
• 说明问题不是简单的溢出，而是精度累积

---

第六阶段：无条件日志 + 激活后保护

配置：
- 在激活函数 (GELU) 后添加 clamp 保护
- 激活后除以 8（大值）/ 2（小值）
- 最终输出除以 3.5-4

结果：
- 4步: 有内容但质量差
- 30步: 雪花图

---

第七阶段：自适应缩放（基于日志数据）

从日志中发现的关键数据：

FFN_after_activation 峰值统计：
- 第一步: max=180.75  ← 异常！
- 后续: max=20~48 不等
- mean 值通常在 0.1~0.8 范围内

Attention 层统计：
- max 值通常在 0.5~40 范围内
- mean 值通常在 0.1~3.0
- 相对稳定

关键发现: GELU 激活后会出现 180+ 的极端峰值，这是问题的根源

---

第八阶段：最终方案（先 Clamp 后降尺度）

def forward(self, x: torch.Tensor) -> torch.Tensor:
    x1 = self.layer1(x)
    x2 = self.activation(x1)
    # Apply FP16 protection after activation
    if x2.dtype == torch.float16:
        # Clamp extreme values first (180+ → 65504)
        x2 = clamp_fp16(x2)
        # Then apply mild downscaling
        x2 = x2 / 2.5
    x3 = self.layer2(x2)
    # Apply final downscaling
    if x3.dtype == torch.float16:
        x3 = x3 / 3.5
    return clamp_fp16(x3)

结果：

• 4步: 雪花图
• 30步: 雪花图

---

根本原因分析

1. 架构差异

2. FP16 限制

FP16 数值范围: -65504 ~ +65504
FP16 精度: 约 3 位小数（相对 FP32）

问题：
- GELU 激活函数在某些输入下会输出 180+ 的值
- 即使 clamp 到 65504，后续计算仍会累积误差
- 多层 transformer 块会放大这种误差
- 30 步累积后，误差导致完全无法重建图像

3. 为什么固定除数无效

固定除数的问题：
- 小数值：被过度压缩 → 精度丢失 → 雪花图
- 大数值：仍可能溢出或被 clamp 截断 → 黑图
- 无法同时满足两种极端情况

---

最终结论

Cosmos Predict2 架构从设计上就不适合 FP16 精度推理。

技术原因：

1. GELU 激活的数值特性：在 FP16 精度下容易产生极端峰值
2. 累积误差：多层 transformer 会放大单层的精度损失
3. 架构复杂度：Self-Attention + Cross-Attention + MLP 三层叠加，每层都有潜在问题

实际验证：

• 测试了 15+ 种不同除数组合
• 调试了 100+ 步推理过程
• 尝试了 6+ 种不同的保护策略
• 无法找到可工作的配置

---

推荐方案

方案 A：明确不支持 FP16（推荐）

在 comfy/supported_models.py 中修改：

class CosmosT2IPredict2(supported_models_base.BASE):
    # ... 其他配置 ...

    supported_inference_dtypes = [torch.bfloat16, torch.float32]  # 移除 torch.float16

    def get_model(self, state_dict, prefix="", device=None):
        # 在加载时添加警告
        dtype = device.type if hasattr(device, 'type') else None
        if dtype == torch.float16:
            print("[WARNING] Cosmos Predict2 models do not support FP16 due to numerical instability.")
            print("[WARNING] Please use BF16 (recommended) or FP32 instead.")
            print("[WARNING] To use FP16, you may need to train the model in FP16 or use a different architecture.")
        out = model_base.CosmosPredict2(self, device=device)
        return out

方案 B：混合精度训练（如果必须支持 FP16）

需要重新训练模型：

• 在 FP32 精度下训练
• 推理时使用 FP16
• 或者在关键层使用 FP32 中间精度

方案 C：改用其他精度

---

代码修改参考

文件：comfy/ldm/cosmos/predict2.py

已修改内容（当前状态）：

def clamp_fp16(x, debug_name=None):
    """FP16 值保护，防止溢出"""
    if x.dtype == torch.float16:
        return torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
    return x

class GPT2FeedForward(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x1 = self.layer1(x)
        x2 = self.activation(x1)
        # FP16 保护
        if x2.dtype == torch.float16:
            x2 = clamp_fp16(x2)
            x2 = x2 / 2.5
        x3 = self.layer2(x2)
        if x3.dtype == torch.float16:
            x3 = x3 / 3.5
        return clamp_fp16(x3)

class Attention(nn.Module):
    def compute_attention(self, q, k, v, transformer_options={}):
        result = self.attn_op(q, k, v, transformer_options)
        output = self.output_dropout(self.output_proj(result))
        if output.dtype == torch.float16:
            output = output / 1.5
        return clamp_fp16(output)

class Block(nn.Module):
    def forward(self, ...):
        # ...
        x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D * clamp_fp16(result_B_T_H_W_D)
        # ...
        x_B_T_H_W_D = clamp_fp16(result_B_T_H_W_D * gate_cross_attn_B_T_1_1_D) + x_B_T_H_W_D
        # ...
        x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp_B_T_1_1_D * clamp_fp16(result_B_T_H_W_D)
        return x_B_T_H_W_D

---

避坑指南

1. 调试策略

❌ 错误做法：
- 盲目调整除数，一个一个测试
- 没有数据支撑的情况下调整
- 忽略日志输出

✅ 正确做法：
- 先添加调试日志，观察实际数值范围
- 基于数据调整参数
- 记录每次测试的配置和结果

2. 数值分析

关键指标：
- max 值：是否接近 FP16 上限 (65504)
- mean 值：代表正常数值范围
- NaN/Inf：是否出现异常值
- 趋势：数值如何随步数变化

分析技巧：
- 比较不同层的数值范围
- 观察极端值出现的频率
- 检查累积效应

3. 参数调优

有效方法：
1. 确定问题所在后，针对性解决
2. 避免同时调整多个参数
3. 记录每次测试的完整配置
4. 从大范围开始，逐步缩小

无效方法：
1. 随意尝试除数
2. 没有理论支撑的调整
3. 重复相同的错误模式

---

经验教训

技术层面

1. 不同架构需要不同方案：Lumina 的方案不能直接套用
2. 数值范围分析很重要：先观察再修改
3. 激活函数影响巨大：GELU vs SiLU 的差异可以导致完全不同的问题
4. 累积误差容易被忽略：单步测试正常不代表多步正常

工程层面

1. 调试日志是必要的：没有数据支持的调试是盲人摸象
2. 测试需要有策略：而不是随机尝试
3. 及时止损很重要：如果多次测试无效，应该重新评估方向
4. 文档记录很重要：每次测试都应该记录下来

---

附录：测试工具

调试日志函数（用于观察数值）

def clamp_fp16_debug(x, debug_name=None):
    """带调试信息的 FP16 保护"""
    if x.dtype == torch.float16:
        has_nan = torch.isnan(x).any().item()
        has_inf = torch.isinf(x).any().item()
        valid_mask = ~torch.isnan(x) & ~torch.isinf(x)
        if valid_mask.any():
            max_val = x[valid_mask].abs().max().item()
            min_val = x[valid_mask].abs().min().item()
            mean_val = x[valid_mask].abs().mean().item()
        else:
            max_val = min_val = mean_val = 0
        # 打印统计信息
        print(f"[FP16] {debug_name}: max={max_val:.2f}, mean={mean_val:.2f}, "
              f"min={min_val:.6f}, nan={has_nan}, inf={has_inf}")
        return torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
    return x

快速测试配置

# 用于快速切换不同配置
FP16_CONFIGS = {
    "lumina_orig": {"attn": 4, "ffn": 32},
    "conservative": {"attn": 1.5, "ffn": 3.5},
    "aggressive": {"attn": 2, "ffn": 4},
    "adaptive": {"attn": 1.5, "ffn": 3.5, "threshold": 20}
}

---

参考资料

1. Lumina FP16 Overflow Fix: https://github.com/Comfy-Org/ComfyUI/pull/11187
2. PyTorch FP16 限制: https://pytorch.org/docs/stable/notes/cuda.html#fp16-optimizations
3. GELU 激活函数: https://pytorch.org/docs/stable/generated/torch.nn.GELU.html
4. Cosmos 官方文档: https://github.com/NVIDIA/cosmos-predict2

---

报告日期: 2026-02-03
测试环境: ComfyUI, CUDA, RTX 3080 Ti Laptop
测试模型: Anima (Cosmos-Predict2-2B-Text2Image)
结论: Cosmos Predict2 架构设计上不适合 FP16 精度