Krea 2: q/k/v dtypes can diverge under gradient checkpointing, breaking backward

On my Ubuntu box with 192 GB ram and a RTX5090 I had problems creating a lora for Krea2. With some debugging with Claude Opus 5 (extra effort) we managed to find the issue. Since I am a creative and not a developer the text below is from Claude, specifically aimed at Ostris as a bug report. . This is my small contribution to this amazing stable diffusion community, I am hugely grateful for everything people like Ostris is doing for all of us.

Summary

With arch: krea2 and gradient_checkpointing: true, training reliably dies in the backward pass within 8-300 steps.

The root cause is that nothing in the Krea 2 attention path forces q, k and v to share a dtype, so under checkpoint recomputation they can disagree and F.scaled_dot_product_attention rejects them.

A two-line fix in attention() resolves it. With the patch, a full 3000-step run at resolution 1024 in fp8 completed without a single failure. Without it, no run survived past step 300.

Environment

• ai-toolkit main @ 561a023
• Ubuntu 24.04.4, Python 3.12.3
• RTX 5090 (sm_120), driver 595.84, CUDA 13.2
• torch 2.12.1+cu130 / torchvision 0.27.1+cu130 (also reproduced on torch 2.13.0+cu130)
• Model krea/Krea-2-Raw, qtype: qfloat8, quantize_te: true, LoRA rank 32, batch size 1, layer_offloading: false

Symptoms

The same underlying problem surfaces in three different ways, non-deterministically.

1. The underlying error, caught in clear text:

File "extensions_built_in/diffusion_models/krea2/src/mmdit.py", line 76, in attention
    x = F.scaled_dot_product_attention(
RuntimeError: Expected query, key, and value to have the same dtype, but got
query.dtype: c10::BFloat16 key.dtype: float and value.dtype: c10::BFloat16

raised from inside torch/utils/checkpoint.py recompute_fn.

2. More commonly, a checkpoint accounting failure:

torch.utils.checkpoint.CheckpointError: A different number of tensors was saved
during the original forward and recomputation.
Number of tensors saved during forward: 100
Number of tensors saved during recomputation: 98

3. Occasionally a hard segfault:

python[4823]: segfault at fffffffffffffff8 ... error 5 in libnvrtc.so.13

All three occur only with gradient_checkpointing: true. With it disabled the model trains fine, memory permitting, which is what pointed at the recomputation path.

Root cause

In extensions_built_in/diffusion_models/krea2/src/mmdit.py:

RMSNorm.forward is dtype-preserving, not dtype-normalising:

t, dtype = x.float(), x.dtype
...
return t.to(dtype)

It computes in fp32 and casts back to whatever dtype came in.

QKNorm then runs q and k through two separate RMSNorm instances and does not touch v at all:

return self.qnorm(q), self.knorm(k), v

ropeapply follows the same pattern: each tensor is restored to its own incoming dtype.

The consequence is that there is no point in the graph where q, k and v are forced to agree. A dtype divergence originating in self.wk(...) propagates untouched all the way to SDPA. This does not show up in the normal forward pass, but under torch.utils.checkpoint recomputation with use_reentrant=False, k can come back as fp32 while q and v are bf16, and SDPA raises.

I suspect, but have not verified, that the 100-vs-98 tensor count mismatch is the same problem seen from the other side. attention() opens an sdpa_kernel context with a backend priority list, so if the original forward selects cuDNN and the recomputation falls back to another backend because k arrived as fp32, the two passes would save a different number of tensors for backward. That would also explain the libnvrtc segfault, since a kernel would then be JIT-compiled on a path the forward pass never took.

Fix

Force the three tensors to a common dtype immediately before SDPA. Patching inside attention() rather than at the call site covers all callers:

     with sdpa_kernel([...], set_priority=True):
+        k = k.to(q.dtype)
+        v = v.to(q.dtype)
         x = F.scaled_dot_product_attention(
             q, k, v, attn_mask=mask, scale=scale, enable_gqa=gqa
         )

Verification

Same dataset (31 images), same config, same machine, repeated across several runs.

Related

Possibly the same root cause as the Krea 2 instability reported in #1032, and adjacent to #954.

Config

job: extension config: name: 0lordis_v3 process:

• type: diffusion_trainer training_folder: /home/funky/ai-toolkit/output sqlite_db_path: /home/funky/ai-toolkit/aitk_db.db device: cuda trigger_word: Nilvara performance_log_every: 10 network: type: lora linear: 32 linear_alpha: 32 lokr_full_rank: true lokr_factor: -1 network_kwargs: ignore_if_contains: [] save: dtype: bf16 save_every: 250 max_step_saves_to_keep: 12 save_format: diffusers push_to_hub: false datasets:
  • folder_path: /home/funky/ai-toolkit/datasets/0lordis_v3 mask_path: null mask_min_value: 0.1 default_caption: '' caption_ext: txt caption_dropout_rate: 0.05 cache_latents_to_disk: true is_reg: false network_weight: 1 resolution:
    • 1024 controls: [] shrink_video_to_frames: true num_frames: 1 flip_x: false flip_y: false num_repeats: 1 train: batch_size: 1 bypass_guidance_embedding: false steps: 3000 gradient_accumulation: 1 train_unet: true train_text_encoder: false gradient_checkpointing: true noise_scheduler: flowmatch optimizer: adamw8bit timestep_type: linear content_or_style: balanced optimizer_params: weight_decay: 0.0001 unload_text_encoder: true cache_text_embeddings: true lr: 0.0001 ema_config: use_ema: false ema_decay: 0.99 skip_first_sample: true force_first_sample: false disable_sampling: true dtype: bf16 diff_output_preservation: false diff_output_preservation_multiplier: 1 diff_output_preservation_class: person switch_boundary_every: 1 loss_type: mse logging: log_every: 1 use_ui_logger: true model: name_or_path: krea/Krea-2-Raw quantize: true qtype: qfloat8 quantize_te: true qtype_te: qfloat8 arch: krea2 low_vram: false model_kwargs: {} compile: false layer_offloading: false layer_offloading_text_encoder_percent: 1 layer_offloading_transformer_percent: 0.5 sample: sampler: flowmatch sample_every: 250 sample_start_step: 0 width: 1024 height: 1024 samples:
    • prompt: photo of [trigger], a man, medium shot, neutral expression, plain background, studio lighting
    • prompt: photo of [trigger], a man smiling, outdoors, natural daylight, shallow depth of field
    • prompt: full body photo of [trigger], a man standing on a city street, overcast day
    • prompt: close up portrait of [trigger], a man, side lighting, dark background neg: '' seed: 42 walk_seed: true guidance_scale: 3.5 sample_steps: 30 num_frames: 1 fps: 1 meta: name: 0lordis_v3 version: '1.0'