update to make more comple friendly

Author: drisspg

float8_experimental/float8_linear.py

@@ -138,23 +138,24 @@ def __init__(self, *args, **kwargs):
         self.recipe = delayed_scaling_recipe
         history_len = self.recipe.history_len
 
-        self.register_always_float32_buffer("fp8_amax_x", torch.tensor(E4M3_MAX_POS))
+        self.register_always_float32_buffer("fp8_amax_x", torch.tensor([E4M3_MAX_POS]))
         self.register_always_float32_buffer(
             "fp8_amax_history_x", torch.zeros(history_len)
         )
-        self.register_always_float32_buffer("fp8_scale_x", torch.tensor(1.0))
-        self.register_always_float32_buffer("fp8_amax_w", torch.tensor(E4M3_MAX_POS))
+        self.register_always_float32_buffer("fp8_scale_x", torch.tensor([1.0]))
+        self.register_always_float32_buffer("fp8_amax_w", torch.tensor([E4M3_MAX_POS]))
         self.register_always_float32_buffer(
             "fp8_amax_history_w", torch.zeros(history_len)
         )
-        self.register_always_float32_buffer("fp8_scale_w", torch.tensor(1.0))
+        self.register_always_float32_buffer("fp8_scale_w", torch.tensor([1.0]))
         self.register_always_float32_buffer(
-            "fp8_amax_dL_dY", torch.tensor(E5M2_MAX_POS)
+            "fp8_amax_dL_dY", torch.tensor([E5M2_MAX_POS])
         )
         self.register_always_float32_buffer(
             "fp8_amax_history_dL_dY", torch.zeros(history_len)
         )
-        self.register_always_float32_buffer("fp8_scale_dL_dY", torch.tensor(1.0))
+        self.register_always_float32_buffer("fp8_scale_dL_dY", torch.tensor([1.0]))
+
         # Whether to emulate the fp8 matmul logic in float32
         self.emulate = False
 

float8_experimental/float8_linear_utils.py

@@ -145,39 +145,37 @@ def sync_float8_amax_and_scale_history(model: torch.nn.Module, fp8_layers=None)
         fp8_layers = get_float8_layers(model)
 
     if dist.is_initialized():
-        fp8_amax_x_tensor = torch.tensor(
-            [child.fp8_amax_x for child in fp8_layers],
-            dtype=torch.float32,
-            device="cuda",
-            requires_grad=False,
-        )
-        fp8_amax_w_tensor = torch.tensor(
-            [child.fp8_amax_w for child in fp8_layers],
-            dtype=torch.float32,
-            device="cuda",
-            requires_grad=False,
-        )
-        fp8_amax_dL_dY_tensor = torch.tensor(
-            [child.fp8_amax_dL_dY for child in fp8_layers],
-            dtype=torch.float32,
-            device="cuda",
-            requires_grad=False,
-        )
-        dist.all_reduce(fp8_amax_x_tensor, op=dist.ReduceOp.MAX)
-        dist.all_reduce(fp8_amax_w_tensor, op=dist.ReduceOp.MAX)
-        dist.all_reduce(fp8_amax_dL_dY_tensor, op=dist.ReduceOp.MAX)
-
+        fp8_amax_x_tensors = [child.fp8_amax_x for child in fp8_layers]
+        fp8_amax_w_tensors = [child.fp8_amax_w for child in fp8_layers]
+        fp8_amax_dL_dY_tensors = [child.fp8_amax_dL_dY for child in fp8_layers]
+
+        assert (
+            len(fp8_amax_x_tensors)
+            == len(fp8_amax_w_tensors)
+            == len(fp8_amax_dL_dY_tensors)
+        ), "Mismatched lengths of amax tensors."
+        if len(fp8_amax_x_tensors) > 0:
+            # Combine all the amax tensors into one tensor and reduce it
+            fp8_amax_x_tensor = torch.cat(fp8_amax_x_tensors)
+            fp8_amax_w_tensor = torch.cat(fp8_amax_w_tensors)
+            fp8_amax_dL_dY_tensor = torch.cat(fp8_amax_dL_dY_tensors)
+
+            dist.all_reduce(fp8_amax_x_tensor, op=dist.ReduceOp.MAX)
+            dist.all_reduce(fp8_amax_w_tensor, op=dist.ReduceOp.MAX)
+            dist.all_reduce(fp8_amax_dL_dY_tensor, op=dist.ReduceOp.MAX)
+
+            # Reassign the reduced amax values to the original tensors
+
+            for idx in range(len(fp8_layers)):
+                child = fp8_layers[idx]
+                child.fp8_amax_x.copy_(fp8_amax_x_tensor[idx].clone())
+                child.fp8_amax_w.copy_(fp8_amax_w_tensor[idx].clone())
+                child.fp8_amax_dL_dY.copy_(fp8_amax_dL_dY_tensor[idx].clone())
+
+    # Itearte over all the layers and update the amax history and scales
     for idx in range(len(fp8_layers)):
         child = fp8_layers[idx]
 
-        #
-        # 1. in distributed contexts, syncs amax values across workers
-        #
-        if dist.is_initialized():
-            child.fp8_amax_x = fp8_amax_x_tensor[idx].clone()
-            child.fp8_amax_w = fp8_amax_w_tensor[idx].clone()
-            child.fp8_amax_dL_dY = fp8_amax_dL_dY_tensor[idx].clone()
-
         #
         # 2. adds the `amax` values to history
         #

float8_experimental/float8_python_api.py

@@ -12,6 +12,8 @@
 
 from typing import Optional, Tuple
 
+import float8_experimental.float8_aten_api  # noqa
+
 import torch
 from float8_experimental.float8_tensor import Float8Tensor
 

float8_experimental/float8_tensor.py

@@ -54,7 +54,9 @@ class FromFloat8ConstrFunc(torch.autograd.Function):
 
     @staticmethod
     def forward(ctx, tensor):
-        return tensor._data.to(tensor._orig_dtype) / tensor._scale
+        return (tensor._data.to(tensor._orig_dtype) / tensor._scale).to(
+            tensor._orig_dtype
+        )
 
     @staticmethod
     def backward(ctx, g):

float8_experimental/float8_utils.py

@@ -23,7 +23,7 @@
 
 @torch.no_grad()
 def amax_to_scale(amax, float8_dtype, orig_dtype):
-    scale = torch.empty((), device=amax.device, dtype=torch.float32)
+    scale = torch.empty((1,), device=amax.device, dtype=torch.float32)
     if float8_dtype == torch.float8_e4m3fn:
         res = E4M3_MAX_POS / torch.clamp(amax, min=EPS)
     else:  # e5m2