bf16 模型在 T4 上跑 fp16 出纯黑图 (fp16 Overflow to NaN on Turing GPUs)

2026-09-24 暂无评论

发生条件

硬件：NVIDIA T4（Colab 免费档的主力卡），架构 Turing，计算能力 7.5
模型：以 bf16 训练的扩散模型 / DiT，例：Anima-2.9B（Cosmos 架构）、Flux、多数 2024 年后的新模型
操作：为了适配老卡，故意把 torch_dtype 从 bf16 降到 fp16

症状

生成的 PNG 全黑，文件大小只有 2 KB 左右（正常该是 500 KB – 2 MB）。
进程不报错，进度条正常走完，日志干干净净，最后 img.save() 也成功。

排查过程

一开始怀疑的方向全错，逐个排除掉的：

怀疑	结论
提示词有问题	不是。换提示词一样黑
显存不够	不是。降分辨率一样黑，而且没有 OOM
模型没下全 / 加载失败	不是。权重文件完整，pipeline 正常实例化
VAE 精度问题	部分相关。Wan VAE 在 fp16 下确实不稳，但不是主因
transformer 数值溢出	是这个

怎么定位到根因（关键手法）

不要靠猜，加探针把中间张量打出来。 分两段查：

第一段——文本编码输出：

emb = pipe.encode_prompt(prompt=PROMPT, negative_prompt=NEG,
                         do_classifier_free_guidance=True,
                         num_images_per_prompt=1, device="cuda",
                         max_sequence_length=512)
for i, e in enumerate(emb):
    f = e.float()
    print("emb[%d] nan=%s inf=%s min=%.3f max=%.3f" % (
        i, torch.isnan(f).any().item(), torch.isinf(f).any().item(),
        f.min().item(), f.max().item()))

实测结果：nan=False min=-0.802 max=0.919 —— 文本编码是干净的。
这一条非常重要，它一次性排除了 tokenizer、文本编码器、LLM adapter 前面所有的环节，把范围缩到"transformer 前向"。

第二段——每步 latent：

def cb(pl, i, t, kw):
    lat = kw["latents"].float()
    print("step %d nan=%s min=%s max=%s" % (
        i + 1, torch.isnan(lat).any().item(), lat.min().item(), lat.max().item()))
    return kw

pipe(..., callback_on_step_end=cb)

实测结果：第 1 步就已经是 nan=True。
不是"跑着跑着崩的"，是第一次前向传播就溢出。

根因

模型权重是 bf16 训练的。为了在 T4 上跑得动，加载时把 torch_dtype 指定为 fp16。

fp16 的指数位比 bf16 少 3 位（fp16 是 5 位指数，bf16 是 8 位）。bf16 靠这 3 位能表示的数值范围，正是训练时依赖的。降到 fp16 后，中间激活值超出 fp16 上限 → 变成 inf → 下一步 inf - inf 或 inf * 0 → NaN → 扩散全程污染 → VAE 解出全黑。

最终解

pipe = AnimaTextToImagePipeline.from_pretrained(
    LOCAL,
    torch_dtype=torch.bfloat16,   # 用模型的原生精度
    trust_remote_code=True,
).to("cuda")
pipe.vae.to(torch.float32)        # Wan VAE 在低精度下也不稳，单独 upcast

改完之后图正常出，std=65.3（黑图是 std=0，所以 std 是个很快的黑图判别指标）。

代价：T4 跑 bf16 很慢

Turing 架构没有原生 bf16 硬件支持。PyTorch 不会报错，会走软件路径算，结果正确但慢。

实测 Anima-2.9B：

精度	640x896 单步耗时
fp16（出黑图，不可用）	快，但结果无效
bf16（正确）	约 11 秒/步

同一张 640x896 / 24 步的图要 4.5 分钟。这是正确的代价，不是可以优化掉的。

下次怎么提前认出来

看计算能力主版本号：

torch.cuda.get_device_capability(0)   # T4 -> (7, 5)

主版本 < 8 就是 Turing 或更老（V100=7.0、T4=7.5、RTX 20 系=7.5），没有原生 bf16。
Ampere（A100=8.0、RTX 30 系=8.6）及以上才有。

一个陷阱：torch.cuda.is_bf16_supported() 在 T4 上可能返回 True。它判断的是"能不能跑得起来"，不是"跑得快不快"。别用这个函数判断该不该用 bf16，用 get_device_capability()。

通用规律：任何"加载成功了、跑完了、没报错，但输出是垃圾"的情况，优先怀疑 dtype 与硬件不匹配，并用上面那套"分段打印中间张量"的办法二分定位是哪一层出的问题。