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 # ##################