make mx recipe name more generic (#1512)

Summary:

Instead of maintaining a mapping in torchtitan with valid mx recipe
names, just pass the string recipe directly to torchao. This way torchao
can iterate on recipes without any changes to torchtitan to use those
recipes. Note that appropriate error messages will be thrown from
torchao if user specifies an invalid config name, so there is no need to
duplicate them in torchtitan.

Test Plan:

```bash
with-proxy CONFIG_FILE="torchtitan/models/llama3/train_configs/debug_model.toml" ./run_train.sh --model.print_after_conversion --training.compile --training.steps 50 --model.converters mx --mx.recipe_name "mxfp8_cublas_rceil"
```

Reviewers:

Subscribers:

Tasks:

Tags:

Author: vkuzo

torchtitan/components/quantization/mx.py

@@ -22,9 +22,6 @@
 
 from .utils import module_filter_fn
 
-# Maps titan recipe names to torchao mx recipe names
-NAME_MAP = {"mxfp8": "mxfp8_cublas"}
-
 
 class MXConverter(ModelConverter):
     """Converts the linear layers of `model` to `MXLinear`."""
@@ -76,7 +73,7 @@ def __init__(self, job_config: JobConfig, parallel_dims: ParallelDims):
         )
 
         mx_job_config: MX = job_config.mx
-        config = MXLinearConfig.from_recipe_name(NAME_MAP[mx_job_config.recipe_name])
+        config = MXLinearConfig.from_recipe_name(mx_job_config.recipe_name)
         config.mxfp8_dim1_cast_kernel_choice = MXFP8Dim1CastKernelChoice[
             mx_job_config.mxfp8_dim1_cast_kernel_choice.upper()
         ]

torchtitan/config/job_config.py

@@ -562,8 +562,11 @@ class MX:
     mxfp8_dim1_cast_kernel_choice: Literal["triton", "cuda", "torch"] = "triton"
     """Temp work around for inductor performance gap"""
 
-    recipe_name: Literal["mxfp8"] = "mxfp8"
-    """If specified, creates float8 config from recipe name"""
+    recipe_name: str = "mxfp8_cublas"
+    """
+    If specified, creates MX config from recipe name. See
+    https://github.com/pytorch/ao/tree/main/torchao/prototype/mx_formats for more information.
+    """
 
     filter_fqns: list[str] = field(default_factory=lambda: ["output"])
     """