feat: log pre-clip gradient total_norm and per-param norm histogram to tensorboard

Author: OutisLi

deepmd/pt/train/training.py

@@ -991,6 +991,13 @@ def step(_step_id: int, task_key: str = "Default") -> None:
                 )
                 loss.backward()
                 if self.gradient_max_norm > 0.0:
+                    # Collect per-parameter gradient norms before clipping.
+                    if self.enable_tensorboard:
+                        pre_clip_named_norms = [
+                            (name, p.grad.detach().norm().item())
+                            for name, p in self.wrapper.named_parameters()
+                            if p.grad is not None
+                        ]
                     # FSDP2 sharded DTensor gradients don't support error_if_nonfinite; use manual isfinite check instead.
                     total_norm = torch.nn.utils.clip_grad_norm_(
                         self.wrapper.parameters(),
@@ -1342,6 +1349,25 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
                     writer.add_scalar(
                         f"{task_key}/{item}", more_loss[item], display_step_id
                     )
+                # === Gradient diagnostics (pre-clip) ===
+                if self.gradient_max_norm > 0.0:
+                    writer.add_scalar(
+                        "grad/total_norm", total_norm.item(), display_step_id
+                    )
+                    norms = torch.tensor(
+                        [gn for _, gn in pre_clip_named_norms],
+                        dtype=torch.float32,
+                        device="cpu",
+                    )
+                    writer.add_histogram("grad/param_norms", norms, display_step_id)
+                    # Log top-10 largest per-parameter gradient norms.
+                    # Shorten name: keep everything after "atomic_model.".
+                    pre_clip_named_norms.sort(key=lambda x: x[1], reverse=True)
+                    for name, gn in pre_clip_named_norms[:10]:
+                        idx = name.find("atomic_model.")
+                        if idx >= 0:
+                            name = name[idx + len("atomic_model.") :]
+                        writer.add_scalar(f"grad_top10/{name}", gn, display_step_id)
 
         self.wrapper.train()
         self.t0 = time.time()