Update

[ghstack-poisoned]

Author: fegin

.github/workflows/integration_test_4gpu.yaml

@@ -39,5 +39,6 @@ jobs:
 
         python -m pip install --force-reinstall --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
         python -m pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/
+        python -m pip install git+https://github.com/pytorch-labs/float8_experimental.git
         mkdir artifacts-to-be-uploaded
         python ./test_runner.py artifacts-to-be-uploaded --ngpu 4

estimation.py

@@ -126,8 +126,8 @@ def loss_fn(pred, labels):
             whole_model = model_cls.from_model_args(model_config)
 
         # apply fp8 linear module swap
-        if job_config.training.fp8_linear:
-            build_fp8_linear(whole_model, job_config)
+        if job_config.training.enable_fp8_linear:
+            build_fp8_linear(whole_model, job_config, parallel_dims.dp_enabled)
 
         # apply PT-D DP/TP parallelisms and activation checkpointing
         model_parts = [whole_model]

test_runner.py

@@ -273,6 +273,39 @@ def build_test_list():
             "fsdp2_mem_tracker",
             ngpu=4,
         ),
+        OverrideDefinitions(
+            [
+                [
+<<<<<<< HEAD
+                    "--training.enable_float8_linear",
+                ]
+            ],
+            "FSDP2 with original dtype",
+            "float8_fsdp2_orig_all_gather",
+            ngpu=4,
+        ),
+        OverrideDefinitions(
+            [
+                [
+                    "--training.enable_float8_linear",
+                    "--training.enable_fsdp_float8_all_gather",
+                ]
+            ],
+            "FSDP2 with float8 all-gather",
+            "fsdp2_float8_all_gather",
+            ngpu=4,
+        ),
+        OverrideDefinitions(
+            [
+                [
+                    "--training.enable_float8_linear",
+                    "--training.enable_fsdp_float8_all_gather",
+                    "--training.precompute_float8_dynamic_scale_for_fsdp",
+                ]
+            ],
+            "FSDP2 with float8 all-gather and precomputed dynamic scales",
+            "fsdp2_float8_all_gather_precompute_dynamic_scales",
+        ),
         OverrideDefinitions(
             [
                 [

torchtitan/config_manager.py

@@ -355,7 +355,7 @@ def __init__(self):
             help="Whether to compile the model",
         )
         self.parser.add_argument(
-            "--training.fp8_linear",
+            "--training.enable_float8_linear",
             action="store_true",
             help="""
                 If true, swaps `torch.nn.Linear` with `Float8Linear` with
@@ -364,6 +364,18 @@ def __init__(self):
                 here: https://github.com/pytorch-labs/float8_experimental
             """,
         )
+        self.parser.add_argument(
+            "--training.enable_fsdp_float8_all_gather",
+            action="store_true",
+            default=False,
+            help="Whether enable float8 all-gather in FSDP",
+        )
+        self.parser.add_argument(
+            "--training.precompute_float8_dynamic_scale_for_fsdp",
+            action="store_true",
+            default=False,
+            help="Whether precompute float8 scales dynamically for FSDP",
+        )
         self.parser.add_argument(
             "--training.gc_freq",
             type=int,

torchtitan/float8_linear.py

@@ -12,31 +12,94 @@
 
 # Note: Performance
 # Float8 experimental is intended to be ran under `torch.compile`` for competitive performance
+import contextlib
+import functools
+from typing import Optional
 
+import torch
 import torch.nn as nn
+from torch._logging import warning_once
 
 from torchtitan.config_manager import JobConfig
 from torchtitan.logging_utils import logger
 
 
-def build_fp8_linear(model: nn.Module, job_config: JobConfig):
+@contextlib.contextmanager
+def set_enable_fsdp_float8_all_gather(enable_fsdp_fp8_all_gather: bool):
+    import float8_experimental.config as config
+
+    prev = config.enable_fsdp_fp8_all_gather
+    torch.distributed.barrier()
+    config.enable_fsdp_fp8_all_gather = enable_fsdp_fp8_all_gather
+    try:
+        yield
+    finally:
+        torch.distributed.barrier()
+        config.enable_fsdp_fp8_all_gather = prev
+
+
+@functools.lru_cache(None)
+def is_sm90_or_later():
+    # Float8 is only supported on H100+ GPUs
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)
+
+
+def maybe_build_fp8_linear(
+    model: nn.Module, job_config: JobConfig, dp_enabled: Optional[bool] = False
+):
     """
     This function converts the linear layers to `Float8Linear`. Note that today,
     only dynamic tensor scaling (the default) is supported.
 
     This will mutate the model inplace.
     """
-    use_fp8_linear = job_config.training.fp8_linear
+    enable_float8_linear = job_config.training.enable_float8_linear
+    if not enable_float8_linear:
+        return
+    if not is_sm90_or_later():
+        warning_once(
+            logger,
+            "Failed to swap to Float8Linear because SM90 or later is not available",
+        )
+        return
     try:
-        from float8_experimental.float8_linear import Float8Linear
+        from float8_experimental.float8_linear import TensorScalingType
         from float8_experimental.float8_linear_utils import (
             swap_linear_with_float8_linear,
         )
+
+        # Mutates the model inplace replacing instances of torch.nn.Linear with Float8Linear
+        enable_fsdp_float8_all_gather = (
+            job_config.training.enable_fsdp_float8_all_gather and dp_enabled
+        )
+        with set_enable_fsdp_float8_all_gather(enable_fsdp_float8_all_gather):
+            swap_linear_with_float8_linear(
+                model, scaling_type_w=TensorScalingType.DYNAMIC
+            )
+        logger.info(
+            f"Swapped to Float8Linear layers with {enable_fsdp_float8_all_gather=}"
+        )
     except ImportError as exc:
         raise ImportError(
             "float8_experimental is not installed. Please install it to use fp8 linear layers."
         ) from exc
-    if use_fp8_linear:
-        # Mutates the model inplace replacing instances of torch.nn.Linear with Float8Linear
-        swap_linear_with_float8_linear(model, Float8Linear)
-        logger.info("Swapped to Float8Linear layers")
+
+
+def maybe_precompute_fp8_dynamic_scale_for_fsdp(
+    model: nn.Module, job_config: JobConfig
+):
+    if not (
+        job_config.training.enable_float8_linear
+        and job_config.training.enable_fsdp_float8_all_gather
+        and job_config.training.precompute_float8_dynamic_scale_for_fsdp
+    ):
+        return
+    if not is_sm90_or_later():
+        warning_once(
+            logger,
+            "Skipped precomputing fp8 scales because SM90 or later is not available",
+        )
+        return
+    from float8_experimental.fsdp_utils import precompute_float8_dynamic_scale_for_fsdp
+
+    precompute_float8_dynamic_scale_for_fsdp(model)

torchtitan/parallelisms/parallelize_llama.py

@@ -125,12 +125,16 @@ def selective_checkpointing_context_fn():
 
 def get_tp_parallel_strategy(
     job_config: JobConfig,
+    model: nn.Module,
 ) -> Tuple[RowwiseParallel, ColwiseParallel, PrepareModuleInput]:
     """Get the parallel strategy for the transformer model.
 
     This function handles the special case of using float8 with tensor parallelism.
     """
-    if job_config.training.fp8_linear == "dynamic":
+    if job_config.training.enable_float8_linear:
+        # TODO(future PR): once float8 configuration supports delayed
+        # scaling, add a check here to enforce supported float8 all-gather
+        # configurations
         from float8_experimental.float8_tensor_parallel import (
             Float8ColwiseParallel,
             Float8RowwiseParallel,
@@ -354,7 +358,7 @@ def apply_tp(
         rowwise_parallel_weight,
         colwise_parallel_weight,
         prepare_module_input,
-    ) = get_tp_parallel_strategy(job_config)
+    ) = get_tp_parallel_strategy(job_config, model)
     loss_parallel = parallel_dims.loss_parallel_enabled
 
     # 1. Parallelize the embedding and shard its outputs (which are the first

train.py

@@ -29,7 +29,10 @@
 from torchtitan.checkpoint import CheckpointManager
 from torchtitan.config_manager import JobConfig
 from torchtitan.datasets import build_hf_data_loader, create_tokenizer
-from torchtitan.float8_linear import build_fp8_linear
+from torchtitan.float8_linear import (
+    maybe_build_fp8_linear,
+    maybe_precompute_fp8_dynamic_scale_for_fsdp,
+)
 from torchtitan.logging_utils import init_logger, logger
 from torchtitan.lr_scheduling import get_lr_schedulers
 from torchtitan.metrics import build_gpu_memory_monitor, build_metric_logger
@@ -249,9 +252,8 @@ def loss_fn(pred, labels):
     with torch.device("meta"):
         whole_model = model_cls.from_model_args(model_config)
 
-    # apply fp8 linear module swap
-    if job_config.training.fp8_linear:
-        build_fp8_linear(whole_model, job_config)
+    # swap to Float8Linear base on fp8 config
+    maybe_build_fp8_linear(whole_model, job_config, parallel_dims.dp_enabled)
 
     # log model size
     model_param_count = get_num_params(whole_model)
@@ -437,6 +439,11 @@ def loss_fn(pred, labels):
             optimizers.step()
             lr_schedulers.step()
 
+            # when fp8 config is on,
+            # calculate float8 dynamic amax/scale for all-parameter for FSDP2
+            # it issues a single all-reduce for all parameters at once for better performance
+            maybe_precompute_fp8_dynamic_scale_for_fsdp(model, job_config)
+
             losses_since_last_log.append(loss)
 
             # log metrics

train_configs/debug_model.toml

@@ -37,7 +37,7 @@ max_norm = 1.0  # grad norm clipping
 steps = 10
 data_parallel_degree = -1
 tensor_parallel_degree = 1
-fp8_linear = false
+enable_float8_linear = false
 compile = false
 dataset = "c4_mini"  # supported datasets: c4_mini (45K), c4 (177M)
 

train_configs/llama2_13b.toml

@@ -33,7 +33,7 @@ max_norm = 1.0  # grad norm clipping
 steps = 1000
 data_parallel_degree = -1
 tensor_parallel_degree = 1
-fp8_linear = false
+enable_float8_linear = false
 compile = false
 dataset = "c4"
 

train_configs/llama2_70b.toml

@@ -33,7 +33,7 @@ max_norm = 1.0  # grad norm clipping
 steps = 1000
 data_parallel_degree = -1
 tensor_parallel_degree = 8  # 8-way TP
-fp8_linear = false
+enable_float8_linear = false
 compile = false
 dataset = "c4"