diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index 984b7b5445..17b5aff61e 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -325,6 +325,7 @@ def _apply_mxfp8_training(transformer, on_event) -> bool: try: from torchao.prototype.mx_formats import MXLinearConfig from torchao.quantization import quantize_ + quantize_( transformer, MXLinearConfig.from_recipe_name("mxfp8_cublas"), @@ -332,9 +333,7 @@ def _apply_mxfp8_training(transformer, on_event) -> bool: ) return True except Exception as exc: # noqa: BLE001 -- mxfp8 is an optimisation, never fatal - _emit( - on_event, "warning", message = f"mxfp8 training unavailable, using bf16 compute: {exc}" - ) + _emit(on_event, "warning", message = f"mxfp8 training unavailable, using bf16 compute: {exc}") return False diff --git a/studio/backend/core/training/diffusion_lora_trainer.py b/studio/backend/core/training/diffusion_lora_trainer.py index 0004b3680f..d22b7f737e 100644 --- a/studio/backend/core/training/diffusion_lora_trainer.py +++ b/studio/backend/core/training/diffusion_lora_trainer.py @@ -360,6 +360,7 @@ def run_diffusion_lora_training( # with a warning event). The U-Net is a dense bf16 base here, the combination that # wrapper compiles under "auto". from core.training.diffusion_dit_trainer import _maybe_compile_transformer + compiled = _maybe_compile_transformer( unet, cfg, False, device, on_event, base_precision = "bf16" ) @@ -550,9 +551,7 @@ def run_diffusion_lora_training( peak_gb = round(torch.cuda.max_memory_allocated() / 1e9, 2) per_step = cfg.train_batch_size * cfg.gradient_accumulation_steps if t_steady is not None and done > 1: - samples_per_second = round( - (done - 1) * per_step / max(now - t_steady, 1e-6), 3 - ) + samples_per_second = round((done - 1) * per_step / max(now - t_steady, 1e-6), 3) else: samples_per_second = round(done * per_step / max(now - t_start, 1e-6), 3) _emit( diff --git a/studio/backend/core/training/diffusion_train_common.py b/studio/backend/core/training/diffusion_train_common.py index dc205ebf5b..c867fde7b1 100644 --- a/studio/backend/core/training/diffusion_train_common.py +++ b/studio/backend/core/training/diffusion_train_common.py @@ -335,9 +335,7 @@ class DiffusionLoraConfig: raise ValueError("compile_transformer must be one of off / on / auto") base_precision = str(self.base_precision or "nf4").strip().lower() if base_precision not in ("nf4", "bf16", "int8", "fp8", "mxfp8", "auto"): - raise ValueError( - "base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto" - ) + raise ValueError("base_precision must be one of nf4 / bf16 / int8 / fp8 / mxfp8 / auto") if base_precision in ("bf16", "int8", "fp8", "mxfp8"): if repo_is_prequantized(self.base_model): raise ValueError( diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py index 1ae037feb1..b2dbc969e4 100644 --- a/studio/backend/tests/test_diffusion_dit_trainer.py +++ b/studio/backend/tests/test_diffusion_dit_trainer.py @@ -114,7 +114,6 @@ def test_family_train_infos_sdxl_supports_compile_without_precision_modes(monkey # ── mxfp8 base precision (DiT dense speed mode) ─────────────────────────────── def _linear(in_features, out_features): import torch.nn as nn - return nn.Linear(in_features, out_features) @@ -172,7 +171,6 @@ def _patch_capability(monkeypatch, capability): # Drive train_precision_modes' GPU probe: pretend CUDA is present at the given tensor # core capability (fp8 needs sm89+, mxfp8 needs sm100+). import torch - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)