From 88efbed72e7d1c002c6a7d5392362db55a68a7ea Mon Sep 17 00:00:00 2001
From: Sora <atri@suzakuintsubaki.com>
Date: Mon, 30 Oct 2023 20:31:22 +0800
Subject: [PATCH] Update train_ms.py. Auto choose gloo backend while using
 Windows platform. (#134)

---
 train_ms.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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()