Update on "Add rowwise scaling to Float8Inference module"

# Summary

# Performance
- Need to investigate the Rowwise dynamic case, I would think this should be faster than TensorWise dynamic
```Shell
Benchmark Results:
+--------------------------+-------------+
| Variant                  |   Time (μs) |
+==========================+=============+
| BF16                     |     2540.56 |
+--------------------------+-------------+
| FP8 Dynamic              |     1512.96 |
+--------------------------+-------------+
| FP8 Static               |     1363.75 |
+--------------------------+-------------+
| FP8 Weight Only          |     2774.22 |
+--------------------------+-------------+
| FP8 Dynamic AxisWise     |     1510.82 |
+--------------------------+-------------+
| FP8 Static AxisWise      |     1438.92 |
+--------------------------+-------------+
| FP8 Weight Only AxisWise |     2762.88 |
+--------------------------+-------------+

Comparison Results:
+--------------------------+-------------+-------------------+---------------+
| Variant                  |   Time (μs) | Speedup vs BF16   |   MAE vs BF16 |
+==========================+=============+===================+===============+
| BF16                     |     2540.56 | 1.00x             |    0          |
+--------------------------+-------------+-------------------+---------------+
| FP8 Dynamic              |     1512.96 | 1.68x             |    0.00543213 |
+--------------------------+-------------+-------------------+---------------+
| FP8 Static               |     1363.75 | 1.86x             |    0.00546265 |
+--------------------------+-------------+-------------------+---------------+
| FP8 Weight Only          |     2774.22 | 0.92x             |    0.00379944 |
+--------------------------+-------------+-------------------+---------------+
| FP8 Dynamic AxisWise     |     1510.82 | 1.68x             |    0.00543213 |
+--------------------------+-------------+-------------------+---------------+
| FP8 Static AxisWise      |     1438.92 | 1.77x             |    0.00546265 |
+--------------------------+-------------+-------------------+---------------+
| FP8 Weight Only AxisWise |     2762.88 | 0.92x             |    0.00379944 |
+--------------------------+-------------+-------------------+---------------+
```

### Numerics

Using this pytorch/ao#446
TensorWise Dynamic scaling:

``` Shell
+------------+--------------------------------------------+
| Task       | Metrics                                    |
+============+============================================+
| winogrande | +-----------------+----------+             |
|            | | acc,none        | 0.735596 |             |
|            | +-----------------+----------+             |
|            | | acc_stderr,none | 0.012395 |             |
|            | +-----------------+----------+             |
+------------+--------------------------------------------+
| wikitext   | +-----------------------------+----------+ |
|            | | bits_per_byte,none          | 0.538637 | |
|            | +-----------------------------+----------+ |
|            | | bits_per_byte_stderr,none   | N/A      | |
|            | +-----------------------------+----------+ |
|            | | byte_perplexity,none        | 1.452600 | |
|            | +-----------------------------+----------+ |
|            | | byte_perplexity_stderr,none | N/A      | |
|            | +-----------------------------+----------+ |
|            | | word_perplexity,none        | 7.363215 | |
|            | +-----------------------------+----------+ |
|            | | word_perplexity_stderr,none | N/A      | |
|            | +-----------------------------+----------+ |
+------------+--------------------------------------------+
```

AxisWise Dynamic Scaling

``` Shell
+------------+--------------------------------------------+
| Task       | Metrics                                    |
+============+============================================+
| winogrande | +-----------------+----------+             |
|            | | acc,none        | 0.735596 |             |
|            | +-----------------+----------+             |
|            | | acc_stderr,none | 0.012395 |             |
|            | +-----------------+----------+             |
+------------+--------------------------------------------+
| wikitext   | +-----------------------------+----------+ |
|            | | bits_per_byte,none          | 0.538637 | |
|            | +-----------------------------+----------+ |
|            | | bits_per_byte_stderr,none   | N/A      | |
|            | +-----------------------------+----------+ |
|            | | byte_perplexity,none        | 1.452600 | |
|            | +-----------------------------+----------+ |
|            | | byte_perplexity_stderr,none | N/A      | |
|            | +-----------------------------+----------+ |
|            | | word_perplexity,none        | 7.363215 | |
|            | +-----------------------------+----------+ |
|            | | word_perplexity_stderr,none | N/A      | |
|            | +-----------------------------+----------+ |
+------------+--------------------------------------------+

```



[ghstack-poisoned]

Author: drisspg

README.md

@@ -2,11 +2,12 @@
 
 This is an early version of a library for accelerating training with float8 in native PyTorch
 according to the recipes laid out in https://arxiv.org/pdf/2209.05433.pdf.
-The codebase strives to stay small, easily hackable, and debuggable with native PyTorch tooling.
-``torch.compile`` is supported out of the box. With ``torch.compile`` on, initial results show
+The codebase strives to stay small, easily hackable, debuggable with native PyTorch tooling,
+and composable with key systems such as autograd, ```torch.compile``` and distributed.
+With ``torch.compile`` on, initial results show
 throughput speedups of up to 1.2x on small scale (8 GPUs) LLaMa pretraining jobs.
 
-:warning: <em>See the [feature tracker](https://github.com/pytorch-labs/float8_experimental/issues/187) for upcoming features. Key features such as weight cast recomputation in backward and large scale distributed support are not ready yet. </em>
+:warning: <em>See the [feature tracker](https://github.com/pytorch-labs/float8_experimental/issues/187) for upcoming features.</em>
 
 :warning: <em>Backwards compatibility is not guaranteed at this point. The codebase is in active development and
 will change rapidly.</em>
@@ -25,7 +26,7 @@ pip install -e .
 pip install -e ".[dev]"
 ```
 
-# User API
+# Single GPU User API
 
 We provide two per-tensor scaling strategies: dynamic and delayed.  See https://arxiv.org/pdf/2209.05433.pdf, Section 4.3 for more details. These strategies are configurable separately for activations (`x`), weights (`w`) and gradients (`dL_dY`).
 
@@ -37,6 +38,7 @@ This is the most accurate recipe as every tensor is scaled dynamically.
 from float8_experimental.float8_linear_utils import (
     swap_linear_with_float8_linear,
 )
+from float8_experimental.fsdp_utils import precompute_float8_dynamic_scale_for_fsdp
 from float8_experimental.float8_linear import Float8Linear
 
 # create model
@@ -51,7 +53,18 @@ model = FSDP(model, use_orig_params=True)
 # optional: enable torch.compile for improved performance
 m = torch.compile(m)
 
-# train/finetune (not shown)
+# toy training loop
+for _ in range(N_ITER):
+    optimizer.zero_grad()
+    y = m(x)
+    y.sum().backward()
+    optimizer.step()
+
+    # specific to fsdp2 + dynamic scaling, when fp8 all-gather is turned on
+    # this method is optional but is highly recommended for performance
+    # it calcuclates scales for all parameters in a single all-reduce
+    precompute_float8_dynamic_scale_for_fsdp(model)
+
 ```
 
 ## float8 linear with delayed scaling
@@ -71,7 +84,7 @@ m = Model(...)
 # convert all `torch.nn.Linear` modules to `Float8Linear`, specifying scaling
 # type
 swap_linear_with_float8_linear(
-    m, 
+    m,
     Float8Linear,
     scaling_type_x=TensorScalingType.DELAYED,
     scaling_type_w=TensorScalingType.DELAYED,
@@ -101,30 +114,32 @@ for _ in range(N_ITER):
     optimizer.step()
 ```
 
-# 🧭 Code Organization
+# Multi GPU User API
 
-* `float8_experimental/float8_linear.py`
-    - `Float8Linear` (main user facing entry point for Float8Linear)
-* `float8_experimental/float8_tensor.py`
-    - `Float8Tensor`, which allows `Float8Linear` to abide by the `x.dtype == x.grad.dtype` restriction
-    - `ScaledMMConfig` defines the semantics for matmul in the forward and backwards pass
+We compose with the `DTensor` based [distributed APIs](https://pytorch.org/docs/stable/distributed.tensor.parallel.html),
+such as FSDP, TP and SP. Please see the [torchtitan](https://github.com/pytorch/torchtitan) repository for e2e examples
+on using `float8_experimental` in a distributed setting.
 
 # Testing
 
 ```bash
 # run single-GPU unit tests
 pytest test/test_base.py
 
-# run a single-GPU integration test on SAM
-pytest test/test_sam.py
-
 # run single-GPU compile tests
 pytest test/test_compile.py
+
+# run single-GPU numerics integration tests
+pytest test/test_numerics_integration.py
+
 # run a two-GPU integration test on FSDP
 ./test/test_fsdp.sh
 
-# run integration tests for TP/SP (outdated)
-./test/test_tp.sh
+# run integration tests on the DTensor TP/SP integration
+./test/test_dtensor.sh
+
+# run integration tests on the FSDP2 integration
+python test/test_fsdp2/test_fsdp2_eager.py
 
 # run all of these tests
 ./test/test_everything.sh

benchmarks/bench_linear_float8.py

@@ -14,11 +14,9 @@
 
 import torch
 import torch.utils.benchmark as benchmark
-from float8_experimental.float8_linear import TensorScalingType
+from float8_experimental.float8_linear import Float8Linear, TensorScalingType
 from float8_experimental.float8_linear_utils import (
-    get_float8_linear,
     linear_requires_sync,
-    LinearType,
     sync_float8_amax_and_scale_history,
 )
 from float8_experimental.float8_tensor import ScaledMMConfig
@@ -69,7 +67,6 @@ class Experiment:
     dtype: torch.dtype
     compiled: bool
     use_fast_accum: bool
-    linear_type: str
     scaling_repr: str
 
     # 3 Times since we are calculating forward backward
@@ -98,7 +95,6 @@ def main(
     n_limit: Optional[int] = None,
     fast_accum_filter: Optional[bool] = None,
     shape_name_filter: Optional[str] = None,
-    linear_type_filter: Optional[str] = None,
     scaling_type_x: str = "delayed",
     scaling_type_w: str = "delayed",
     scaling_type_dL_dY: str = "delayed",
@@ -123,44 +119,28 @@ def main(
         use_fast_accum = [fast_accum_filter]
     else:
         use_fast_accum = [True, False]
-    if linear_type_filter is not None:
-        linear_types = [linear_type_filter]
-    else:
-        linear_types = ["delayed", "dynamic"]
     if shape_name_filter is not None:
         k = shape_name_filter
         name_to_shapes_70b = {k: name_to_shapes_70b[k]}
     experiment_list: List[Experiment] = []
     dtype = torch.bfloat16
-    for idx, (fast_accum, (name, (K, N)), linear_type) in enumerate(
-        tqdm(list(product(use_fast_accum, name_to_shapes_70b.items(), linear_types)))
+    for idx, (fast_accum, (name, (K, N))) in enumerate(
+        tqdm(list(product(use_fast_accum, name_to_shapes_70b.items())))
     ):
         if n_limit is not None and idx >= n_limit:
             break
         linear_ref = torch.nn.Linear(K, N, bias=input_bias).to(
             device=device, dtype=dtype
         )
-        linear_type_enum = (
-            LinearType.DELAYED if linear_type == "delayed" else LinearType.DYNAMIC
-        )
 
-        if linear_type == "delayed":
-            linear_float8 = get_float8_linear(
-                linear_type_enum,
-                copy.deepcopy(linear_ref),
-                emulate=False,
-                scaling_type_x=scaling_type_x,
-                scaling_type_w=scaling_type_w,
-                scaling_type_dL_dY=scaling_type_dL_dY,
-            )
-            scaling_repr = linear_float8.scaling_repr()
-        else:
-            linear_float8 = get_float8_linear(
-                linear_type_enum,
-                copy.deepcopy(linear_ref),
-                emulate=False,
-            )
-            scaling_repr = None
+        linear_float8 = Float8Linear.from_float(
+            copy.deepcopy(linear_ref),
+            emulate=False,
+            scaling_type_x=scaling_type_x,
+            scaling_type_w=scaling_type_w,
+            scaling_type_dL_dY=scaling_type_dL_dY,
+        )
+        scaling_repr = linear_float8.scaling_repr()
 
         if fast_accum:
             linear_float8.forward_config = ScaledMMConfig(False, True, False)
@@ -172,19 +152,10 @@ def main(
         input_tensor = torch.randn(M, K, device=device, dtype=dtype, requires_grad=True)
         ref_forw_backward = lambda: linear_ref(input_tensor).sum().backward()
 
-        if linear_type_enum == LinearType.DELAYED:
-
-            def float8_forw_backward():
-                if linear_requires_sync(
-                    linear_type_enum, scaling_type_x, scaling_type_w, scaling_type_dL_dY
-                ):
-                    sync_float8_amax_and_scale_history(linear_float8)
-                linear_float8(input_tensor).sum().backward()
-
-        else:
-
-            def float8_forw_backward():
-                linear_float8(input_tensor).sum().backward()
+        def float8_forw_backward():
+            if linear_requires_sync(scaling_type_x, scaling_type_w, scaling_type_dL_dY):
+                sync_float8_amax_and_scale_history(linear_float8)
+            linear_float8(input_tensor).sum().backward()
 
         def n_times(n, fn, *args, **kwargs):
             def wrapper(*args, **kwargs):
@@ -224,7 +195,6 @@ def wrapper(*args, **kwargs):
             dtype,
             compile,
             use_fast_accum=fast_accum,
-            linear_type=linear_type,
             scaling_repr=scaling_repr,
         )
         print(experiment)
@@ -237,7 +207,6 @@ def wrapper(*args, **kwargs):
         "M",
         "K",
         "N",
-        "linear_type",
         "scaling_repr",
         "ref_dtype",
         "compiled",
@@ -257,7 +226,6 @@ def wrapper(*args, **kwargs):
                 experiment.shape[0],
                 experiment.shape[1],
                 experiment.shape[2],
-                experiment.linear_type,
                 experiment.scaling_repr,
                 experiment.dtype,
                 experiment.compiled,
@@ -287,7 +255,6 @@ def wrapper(*args, **kwargs):
         [
             "name",
             "shape",
-            "linear_type",
             "scaling_repr",
             "compiled",
             "use_fast_accum",
@@ -311,7 +278,6 @@ def invoke_main() -> None:
     parser.add_argument("-n", "--n_limit", type=int, required=False)
     parser.add_argument("--fast_accum_filter", type=bool, required=False)
     parser.add_argument("--shape_name_filter", type=str, required=False)
-    parser.add_argument("--linear_type_filter", type=str, required=False)
     parser.add_argument("--scaling_type_x", type=str, required=False)
     parser.add_argument("--scaling_type_w", type=str, required=False)
     parser.add_argument("--scaling_type_dL_dY", type=str, required=False)
@@ -330,7 +296,6 @@ def invoke_main() -> None:
         args.n_limit,
         args.fast_accum_filter,
         args.shape_name_filter,
-        args.linear_type_filter,
         **kwargs,
     )
 

benchmarks/bench_multi_gpu.py

@@ -14,7 +14,7 @@
 import torch.multiprocessing as mp
 import torch.nn as nn
 import torch.utils.benchmark as benchmark
-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,
     sync_float8_amax_and_scale_history,
@@ -65,7 +65,13 @@ def get_model(K, N, is_fp8, base_dtype=torch.float32):
         modules.append(nn.ReLU())
     m = nn.Sequential(*modules)
     if is_fp8:
-        swap_linear_with_float8_linear(m, Float8Linear, emulate=False)
+        swap_linear_with_float8_linear(
+            m,
+            emulate=False,
+            scaling_type_x=TensorScalingType.DELAYED,
+            scaling_type_w=TensorScalingType.DELAYED,
+            scaling_type_dL_dY=TensorScalingType.DELAYED,
+        )
     return m
 
 

benchmarks/profile_linear_float8.py

@@ -18,11 +18,9 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from float8_experimental.float8_dynamic_linear import Float8DynamicLinear
-from float8_experimental.float8_linear import Float8Linear, TensorScalingType
+from float8_experimental.float8_linear import TensorScalingType
 from float8_experimental.float8_linear_utils import (
     linear_requires_sync,
-    LinearType,
     swap_linear_with_float8_linear,
     sync_float8_amax_and_scale_history,
 )
@@ -206,19 +204,25 @@ def profile_function(
 def main(
     profile_path_prefix: Path,
     compile: bool = True,
-    linear_type: str = "dynamic",
-    scaling_type_x: str = "delayed",
-    scaling_type_w: str = "delayed",
-    scaling_type_dL_dY: str = "delayed",
+    scaling_type_x: str = "dynamic",
+    scaling_type_w: str = "dynamic",
+    scaling_type_dL_dY: str = "dynamic",
     model_type: str = "linear",
     dtype_filter: str = "both",
 ):
     assert model_type in ("linear", "ln_linear", "norm_ffn_norm"), "unsupported"
     assert dtype_filter in ("both", "float8", "bfloat16")
 
-    print(f"Compile is set to          | {compile}")
-    print(f"Using Linear type:         | {linear_type}")
-    print(f"model_type is set to       | {model_type}")
+    scaling_type_x = TensorScalingType(scaling_type_x)
+    scaling_type_w = TensorScalingType(scaling_type_w)
+    scaling_type_dL_dY = TensorScalingType(scaling_type_dL_dY)
+    scaling_repr = "_".join(
+        [s.short_str() for s in (scaling_type_x, scaling_type_w, scaling_type_dL_dY)]
+    )
+
+    print(f"Compile is set to       | {compile}")
+    print(f"model_type is set to    | {model_type}")
+    print(f"scaling_repr is set to  | {scaling_repr}")
 
     device = "cuda"
     ref_dtype = torch.bfloat16
@@ -249,21 +253,14 @@ def main(
 
     m_ref = m_ref.to(device).to(ref_dtype)
 
-    linear_type = LinearType[linear_type.upper()]
-    linear_cls = (
-        Float8Linear if linear_type is LinearType.DELAYED else Float8DynamicLinear
-    )
-    extra_kwargs = {}
-    scaling_type_x = TensorScalingType(scaling_type_x)
-    scaling_type_w = TensorScalingType(scaling_type_w)
-    scaling_type_dL_dY = TensorScalingType(scaling_type_dL_dY)
-    if linear_type is LinearType.DELAYED:
-        extra_kwargs["scaling_type_x"] = scaling_type_x
-        extra_kwargs["scaling_type_w"] = scaling_type_w
-        extra_kwargs["scaling_type_dL_dY"] = scaling_type_dL_dY
+    extra_kwargs = {
+        "scaling_type_x": scaling_type_x,
+        "scaling_type_w": scaling_type_w,
+        "scaling_type_dL_dY": scaling_type_dL_dY,
+    }
 
     m_float8 = copy.deepcopy(m_ref)
-    swap_linear_with_float8_linear(m_float8, linear_cls, **extra_kwargs)
+    swap_linear_with_float8_linear(m_float8, **extra_kwargs)
 
     def ref_forw_backward(x):
         out = m_ref(x)
@@ -281,9 +278,7 @@ def float8_forw_backward_wrapper(x):
         # inspection of the fw+bw torch.compile without the scale
         # syncing code
         # TODO(future): make this better
-        if linear_requires_sync(
-            linear_type, scaling_type_x, scaling_type_w, scaling_type_dL_dY
-        ):
+        if linear_requires_sync(scaling_type_x, scaling_type_w, scaling_type_dL_dY):
             with record_function("scale_amax_and_scales"):
                 sync_amax_history(m_float8)
         out = float8_forw(x)
@@ -345,7 +340,9 @@ def float8_forw_backward_wrapper(x):
         if dtype_filter != "bfloat16":
             # Profile Float8 Model
             print("profiling float8")
-            float8_suffix = f"_{model_type}_float8_compile_{compile}_{linear_type}.json"
+            float8_suffix = (
+                f"_{model_type}_float8_compile_{compile}_{scaling_repr}.json"
+            )
             float8_path = profile_path_prefix + float8_suffix
             profile_config = ProfileConfig(
                 float8_path,