test(pt_expt): shrink change-bias water dataset to 5 frames (deepmodeling#5467)

## Summary

`TestChangeBias` is the dominant memory hog in `Test Python` shard `(10,
3.13)` of the CI matrix — by itself it peaks at **~5 GB RSS**, leaving
so little headroom under the 7 GB GitHub-hosted runner that the shard
intermittently loses communication with the GitHub Actions server. This
causes the recurring `runner lost communication` failure that has
affected many recent PRs (deepmodeling#5446, deepmodeling#5448, deepmodeling#5450, deepmodeling#5455, deepmodeling#5456, …).

This PR shrinks the change-bias test dataset from 80 frames to 5 frames,
dropping the class's peak RSS to **~1.7 GB** while keeping all 9 tests
passing — including the strict `atol=1e-10` `pt2_pte_consistency` check.

## How I located it

Local reproduction of shard `(10, 3.13)` (using the same
`.test_durations` cache CI uses, so identical test partitioning):

| Test profiled in isolation | Peak RSS |
|---|---|
| **`test_change_bias_frozen_pte`** | **5.04 GB** ← outlier |
| `TestDeepEvalEnerPt2` class | 1.43 GB |
| `TestDeepEvalEnerAparamPt2` class | 1.44 GB |
| `TestSpinInference::test_get_use_spin` | 1.41 GB |
| `test_finetune_from_pt2_use_pretrain_script` | 1.41 GB |
| `test_training_loop_compiled` | 1.32 GB |
| `test_export_pipeline` | 1.57 GB |
| `test_descriptor_shape_dpa1` | 1.34 GB |

Then phase-by-phase RSS profiling inside `dp change-bias` showed the 4.3
GB jump happens entirely inside `compute_output_stats` →
`_compute_model_predict`. Scaling experiment confirms it: peak grows
**linearly at ~50 MB per frame** of input data.

| nbatches | Peak RSS | per-frame |
|---|---|---|
| 1  | 567 MB | — |
| 5  | 781 MB | +43 MB |
| 20 | 1583 MB | +53 MB |
| 80 | 4797 MB | +53 MB |

That's a leak in the `torch.no_grad()`-wrapped `forward_common_atomic`
somewhere — separate from autograd. The water example has 80 frames at
batch_size=1, so the CLI default `nbatches = min(data.get_nbatches()) =
80` triggers all 80 forwards in one go.

## Why I can't just pass `-n 5`

`_load_batch_set` shuffles when it loads the set. If `nbatches <
total_frames`, the loop samples a random subset — and the two calls in
`test_change_bias_pt2_pte_consistency` (running in the **same Python
process** via `main(cmds)`, with `dp_random`'s state advancing between
calls) would see **different** subsets → different biases → the
`atol=1e-10` assertion fails.

`nbatches == total_frames` makes the forward enumerate **every** frame
regardless of shuffle order, so the aggregate bias is invariant under
shuffle. Determinism is preserved.

## The fix

Build a 5-frame subset of `examples/water/data/data_0` in
`TestChangeBias.setUpClass` and point both the trainer config and the
change-bias `-s` argument at it. `nbatches` then resolves to 5 (= the
new dataset size = full enumeration), and all 9 tests pass with peak RSS
at ~1.7 GB.

## Test plan

- [x] All 9 tests in `TestChangeBias` pass locally (CPU fp64):
  - `test_change_bias_with_data`
  - `test_change_bias_with_data_sys_file`
  - `test_change_bias_with_user_defined`
  - `test_change_bias_frozen_pte`
  - `test_change_bias_frozen_pt2`
  - `test_change_bias_frozen_pt2_user_defined`
- `test_change_bias_pt2_pte_consistency` (atol=1e-10 — the
determinism-sensitive one)
  - `test_change_bias_pte_preserves_model_def_script`
  - `test_change_bias_pt2_preserves_model_def_script`
- [x] Peak RSS measurement (kernel `ru_maxrss`):
  - Before: 5.66 GB
  - After: **1.62 GB** (single test) / **1.75 GB** (whole class)
- [ ] CI shard `(10, 3.13)` confirms no more `runner lost communication`
on this branch (pending CI run)

## Known limitations

- **The underlying ~50 MB/frame leak in `forward_common_atomic` remains
as a production bug.** Users running `dp change-bias` on real datasets
with thousands of frames will see multi-GB RSS growth. Worth a separate
follow-up to find and patch the leak.
- The 5-frame number is somewhat arbitrary. It's chosen as the smallest
value that (a) keeps the assertion logic working, (b) leaves room for
the least-squares regression (need ≥ ntypes = 2 frames).
- `_make_subset_dataset` only handles `set.000`; multi-set datasets
would need extension. Not needed for water/data_0.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Tests**
* Reduced the dataset used by a change-bias test to a small, fixed
subset (5 frames) so test runs use far less disk space and complete
faster.
* Test setup now builds and points to the truncated dataset for all
related invocations, lowering resource overhead during CI and local
testing.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/deepmodeling/deepmd-kit/pull/5467?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

source/tests/pt_expt/test_change_bias.py

@@ -118,19 +118,68 @@ def _make_config(data_dir: str) -> dict:
     }
 
 
+def _make_subset_dataset(src_system: str, dst_system: str, n_frames: int) -> None:
+    """Copy ``type{,_map}.raw`` and the first ``n_frames`` of every ``.npy``
+    in ``set.000`` from ``src_system`` to ``dst_system``.
+
+    Used by ``TestChangeBias`` to shrink the water/data_0 example (80
+    frames) down to a tiny subset so that ``dp change-bias`` enumerates
+    over only ``n_frames`` frames.  Why this matters: the in-process
+    ``main(cmds)`` path runs the model forward over ``nbatches`` frames
+    via ``compute_output_stats``, and each frame leaks ~50 MB into
+    torch's caching allocator.  At ``n_frames=80`` (the default,
+    ``min(data.get_nbatches()) = 80``) peak RSS hits ~5 GB which OOMs
+    the 7 GB GitHub-hosted CI runner.  Shrinking to ``n_frames=5`` keeps
+    peak at ~800 MB while preserving **determinism**: the test
+    ``test_change_bias_pt2_pte_consistency`` asserts ``atol=1e-10``
+    between two .pte and .pt2 calls in the same process, which requires
+    every frame to be seen on each call regardless of the
+    shuffle-based ``_load_batch_set`` order.  ``nbatches == total
+    frames`` makes the forward enumerate every frame and so the
+    aggregate bias is invariant under shuffle.
+    """
+    src_set = os.path.join(src_system, "set.000")
+    dst_set = os.path.join(dst_system, "set.000")
+    os.makedirs(dst_set, exist_ok=True)
+    for raw in ("type.raw", "type_map.raw"):
+        src = os.path.join(src_system, raw)
+        if os.path.isfile(src):
+            shutil.copyfile(src, os.path.join(dst_system, raw))
+    for fname in os.listdir(src_set):
+        if not fname.endswith(".npy"):
+            continue
+        arr = np.load(os.path.join(src_set, fname))
+        np.save(os.path.join(dst_set, fname), arr[:n_frames])
+
+
 class TestChangeBias(unittest.TestCase):
     """Test dp change-bias for the pt_expt backend."""
 
     @classmethod
     def setUpClass(cls) -> None:
-        data_dir = os.path.join(EXAMPLE_DIR, "data")
-        if not os.path.isdir(data_dir):
-            raise unittest.SkipTest(f"Example data not found: {data_dir}")
+        full_data_dir = os.path.join(EXAMPLE_DIR, "data")
+        if not os.path.isdir(full_data_dir):
+            raise unittest.SkipTest(f"Example data not found: {full_data_dir}")
+        cls.tmpdir = tempfile.mkdtemp()
+        cls.old_cwd = os.getcwd()
+
+        # Shrink the water example dataset (80 frames) to a 5-frame
+        # subset.  ``dp change-bias`` defaults to enumerating every
+        # frame (``nbatches = min(data.get_nbatches())``), and each
+        # frame's forward pass leaks ~50 MB into torch's allocator; at
+        # 80 frames peak RSS pushes the 7 GB CI runner into OOM.  See
+        # the docstring of ``_make_subset_dataset`` for why we keep
+        # full enumeration (determinism) but shrink the dataset.
+        data_dir = os.path.join(cls.tmpdir, "data")
+        os.makedirs(data_dir, exist_ok=True)
+        _make_subset_dataset(
+            src_system=os.path.join(full_data_dir, "data_0"),
+            dst_system=os.path.join(data_dir, "data_0"),
+            n_frames=5,
+        )
         cls.data_dir = data_dir
         cls.data_file = [os.path.join(data_dir, "data_0")]
 
-        cls.tmpdir = tempfile.mkdtemp()
-        cls.old_cwd = os.getcwd()
         os.chdir(cls.tmpdir)
 
         # Build & train 1-step model