diff --git a/train_ms.py b/train_ms.py index 4e9025d..8be0e8e 100644 --- a/train_ms.py +++ b/train_ms.py @@ -1,5 +1,6 @@ # flake8: noqa: E402 +import platform import os import torch from torch.nn import functional as F @@ -53,8 +54,11 @@ def run(): os.environ[env_name] = str(env_value) # 多卡训练设置 + backend = "nccl" + if platform.system() == "Windows": + backend = "gloo" dist.init_process_group( - backend="nccl", + backend=backend, init_method="env://", # If Windows,switch to gloo backend. ) # Use torchrun instead of mp.spawn rank = dist.get_rank()