Use gloo as part of DeviceGPU's process group backend

composer/devices/device_gpu.py

@@ -12,7 +12,9 @@
 import torch.backends.cudnn
 import torch.cuda
 import torch.cuda.amp
+import torch.distributed as torch_dist
 import torch.utils.data
+from packaging import version
 
 from composer.devices.device import Device
 from composer.utils import dist
@@ -42,6 +44,9 @@ def __init__(
     ):
         if not torch.cuda.is_available():
             raise ValueError('DeviceGPU cannot be created as torch.cuda is not available.')
+        if torch_dist.is_gloo_available() and version.parse(torch.__version__) >= version.parse('2.3.0'):
+            # Composer checkpoint load / save from before torch 2.3.0 is not compatible with gloo + nccl backends.
+            DeviceGPU.dist_backend = 'cuda:nccl,cpu:gloo'
         if device_id is None:
             device_id = dist.get_local_rank()
         self._device = torch.device(f'cuda:{device_id}')

tests/checkpoint/test_state_dict.py

@@ -6,6 +6,7 @@
 
 import pytest
 import torch
+import torch.distributed as torch_dist
 from packaging import version
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
@@ -446,7 +447,10 @@ def test_get_metadata_sharded_model(model_type: str, tensor_type: str, world_siz
         assert 'model_name' in metadata_sd
 
     assert 'dist_backend' in metadata_sd
-    assert metadata_sd['dist_backend'] == 'nccl'
+    if torch_dist.is_gloo_available() and version.parse(torch.__version__) >= version.parse('2.3.0'):
+        assert metadata_sd['dist_backend'] == 'cuda:nccl,cpu:gloo'
+    else:
+        assert metadata_sd['dist_backend'] == 'nccl'
 
 
 @pytest.mark.filterwarnings('ignore:SWA has')