Comfy-Org / Comfy-Org/ComfyUI

Apple Silicon (MPS): random all-NaN attention -> black videos with bf16 models (LTX-2.x); two suggested fixes

Open Beginner friendly
#15,804 0 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

## Summary

On Apple Silicon (MPS, macOS ≥ 14.5), video generation with **bf16** checkpoints (all current LTX-2.x models) intermittently collapses mid-sampling: a single transformer forward returns all-NaN hidden states, after which every remaining sampler step and both VAE decodes stay NaN — the saved video is solid black (audio waveform included).

Observed failure rates on one machine (Apple M5 Max, macOS 26.6.2, torch 2.13.0, ComfyUI v0.33.3):

- LTX-2.5 t2v workflow: **19 of 26 runs** produced all-black videos
- LTX-2.5 i2v workflow: **8 of 8 runs** before mitigation

Same seed reproduces bit-identically; different seeds fail stochastically — the failure depends on memory contents, not on model math.

Instrumenting the LTXAV transformer with per-stage finite checks shows the first non-finite values appear as the **entire output of a single attention op while its inputs are still finite** (e.g. `audio_to_video_attn` returning `nan=14417920/14417920`), consistent with garbage being read from uninitialized memory inside attention.

## Root cause 1: sub-quadratic attention relies on `beta=0` semantics that MPS does not honor

`comfy/ldm/modules/sub_quadratic_attention.py` computes attention scores as:

```python
attn_weights = torch.baddbmm(
torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
query, key_t, alpha=scale, beta=0,
)
```

PyTorch documents that with `beta=0` the input tensor is ignored ("nan and infinity in it will not be propagated"). MPS violates this today:

```python
import torch
q = torch.randn(1, 64, 128, device="mps", dtype=torch.bfloat16)
k = torch.randn(1, 128, 3520, device="mps", dtype=torch.bfloat16)
nan_in = torch.full((1, 1, 1), float("nan"), device="mps", dtype=torch.bfloat16)
out = torch.baddbmm(nan_in, q, k, alpha=0.088, beta=0)
print(int(torch.isnan(out).sum()), "/", out.numel()) # MPS: 225280/225280 ; CPU: 0
```

Since `torch.empty(1, 1, 1)` may contain NaN bit patterns, any run can get its whole score matrix poisoned by one broadcast scalar.

**Suggested fix** (verified locally): initialize with `zeros` instead of `empty` — 4 occurrences (`_summarize_chunk` lines ~71/79 and `_get_attention_scores_no_kv_chunking` lines ~152/160). No measurable cost.

## Root cause 2: macOS fp32-attention workaround covers fp16 but not bf16

`comfy/model_management.py::force_upcast_attention_dtype()` already works around a known macOS black-image bug:

```python
if macos_version is not None and ((14, 5) <= macos_version):
upcast = True # black image bug on recent versions of macOS ...
if upcast:
return {torch.float16: torch.float32}
```

`bfloat16` is missing from the map, so fully-bf16 pipelines (all current LTX-2.x checkpoints) bypass the protection and keep running attention through the fragile path above.

**Suggested fix**: `{torch.float16: torch.float32, torch.bfloat16: torch.float32}`.

## Verification after applying both fixes locally

- 3 consecutive seeds that previously collapsed generated valid videos (~1 MB vs ~120 KB black files); step-level NaN probes report zero non-finite values across base sampling (8 steps), latent upsampler, refine stage (3 steps) and both VAE decodes.
- Workflow details: two-stage ManualSigmas + euler_ancestral + DualCFGGuider [1,1], joint audio-video latent; failure signature was identical for t2v and i2v, confirming the shared attention path as the culprit rather than anything i2v-specific.

Contributor guide

Open the contributing guide

Research direction

Inspect comfy/ldm/modules/sub_quadratic_attention.py at the four baddbmm initializations in _summarize_chunk and _get_attention_scores_no_kv_chunking, then review force_upcast_attention_dtype() in comfy/model_management.py. Verify the changes on Apple Silicon with bf16 LTX-2.x workflows, confirming that sampling and both VAE decodes produce finite values and valid non-black videos.

Written by the indexing model from the issue text.

Assessment

Tech stack
macos, python, pytorch
Domain
ai, machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
82/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.