update release_grads (#8834)

Author: DesmonDay

paddlenlp/trainer/trainer.py

@@ -1062,11 +1062,12 @@ def fused_allreduce_gradients_no_sync(paramlist, hcg):
  if optimizer_was_run:
  self.lr_scheduler.step()
 
- if enable_release_grads and args.pipeline_parallel_degree > 1:
+ if args.release_grads or enable_release_grads:
  self.optimizer.clear_grad(set_to_zero=False)
- for _, buffers in model._chunk_2_comm_buffers.items():
- for buffer in buffers:
- buffer._clear_grad_storage()
+ if args.pipeline_parallel_degree > 1:
+ for _, buffers in model._chunk_2_comm_buffers.items():
+ for buffer in buffers:
+ buffer._clear_grad_storage()
  else:
  self.optimizer.clear_grad()
 

paddlenlp/trainer/training_args.py

@@ -344,6 +344,8 @@ class TrainingArguments:
  Whether skip profile timer, timer will record time usage of forward/ backward/ step, etc.
  distributed_dataloader (`bool`, *optional*):
  Whether to use distributed dataloader. Default is `False`.
+ release_grads (`bool`, *optional*):
+ Whether to release gradients during training. Default is `False`.
  """
 
  output_dir: str = field(
@@ -791,6 +793,9 @@ class TrainingArguments:
  default=False,
  metadata={"help": "Enable MoE (Mixture of Experts) expert parallel training"},
  )
+ release_grads: Optional[bool] = field(
+ default=False, metadata={"help": "Whether to release gradients during training. Default is `False`."}
+ )
 
  def __post_init__(self):
  env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))