Root cause analysis: Gemma4 sliding-window BF16 cuDNN SDPA NaN on H100

Summary

Follow-up to #2208 and its workaround
in #3141:

• Problem: Gemma4 BF16 training with attn_implementation="sdpa" produces NaN gradients and loss on H100,
  but not A100.
• Debug result: Backend controls and first-nonfinite capture strongly isolate the failure to the masked
  head_dim=256 sliding-window cuDNN SDPA backward path on SM90. The head_dim=512 full-attention SDPA
  EfficientAttention path is not the source of the observed NaN.
• Proposed fix: Address the root cause by deprioritizing cuDNN for the affected sliding-window path. The
  existing global FP32 SDPA upcast also fixes the observed NaN, but is a broader workaround that changes
  precision and dispatch for all Gemma4 SDPA calls. For supported workloads, prefer
  attn_implementation="ffpa" so sliding-window layers use FlexAttention BlockMask objects while eligible
  head_dim=512 full-attention layers use FFPA.

Reproduction

Environment:

• nvcr.io/nvidia/nemo-automodel:26.06
• 2× H100
• BF16
• local_batch_size=2

torchrun --nproc-per-node=2 \
  examples/vlm_finetune/finetune.py \
  -c examples/vlm_finetune/gemma4/gemma4_2b.yaml \
  --step_scheduler.local_batch_size 2 \
  --model.attn_implementation sdpa

Observed behavior:

• grad_norm becomes NaN at approximately step 22.
• Loss becomes NaN at approximately step 23.
• The same container and command are stable on A100.

See the original H100/A100 reproduction (https://github.com/NVIDIA-NeMo/Automodel/issues/2208#issuecomment-5030462469).

Why cuDNN SDPA is the affected path

On the affected H100 stack, the default SDPA priority is
cuDNN -> FlashAttention -> EfficientAttention -> Math, and the dispatcher selects the first eligible backend.

Gemma4 dispatch before the SDPA change reflected by Transformers #47042

Transformers #47042 updates an integration-test
expectation after #46960 changed the head_dim=512
GQA path from a Math fallback to EfficientAttention. Before that dispatch change, the observed operators were:

Backend	Layer type	head_dim	Mask object	Dispatched aten/HOP op	Actual CUDA kernel
SDPA	sliding	256	4D bool Tensor	aten::_scaled_dot_product_cudnn_attention	cudnn_generated_fort_native_sdpa_sm90_flash_{fprop,bprop}_wgmma_f16
SDPA	full (no padding)	512	None (is_causal fast path)	Math — no fused op	aten::bmm ×6 + aten::_softmax (materializes S×S)
SDPA	full (with padding)	512	4D bool Tensor	aten::_scaled_dot_product_efficient_attention	CUTLASS memory-efficient fmha_cutlassF/B_bf16_aligned_*_sm80
SDPA	sliding (with padding)	256	4D bool Tensor	aten::_scaled_dot_product_cudnn_attention	Same cuDNN flash kernel as above
Flex	sliding / full	256 / 512	BlockMask	torch.ops.higher_order.flex_attention (compiled)	Inductor triton_tem_fused_flex_attention*

The masked head_dim=256 sliding layers are the only rows that dispatch to the cuDNN SM90 forward/backward
kernels. The controlled training runs below strongly support this path as the first source of non-finite values
in the affected stack.

Debug and validation results

The validation used the exact 26.06 stack (PyTorch 2.12.0a0+0291f960b6, CUDA 13.2, cuDNN 9.21), 2 GPUs,
BF16, local batch size 2, and model revision 3e22461f65e89153144f8adb70e3b8c2cc9845a7.

GPU / run	Sliding d256 path	Result
H100 baseline	Default cuDNN	grad norm NaN at step 22; loss NaN at step 23
H100 forced cuDNN	Force cuDNN for sliding d256	Reproduced the same step 22/23 failure
H100 no cuDNN	cuDNN excluded; EfficientAttention selected	28/28 steps finite; all 22,416 attention tensor checks finite
H100 cuDNN last	cuDNN deprioritized; EfficientAttention selected	28/28 steps finite; all 22,416 attention tensor checks finite
A100 baseline	Default EfficientAttention	28/28 steps finite; all 22,416 attention tensor checks finite

Attention tensor checks are individual torch.isfinite(tensor).all() results recorded for Q/K/V/output and
dO/dQ/dK/dV on sliding d256 calls during the instrumented window; the 22,416 checks cover both ranks.

The first failing edge was rank1:layer33:call45: finite dO entered the sliding attention backward and produced
non-finite dQ/dK/dV. Operator traces mapped sliding d256 to cuDNN SDPA forward/backward on H100, but to
EfficientAttention in both the H100 no-cuDNN controls and A100 baseline.

This strongly supports an SM90 Gemma4 sliding-window BF16 cuDNN SDPA backward numerical issue in this
historical stack.

Proposed resolution

1. Deprioritize or disable cuDNN SDPA specifically for affected Gemma4 training calls on Hopper.

   Prefer a scoped torch.nn.attention.sdpa_kernel(...) priority containing EfficientAttention and Math
   fallback over setting a process-wide environment variable or replacing the global "sdpa" registry entry.

2. Prefer the existing FFPA backend for supported Gemma4 recipes:

   model:
     attn_implementation: ffpa
   

   AutoModel’s FFPA integration already converts sliding/non-causal mask functions into FlexAttention
   BlockMask objects and routes them directly to FlexAttention. Eligible head_dim=512 full-attention layers
   use FFPA. Discussion #2928 reports that
   FlexAttention was noticeably faster than SDPA for these sliding-window layers in its end-to-end Gemma4
   tests. See PR #2436 and
   ffpa_attention.py.

3. Draft PR #2984 adds FlashAttention-2 as an alternative
   to FlexAttention for supported Gemma4 sliding-window context-parallel ring layers. This route also bypasses
   the affected cuDNN SDPA kernel and avoids this issue.

4. Keep FP32 SDPA as a fallback only if BF16 EfficientAttention/FlexAttention still reproduces the NaN.