Used functional all-reduce for amax reduction

ghstack-source-id: 256aebe
Pull Request resolved: #219

Author: Andrew Gu

float8_experimental/float8_utils.py

@@ -75,7 +75,13 @@ def tensor_to_amax(x, distributed_reduction=False):
     # If the user did not ask for it, assume that it will
     # happen elsewhere.
     if distributed_reduction and dist.is_initialized():
-        dist.all_reduce(amax, op=dist.ReduceOp.MAX)
+        # TODO: Dynamo rewriting synchronous in-place collectives fails:
+        # https://github.com/pytorch/pytorch/issues/120082
+        # Use functional all-reduce to avoid graph breaking.
+        amax = dist._functional_collectives.all_reduce(
+            amax, "MAX", list(range(dist.get_world_size()))
+        )
+        # dist.all_reduce(amax, op=dist.ReduceOp.MAX)
 
     return amax