diff --git a/train/comms/pt/pytorch_backend_utils.py b/train/comms/pt/pytorch_backend_utils.py index 1a59b766..5a4570c2 100644 --- a/train/comms/pt/pytorch_backend_utils.py +++ b/train/comms/pt/pytorch_backend_utils.py @@ -31,7 +31,7 @@ class CollectiveArgsMixin: supportedDevices = ["cpu", "cuda", "rocm", "tpu"] supportedC10dBackends = ["nccl", "gloo", "mpi", "ucc", "xla", "mccl"] -supportedTorchcommsBackends = ["rcclx", "ncclx", "rccl", "mccl"] +supportedTorchcommsBackends = ["ncclx", "mccl"] supportedCollectives = [ "reduce", "all_reduce", diff --git a/train/comms/pt/pytorch_torchcomms_backend.py b/train/comms/pt/pytorch_torchcomms_backend.py index e6019381..4d8b428f 100644 --- a/train/comms/pt/pytorch_torchcomms_backend.py +++ b/train/comms/pt/pytorch_torchcomms_backend.py @@ -17,7 +17,7 @@ logger = logging.getLogger(__name__) -_VALID_BACKENDS = ["nccl", "ncclx", "rcclx", "mccl"] +_VALID_BACKENDS = ["nccl", "ncclx", "mccl"] class PyTorchTorchcommsBackend(backendFunctions): @@ -742,6 +742,11 @@ def initialize_backend( self, master_ip, master_port, backend="ncclx", eager_mode=False ): """Initialize the torchcomms backend""" + if torch.version.hip is not None: + raise ValueError( + "The pytorch-torchcomms stack does not support AMD GPUs. " + "Use --nw-stack pytorch-dist for AMD benchmarks." + ) if backend not in _VALID_BACKENDS: raise ValueError( f"Invalid backend {backend}. Valid backends are {_VALID_BACKENDS}"