ThunderFX in BF16 mode shows a larger divergence from eager FP64 results compared to the divergence between eager BF16 and eager FP64

[kiya00]

Note: If you have a model or program that is not supported yet but should be, please use the program coverage template.

🐛 Bug

While adding the coverage test for HF Diffusers (PR #2141), I noticed that ThunderFX in BF16 showed a significant divergence compared to Eager BF16. To investigate further, I tested various dtypes (FP64, FP32, FP16, BF16) across different backends (Eager, ThunderFX, and torch.compile) and compared their outputs to Eager FP64. The results showed that ThunderFX generally stays within an acceptable range of divergence—except for gradient outputs in FP16 and BF16, where the discrepancies were noticeably larger.

The results of runwayml/stable-diffusion-v1-5:
BF16 vs eager FP64:

ThunderFX with default executors
Forward pass comparison: FAILED
Error: Tensor-likes are not close!

Mismatched elements: 1023 / 1024 (99.9%)
Greatest absolute difference: 0.02505786082606219 at index (0, 3, 9, 3) (up to 1e-05 allowed)
Greatest relative difference: 15.79343976974047 at index (0, 0, 6, 2) (up to 1e-05 allowed)
Computing gradients for torch eager float64...
Computing gradients for thunder torch.bfloat16...
Gradient comparison: FAILED
Error: Tensor-likes are not close!

Mismatched elements: 11520 / 11520 (100.0%)
Greatest absolute difference: **52.038735867250054** at index (145, 0, 1, 0) (up to 1e-05 allowed)
Greatest relative difference: **118164.35679032684** at index (162, 0, 0, 2) (up to 1e-05 allowed)

The failure occurred for item [0]

eager

Mismatched elements: 1022 / 1024 (99.8%)
Greatest absolute difference: 0.031462447813257244 at index (0, 3, 12, 14) (up to 1e-05 allowed)
Greatest relative difference: 5.43559016117814 at index (0, 0, 6, 4) (up to 1e-05 allowed)
Computing gradients for torch eager float64...
Computing gradients for eager torch.bfloat16...
Gradient comparison: FAILED
Error: Tensor-likes are not close!

Mismatched elements: 11518 / 11520 (100.0%)
Greatest absolute difference: **0.23972053767860058** at index (127, 3, 1, 1) (up to 1e-05 allowed)
Greatest relative difference: **532.5203075118578** at index (162, 0, 0, 2) (up to 1e-05 allowed)

torch.compile

Mismatched elements: 1023 / 1024 (99.9%)
Greatest absolute difference: 0.02224970804979243 at index (0, 3, 1, 14) (up to 1e-05 allowed)
Greatest relative difference: 7.853684232346166 at index (0, 0, 6, 2) (up to 1e-05 allowed)
Computing gradients for torch eager float64...
Computing gradients for inductor torch.bfloat16...
Gradient comparison: FAILED
Error: Tensor-likes are not close!

Mismatched elements: 11515 / 11520 (100.0%)
Greatest absolute difference: **0.1986014106462043** at index (202, 0, 0, 1) (up to 1e-05 allowed)
Greatest relative difference: **328.6249694895744** at index (239, 3, 0, 2) (up to 1e-05 allowed)

ThunderFX with executors=[thunder.pytorch_executor]

Mismatched elements: 1024 / 1024 (100.0%)
Greatest absolute difference: 0.02503035742302913 at index (0, 1, 11, 1) (up to 1e-05 allowed)
Greatest relative difference: 14.88029205458141 at index (0, 0, 6, 2) (up to 1e-05 allowed)
Computing gradients for torch eager float64...
Computing gradients for thunder torch.bfloat16...
Gradient comparison: FAILED
Error: Tensor-likes are not close!

Mismatched elements: 11517 / 11520 (100.0%)
Greatest absolute difference: **0.25748365664500517** at index (122, 0, 1, 2) (up to 1e-05 allowed)
Greatest relative difference: **387.0515917837159** at index (162, 0, 0, 2) (up to 1e-05 allowed)

To Reproduce

1. Uses container pytorchlightning/lightning-thunder:ubuntu24.04-cuda12.6.3-cudnn-fe1.10.0-py3.10-pt_2.7.0-dev
2. Modifies the bitsandbytes version and adds diffusers, then installs it:

--- a/requirements/test.txt
+++ b/requirements/test.txt
@@ -17,7 +17,9 @@ absl-py # thunder/benchmarks/test_benchmark_litgpt.py
 pandas # thunder/benchmarks/test_benchmark_litgpt.py
 xlsxwriter # thunder/benchmarks/test_benchmark_litgpt.py
 jsonargparse # thunder/benchmarks/benchmark_litgpt.py
-bitsandbytes==0.42.0  # fixed version!
+# bitsandbytes==0.42.0  # fixed version!
+bitsandbytes>=0.42,<0.43; sys_platform=='darwin'
+bitsandbytes>=0.45.2,<0.45.5; sys_platform!='darwin'
 transformers==4.50.3 # for test_networks.py
 diffusers==0.33.0 # for test_networks.py

3. python compare_acc64.py --backend thunder --dtype bfloat16:

import torch
from thunder.dynamo import thunderfx
from diffusers import UNet2DConditionModel

# torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
torch.set_float32_matmul_precision('high')
torch.manual_seed(0)
import argparse

# Parse command line arguments for backend selection
parser = argparse.ArgumentParser(description='UNet2D accuracy test with different backends')
parser.add_argument('--backend', type=str, default='thunder', choices=['thunder', 'inductor', 'eager'],
                    help='Backend to use for model compilation (default: thunder)')
parser.add_argument('--dtype', type=str, default='bfloat16', choices=['float16', 'bfloat16', 'float32', 'float64'],
                    help='Data type to use for model (default: bfloat16)')
parser.add_argument('--model-id', type=str, default='runwayml/stable-diffusion-v1-5',
                    help='Model ID to use for testing (default: runwayml/stable-diffusion-v1-5)')
args = parser.parse_args()
backend = args.backend
if args.dtype == "float16":
    dtype = torch.float16
elif args.dtype == "bfloat16":
    dtype = torch.bfloat16
elif args.dtype == "float32":
    dtype = torch.float32
elif args.dtype == "float64":
    dtype = torch.float64
else:
    raise ValueError(f"Invalid dtype: {args.dtype}")
model_id = args.model_id

print(f"Using backend: {backend}")
print(f"Using dtype: {dtype}")
print(f"Using model ID: {model_id}")

# model_id = "runwayml/stable-diffusion-v1-5" #"stabilityai/stable-diffusion-2-1"

try:
    unet_bf16 = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", torch_dtype=dtype)
    unet_f64 = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", torch_dtype=torch.float64)
except OSError:
    unet_bf16 = UNet2DConditionModel.from_pretrained(
        model_id, subfolder="unet", use_safetensors=False, torch_dtype=dtype
    )
    unet_f64 = UNet2DConditionModel.from_pretrained(
        model_id, subfolder="unet", use_safetensors=False, torch_dtype=torch.float64
    )

config = unet_bf16.config
in_channels = config.in_channels
sample_size = 16 # config.sample_size
cross_attention_dim = config.cross_attention_dim
addition_embed_type = config.addition_embed_type

batch_size = 1
seq_length = 8

if "xl" in model_id:
    time_ids_dim = 6
    text_embeds_dim = 1280
    if "refiner" in model_id:
        time_ids_dim = 2
        text_embeds_dim = 2048
else:
    time_ids_dim = None
    text_embeds_dim = None

input_shape = (batch_size, in_channels, sample_size, sample_size)
hidden_states_shape = (batch_size, seq_length, cross_attention_dim)

# Setup models
unet_f64 = unet_f64.to("cuda", dtype=torch.float64).requires_grad_(True)
unet_bf16 = unet_bf16.to("cuda", dtype=dtype).requires_grad_(True)
if backend == "thunder":
    import thunder
    compiled_model = thunderfx(unet_bf16)
    # compiled_model = thunderfx(unet_bf16, executors=[thunder.pytorch_executor])
elif backend == "inductor":
    compiled_model = torch.compile(unet_bf16)
elif backend == "eager":
    compiled_model = unet_bf16
else:
    raise ValueError(f"Invalid backend: {backend}")

def make_inputs(dtype=torch.bfloat16):
    added_cond_kwargs = {}
    with torch.device("cuda"):
        input = torch.randn(input_shape, dtype=dtype)
        hidden_states = torch.randn(hidden_states_shape, dtype=dtype)
        timestep = torch.ones(batch_size, dtype=torch.long)
        if addition_embed_type is not None:
            assert text_embeds_dim is not None and time_ids_dim is not None
            time_ids_shape = (batch_size, time_ids_dim)
            text_embeds_shape = (batch_size, text_embeds_dim)
            added_cond_kwargs["time_ids"] = torch.randn(time_ids_shape, device="cuda", dtype=dtype)
            added_cond_kwargs["text_embeds"] = torch.randn(text_embeds_shape, device="cuda", dtype=dtype)
    return (input, timestep, hidden_states), {"added_cond_kwargs": added_cond_kwargs}

# Create identical inputs, then convert to appropriate dtypes
base_args, base_kwargs = make_inputs(torch.float64)

# Use the same inputs for both models, just convert dtype
f64_args = base_args
f64_kwargs = base_kwargs

# Convert to bfloat16 for the compiled model
bf16_args = tuple(arg.to(dtype) for arg in base_args)
bf16_kwargs = {}
for key, value in base_kwargs.items():
    if key == "added_cond_kwargs":
        bf16_kwargs[key] = {k: v.to(dtype) for k, v in value.items()}
    else:
        bf16_kwargs[key] = value

# Run inference
print("Running torch eager float64...")
ref_output = unet_f64(*f64_args, **f64_kwargs)

print(f"Running {backend} {dtype}...")
compiled_output = compiled_model(*bf16_args, **bf16_kwargs)

ref_output = ref_output.sample
compiled_output = compiled_output.sample

# Convert compiled output to float64 for comparison
compiled_output_f64 = compiled_output.to(torch.float64)

# Compare outputs
try:
    torch.testing.assert_close(compiled_output_f64, ref_output, rtol=1e-5, atol=1e-5)
except Exception as e:
    print(f"Forward pass comparison: FAILED")
    print(f"Error: {e}")
else:
    print("Forward pass comparison: PASSED")

# Gradient comparison
# loss_grad_f64 = torch.randn_like(ref_output)
# loss_grad_bf16 = loss_grad_f64.to(dtype)

loss_grad_bf16 = torch.randn_like(compiled_output)
loss_grad_f64 = loss_grad_bf16.to(torch.float64)

print("Computing gradients for torch eager float64...")
grads_ref = torch.autograd.grad(ref_output, unet_f64.parameters(), grad_outputs=loss_grad_f64)

print(f"Computing gradients for {backend} {dtype}...")
grads_compiled = torch.autograd.grad(compiled_output, unet_bf16.parameters(), grad_outputs=loss_grad_bf16)

# Convert compiled gradients to float64 for comparison
grads_compiled_f64 = [g.to(torch.float64) for g in grads_compiled]

try:
    # Compare gradients with relaxed tolerance due to precision differences
    torch.testing.assert_close(grads_compiled_f64, grads_ref, rtol=1e-5, atol=1e-5)
except Exception as e:
    print(f"Gradient comparison: FAILED")
    print(f"Error: {e}")
else:
    print("Gradient comparison: PASSED")

[kiya00]

Some updates:
By spliting the model, found the minimum reproduce pattern as follows. The preliminary analysis suggests that the issue is caused by precision errors in Thunder's cudnn SDPA backward computation. When the cudnn_executor is disabled, the precision returns to normal.

from math import inf
from math import nan
NoneType = type(None)
import torch
from torch import device
import torch.fx._pytree as fx_pytree
import torch.utils._pytree as pytree
import thunder
torch.manual_seed(0)


class DynamoModule(torch.nn.Module):
    def forward(self, query_1, key_1, value_1, l_self_modules_to_out_modules_0_parameters_weight_: "f16[768, 768][768, 1]cuda:0", l_self_modules_to_out_modules_0_parameters_bias_: "f16[768][1]cuda:0"):
         # File: /usr/local/lib/python3.10/dist-packages/diffusers/models/attention_processor.py:3313 in __call__, code: hidden_states = F.scaled_dot_product_attention(
        hidden_states: "f16[2, 8, 64, 96][49152, 96, 768, 1]cuda:0" = torch._C._nn.scaled_dot_product_attention(query_1, key_1, value_1, attn_mask = None, dropout_p = 0.0, is_causal = False);  query_1 = key_1 = value_1 = None
         # File: /usr/local/lib/python3.10/dist-packages/diffusers/models/attention_processor.py:3317 in __call__, code: hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        transpose_3: "f16[2, 64, 8, 96][49152, 768, 96, 1]cuda:0" = hidden_states.transpose(1, 2);  hidden_states = None
        hidden_states_1: "f16[2, 64, 768][49152, 768, 1]cuda:0" = transpose_3.reshape(2, -1, 768);  transpose_3 = None
        
         # File: /usr/local/lib/python3.10/dist-packages/diffusers/models/attention_processor.py:3321 in __call__, code: hidden_states = attn.to_out[0](hidden_states)
        hidden_states_3: "f16[2, 64, 768][49152, 768, 1]cuda:0" = torch._C._nn.linear(hidden_states_1, l_self_modules_to_out_modules_0_parameters_weight_, l_self_modules_to_out_modules_0_parameters_bias_);  hidden_states_1 = l_self_modules_to_out_modules_0_parameters_weight_ = l_self_modules_to_out_modules_0_parameters_bias_ = None
        
        return hidden_states_3

        
def test_graph0_thunder_0():
    inputs = [
        torch.testing.make_tensor((2, 64, 768), dtype=torch.float16,  device='cuda:0', requires_grad=True, low=-5.1171875, high=4.26953125,).as_strided(size=(2, 8, 64, 96), stride=(49152, 96, 768, 1)),
        torch.testing.make_tensor((2, 64, 768), dtype=torch.float16,  device='cuda:0', requires_grad=True, low=-5.1171875, high=4.26953125,).as_strided(size=(2, 8, 64, 96), stride=(49152, 96, 768, 1)),
        torch.testing.make_tensor((2, 64, 768), dtype=torch.float16,  device='cuda:0', requires_grad=True, low=-5.1171875, high=4.26953125,).as_strided(size=(2, 8, 64, 96), stride=(49152, 96, 768, 1)),
        torch.testing.make_tensor((768, 768), dtype=torch.float16,  device='cuda:0', requires_grad=True, low=-0.03607177734375, high=0.03607177734375,),
        torch.testing.make_tensor((768,), dtype=torch.float16,  device='cuda:0', requires_grad=True, low=-0.036041259765625, high=0.036041259765625,),
    ]

    model = DynamoModule()

    # Accuracy check passes
    # compiled_model = thunder.jit(model,executors=[thunder.sdpa_executor, thunder.torchcompile_xentropy_executor, thunder.pytorch_executor],fusion_type='dataflow',transforms=[thunder.transforms.extraction_only_prologue_transform.ExtractionOnlyPrologueTransform()],thunderfx_disable_split_autograd=True,)
    # compiled_model = torch.compile(model)
    # Accuracy check fails
    compiled_model = thunder.jit(model,executors=[thunder.cudnn_executor, thunder.sdpa_executor, thunder.torchcompile_xentropy_executor, thunder.pytorch_executor],fusion_type='dataflow',transforms=[thunder.transforms.extraction_only_prologue_transform.ExtractionOnlyPrologueTransform()],thunderfx_disable_split_autograd=True,)

    from thunder.dynamo.report import run_forward_backward
    fwd_result, grads = run_forward_backward(compiled_model, *inputs)
    eager_fwd_result, eager_grads = run_forward_backward(model, *inputs)


    #trcs = thunder.last_traces(compiled_model)
    #print(trcs[-1])
    #bwd_trcs = thunder.last_backward_traces(compiled_model)
    #print(bwd_trcs[-1])

    for i, (fwd_res, eager_fwd_res) in enumerate(zip(fwd_result, eager_fwd_result)):
        if not torch.allclose(fwd_res, eager_fwd_res, rtol=1e-2, atol=1e-2):
            max_diff = torch.max(torch.abs(fwd_res - eager_fwd_res))
            print(f"fwd Outputs differ by {max_diff}, {(fwd_res - eager_fwd_res).abs().max()}")
    
    # Compare gradients
    for i, (grad, eager_grad) in enumerate(zip(grads, eager_grads)):
        if not torch.allclose(grad, eager_grad, rtol=1e-2, atol=1e-2, equal_nan=True):
            max_diff = torch.max(torch.abs(grad - eager_grad))
            print(f"Gradient {i} differs by {max_diff}, {(grad - eager_grad).abs().max()}")

    torch.testing.assert_close(fwd_result, eager_fwd_result,rtol=1e-2, atol=1e-2)
    torch.testing.assert_close(grads, eager_grads,rtol=1e-2, atol=1e-2, equal_nan=True)


if __name__ == "__main__":
    test_graph0_thunder_0()


"""
Environment information get from `torch.utils.collect_env.get_pretty_env_info()`:
CUDA devices:
  0: NVIDIA RTX 6000 Ada Generation
  1: NVIDIA RTX 6000 Ada Generation
CUDA version: 12.6
numpy==2.2.6
nvfuser_cu126_torch27==0.2.27.dev20250520
nvidia-cublas-cu12==12.6.4.1
nvidia-cuda-cupti-cu12==12.6.80
nvidia-cuda-nvrtc-cu12==12.6.77
nvidia-cuda-runtime-cu12==12.6.77
nvidia-cudnn-cu12==9.5.1.17
nvidia-cudnn-frontend==1.10.0
nvidia-cufft-cu12==11.3.0.4
nvidia-curand-cu12==10.3.7.77
nvidia-cusolver-cu12==11.7.1.2
nvidia-cusparse-cu12==12.5.4.2
nvidia-cusparselt-cu12==0.6.3
nvidia-nccl-cu12==2.26.2
nvidia-nvjitlink-cu12==12.6.85
nvidia-nvtx-cu12==12.6.77
optree==0.15.0
pytorch-lightning==2.5.1.post0
torch==2.7.0+cu126
torchmetrics==1.7.1
triton==3.3.0

Versions of Thunder related libraries:
lightning-thunder==0.2.4.dev0
nvfuser_cu126_torch27==0.2.27.dev20250520

Split Information:
The original graph is not split, and is entirely run by Thunder.

"""

[kiya00]

update:
The cause appears to be that cudnn_sdpa_bwd does not correctly handle strided input tensors.
This is the minimal repro for the accuracy issue:
inputs file: https://drive.google.com/file/d/10bRAbJK5MWoLmfyhMM6BtlCdK3MJAxDT/view?usp=sharing

diff --git a/thunder/__init__.py b/thunder/__init__.py
index 44cb5989..9e546a2f 100644
--- a/thunder/__init__.py
+++ b/thunder/__init__.py
@@ -550,7 +550,7 @@ def jit(
                 requires_grad = any(isinstance(arg, tensor_cls) and arg.requires_grad for arg in computation_trc.args)
             else:
                 requires_grad = False
-
+            requires_grad = False
             if requires_grad:
                 # Currently split_forward_backward also includes
                 # transform_for_execution and various sorting of symbols,

from thunder.executors.cudnnex import cudnn_sdpa_bwd
import thunder
import torch
def func(*args):
    # (grad_for_query_1, grad_for_key_1, grad_for_value_1) = cudnn_sdpa_bwd(bw_t56, query_1, key_1, value_1, None, 0.0, False, t23, t24, t25, t26, scale=None, cat_grad_qkv=False)
    (grad_for_query_1, grad_for_key_1, grad_for_value_1) = cudnn_sdpa_bwd(*args, scale=None, cat_grad_qkv=False)
    return (grad_for_query_1, grad_for_key_1, grad_for_value_1)

from thunder.executors.sdpaex import sdpfa_bwd
def func2(*args):
    # (grad_for_query_1, grad_for_key_1, grad_for_value_1) = sdpfa_bwd(bw_t59, query_1, key_1, value_1, t23, t24, t25, t26, 64, 64, 0.0, False, t27, t28, scale=None)
    (grad_for_query_1, grad_for_key_1, grad_for_value_1) = sdpfa_bwd(*args, scale=None)
    return (grad_for_query_1, grad_for_key_1, grad_for_value_1)


inps = torch.load("./inputs.pt")
compiled_model = thunder.jit(func,thunderfx_disable_split_autograd=True)
out = compiled_model(*inps)

contiguous_inps = [inp.contiguous() if isinstance(inp, torch.Tensor) else inp for inp in inps]
compiled_model2 = thunder.jit(func,thunderfx_disable_split_autograd=True)
out2 = compiled_model2(*contiguous_inps)

inps2 = torch.load("./inputs.pt")
inps2 = inps2[0:4] + inps2[7:]
t27=torch.tensor([0, 0], device='cuda:0', dtype=torch.int64) 
t28=torch.tensor(0, device='cuda:0', dtype=torch.int64)
inps2.extend([64, 64, 0.0, False, t27, t28])
compiled_model2 = thunder.jit(func2)
expected_out = compiled_model2(*inps2)

# compare the strided case
print("strided case")
for o, eo in zip(out, expected_out):
    max_diff = torch.max(torch.abs(o - eo))
    print(max_diff)

# compare the contiguous case
print("contiguous case")
for o, eo in zip(out2, expected_out):
    max_diff = torch.max(torch.abs(o - eo))
    print(max_diff)

outputs:

strided case
tensor(14.5156, device='cuda:0', dtype=torch.float16)
tensor(14.7109, device='cuda:0', dtype=torch.float16)
tensor(8.9609, device='cuda:0', dtype=torch.float16)
contiguous case
tensor(0.0078, device='cuda:0', dtype=torch.float16)
tensor(0.0078, device='cuda:0', dtype=torch.float16)
tensor(0.0039, device='cuda:0', dtype=torch.float16)

[kiya00]

I noticed there's a comment

lightning-thunder/thunder/executors/cudnnex.py

Lines 244 to 253 in 62d9535

	def _sdpa_enforce_input_tensor_contiguity(a: torch.Tensor) -> torch.Tensor:
	if a.stride(-1) == 1:
	# TODO(vedaanta-nvidia): there's an inconsistency between
	# _transform_sdpa_inputs and this function, leading to a potential bug.
	# _transform_sdpa_inputs always creates contiguous strides, but code
	# here creates a partially contiguous stride when the last dimension is
	# contiguous but other dimensions are not.
	return a
	else:
	return a.contiguous()

If the contiguous tensor is simply returned, there's no accuracy issue:

diff --git a/thunder/executors/cudnnex.py b/thunder/executors/cudnnex.py
index 447ed6ec..efc1e393 100644
--- a/thunder/executors/cudnnex.py
+++ b/thunder/executors/cudnnex.py
@@ -248,7 +248,8 @@ def _sdpa_enforce_input_tensor_contiguity(a: torch.Tensor) -> torch.Tensor:
         # _transform_sdpa_inputs always creates contiguous strides, but code
         # here creates a partially contiguous stride when the last dimension is
         # contiguous but other dimensions are not.
-        return a
+        # return a
+        return a.contiguous()
     else:
         return a.contiguous()
 
@@ -646,6 +647,7 @@ def _cudnn_sdpa_bwd_impl(
     query = _sdpa_enforce_input_tensor_contiguity(query)
     key = _sdpa_enforce_input_tensor_contiguity(key)
     value = _sdpa_enforce_input_tensor_contiguity(value)
+    grad_out = _sdpa_enforce_input_tensor_contiguity(grad_out)
 
     # When cat_grad_qkv is on, allocate dQKV and make dQ, dK, and dV
     # slices of that. Otherwise, allocate them individually.

cc: @IvanYashchuk @mruberry @vedaanta

[IvanYashchuk]

@vedaanta, do you know what the stride requirements are for cuDNN's backward SDPA? Should we always pass only contiguous inputs?