From ef9a592d1431f10d1d341d4b715dd75d3017ccf9 Mon Sep 17 00:00:00 2001
From: Yu Chengzhang <iomgaaycz@gmail.com>
Date: Tue, 24 Sep 2024 11:43:30 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E4=BA=86wandb=E7=9A=84bug?=
 =?UTF-8?q?=EF=BC=8C=E9=81=BF=E5=85=8D=E4=BA=86=E5=A4=9A=E6=AC=A1=E4=BA=A7?=
 =?UTF-8?q?=E7=94=9F=E9=A1=B9=E7=9B=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 1-pretrain.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/1-pretrain.py b/1-pretrain.py
index 8560126..56c937d 100644
--- a/1-pretrain.py
+++ b/1-pretrain.py
@@ -146,11 +146,6 @@ if __name__ == "__main__":
     use_wandb = True #是否使用wandb
     wandb_project = "MiniMind-Pretrain"
     wandb_run_name = f"MiniMind-Pretrain-Epoch-{epochs}-BatchSize-{batch_size}-LearningRate-{learning_rate}"
-    if use_wandb:
-        import wandb
-        wandb.init(project=wandb_project, name=wandb_run_name)
-    else:
-        wandb = None
 
 
     ctx = (
@@ -163,6 +158,12 @@ if __name__ == "__main__":
     if ddp:
         init_distributed_mode()
         device = torch.device(DEVICE)
+
+    if use_wandb and (not ddp or ddp_local_rank == 0):
+        import wandb
+        wandb.init(project=wandb_project, name=wandb_run_name)
+    else:
+        wandb = None
     # -----------------------------------------------------------------------------
 
     # -----init dataloader------