update sharding release grads

DesmonDay · Nov 19, 2024 · a44b4bb · a44b4bb
1 parent f21329c
commit a44b4bb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/paddlenlp/trainer/training_args.py b/paddlenlp/trainer/training_args.py
@@ -1346,7 +1346,7 @@ def is_segment_parallel_supported():
                             strategy.hybrid_configs["sharding_configs"].split_param = True
                             assert self.amp_master_grad, "Currently sharding stage1 v2 only support amp_master_grad"
 
-                        if "enable_release_grads" in sharding_parallel_config:
+                        if "enable_release_grads" in sharding_parallel_config or self.release_grads:
                             strategy.hybrid_configs["sharding_configs"].release_gradients = True
 
                         if self.pipeline_parallel_degree == 1: