Fp8 integration

examples/complete_cv_example.py

@@ -272,7 +272,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/complete_nlp_example.py

@@ -109,9 +109,17 @@ def tokenize_function(examples):
 
     def collate_fn(examples):
         # On TPU it's best to pad everything to the same length or training will be very slow.
-        if accelerator.distributed_type == DistributedType.TPU:
-            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
-        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples, padding="longest", max_length=128, pad_to_multiple_of=pad_to_multiple_of, return_tensors="pt"
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -251,7 +259,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/cv_example.py

@@ -190,7 +190,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

examples/nlp_example.py

@@ -79,9 +79,17 @@ def tokenize_function(examples):
 
     def collate_fn(examples):
         # On TPU it's best to pad everything to the same length or training will be very slow.
-        if accelerator.distributed_type == DistributedType.TPU:
-            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
-        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+        # When using mixed precision we want round multiples of 8/16
+        if accelerator.mixed_precision == "fp8":
+            pad_to_multiple_of = 16
+        elif accelerator.mixed_precision != "no":
+            pad_to_multiple_of = 8
+        else:
+            pad_to_multiple_of = None
+
+        return tokenizer.pad(
+            examples, padding="longest", max_length=128, pad_to_multiple_of=pad_to_multiple_of, return_tensors="pt"
+        )
 
     # Instantiate dataloaders.
     train_dataloader = DataLoader(
@@ -120,7 +128,6 @@ def training_function(config, args):
     # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
     # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
     model = model.to(accelerator.device)
-
     # Instantiate optimizer
     optimizer = AdamW(params=model.parameters(), lr=lr)
 
@@ -134,6 +141,7 @@ def training_function(config, args):
     # Prepare everything
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
     # prepare method.
+
     model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
         model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
     )
@@ -177,7 +185,7 @@ def main():
         "--mixed_precision",
         type=str,
         default=None,
-        choices=["no", "fp16", "bf16"],
+        choices=["no", "fp16", "bf16", "fp8"],
         help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU.",

src/accelerate/accelerator.py

@@ -39,6 +39,7 @@
     DistributedDataParallelKwargs,
     DistributedType,
     DynamoBackend,
+    FP8RecipeKwargs,
     FullyShardedDataParallelPlugin,
     GradScalerKwargs,
     InitProcessGroupKwargs,
@@ -49,12 +50,15 @@
     ProjectConfiguration,
     RNGType,
     compare_versions,
+    convert_model,
     convert_outputs_to_fp32,
     extract_model_from_parallel,
     gather,
     get_pretty_name,
+    has_transformer_engine_layers,
     is_bf16_available,
     is_deepspeed_available,
+    is_fp8_available,
     is_megatron_lm_available,
     is_torch_version,
     is_tpu_available,
@@ -79,6 +83,11 @@
         DummyScheduler,
     )
 
+if is_fp8_available():
+    import transformer_engine.common.recipe as te_recipe
+    from transformer_engine.pytorch import fp8_autocast
+
+
 if is_megatron_lm_available():
     from .utils import (
         MegatronEngine,
@@ -123,10 +132,11 @@ class Accelerator:
             round multiple of the `num_processes` you are using. If `False`, actual batch size used will be the one set
             in your script multiplied by the number of processes.
         mixed_precision (`str`, *optional*):
-            Whether or not to use mixed precision training (fp16 or bfloat16). Choose from 'no','fp16','bf16'. Will
-            default to the value in the environment variable `ACCELERATE_MIXED_PRECISION`, which will use the default
-            value in the accelerate config of the current system or the flag passed with the `accelerate.launch`
-            command. 'fp16' requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+            Whether or not to use mixed precision training (fp16 or bfloat16). Choose from 'no','fp16','bf16 or 'fp8'.
+            Will default to the value in the environment variable `ACCELERATE_MIXED_PRECISION`, which will use the
+            default value in the accelerate config of the current system or the flag passed with the
+            `accelerate.launch` command. 'fp16' requires pytorch 1.6 or higher. 'bf16' requires pytorch 1.10 or higher.
+            'fp8' requires the installation of transformers-engine.
         gradient_accumulation_steps (`int`, *optional*, default to 1):
             The number of steps that should pass before gradients are accumulated. A number > 1 should be combined with
             `Accelerator.accumulate`. If not passed, will default to the value in the environment variable
@@ -298,6 +308,7 @@ def __init__(
         self.ddp_handler = None
         self.scaler_handler = None
         self.init_handler = None
+        self.fp8_recipe_handler = None
         if kwargs_handlers is not None:
             for handler in kwargs_handlers:
                 assert isinstance(
@@ -318,6 +329,11 @@ def __init__(
                         raise ValueError("You can only pass one `InitProcessGroupKwargs` in `kwargs_handler`.")
                     else:
                         self.init_handler = handler
+                elif isinstance(handler, FP8RecipeKwargs):
+                    if self.fp8_recipe_handler is not None:
+                        raise ValueError("You can only pass one `FP8RecipeKwargs` in `kwargs_handler`.")
+                    else:
+                        self.fp8_recipe_handler = handler
 
         kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {}
         self.state = AcceleratorState(
@@ -1046,7 +1062,7 @@ def prepare(self, *args, device_placement=None):
 
         # If we're dealing with device placement, this deals with that by...
         tpu_should_fix_optimizer = self.device_placement and self.distributed_type == DistributedType.TPU
-        if tpu_should_fix_optimizer:
+        if tpu_should_fix_optimizer or self.mixed_precision == "fp8":
             # 1. grabbing old model parameters
             old_named_params = self._get_named_parameters(*args)
 
@@ -1060,7 +1076,7 @@ def prepare(self, *args, device_placement=None):
             )
             result = tuple(self._prepare_one(obj, device_placement=d) for obj, d in zip(result, device_placement))
 
-        if tpu_should_fix_optimizer:
+        if tpu_should_fix_optimizer or self.mixed_precision == "fp8":
             # 2. grabbing new model parameters
             new_named_params = self._get_named_parameters(*result)
             # 3. building a map from the first to the second
@@ -1144,6 +1160,25 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
             else:
                 model.forward = torch.cuda.amp.autocast()(model.forward)
             model.forward = convert_outputs_to_fp32(model.forward)
+        elif self.mixed_precision == "fp8":
+            if not has_transformer_engine_layers(model):
+                with torch.no_grad():
+                    convert_model(model)
+                model._converted_to_transformer_engine = True
+            model._original_forward = model.forward
+
+            kwargs = self.fp8_recipe_handler.to_kwargs() if self.fp8_recipe_handler is not None else {}
+            if "fp8_format" in kwargs:
+                kwargs["fp8_format"] = getattr(te_recipe.Format, kwargs["fp8_format"])
+            fp8_recipe = te_recipe.DelayedScaling(**kwargs)
+            fp8_enabled = torch.cuda.get_device_capability()[0] >= 9
+            if not fp8_enabled:
+                logger.warn(
+                    f"The current device has compute capability of {torch.cuda.get_device_capability()} which is "
+                    "insufficient for FP8 mixed precision training (requires a GPU Hopper or higher, compute "
+                    "capability of 9 or higher). Will using FP16 instead."
+                )
+            model.forward = fp8_autocast(enabled=fp8_enabled, fp8_recipe=fp8_recipe)(model.forward)
         if self.distributed_type == DistributedType.TPU and self.state.fork_launched:
             model = xmp.MpModelWrapper(model).to(self.device)
         return model

src/accelerate/commands/config/cluster.py

@@ -485,7 +485,7 @@ def get_cluster_input():
         else:
             mixed_precision = _ask_options(
                 "Do you wish to use FP16 or BF16 (mixed precision)?",
-                ["no", "fp16", "bf16"],
+                ["no", "fp16", "bf16", "fp8"],
                 _convert_mixed_precision,
             )
 

src/accelerate/commands/config/config_utils.py

@@ -76,7 +76,7 @@ def _convert_dynamo_backend(value):
 
 def _convert_mixed_precision(value):
     value = int(value)
-    return PrecisionType(["no", "fp16", "bf16"][value])
+    return PrecisionType(["no", "fp16", "bf16", "fp8"][value])
 
 
 def _convert_sagemaker_distributed_mode(value):

src/accelerate/utils/__init__.py

@@ -5,6 +5,7 @@
     DistributedDataParallelKwargs,
     DistributedType,
     DynamoBackend,
+    FP8RecipeKwargs,
     FullyShardedDataParallelPlugin,
     GradScalerKwargs,
     InitProcessGroupKwargs,
@@ -28,6 +29,7 @@
     is_comet_ml_available,
     is_datasets_available,
     is_deepspeed_available,
+    is_fp8_available,
     is_megatron_lm_available,
     is_mlflow_available,
     is_mps_available,
@@ -130,3 +132,4 @@
 from .random import set_seed, synchronize_rng_state, synchronize_rng_states
 from .torch_xla import install_xla
 from .tqdm import tqdm
+from .transformer_engine import convert_model, has_transformer_engine_layers

src/accelerate/utils/dataclasses.py

@@ -27,7 +27,7 @@
 from dataclasses import dataclass, field
 from datetime import timedelta
 from distutils.util import strtobool
-from typing import Any, Callable, Dict, Iterable, List, Optional
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple
 
 import torch
 
@@ -141,6 +141,38 @@ class InitProcessGroupKwargs(KwargsHandler):
     timeout: timedelta = timedelta(seconds=1800)
 
 
+@dataclass
+class FP8RecipeKwargs(KwargsHandler):
+    """
+    Use this object in your [`Accelerator`] to customize the initialization of the recipe for FP8 mixed precision
+    training. Please refer to the documentation of this
+    [class](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/common.html#transformer_engine.common.recipe.DelayedScaling)
+    for more information on each argument.
+
+    ```python
+    from accelerate import Accelerator
+    from accelerate.utils import FP8RecipeKwargs
+
+    kwargs = FP8RecipeKwargs(fp8_format="HYBRID")
+    accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=[kwargs])
+    ```
+    """
+
+    margin: int = 0
+    interval: int = 1
+    fp8_format: str = "E4M3"
+    amax_history_len: int = 1
+    amax_compute_algo: str = "most_recent"
+    override_linear_precision: Tuple[bool, bool, bool] = (False, False, False)
+
+    def __post_init__(self):
+        self.fp8_format = self.fp8_format.upper()
+        if self.fp8_format not in ["E4M3", "HYBRID"]:
+            raise ValueError("`fp8_format` must be 'E4M3' or 'HYBRID'.")
+        if self.amax_compute_algo not in ["max", "most_recent"]:
+            raise ValueError("`amax_compute_algo` must be 'max' or 'most_recent'")
+
+
 class DistributedType(str, enum.Enum):
     """
     Represents a type of distributed environment.
@@ -294,6 +326,7 @@ class PrecisionType(BaseEnum):
     """
 
     NO = "no"
+    FP8 = "fp8"
     FP16 = "fp16"
     BF16 = "bf16"
 

src/accelerate/utils/imports.py

@@ -56,6 +56,10 @@ def is_apex_available():
     return importlib.util.find_spec("apex") is not None
 
 
+def is_fp8_available():
+    return importlib.util.find_spec("transformer_engine") is not None
+
+
 @lru_cache()
 def is_tpu_available(check_device=True):
     "Checks if `torch_xla` is installed and potentially if a TPU is in the environment"

src/accelerate/utils/other.py

@@ -21,6 +21,7 @@
 from ..state import PartialState
 from .dataclasses import DistributedType
 from .imports import is_deepspeed_available, is_tpu_available
+from .transformer_engine import convert_model
 
 
 if is_deepspeed_available():
@@ -59,6 +60,8 @@ def extract_model_from_parallel(model, keep_fp32_wrapper: bool = True):
                 if forward == original_forward:
                     break
             model.forward = forward
+        if getattr(model, "_converted_to_transformer_engine", False):
+            convert_model(model, to_transformer_engine=False)
     return model