fix sync condition (hpcaitech#6000)

TongLi3701 · Aug 14, 2024 · ceb1e26 · ceb1e26
1 parent ed97d3a
commit ceb1e26
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py
@@ -1326,8 +1326,10 @@ def execute_pipeline(
             )
 
         # run with gradients accumulation
-        if model.require_grad_sync == False or (
-            isinstance(optimizer, HybridParallelZeroOptimizer) and optimizer.require_grad_sync == False
+        if (
+            model.require_grad_sync == False
+            or (isinstance(optimizer, HybridParallelZeroOptimizer) and optimizer.require_grad_sync == False)
+            or not torch.is_grad_enabled()
         ):
             return outputs