Bug: Inference crashes on Hopper GPUs (H20/H100) with SIGFPE and dtype mismatch errors

[jyx13121802323]

Summary

When running batch inference (batch_infer.py) on NVIDIA Hopper architecture GPUs (H20, H100, H200), two critical bugs cause the process to crash:

1. time_embedding.float() in-place mutation — permanently converts time_embedding weights to float32, causing mat1 and mat2 must have the same dtype, but got Float and BFloat16 on subsequent calls.
2. NVIDIA cuBLAS SIGFPE on Hopper GPUs — bfloat16 GEMM in cublasLtMatmul triggers a C-level Floating point exception (core dumped) crash due to a known cuBLAS bug (fixed in cuBLAS 12.4.5.8).

Environment

• GPU: NVIDIA H20 (Hopper, compute capability 9.0)
• PyTorch: 2.4.1+cu121
• CUDA: 12.1
• nvidia-cublas-cu12: 12.1.3.1
• Model dtype: bfloat16

Bug 1: time_embedding.float() In-place Weight Pollution

Error Message

RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16

Root Cause

In diffsynth/pipelines/wan_video_new.py, the model_fn_wans2v function contained debug code (around lines 353-398) that called:

with torch.autocast("cuda", enabled=False):
    t = dit.time_embedding.float()(emb.float())

nn.Module.float() is an in-place operation — it permanently converts all parameters of the module to float32. This is NOT equivalent to creating a temporary copy. After this call, dit.time_embedding's weights are permanently float32, while the rest of the model remains bfloat16.

When the code subsequently calls:

t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep))

The sinusoidal_embedding_1d returns a bfloat16 tensor (matching timestep's dtype), but time_embedding's weights are now float32 → dtype mismatch.

Fix

1. Remove the debug code that contains .float() in-place calls on modules.
2. Ensure time_embedding is called with input that matches its weight dtype:

emb = sinusoidal_embedding_1d(dit.freq_dim, timestep)
t = dit.time_embedding(emb.to(dit.time_embedding[0].weight.dtype))

---

Bug 2: cuBLAS SIGFPE on Hopper Architecture GPUs

Error Message

Fatal Python error: Floating point exception

Thread 0x00007f43056ff640 (most recent call first):
  File ".../threading.py", line 324 in wait
  ...
Current thread 0x00007f457a28c740 (most recent call first):
  File ".../torch/nn/modules/linear.py", line 117 in forward
  File ".../torch/nn/modules/container.py", line 219 in forward
  File "diffsynth/pipelines/wan_video_new.py", line 2354 in model_fn_wans2v

Floating point exception (core dumped)

Root Cause

This is a known NVIDIA cuBLAS bug affecting Hopper architecture GPUs (H20, H100, H200, etc.):

cublasLtMatmul() and cublasLtMatmulAlgoGetHeuristic() could have resulted in floating point exceptions (FPE) on some Hopper-based GPUs, including Multi-Instance GPU (MIG). The issue was introduced in cuBLAS 11.8.
— NVIDIA CUDA Toolkit 12.4 Update 1 Release Notes

The bug occurs because cuBLASLt's internal heuristic algorithm selection contains an integer division by zero for certain GEMM sizes when using bfloat16 on Hopper GPUs. This triggers a SIGFPE signal at the C level, which Python cannot catch — the process immediately crashes.

Related issues:

• pytorch/pytorch#133548 — FPE on H20 with nn.Linear
• pytorch/pytorch#152017 — FPE on H20 with specific matmul sizes
• vllm-project/vllm#4392 — FPE on H20, fix with nvidia-cublas-cu12==12.4.5.8
• aigc-apps/EasyAnimate#32 — bf16 causes FPE on H20

In our case, the crash happens at dit.time_projection(t) in model_fn_wans2v (line 2354). The time_projection layer is nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) — the large GEMM size (5120 → 30720 for the 14B model) triggers the cuBLAS bug.

Fix (Two Options)

Option A (Recommended): Upgrade nvidia-cublas-cu12

pip install nvidia-cublas-cu12==12.4.5.8

This is NVIDIA's official fix. The pip package is forward-compatible and works with PyTorch 2.4.1+cu121.

Option B (Code-level workaround): Use float32 for time_embedding/time_projection computation

Add a helper function that computes nn.Sequential layers in float32, avoiding the bfloat16 GEMM that triggers the cuBLAS bug:

import torch.nn as nn

def _safe_sequential_forward(module: nn.Sequential, x: torch.Tensor) -> torch.Tensor:
    """Forward through nn.Sequential in float32 to avoid cuBLAS SIGFPE on Hopper GPUs.

    NVIDIA cuBLAS <= 12.4.0 has a known bug on Hopper architecture GPUs (H20/H100/H200)
    where bfloat16 GEMM in cublasLtMatmul triggers SIGFPE due to internal integer
    division by zero. This function computes the Sequential layers in float32 as a
    workaround, then casts the result back to the original input dtype.

    This is safe because time_embedding/time_projection are relatively small layers
    and the float32 overhead is negligible compared to the DiT backbone.
    """
    original_dtype = x.dtype
    x = x.float()
    for layer in module:
        if isinstance(layer, nn.Linear):
            x = torch.nn.functional.linear(
                x, layer.weight.float(),
                layer.bias.float() if layer.bias is not None else None
            )
        else:
            x = layer(x)
    return x.to(original_dtype)

Then replace all dit.time_embedding(...) and dit.time_projection(...) calls in both model_fn_wan_video and model_fn_wans2v:

# Before (crashes on Hopper):
t = dit.time_embedding(sinusoidal_embedding_1d(dit.freq_dim, timestep))
t_mod = dit.time_projection(t).unflatten(1, (6, dit.dim))

# After (safe on Hopper):
emb = sinusoidal_embedding_1d(dit.freq_dim, timestep)
t = _safe_sequential_forward(dit.time_embedding, emb)
t_mod = _safe_sequential_forward(dit.time_projection, t).unflatten(1, (6, dit.dim))

The code change covers 3 call sites:

1. model_fn_wan_video TI2V-5B path (seperated_timestep + fuse_vae_embedding_in_latents)
2. model_fn_wan_video T2V path (scalar timestep)
3. model_fn_wans2v S2V path (and depth_timestep branch)

Note: This code workaround only protects time_embedding and time_projection. If SIGFPE occurs in other Linear layers after upgrading cuBLAS is not possible, a more comprehensive approach (e.g., using float16 as torch_dtype instead of bfloat16) may be needed.

---

Recommended Complete Fix

Apply both fixes together:

1. Remove debug code with in-place .float() calls (Bug 1)
2. Upgrade nvidia-cublas-cu12 to >=12.4.5.8 (Bug 2, primary)
3. Add _safe_sequential_forward as a code-level safeguard (Bug 2, backup)

This ensures the inference works reliably on both Hopper and non-Hopper GPUs, regardless of the cuBLAS version installed.

[luoxyhappy]

Thank you for pointing this out. Feel free to submit a PR, and I’ll review and merge it.

[jyx13121802323]

Thank you for pointing this out. Feel free to submit a PR, and I’ll review and merge it.

ok, I'll submit it while i'm free