Add support for MXFP8 recipe in accelerate

Author: pstjohn

src/accelerate/test_utils/__init__.py

@@ -53,6 +53,7 @@
     require_torchvision,
     require_tpu,
     require_transformer_engine,
+    require_transformer_engine_mxfp8,
     require_xpu,
     run_first,
     skip,

src/accelerate/test_utils/testing.py

@@ -71,6 +71,7 @@
     is_torchvision_available,
     is_trackio_available,
     is_transformer_engine_available,
+    is_transformer_engine_mxfp8_available,
     is_transformers_available,
     is_triton_available,
     is_wandb_available,
@@ -540,6 +541,13 @@ def require_transformer_engine(test_case):
     return unittest.skipUnless(is_transformer_engine_available(), "test requires transformers engine")(test_case)
 
 
+def require_transformer_engine_mxfp8(test_case):
+    """
+    Decorator marking a test that requires transformers engine FP8 block scaling available. These tests are skipped when transformers
+    engine FP8 block scaling isn't available
+    """
+    return unittest.skipUnless(is_transformer_engine_mxfp8_available(), "test requires transformers engine FP8 block scaling")(test_case)
+
 def require_torchao(test_case):
     """
     Decorator marking a test that requires torchao installed. These tests are skipped when torchao isn't installed

src/accelerate/utils/__init__.py

@@ -131,6 +131,7 @@
     is_torchvision_available,
     is_trackio_available,
     is_transformer_engine_available,
+    is_transformer_engine_mxfp8_available,
     is_transformers_available,
     is_triton_available,
     is_wandb_available,

src/accelerate/utils/dataclasses.py

@@ -359,6 +359,7 @@ class TERecipeKwargs(KwargsHandler):
     amax_history_len: int = None
     amax_compute_algo: AmaxComputeAlgorithm = None
     override_linear_precision: tuple[bool, bool, bool] = None
+    use_mxfp8_block_scaling: bool = None
 
     def __post_init__(self):
         env_prefix = "ACCELERATE_FP8_"
@@ -387,6 +388,8 @@ def __post_init__(self):
             dgrad = parse_flag_from_env(env_prefix + "OVERRIDE_DGRAD")
             wgrad = parse_flag_from_env(env_prefix + "OVERRIDE_WGRAD")
             self.override_linear_precision = (fprop, dgrad, wgrad)
+        if self.use_mxfp8_block_scaling is None:
+            self.use_mxfp8_block_scaling = parse_flag_from_env(env_prefix + "USE_MXFP8_BLOCK_SCALING")
 
 
 @dataclass

src/accelerate/utils/imports.py

@@ -114,6 +114,12 @@ def is_transformer_engine_available():
         return _is_package_available("transformer_engine", "transformer-engine")
 
 
+def is_transformer_engine_mxfp8_available():
+    if _is_package_available("transformer_engine", "transformer-engine"):
+        import transformer_engine.pytorch as te
+        return te.fp8.check_mxfp8_support()[0]
+    return False
+
 def is_lomo_available():
     return _is_package_available("lomo_optim")
 

src/accelerate/utils/transformer_engine.py

@@ -146,14 +146,32 @@ def apply_fp8_autowrap(model, fp8_recipe_handler):
 
     if is_hpu_available():
         import intel_transformer_engine.recipe as te_recipe
+        is_fp8_block_scaling_available = False
+        message = "MXFP8 block scaling is not available on HPU."
+
     else:
         import transformer_engine.common.recipe as te_recipe
+        import transformer_engine.pytorch as te
+        is_fp8_block_scaling_available, message = te.fp8.check_mxfp8_support()
 
     kwargs = fp8_recipe_handler.to_kwargs() if fp8_recipe_handler is not None else {}
     if "fp8_format" in kwargs:
         kwargs["fp8_format"] = getattr(te_recipe.Format, kwargs["fp8_format"])
     use_during_eval = kwargs.pop("use_autocast_during_eval", False)
-    fp8_recipe = te_recipe.DelayedScaling(**kwargs)
+    use_mxfp8_block_scaling = kwargs.pop("use_mxfp8_block_scaling", False)
+
+    if use_mxfp8_block_scaling and not is_fp8_block_scaling_available:
+        raise ValueError(f"MXFP8 block scaling is not available: {message}")
+
+    if use_mxfp8_block_scaling:
+        if "amax_compute_algo" in kwargs:
+            raise ValueError("`amax_compute_algo` is not supported for MXFP8 block scaling.")
+        if "amax_history_len" in kwargs:
+            raise ValueError("`amax_history_len` is not supported for MXFP8 block scaling.")
+        fp8_recipe = te_recipe.MXFP8BlockScaling(**kwargs)
+    else:
+        fp8_recipe = te_recipe.DelayedScaling(**kwargs)
+
     new_forward = contextual_fp8_autocast(model.forward, fp8_recipe, use_during_eval)
 
     if hasattr(model.forward, "__func__"):

tests/test_fp8.py

@@ -31,6 +31,7 @@
     require_multi_device,
     require_torchao,
     require_transformer_engine,
+    require_transformer_engine_mxfp8,
     run_first,
 )
 from accelerate.test_utils.testing import require_deepspeed, run_command
@@ -109,6 +110,27 @@ def test_can_prepare_model_single_gpu_from_config(self):
             command += ["-m", "tests.test_fp8", "--test_te", "--from_config"]
             run_command(command)
 
+
+    @require_transformer_engine_mxfp8
+    def test_can_prepare_model_with_mxfp8_block_scaling(self):
+        with tempfile.TemporaryDirectory() as dir_name:
+            config_file = Path(dir_name) / "config.yaml"
+            config_file.write_text(
+                textwrap.dedent(
+                    """
+                    distributed_type: "NO"
+                    num_processes: 1
+                    mixed_precision: fp8
+                    fp8_config:
+                      backend: TE
+                      use_mxfp8_block_scaling: true
+                    """
+                )
+            )
+            command = get_launch_command(config_file=str(config_file), monitor_interval=0.1)
+            command += ["-m", "tests.test_fp8", "--test_te", "--from_config"]
+            run_command(command)
+
     @require_multi_device
     def test_can_prepare_model_multi_gpu(self):
         command = get_launch_command(num_processes=2, monitor_interval=0.1)