Merge branch 'main' of https://github.com/NVIDIA/TransformerEngine

transformer_engine/pytorch/utils.py

@@ -411,3 +411,16 @@ def nvtx_range_pop(msg: Optional[str] = None) -> None:
 
     # Pop NVTX range
     torch.cuda.nvtx.range_pop()
+
+
+def canonicalize_process_group(
+    group: Optional[torch.distributed.ProcessGroup],
+) -> torch.distributed.ProcessGroup:
+    """Convert to PyTorch process group
+
+    If `None`, returns default process group.
+
+    """
+    if group is None:
+        return torch.distributed.distributed_c10d._get_default_group()
+    return group