make mxfp8 dim1 cast kernel configurable

stack-info: PR: #1427, branch: danielvegamyhre/stack/1

Author: danielvegamyhre

torchtitan/components/quantization/mx.py

@@ -40,10 +40,12 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
                 "torchao is not installed. Please install it to use MXFP8 linear layers."
             )
         torchao_version = version("torchao")
-        mxfp8_min_version = "0.11.0"
-        if torchao_version < mxfp8_min_version:
+
+        # Last torchao release was 0.12.0, so nightly build starts with 0.13.0+git...
+        is_nightly_build = torchao_version.startswith("0.13.0")
+        if not is_nightly_build:
             raise ImportError(
-                f"torchao version {torchao_version} is too old, please install torchao {mxfp8_min_version} or later and try again"
+                f"torchao version {torchao_version} is too old, please install torchao nightly build and try again"
             )
 
         # Can be removed if we enable the emulated versions
@@ -56,12 +58,17 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
         self.filter_fqns = mx_job_config.filter_fqns
 
         # Configure MXFP8
-        from torchao.prototype.mx_formats.config import MXLinearConfig
+        from torchao.prototype.mx_formats.config import (
+            MXFP8Dim1CastKernelChoice,
+            MXLinearConfig,
+        )
 
         config = MXLinearConfig.from_recipe_name(NAME_MAP[mx_job_config.recipe_name])
-        config.use_fp8_dim1_cast_triton_kernel = (
-            mx_job_config.use_fp8_dim1_cast_triton_kernel
-        )
+
+        # String to enum
+        config.mxfp8_dim1_cast_kernel_choice = MXFP8Dim1CastKernelChoice[
+            mx_job_config.mxfp8_dim1_cast_kernel_choice.upper()
+        ]
         self.config = config
 
         logger.info(f"Float8 training active with recipe {mx_job_config.recipe_name}")

torchtitan/config/job_config.py

@@ -534,7 +534,7 @@ class Float8:
 
 @dataclass
 class MX:
-    use_fp8_dim1_cast_triton_kernel: bool = True
+    mxfp8_dim1_cast_kernel_choice: Literal["triton", "cuda", "torch"] = "triton"
     """Temp work around for inductor performance gap"""
 
     recipe_name: Literal["mxfp8"] = "mxfp8"