Validate the training config before importing diffusers and pin the arbiter test's device

The fp16-on-bf16-family refusal in run_dit_lora_training now fires before the heavy
imports, so a host without diffusers gets the real validation error instead of
ModuleNotFoundError. test_in_progress_returns_409_after_validation_passes pins the
resolved device to cuda because the load route only takes the GPU arbiter for non-CPU
loads, which made the ownership assert host-dependent.

studio/backend/core/training/diffusion_dit_trainer.py

@@ -487,6 +487,21 @@ def run_dit_lora_training(
     should_stop: Optional[StopCb] = None,
 ) -> str:
     """Train a flow-matching DiT LoRA (FLUX.1-dev / Qwen-Image / Z-Image) and export it."""
+    cfg = config.normalized()
+    spec = _SPECS.get(cfg.resolved_family)
+    if spec is None:
+        raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
+
+    # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
+    # explicitly asks for fp16 on a bf16-only family is refused rather than silently
+    # upgraded, so the choice is never misrepresented. Validation runs before the heavy
+    # imports so a host without diffusers still sees the real error.
+    if cfg.mixed_precision == "fp16" and spec.force_bf16:
+        raise ValueError(
+            f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
+            f"embedder internals. Set mixed precision to bf16."
+        )
+
     import torch
     import torch.nn.functional as F
     from diffusers import FlowMatchEulerDiscreteScheduler
@@ -494,11 +509,6 @@ def run_dit_lora_training(
     from peft import LoraConfig
     from peft.utils import get_peft_model_state_dict
 
-    cfg = config.normalized()
-    spec = _SPECS.get(cfg.resolved_family)
-    if spec is None:
-        raise ValueError(f"No DiT trainer for family {cfg.resolved_family!r}")
-
     rng = random.Random(cfg.seed)
     torch.manual_seed(cfg.seed)
 
@@ -515,15 +525,6 @@ def run_dit_lora_training(
             save_on_stop = False
         return True
 
-    # DiT families train in bf16 (Z-Image/Qwen require it; FLUX prefers it). A caller that
-    # explicitly asks for fp16 on a bf16-only family is refused rather than silently
-    # upgraded, so the choice is never misrepresented.
-    if cfg.mixed_precision == "fp16" and spec.force_bf16:
-        raise ValueError(
-            f"{spec.family} LoRA training requires bf16: fp16 overflows its fp32 RoPE / "
-            f"embedder internals. Set mixed precision to bf16."
-        )
-
     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).

studio/backend/tests/test_diffusion_routes.py

@@ -659,6 +659,17 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch):
     backend = _FakeBackend()
     backend.begin_load = _busy
     monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
+    # Pin the resolved device to cuda: the route only takes the arbiter for non-CPU
+    # loads, so on a CPU-only host the ownership assert below would never hold.
+    import types as _types
+
+    import core.inference.diffusion_device as devmod
+
+    monkeypatch.setattr(
+        devmod,
+        "resolve_diffusion_device_target",
+        lambda: _types.SimpleNamespace(device = "cuda"),
+    )
     resp = client.post(
         "/api/inference/images/load",
         json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "q.gguf"},