update for dpa3_dynamic debug code

HydrogenSulfate · HydrogenSulfate · commit d2a46b72638b · 2025-07-21T01:06:13.000+08:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -65,13 +65,13 @@ repos:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$)
   # markdown, yaml, CSS, javascript
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: v4.0.0-alpha.8
-    hooks:
-      - id: prettier
-        types_or: [markdown, yaml, css]
-        # workflow files cannot be modified by pre-commit.ci
-        exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
+  # - repo: https://github.com/pre-commit/mirrors-prettier
+  #   rev: v4.0.0-alpha.8
+  #   hooks:
+  #     - id: prettier
+  #       types_or: [markdown, yaml, css]
+  #       # workflow files cannot be modified by pre-commit.ci
+  #       exclude: ^(source/3rdparty|\.github/workflows|\.clang-format)
   # Shell
   - repo: https://github.com/scop/pre-commit-shfmt
     rev: v3.12.0-1
@@ -83,25 +83,25 @@ repos:
     hooks:
       - id: cmake-format
       #- id: cmake-lint
-  - repo: https://github.com/njzjz/mirrors-bibtex-tidy
-    rev: v1.13.0
-    hooks:
-      - id: bibtex-tidy
-        args:
-          - --curly
-          - --numeric
-          - --align=13
-          - --blank-lines
-          # disable sort: the order of keys and fields has explict meanings
-          #- --sort=key
-          - --duplicates=key,doi,citation,abstract
-          - --merge=combine
-          #- --sort-fields
-          #- --strip-comments
-          - --trailing-commas
-          - --encode-urls
-          - --remove-empty-fields
-          - --wrap=80
+  # - repo: https://github.com/njzjz/mirrors-bibtex-tidy
+  #   rev: v1.13.0
+  #   hooks:
+  #     - id: bibtex-tidy
+  #       args:
+  #         - --curly
+  #         - --numeric
+  #         - --align=13
+  #         - --blank-lines
+  #         # disable sort: the order of keys and fields has explict meanings
+  #         #- --sort=key
+  #         - --duplicates=key,doi,citation,abstract
+  #         - --merge=combine
+  #         #- --sort-fields
+  #         #- --strip-comments
+  #         - --trailing-commas
+  #         - --encode-urls
+  #         - --remove-empty-fields
+  #         - --wrap=80
   # license header
   - repo: https://github.com/Lucas-C/pre-commit-hooks
     rev: v1.5.5
diff --git a/deepmd/pd/model/descriptor/repflows.py b/deepmd/pd/model/descriptor/repflows.py
@@ -5,6 +5,7 @@
     Union,
 )
 
+import numpy as np
 import paddle
 
 from deepmd.dpmodel.utils.seed import (
@@ -493,7 +494,7 @@ def forward(
         # nf x nloc x a_nnei x a_nnei
         # 1 - 1e-6 for paddle.acos stability
         cosine_ij = paddle.matmul(normalized_diff_i, normalized_diff_j) * (1 - 1e-6)
-        angle_input = cosine_ij.unsqueeze(-1) / (paddle.pi**0.5)
+        angle_input = cosine_ij.unsqueeze(-1) / (np.pi**0.5)
 
         if not parallel_mode and self.use_loc_mapping:
             assert mapping is not None
diff --git a/deepmd/pd/model/model/transform_output.py b/deepmd/pd/model/model/transform_output.py
@@ -27,6 +27,7 @@ def atomic_virial_corr(
     ce = coord * atom_energy
     sumce0, sumce1, sumce2 = paddle.split(paddle.sum(ce, axis=1), [1, 1, 1], axis=-1)
     # faked_grad = paddle.ones_like(sumce0)
+    raise
     extended_virial_corr0 = paddle.autograd.grad(
         [sumce0],
         [extended_coord],
@@ -93,6 +94,9 @@ def task_deriv_one(
         )
     else:
         extended_virial = None
+    print(
+        f"extended_force: {extended_force.min().item():.10f} {extended_force.max().item():.10f} {extended_force.mean().item():.10f} {extended_force.std().item():.10f}"
+    )
     return extended_force, extended_virial
 
 
diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py
@@ -760,6 +760,38 @@ def step(_step_id, task_key="Default") -> None:
                     pref_lr = _lr.start_lr
                 else:
                     pref_lr = cur_lr
+                self.wrapper.load_state_dict(paddle.load("./wrapper_dict.pd"))
+                print("model loaded")
+                inp = np.load("./input_dict.npz", allow_pickle=True)
+                for k, v in inp.items():
+                    if isinstance(v, np.ndarray):
+                        # print(k, type(v), v.shape, v.dtype)
+                        try:
+                            input_dict[k] = paddle.to_tensor(v)
+                            # print(k)
+                        except Exception:
+                            pass
+                        if isinstance(input_dict[k], paddle.Tensor):
+                            input_dict[k] = input_dict[k].cuda()
+                print("input_dict loaded")
+                lab = np.load("./label_dict.npz", allow_pickle=True)
+                for k, v in lab.items():
+                    if isinstance(v, np.ndarray):
+                        # print(k, type(v), v.shape, v.dtype)
+                        try:
+                            label_dict[k] = paddle.to_tensor(v)
+                            # print(k)
+                        except Exception:
+                            pass
+                        if isinstance(label_dict[k], paddle.Tensor):
+                            label_dict[k] = label_dict[k].cuda()
+                print("label_dict loaded")
+                model_pred, loss, more_loss = self.wrapper(
+                    **input_dict, cur_lr=pref_lr, label=label_dict, task_key=task_key
+                )
+                print({k: float(v) for k, v in more_loss.items()})
+                print(f"{loss.item():.10f}")
+                exit()
 
                 # disable synchronization in forward-backward manually
                 # as derivatives exist in model forward
@@ -779,7 +811,7 @@ def step(_step_id, task_key="Default") -> None:
 
                     with nvprof_context(enable_profiling, "Backward pass"):
                         loss.backward()
-
+                exit()
                 # 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:
diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py
@@ -173,6 +173,7 @@ def forward(
             model_pred = self.model[task_key](**input_dict)
             return model_pred, None, None
         else:
+            print(self.loss)
             natoms = atype.shape[-1]
             model_pred, loss, more_loss = self.loss[task_key](
                 input_dict,
diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py
@@ -96,6 +96,9 @@ def task_deriv_one(
         extended_virial = extended_virial.view(list(extended_virial.shape[:-2]) + [9])  # noqa:RUF005
     else:
         extended_virial = None
+    print(
+        f"extended_force: {extended_force.min().item():.10f} {extended_force.max().item():.10f} {extended_force.mean().item():.10f} {extended_force.std().item():.10f}"
+    )
     return extended_force, extended_virial
 
 
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -725,9 +725,64 @@ def step(_step_id, task_key="Default") -> None:
                     pref_lr = _lr.start_lr
                 else:
                     pref_lr = cur_lr
+
+                # save
+                # torch.save(self.wrapper.state_dict(), "wrapper_dict.pt")
+                # import paddle
+                # psd = {}
+                # for k, v in self.wrapper.state_dict().items():
+                #     if isinstance(v, torch.Tensor):
+                #         psd[k] = paddle.from_dlpack(v.detach())
+                #     else:
+                #         psd[k] = v
+                # paddle.save(psd, "wrapper_dict.pd")
+                # inp = {}
+                # for k, v in input_dict.items():
+                #     if isinstance(v, torch.Tensor):
+                #         inp[k] = v.detach().cpu().numpy()
+                #     else:
+                #         inp[k] = v
+                # np.savez("./input_dict.npz", **inp)
+                # lab = {}
+                # for k, v in label_dict.items():
+                #     if isinstance(v, torch.Tensor):
+                #         lab[k] = v.detach().cpu().numpy()
+                #     else:
+                #         lab[k] = v
+                # np.savez("./label_dict.npz", **lab)
+
+                # load
+                self.wrapper.load_state_dict(torch.load("./wrapper_dict.pt"))
+                print("model loaded")
+                inp = np.load("./input_dict.npz", allow_pickle=True)
+                for k, v in inp.items():
+                    if isinstance(v, np.ndarray):
+                        # print(k, type(v), v.shape, v.dtype)
+                        try:
+                            input_dict[k] = torch.tensor(v)
+                        except TypeError:
+                            pass
+                        if isinstance(input_dict[k], torch.Tensor):
+                            input_dict[k] = input_dict[k].cuda()
+                print("input_dict loaded")
+                lab = np.load("./label_dict.npz", allow_pickle=True)
+                for k, v in lab.items():
+                    if isinstance(v, np.ndarray):
+                        # print(k, type(v), v.shape, v.dtype)
+                        try:
+                            label_dict[k] = torch.tensor(v)
+                        except TypeError:
+                            pass
+                        if isinstance(label_dict[k], torch.Tensor):
+                            label_dict[k] = label_dict[k].cuda()
+                print("label_dict loaded")
+
                 model_pred, loss, more_loss = self.wrapper(
                     **input_dict, cur_lr=pref_lr, label=label_dict, task_key=task_key
                 )
+                print({k: float(v) for k, v in more_loss.items()})
+                print(f"{loss.item():.10f}")
+                exit()
                 loss.backward()
                 if self.gradient_max_norm > 0.0:
                     torch.nn.utils.clip_grad_norm_(
diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py
@@ -175,6 +175,7 @@ def forward(
             return model_pred, None, None
         else:
             natoms = atype.shape[-1]
+            print(self.loss)
             model_pred, loss, more_loss = self.loss[task_key](
                 input_dict,
                 self.model[task_key],
