enable float types in pytorch for non comptue comms

pytorch-labs · May 17, 2024 · 02a80f9 · 02a80f9
1 parent 6891cbe
commit 02a80f9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/float8_experimental/float8_ops.py b/float8_experimental/float8_ops.py
@@ -235,11 +235,8 @@ def allgather_fp8(aten_op, args, kwargs=None):
     ), f"expecting a Float8Tensor for allgather but found {type(fp8_input)}"
 
     fp8_data = fp8_input._data
-    fp8_data = fp8_data.view(torch.uint8)
     fp8_data = fp8_data.contiguous()
     fp8_out = aten_op(fp8_data, *args[1:], **kwargs)
-    fp8_out = torch.ops._c10d_functional.wait_tensor(fp8_out)
-    fp8_out = fp8_out.view(fp8_input._data.dtype)
     return Float8Tensor(
         fp8_out, fp8_input._scale, fp8_input._orig_dtype, fp8_input._mm_config
     )

diff --git a/test/test_dtensor.py b/test/test_dtensor.py
@@ -246,3 +246,5 @@ def test_fp8_mlp_tensor_parallelism_compile(mesh: DeviceMesh, size=16):
         except Exception as e:
             print(f"Test {test.__name__} failed with error: {e}")
             raise e
+
+    torch.distributed.destroy_process_group()