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