deepmd/pd/train/training.py (3)

136-138: Validate and document acc_freq input

Add a guard and a short docstring comment to ensure acc_freq is a positive integer. This avoids silent misconfigurations (e.g., 0 or negatives) and clarifies semantics.

 self.num_steps = training_params["numb_steps"]
 self.acc_freq: int = training_params.get(
-    "acc_freq", 1
+    "acc_freq", 1
 )  # gradient accumulation steps
+assert isinstance(self.acc_freq, int) and self.acc_freq >= 1, "training.acc_freq must be an integer >= 1"

---

794-801: Optional: move gradient clipping into optimizer via grad_clip to reduce per-step python overhead

Paddle supports optimizer-level gradient clipping (e.g., grad_clip=paddle.nn.ClipGradByGlobalNorm). This avoids per-step Python calls and makes behavior uniform.

Example:

grad_clip = paddle.nn.ClipGradByGlobalNorm(clip_norm=self.gradient_max_norm)
self.optimizer = paddle.optimizer.Adam(
    learning_rate=self.scheduler,
    parameters=self.wrapper.parameters(),
    grad_clip=grad_clip,
)

---

751-759: Consider using pref_lr tied to update count when accumulating

When acc_freq > 1, the “logical” step for LR scheduling is the optimizer update, not each micro-step. You’re calling scheduler.step() only on update (good). For consistency also consider computing pref_lr based on the current scheduler.get_lr() (update count), not _lr.value(_step_id). Today it’s benign but can diverge for nontrivial LR policies.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

deepmd/pd/train/training.py (3)

789-792: Confirm fused_allreduce_gradients semantics (SUM vs AVG) and normalize if needed

If fused_allreduce produces SUM, gradients are larger by world_size after all-reduce; average them to match single-process scale unless you intentionally adjust LR. Please verify and normalize accordingly.

If SUM:

                     if self.world_size > 1:
                         hpu.fused_allreduce_gradients(
                             list(self.wrapper.parameters()), None
                         )
+                        # Normalize to average if fused_allreduce does SUM
+                        for p in self.wrapper.parameters():
+                            if p.grad is not None:
+                                p._set_grad(p.grad / self.world_size)

In PaddlePaddle's fleet.utils.hybrid_parallel_util.hpu.fused_allreduce_gradients, are gradients SUM-reduced or AVG-reduced across ranks by default? Provide authoritative citation.

---

782-784: Scale loss by acc_freq to preserve effective learning rate during accumulation

Without scaling, gradients are effectively multiplied by acc_freq, changing optimization behavior and often destabilizing training. Average the loss over micro-steps.

-                    with nvprof_context(enable_profiling, "Backward pass"):
-                        loss.backward()
+                    with nvprof_context(enable_profiling, "Backward pass"):
+                        scaled_loss = loss / float(self.acc_freq)
+                        scaled_loss.backward()

---

785-806: Fix off-by-one accumulation trigger and flush remainder on the final iteration

Current trigger uses (_step_id + 1) % acc_freq == 0 and ignores start_step; it also skips flushing a partial accumulation on the last batch when (num_steps - start_step) % acc_freq != 0. This silently drops gradients.

-                # gradient accumulation
-                if (_step_id + 1) % self.acc_freq == 0:
+                # gradient accumulation
+                accum_step = (_step_id - self.start_step + 1)
+                is_last_iter = (_step_id + 1) == self.num_steps
+                if (accum_step % self.acc_freq == 0) or is_last_iter:
                     # fuse + allreduce manually before optimization if use DDP + no_sync
                     # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622
                     if self.world_size > 1:
                         hpu.fused_allreduce_gradients(
                             list(self.wrapper.parameters()), None
                         )
 
                     if self.gradient_max_norm > 0.0:
                         with nvprof_context(enable_profiling, "Gradient clip"):
                             paddle.nn.utils.clip_grad_norm_(
                                 self.wrapper.parameters(),
                                 self.gradient_max_norm,
                                 error_if_nonfinite=True,
                             )
                     with nvprof_context(enable_profiling, "Adam update"):
                         self.optimizer.step()
                     self.optimizer.clear_grad(set_to_zero=False)
                     self.scheduler.step()

source/tests/pd/test_training.py (1)

161-172: Make the accumulation test actually accumulate and avoid display-time interference

With numb_steps=1 and acc_freq=4, no accumulation occurs, and step-1 display will zero grads in current implementation. Run at least acc_freq steps, disable in-training display, and optionally assert prim is enabled.

     def setUp(self) -> None:
@@
         self.config["model"] = deepcopy(model_se_e2_a)
-        self.config["training"]["numb_steps"] = 1
+        # Use acc_freq steps so we actually accumulate before an optimizer step
+        self.config["training"]["numb_steps"] = 4
         self.config["training"]["save_freq"] = 1
         self.config["training"]["acc_freq"] = 4
+        # Avoid display-time eval that may touch grads during accumulation
+        self.config["training"]["disp_training"] = False
         enable_prim(True)
+        # Optional sanity check
+        # assert paddle.framework.core._is_eager_prim_enabled(), \
+        #     "Eager prim should be enabled for gradient accumulation tests"

deepmd/pd/train/training.py (1)

811-816: Optional: only run heavy display/validation at accumulation boundaries

To avoid extra compute and any future risk of interfering with in-flight accumulation, consider logging only when you step the optimizer or on the very first batch.

Example gate:

-            if self.display_in_training and (
-                display_step_id % self.disp_freq == 0 or display_step_id == 1
-            ):
+            if self.display_in_training and (
+                display_step_id % self.disp_freq == 0 or display_step_id == 1
+            ) and (
+                ((_step_id + 1 - self.start_step) % self.acc_freq == 0) or ((_step_id + 1) == self.num_steps)
+            ):

source/tests/pd/test_training.py (1)

161-172: Optional: add a remainder-flush test case

To cover the final-step flush logic, consider setting numb_steps to acc_freq+1 (e.g., 5 with acc_freq=4) and asserting one optimizer update for the remainder. I can provide a concrete test if desired.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

deepmd/utils/argcheck.py (1)

3171-3171: Clarify acc_freq semantics and constraints in docs

Explicitly state meaning, constraint (>=1), and effective batch size to avoid user confusion.

-    doc_acc_freq = "The accumulation steps for the gradients."
+    doc_acc_freq = (
+        "Number of mini-batches to accumulate gradients before applying optimizer "
+        "and scheduler updates. Must be >= 1. Effective batch size = batch_size * acc_freq."
+    )

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.