From d6795ed077e42c2984d8e75e7c93cc682b26aa81 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:27:58 +0000 Subject: [PATCH 1/2] Restore pre-Ampere bf16 fail-fast in the DiT trainer The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype error instead of a clear message. Restores parity with the SDXL trainer. --- studio/backend/core/training/diffusion_dit_trainer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index af443059aa..c27f1f6fef 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -812,6 +812,13 @@ def run_dit_lora_training( device = "cuda" if torch.cuda.is_available() else "cpu" # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is # unsupported for real runs but keeps import/unit tests architecture-agnostic). + # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the run + # would otherwise die deep in model load with an opaque dtype error. + if device == "cuda" and not torch.cuda.is_bf16_supported(): + raise ValueError( + "This trainer requires a bfloat16-capable GPU (Ampere or newer); " + "this CUDA device does not support bf16." + ) weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32 _assert_trusted_base_model(cfg.base_model) From 7a05d8655b8b2e29676b7e7fdbbb62802f8a36f5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:30:58 +0000 Subject: [PATCH 2/2] Stream every DiT through group offload, not just the primary transformer A dual-DiT pipeline (Ideogram 4's unconditional tower) placed its second denoiser resident under the group tier, which defeats the tier since the pair rarely fits where one alone did not. Stream transformer_2 and unconditional_transformer alongside the transformer and keep only the smaller companions resident. --- studio/backend/core/inference/diffusion_memory.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index c56d1592ce..4c4a3cca2c 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -502,6 +502,16 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: import torch from diffusers.hooks import apply_group_offloading + # A dual-DiT pipeline (e.g. Ideogram 4's unconditional tower) carries a second + # denoiser as large as the first; leaving it resident would defeat this tier + # (the pair rarely fits where one alone did not). Stream every DiT and keep + # only the genuinely smaller companions resident. + streamed: dict[str, Any] = {"transformer": transformer} + for extra in ("transformer_2", "unconditional_transformer"): + module = getattr(pipe, extra, None) + if isinstance(module, torch.nn.Module): + streamed[extra] = module + onload = torch.device(device) use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA gkwargs: dict[str, Any] = { @@ -531,11 +541,12 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: # load-time crash. The streamed transformer manages its own placement via the # offloading hooks applied next. for name, comp in getattr(pipe, "components", {}).items(): - if name == "transformer": + if name in streamed: continue if isinstance(comp, torch.nn.Module): comp.to(onload) - apply_group_offloading(transformer, **gkwargs) + for module in streamed.values(): + apply_group_offloading(module, **gkwargs) return True except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if logger is not None: