From b8da49118dc79db686bdd2e1de29f9b0f9c44a94 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 23:55:41 +0000 Subject: [PATCH] Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16 The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 / RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported() and drop to fp16 there. --- studio/backend/core/training/diffusion_lora_trainer.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index f798496f43..ee095e63d2 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -261,9 +261,12 @@ def run_diffusion_lora_training( torch.manual_seed(cfg.seed) device = "cuda" if torch.cuda.is_available() else "cpu" - weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[ - cfg.mixed_precision if device == "cuda" else "no" - ] + precision = cfg.mixed_precision if device == "cuda" else "no" + if precision == "bf16" and device == "cuda" and not torch.cuda.is_bf16_supported(): + # The default is bf16, but pre-Ampere GPUs (T4 / V100 / RTX 20xx) have no + # bf16 compute; fall back to fp16 there instead of failing at load/forward. + precision = "fp16" + weight_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "no": torch.float32}[precision] pairs = discover_image_caption_pairs( cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column