enable float types in pytorch for non comptue comms

Author: drisspg

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
     )

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()