Support the torchao 0.17 mxfp8 recipe API

torchao 0.17 removed MXLinearConfig from prototype.mx_formats in favour of
MXFP8TrainingOpConfig.from_recipe shared with MoE training. _mxfp8_training_config
tries the 0.16 API first and falls back to the 0.17 one; both feed quantize_.
mxfp8 still degrades to bf16 with a warning when neither import resolves

studio/backend/core/training/diffusion_dit_trainer.py

@@ -314,6 +314,22 @@ def _mx_module_filter(mod, fqn: str) -> bool:
     return mod.in_features % 32 == 0 and mod.out_features % 32 == 0
 
 
+def _mxfp8_training_config():
+    """The torchao MX training config across the prototype API's revisions: torchao 0.16
+    ships ``MXLinearConfig`` in ``prototype.mx_formats``; 0.17 removed it in favour of the
+    ``MXFP8TrainingOpConfig`` recipe API shared with MoE training. Both feed ``quantize_``.
+    Raises ImportError when neither API exists (mxfp8 then falls back to bf16)."""
+    try:
+        from torchao.prototype.mx_formats import MXLinearConfig
+        return MXLinearConfig.from_recipe_name("mxfp8_cublas")
+    except ImportError:
+        from torchao.prototype.moe_training.config import (
+            MXFP8TrainingOpConfig,
+            MXFP8TrainingRecipe,
+        )
+        return MXFP8TrainingOpConfig.from_recipe(MXFP8TrainingRecipe.MXFP8_RCEIL)
+
+
 def _apply_mxfp8_training(transformer, on_event) -> bool:
     """Swap the frozen base linears to torchao MX float8 training compute (mxfp8, the
     Blackwell-native block-scaled format; the swap is in place and the weights stay bf16
@@ -323,12 +339,11 @@ def _apply_mxfp8_training(transformer, on_event) -> bool:
     or batch), which is why it stays an explicit opt-in rather than an "auto" pick.
     Never fatal: on any failure the run continues in bf16 with a warning."""
     try:
-        from torchao.prototype.mx_formats import MXLinearConfig
         from torchao.quantization import quantize_
 
         quantize_(
             transformer,
-            MXLinearConfig.from_recipe_name("mxfp8_cublas"),
+            _mxfp8_training_config(),
             filter_fn = _mx_module_filter,
         )
         return True

studio/backend/tests/test_diffusion_dit_trainer.py

@@ -156,9 +156,11 @@ def test_should_compile_auto_mxfp8_on_cuda():
 
 
 def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
-    # An unavailable torchao MX path must never be fatal: force the import to raise, then
-    # assert the helper returns False and emits exactly one warning naming mxfp8.
+    # An unavailable torchao MX path must never be fatal: force both API revisions'
+    # imports to raise, then assert the helper returns False and emits exactly one
+    # warning naming mxfp8.
     monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None)
+    monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", None)
     events = []
     ok = _apply_mxfp8_training(object(), lambda e: events.append(e))
     assert ok is False
@@ -167,6 +169,38 @@ def test_apply_mxfp8_training_failure_falls_back_with_warning(monkeypatch):
     assert "mxfp8" in warnings[0]["message"]
 
 
+def test_mxfp8_training_config_falls_back_to_the_torchao_0_17_api(monkeypatch):
+    # torchao 0.17 removed prototype.mx_formats.MXLinearConfig in favour of the
+    # MXFP8TrainingOpConfig recipe API; the config helper must fall back to it so the
+    # advertised mxfp8 mode keeps engaging on those installs instead of silently
+    # training dense bf16.
+    from types import SimpleNamespace
+
+    from core.training.diffusion_dit_trainer import _mxfp8_training_config
+
+    calls = {}
+
+    class _Recipe:
+        MXFP8_RCEIL = "mxfp8_rceil"
+
+    class _OpConfig:
+        @staticmethod
+        def from_recipe(recipe):
+            calls["recipe"] = recipe
+            return "cfg-0.17"
+
+    fake_config = SimpleNamespace(
+        MXFP8TrainingOpConfig = _OpConfig, MXFP8TrainingRecipe = _Recipe
+    )
+    monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", None)
+    monkeypatch.setitem(
+        sys.modules, "torchao.prototype.moe_training", SimpleNamespace(config = fake_config)
+    )
+    monkeypatch.setitem(sys.modules, "torchao.prototype.moe_training.config", fake_config)
+    assert _mxfp8_training_config() == "cfg-0.17"
+    assert calls["recipe"] == _Recipe.MXFP8_RCEIL
+
+
 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+).