From aeee9fd2a9e908643cf4c8d021ed8618dac26dd1 Mon Sep 17 00:00:00 2001
From: ytxiong <45058324+yingtongxiong@users.noreply.github.com>
Date: Thu, 26 Oct 2023 17:33:00 +0800
Subject: [PATCH] fix broadcast synchronize() (#450)

---
 internlm/solver/optimizer/hybrid_zero_optim.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/internlm/solver/optimizer/hybrid_zero_optim.py b/internlm/solver/optimizer/hybrid_zero_optim.py
index 2817258..aa058c7 100644
--- a/internlm/solver/optimizer/hybrid_zero_optim.py
+++ b/internlm/solver/optimizer/hybrid_zero_optim.py
@@ -809,6 +809,8 @@ class HybridZeroOptimizer(BaseOptimizer):
         for handle in handles:
             handle.wait()
 
+        torch.cuda().synchronize()
+
     ##################
     # FP16 Utilities #
     ##################