From 07128c784a10d69139ebf6835a70c936d2d962e1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 3 Jul 2026 08:46:16 +0000 Subject: [PATCH 1/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../core/training/diffusion_dit_trainer.py | 49 +++++++++++++++---- .../core/training/diffusion_lora_trainer.py | 10 +++- .../core/training/diffusion_train_common.py | 11 ++--- .../tests/test_diffusion_train_perf.py | 8 +-- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 4da16d764a..047583f881 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -258,7 +258,12 @@ def _flux_encode_latent_stats(vae, pixel_values): return (dist.mean - vae.config.shift_factor) * scale, dist.std * scale -def _flux_collate(entries, device, weight_dtype, pad_to = None): +def _flux_collate( + entries, + device, + weight_dtype, + pad_to = None, +): import torch # FLUX embeds are fixed-length (encode_prompt pads to max_sequence_length), so a plain @@ -381,7 +386,12 @@ def _qwen_encode_latent_stats(vae, pixel_values): return (dist.mean - mean) / std, dist.std / std -def _qwen_collate(entries, device, weight_dtype, pad_to = None): +def _qwen_collate( + entries, + device, + weight_dtype, + pad_to = None, +): import torch import torch.nn.functional as F @@ -482,7 +492,12 @@ def _zimage_encode_latent_stats(vae, pixel_values): return _zimage_encode_latents(vae, pixel_values), None -def _zimage_collate(entries, device, weight_dtype, pad_to = None): +def _zimage_collate( + entries, + device, + weight_dtype, + pad_to = None, +): caps = [e[0].to(device = device, dtype = weight_dtype) for e in entries] return (caps,) @@ -647,7 +662,7 @@ def _build_latent_cache(spec, vae, image_paths, cfg, device, weight_dtype, on_ev total = len(image_paths) for i, path in enumerate(image_paths): variants = [] - for (u_left, u_top, flip) in plan[i]: + for u_left, u_top, flip in plan[i]: px = ( _load_pixel_tensor_planned( path, cfg.resolution, cfg.center_crop, u_left, u_top, flip @@ -706,7 +721,9 @@ def _maybe_compile_transformer(transformer, cfg, base_is_bnb, device, on_event) fn = getattr(transformer, "compile_repeated_blocks", None) if not callable(fn): - _emit(on_event, "warning", message = "torch.compile unavailable for this model; running eager.") + _emit( + on_event, "warning", message = "torch.compile unavailable for this model; running eager." + ) return False try: dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) @@ -792,7 +809,14 @@ def run_dit_lora_training( perf_snap = _apply_perf_flags(cfg, device) try: return _train_dit( - cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_stop, + cfg, + spec, + pairs, + rng, + device, + weight_dtype, + on_event, + _check_stop, lambda: save_on_stop, ) finally: @@ -841,8 +865,12 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto ) if latent_cache is None: # stopped during the cache build; nothing trained yet _emit( - on_event, "complete", output_dir = str(out_dir), lora_path = None, - stopped = True, steps_run = 0, + on_event, + "complete", + output_dir = str(out_dir), + lora_path = None, + stopped = True, + steps_run = 0, ) return str(out_dir) try: @@ -952,7 +980,9 @@ def _train_dit(cfg, spec, pairs, rng, device, weight_dtype, on_event, _check_sto noisy = (1.0 - sigmas) * latents + sigmas * noise embeds = spec.collate( - [caption_embeds[captions[i]] for i in idxs], device, weight_dtype, + [caption_embeds[captions[i]] for i in idxs], + device, + weight_dtype, pad_to = qwen_pad_to, ) with autocast: @@ -1027,6 +1057,7 @@ def _make_optimizer(params, lr): regression for LoRA -- else torch AdamW, fused on CUDA (with a fallback when this build/device lacks the fused kernel).""" import torch + try: import bitsandbytes as bnb return bnb.optim.AdamW8bit(params, lr = lr) diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index ceb3802ea7..7c7ca79d51 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -201,7 +201,7 @@ def _build_sdxl_latent_cache( total = len(image_paths) for i, path in enumerate(image_paths): variants = [] - for (u_left, u_top, flip) in plan[i]: + for u_left, u_top, flip in plan[i]: tensor, time_ids = _load_image_tensor_planned( path, cfg.resolution, cfg.center_crop, u_left, u_top, flip ) @@ -390,7 +390,13 @@ def run_diffusion_lora_training( latent_cache = None if use_cache: latent_cache = _build_sdxl_latent_cache( - vae, vae_scale, [p for p, _ in pairs], cfg, device, weight_dtype, on_event, + vae, + vae_scale, + [p for p, _ in pairs], + cfg, + device, + weight_dtype, + on_event, _check_stop, ) if latent_cache is None: # stopped during the cache build; nothing trained yet diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index f1989f6305..5b035cd57e 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -376,11 +376,7 @@ def _emit(on_event: Optional[EventCb], type_: str, **kw: Any) -> None: def _plan_cache_variants( - num_images: int, - cache_variants: int, - center_crop: bool, - random_flip: bool, - seed: int, + num_images: int, cache_variants: int, center_crop: bool, random_flip: bool, seed: int ) -> list[list[tuple[float, float, bool]]]: """Seed-deterministic crop/flip plan for the latent cache: per image, up to ``cache_variants`` draws of (u_left, u_top, flip) with the crop as unit fractions the @@ -405,7 +401,9 @@ def _plan_cache_variants( def _apply_perf_flags( - cfg: "DiffusionLoraConfig", device: str, cudnn_benchmark: bool = False + cfg: "DiffusionLoraConfig", + device: str, + cudnn_benchmark: bool = False, ) -> dict: """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision (under ``cfg.enable_tf32``), plus cudnn autotuning when the caller opts in. Autotune is @@ -444,7 +442,6 @@ def _restore_perf_flags(snap: Optional[dict]) -> None: if snap.get("matmul_precision"): try: import torch - torch.set_float32_matmul_precision(snap["matmul_precision"]) except Exception: # noqa: BLE001 -- best-effort restore pass diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index b8d805ead2..ecfe966b5f 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -82,9 +82,7 @@ def test_plan_cache_variants_deterministic_and_deduped(): # ── per-family collate fns ──────────────────────────────────────────────────── def test_flux_collate_shapes(): # FLUX embeds are fixed length: 3 entries batch by a plain cat; text_ids are shared. - entries = [ - (torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3) - ] + entries = [(torch.randn(1, 512, 32), torch.randn(1, 16), torch.randn(512, 3)) for _ in range(3)] pe, pooled, text_ids = _flux_collate(entries, "cpu", torch.float32) assert pe.shape == (3, 512, 32) assert pooled.shape == (3, 16) @@ -239,9 +237,7 @@ def test_service_stop_save_flag(): # ── preparing / warning events + stopped completion messages ────────────────── def test_apply_event_preparing_and_warning(): svc = DiffusionTrainingService() - svc._apply_event( - {"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8} - ) + svc._apply_event({"type": "preparing", "stage": "cache_latents", "done": 4, "total": 8}) st = svc.status() assert st["status"] == "running" assert st["in_model_load"] is True From fabd930c39bdda3aa661590cbcab285bc82209aa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 3 Jul 2026 09:39:43 +0000 Subject: [PATCH 2/2] Clear TF32 flags when enable_tf32 is off so the opt-out is strict fp32 --- .../core/training/diffusion_train_common.py | 9 ++++++- .../tests/test_diffusion_train_perf.py | 27 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index 5b035cd57e..234bbdd0c4 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -406,7 +406,8 @@ def _apply_perf_flags( cudnn_benchmark: bool = False, ) -> dict: """Set the run-scoped torch backend knobs: TF32 matmuls + high fp32 matmul precision - (under ``cfg.enable_tf32``), plus cudnn autotuning when the caller opts in. Autotune is + when ``cfg.enable_tf32`` is on, strict fp32 (all TF32 flags cleared) when it is off, + plus cudnn autotuning when the caller opts in. Autotune is for the conv-heavy SDXL U-Net only: measured on B200, it DOUBLES peak VRAM (fp32 VAE conv workspaces) while the DiT loop -- pure matmuls once the latent cache is built -- gains nothing from it. Returns a snapshot for ``_restore_perf_flags``. Best-effort: @@ -424,6 +425,12 @@ def _apply_perf_flags( torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision("high") + else: + # The opt-out is a strict-fp32 A/B mode, so actively clear the flags rather + # than inherit ambient state (cudnn TF32 defaults to ON in torch). + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") if cudnn_benchmark: torch.backends.cudnn.benchmark = True except Exception: # noqa: BLE001 -- perf flags are never fatal diff --git a/studio/backend/tests/test_diffusion_train_perf.py b/studio/backend/tests/test_diffusion_train_perf.py index ecfe966b5f..4339fbb0e7 100644 --- a/studio/backend/tests/test_diffusion_train_perf.py +++ b/studio/backend/tests/test_diffusion_train_perf.py @@ -328,3 +328,30 @@ def test_perf_flags_cpu_roundtrip(): snap = _apply_perf_flags(_cfg(), "cpu") assert isinstance(snap, dict) _restore_perf_flags(snap) # no exception + + +def test_perf_flags_tf32_off_clears_flags(): + # enable_tf32=False is the strict-fp32 A/B mode: it must actively clear the TF32 flags + # (cudnn TF32 defaults ON in torch) rather than inherit ambient state, and restore must + # put the ambient values back. The flag attributes are plain Python state, present and + # settable on CPU-only torch builds, so this runs without a GPU. + import torch + + before = ( + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) + snap = _apply_perf_flags(_cfg(enable_tf32 = False), "cuda") + try: + assert torch.backends.cuda.matmul.allow_tf32 is False + assert torch.backends.cudnn.allow_tf32 is False + assert torch.get_float32_matmul_precision() == "highest" + finally: + _restore_perf_flags(snap) + after = ( + torch.backends.cuda.matmul.allow_tf32, + torch.backends.cudnn.allow_tf32, + torch.get_float32_matmul_precision(), + ) + assert after == before