perf: bound memory in model testing and output statistics - #5944
Conversation
Testing read every frame of a system before looking at any of them. On an LMDB of 38.8 million frames that meant decoding the whole 54 GB set into memory, and `-n` was applied only afterwards, so the default run of 100 frames paid for all of them while the GPU sat idle. Frames are now chosen before anything is read -- atom counts come from the LMDB metadata, so grouping, shuffling and truncation are index arithmetic -- and a system is walked in chunks bounded by atom count, which also lifts the ceiling on how large a dataset can be tested. Chunking is exact rather than an approximation. An MAE is the mean of the absolute errors and an RMSE the root of the mean of their squares, so both are recovered from partial results weighted by the number of elements each was taken over. `merge_weighted_errors` performs that combination, and `weighted_average`, which already combined systems the same way, is now expressed in terms of it. The five test routines had drifted into five shapes of one thing. They are now a single skeleton -- declare the labels, evaluate chunk by chunk, combine, report -- with a tester per model class supplying only what differs. Choosing a tester replaces two parallel isinstance chains, one selecting the routine and one its printer, because the report of a tester drives the per-system and the run-level table alike. That machinery lives in `deepmd/infer/model_test`, beside `model_devi`, which is the same kind of backend-independent analysis driven by a command-line entry point. `deepmd/entrypoints/test.py` keeps system discovery and the run-level report, which returns it to the size of every other entry point. A spin model now reports the virial and the stress. Its magnetic degrees of freedom reach the virial only through the virtual atoms, whose displacement the model already removes, so the virial is with respect to the real atomic positions as for any other energy model; the quantity was computed and then discarded. Deriving the spin tester from the energy one leaves it overriding the force alone, and the exclusions vanish with the branches that carried them. Two defects surface in the same code. A system split into sub-groups, as a mixed-nloc LMDB is, had each group overwrite the detail file of the last, because the append flag tracked the system instead of the group. And `test_wfc` was unreachable while the dispatch still carried its printer, so a wave-function model would have failed on an unbound name. Verified against the previous implementation over the same frames: 144 metrics agree to 8.3e-07, below the 9.98e-07 spread between two runs of one unchanged implementation, which is nondeterministic on this GPU. Chunked and unchunked evaluation agree to 4.6e-07. Tensor detail headers are identical to the previous ones for both model classes, atomic and not, over a range of selected atom counts. The pt, pt_expt and dpmodel suites pass, and a JAX model was trained, frozen and tested end to end. The TensorFlow backend is not built here, so its path was reviewed rather than run.
…ares The output-bias forward wrapper is the only caller of an atomic model that starts from raw coordinates, so it constructs the neighbor input itself. It always built a fixed-capacity neighbor list sized by `get_sel()`, even though the model already declares through `uses_graph_lower()` which representation it consumes, and already implements both `forward_common_atomic` and `forward_common_atomic_graph`. Every caller that starts from an extended input honours that declaration; this one did not. A graph-native model reports no finite neighbor capacity, so sizing a dense list from `get_sel()` is not merely wasteful there: the allocation is unbounded and the index array alone reaches tens of gigabytes on a few hundred atoms. Fine-tuning ran out of memory at `change_out_bias`, which drives this wrapper over sampled training frames. The wrapper now builds a carry-all `NeighborGraph` when the model is graph-native and keeps the neighbor list otherwise. Pair exclusion stays a build-time transform on both routes, folded into `edge_mask` by the graph builder. The graph route works on a flat node axis, so its result is restored to the per-frame layout the dense route returns and the contract with `compute_output_stats` is unchanged. The routing tests a capability, not a descriptor, so any graph-native model is covered. Backends whose atomic models never report `uses_graph_lower` keep their own dense wrappers untouched.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughThe change adds graph-native atomic forwarding and replaces legacy model testing with centralized, chunked testers. LMDB data loads lazily with frame limits. The ChangesGraph-native atomic forwarding
Modular model testing
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant DPTest
participant LmdbTestData
participant ModelTester
participant ModelEvaluator
DPTest->>LmdbTestData: load retained frames with max_frames
DPTest->>ModelTester: build_tester and run
LmdbTestData->>ModelTester: yield atom-count chunks
ModelTester->>ModelEvaluator: evaluate_chunk
ModelEvaluator-->>ModelTester: weighted errors and details
ModelTester-->>DPTest: merged errors
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (7)
deepmd/infer/model_test/property.py (1)
41-48: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider requiring the atomic label when
--atomicis set.
atom_{var_name}is registered withmust=False. If the label file is absent,DeepmdDatasupplies a default-filled array.evaluate_chunkthen compares the prediction against zeros and reportsmae_apropertyandrmse_apropertyas if a label existed.DosTesterregistersatom_doswithmust=Truefor the same situation. Align the two, or gate the atomic metrics onfind_atom_{var_name}.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/infer/model_test/property.py` around lines 41 - 48, The atomic label registration in the property test’s atomic branch should require an actual label when --atomic is enabled. Update the `data.add` call for `atom_{var_name}` to match `DosTester`’s required-label behavior, or alternatively ensure `evaluate_chunk` only computes atomic metrics when `find_atom_{var_name}` confirms the label exists.deepmd/dpmodel/utils/lmdb_data.py (1)
2430-2435: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse the project random generator for test shuffling.
_select_framescreates an unseedednp.random.default_rng().DeepmdData._shuffle_datausesdeepmd.utils.random, which the seed configuration controls. As written, LMDB test frame selection is not reproducible across runs, even with a fixed seed.♻️ Proposed change to reuse the shared generator
- rng = np.random.default_rng() groups: dict[int, list[int]] = {} for begin, end in pairwise(starts): indices = order[begin:end] if shuffle_test: - indices = rng.permutation(indices) + indices = indices.copy() + dp_random.shuffle(indices)Add the import:
from deepmd.utils import random as dp_random🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/utils/lmdb_data.py` around lines 2430 - 2435, Update _select_frames to use the project-controlled generator from deepmd.utils.random (import it as dp_random) instead of creating an unseeded np.random.default_rng(), and call its permutation operation for test-frame shuffling so configured seeds remain reproducible.deepmd/infer/model_test/base.py (2)
266-268: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove or correct the section banners.
base.pyno longer holds the energy or tensor testers; they live indeepmd/infer/model_test/ener.pyanddeepmd/infer/model_test/tensor.py. The "Energy models" banner now sits above_write_per_frame_details, which the DOS and property testers use. The "Tensor models" banner closes the file with no content under it.Also applies to: 308-310
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/infer/model_test/base.py` around lines 266 - 268, Remove or relocate the obsolete “Energy models” and “Tensor models” section banners in base.py so they no longer label unrelated helpers or an empty section. Keep the actual energy and tensor tester organization represented by the corresponding symbols in ener.py and tensor.py.
62-62: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicate logger assignment.
logis already defined at line 33. Line 62 repeats the same assignment.♻️ Proposed cleanup
-log = logging.getLogger(__name__) - - def save_txt_file(Keep the definition at line 33.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/infer/model_test/base.py` at line 62, Remove the duplicate logging.getLogger(__name__) assignment near the later location in base.py, keeping the existing log definition established earlier in the module.deepmd/entrypoints/test.py (1)
205-212: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe reported system count can disagree with the merged results.
Line 211 logs
len(all_sys). For a mixed-nloc LMDB system,err_collholds one entry per nloc group, as the comment at lines 205-206 states. The run-level header then understates the number of merged results. Consider logginglen(err_coll)as well.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/entrypoints/test.py` around lines 205 - 212, Update the weighted-average summary near tester.log_errors to report the merged result count using len(err_coll), alongside or instead of the all_sys count, so mixed-nloc LMDB runs accurately reflect the entries being merged.source/tests/pt/test_dp_test.py (1)
427-451: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for a multi-chunk run.
Both tests use
numb_test=1on a single-frame system, sorunexecutes exactly one chunk. The chunked aggregation and the detail-file appending across chunks stay untested.test_chunk_atoms()readsDP_TEST_CHUNK_ATOMS, so a test can force several chunks over a multi-frame system and then compare the merged MAE and RMSE against a single-chunk run.Run a single case with
pytest source/tests/pt/test_dp_test.py::TestDPTestStress::test_stress -v.I can draft the multi-chunk test if you want it.
As per coding guidelines: "Use
pytestfor testing single test cases ... instead of full test suite (60+ minutes)".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/test_dp_test.py` around lines 427 - 451, The stress test currently covers only a single chunk; extend TestDPTestStress.test_stress to use a multi-frame system and force multiple chunks through the DP_TEST_CHUNK_ATOMS setting consumed by test_chunk_atoms(). Compare the chunked run’s merged detail arrays and reported mae_s/rmse_s against an equivalent single-chunk run, preserving the existing stress/virial and volume validations.Source: Coding guidelines
source/tests/pt/test_weighted_avg.py (1)
227-230: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWrite the virial label into a temporary copy of the system.
np.savewritesvirial.npyinto the tracked test-data directoryNiO/data/single/set.000.tearDownremoves it, but an interrupted or crashed run leaves the file behind.test_dp_test_ener_with_spinin the same class then asserts thatmae_vis absent and fails. The two tests are coupled through shared on-disk state.
TestDPTestForceWeight._prepare_weighted_systemandTestDPTestStress._prepare_virial_systeminsource/tests/pt/test_dp_test.pycopy the system withshutil.copytreeinto a temp directory first. Use the same approach here.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/test_weighted_avg.py` around lines 227 - 230, Update the test setup around the np.save call to create and use a temporary copy of the system via shutil.copytree, matching TestDPTestForceWeight._prepare_weighted_system and TestDPTestStress._prepare_virial_system. Save virial.npy only in that copied system and pass the temporary path through the existing test flow, leaving tracked test data untouched.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/atomic_model/base_atomic_model.py`:
- Around line 869-875: Update the call to forward_common_atomic_graph in the
surrounding atomic forward branch to reshape aparam from (nframes, nloc, nda)
into (-1, nda) before passing it, matching the existing flattening of atype
while preserving the required per-atom parameter ordering.
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 2823-2833: Update LmdbTestDataNlocView.iter_test to honor
self._frame_indices, matching LmdbTestDataNlocView.get_test, while preserving
chunk_atoms, numb_test, and self._nloc behavior. Delegate iteration through the
selected frame-index subgroup so ModelTester.run tests only the view’s frames
and avoids mixing availability groups.
In `@deepmd/infer/model_test/__init__.py`:
- Around line 92-111: Update build_tester to recognize DeepWFC before the
fallback RuntimeError and return the appropriate WFC/tensor tester for it.
Preserve the existing dispatch behavior for DeepPot, DeepDOS, DeepProperty,
DeepGlobalPolar, DeepPolar, and DeepDipole.
In `@deepmd/infer/model_test/base.py`:
- Around line 291-305: Update the per-frame output logic in the detail-file loop
around save_txt_file so each frame file is written fresh by disabling
append_detail for these files. Also include a system-unique component in the
filename derived from context.system, while preserving the existing frame and
suffix information, so multi-system runs cannot overwrite or combine unrelated
frames.
---
Nitpick comments:
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 2430-2435: Update _select_frames to use the project-controlled
generator from deepmd.utils.random (import it as dp_random) instead of creating
an unseeded np.random.default_rng(), and call its permutation operation for
test-frame shuffling so configured seeds remain reproducible.
In `@deepmd/entrypoints/test.py`:
- Around line 205-212: Update the weighted-average summary near
tester.log_errors to report the merged result count using len(err_coll),
alongside or instead of the all_sys count, so mixed-nloc LMDB runs accurately
reflect the entries being merged.
In `@deepmd/infer/model_test/base.py`:
- Around line 266-268: Remove or relocate the obsolete “Energy models” and
“Tensor models” section banners in base.py so they no longer label unrelated
helpers or an empty section. Keep the actual energy and tensor tester
organization represented by the corresponding symbols in ener.py and tensor.py.
- Line 62: Remove the duplicate logging.getLogger(__name__) assignment near the
later location in base.py, keeping the existing log definition established
earlier in the module.
In `@deepmd/infer/model_test/property.py`:
- Around line 41-48: The atomic label registration in the property test’s atomic
branch should require an actual label when --atomic is enabled. Update the
`data.add` call for `atom_{var_name}` to match `DosTester`’s required-label
behavior, or alternatively ensure `evaluate_chunk` only computes atomic metrics
when `find_atom_{var_name}` confirms the label exists.
In `@source/tests/pt/test_dp_test.py`:
- Around line 427-451: The stress test currently covers only a single chunk;
extend TestDPTestStress.test_stress to use a multi-frame system and force
multiple chunks through the DP_TEST_CHUNK_ATOMS setting consumed by
test_chunk_atoms(). Compare the chunked run’s merged detail arrays and reported
mae_s/rmse_s against an equivalent single-chunk run, preserving the existing
stress/virial and volume validations.
In `@source/tests/pt/test_weighted_avg.py`:
- Around line 227-230: Update the test setup around the np.save call to create
and use a temporary copy of the system via shutil.copytree, matching
TestDPTestForceWeight._prepare_weighted_system and
TestDPTestStress._prepare_virial_system. Save virial.npy only in that copied
system and pass the temporary path through the existing test flow, leaving
tracked test data untouched.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: cfe48005-66e0-4b4e-b1ea-d47d1d0517d6
📒 Files selected for processing (14)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/entrypoints/test.pydeepmd/infer/model_test/__init__.pydeepmd/infer/model_test/base.pydeepmd/infer/model_test/dos.pydeepmd/infer/model_test/ener.pydeepmd/infer/model_test/property.pydeepmd/infer/model_test/tensor.pydeepmd/utils/data.pydeepmd/utils/weight_avg.pysource/tests/common/test_dp_test_ener_split.pysource/tests/pt/test_dp_test.pysource/tests/pt/test_weighted_avg.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5944 +/- ##
==========================================
- Coverage 79.47% 79.27% -0.20%
==========================================
Files 1072 1079 +7
Lines 125055 125466 +411
Branches 4541 4591 +50
==========================================
+ Hits 99388 99466 +78
- Misses 24043 24348 +305
- Partials 1624 1652 +28 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The design here is good and I want to say that first, because most of what follows is about things lost in the move rather than the idea. I checked the aggregation math, since everything else rests on it: MAE accumulates sum(e_i * w_i) over sum(w_i) and RMSE accumulates sum(e_i^2 * w_i) before the root, so chunked evaluation is exactly equal to one-shot, and a merged (error, weight) really is reusable as a further partial — the merge_weighted_errors docstring's claim holds. Frame order and numb_test semantics are preserved too: iter_test serves the same prefix get_test() did, with shuffle_test still applied once inside _load_test_set. Of the three incidental fixes claimed, two are real and correct — spin models now report virial/stress, and append_detail = bool(err_coll) genuinely fixes LMDB nloc-subgroups overwriting each other's detail files.
The problem is that a 1470-line move carried three regressions across, and the feature the PR exists to add has no test to have caught them. Four blocking comments inline, one non-blocking.
The third claimed fix, "restore reachable wave-function model dispatch", is inverted rather than wrong in a way that matters: build_tester has no DeepWFC branch and test_wfc was deleted, so WFC testing now raises RuntimeError. That is fine as an outcome — the old dispatcher had no reachable DeepWFC branch either, so test_wfc was already dead code and the explicit error beats the old UnboundLocalError — but the body describes it as a restoration when it is a removal. Worth correcting the sentence.
Two smaller things I am not asking you to change, recorded so they are not lost. The ModelTester and deepmd/infer/model_test/__init__.py docstrings say chunking is "what makes a dataset larger than memory testable", stated across all backends; that is true only for the lazy LmdbTestData path, and data.py's own iter_test docstring is honest that "a set is loaded as a whole, so chunking here bounds what a consumer holds at once rather than what is read". The two contradict each other and the narrower one is right. Separately, LmdbTestDataNlocView.iter_test delegates on nloc alone while get_test honours self._frame_indices; no in-tree caller passes frame_indices today so it is latent, but the two accessors now disagree and the next caller gets the wrong frames silently.
One point I considered raising and decided against: routing on uses_graph_lower() means output bias for dpa1/dpa2 now comes from the sel-free graph builder rather than the get_sel()-truncated dense list, so the value changes for existing models. dpa1's own docstring already discloses that divergence and this only affects bias initialisation, so I do not think it blocks — but gating on the absence of a finite sel, which is the actual motivation, would keep the change to the models that need it.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_lmdb_data.py (1)
790-806: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the expected
dp_randompermutation.Both instances can have equal
nloc_groupsif shuffling is disabled or uses another deterministic generator. Build the expected order from unshuffled groups withdp_random.shuffle()after the same seed. Then compare it withfirst.nloc_groups.Proposed test update
def test_test_data_shuffle_uses_global_seed(self): """The CLI random seed makes LMDB test-frame shuffling reproducible.""" self.addCleanup(dp_random.seed, None) + unshuffled = LmdbTestData( + self._lmdb_path, + type_map=self._type_map, + shuffle_test=False, + ) + expected = { + nloc: np.asarray(indices) + for nloc, indices in unshuffled.nloc_groups.items() + } dp_random.seed(123) + for indices in expected.values(): + dp_random.shuffle(indices) dp_random.seed(123) first = LmdbTestData( self._lmdb_path, type_map=self._type_map, shuffle_test=True, @@ self.assertEqual(first.nloc_groups, second.nloc_groups) + self.assertEqual( + first.nloc_groups, + {nloc: indices.tolist() for nloc, indices in expected.items()}, + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_lmdb_data.py` around lines 790 - 806, Update test_test_data_shuffle_uses_global_seed to construct the unshuffled nloc_groups, seed dp_random with 123, and apply dp_random.shuffle() to derive the expected permutation; then compare expected order with first.nloc_groups while retaining the reproducibility assertion for the second instance.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@source/tests/common/dpmodel/test_lmdb_data.py`:
- Around line 790-806: Update test_test_data_shuffle_uses_global_seed to
construct the unshuffled nloc_groups, seed dp_random with 123, and apply
dp_random.shuffle() to derive the expected permutation; then compare expected
order with first.nloc_groups while retaining the reproducibility assertion for
the second instance.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 23c06d3e-f73b-4eee-a8c4-e6f589df770e
📒 Files selected for processing (15)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/entrypoints/test.pydeepmd/infer/model_test/__init__.pydeepmd/infer/model_test/base.pydeepmd/infer/model_test/dos.pydeepmd/infer/model_test/ener.pydeepmd/infer/model_test/property.pydeepmd/infer/model_test/tensor.pydeepmd/utils/data.pysource/tests/common/dpmodel/test_dpa4_call_graph.pysource/tests/common/dpmodel/test_lmdb_data.pysource/tests/common/test_dp_test_ener_split.pysource/tests/pt/test_dp_test.pysource/tests/pt/test_weighted_avg.py
💤 Files with no reviewable changes (2)
- deepmd/infer/model_test/ener.py
- deepmd/infer/model_test/dos.py
🚧 Files skipped from review as they are similar to previous changes (8)
- deepmd/infer/model_test/init.py
- deepmd/entrypoints/test.py
- deepmd/dpmodel/atomic_model/base_atomic_model.py
- source/tests/pt/test_weighted_avg.py
- deepmd/utils/data.py
- deepmd/infer/model_test/base.py
- deepmd/dpmodel/utils/lmdb_data.py
- deepmd/infer/model_test/tensor.py
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 16 out of 16 changed files in this pull request and generated 1 comment.
Suppressed comments (3)
deepmd/dpmodel/utils/lmdb_data.py:2498
_frame_has_datatreats msgpack-encoded arrays as present, but it does not treat a list of encoded arrays as present (even though_decode_valuesupports that form). If LMDB frames store some labels as lists of encoded arrays, availability grouping and required-label enforcement can incorrectly mark these fields as missing. Consider extending the check to returnTruewhenvalueis a non-empty list whose first element satisfies_is_encoded_array.
find_key = f"find_{key}"
if find_key in frame:
return bool(float(np.asarray(_decode_value(frame[find_key])).item()))
value = frame.get(key)
if _is_encoded_array(value):
return True
return isinstance(value, (np.ndarray, np.generic, int, float, bool))
deepmd/infer/model_test/ener.py:564
- Virial/stress detail output is currently driven by
reference_virial/prediction_virialbeing non-None, but this call passes virial arrays unconditionally even whenreports_virialis false (e.g., no virial label, or non-PBC). That can produce.v.out(and possibly.s.out) files that look valid while the run intentionally did not compute/record virial metrics. To keep detail outputs consistent with the reported metrics, passreference_virial=Noneandprediction_virial=None(and similarly stress) whenreports_virialis false.
if context.detail_path is not None:
_write_energy_test_details(
detail_path=context.detail_path,
system=context.system,
natoms=natoms,
append_detail=context.append_detail,
reference_energy=test_data["energy"],
prediction_energy=energy,
reference_force=test_data["force"],
prediction_force=force,
reference_virial=test_data["virial"],
prediction_virial=virial,
reference_stress=reference_stress,
prediction_stress=prediction_stress,
out_put_spin=not self.reports_plain_force,
reference_force_real=force_details.reference_real,
prediction_force_real=force_details.prediction_real,
reference_force_magnetic=force_details.reference_magnetic,
prediction_force_magnetic=force_details.prediction_magnetic,
reference_hessian=test_data["hessian"] if dp.has_hessian else None,
prediction_hessian=optional_outputs.hessian if dp.has_hessian else None,
)
deepmd/utils/weight_avg.py:47
- The raised
RuntimeError(\"unknown error type\")doesn’t include the offending quantity name, which makes debugging callers harder (especially now that errors can be merged across chunks and systems). Consider includingkkin the exception message (e.g., mention the key and expected prefixes).
for kk, (ee, ss) in err.items():
if kk.startswith("mae"):
sum_err[kk] += ee * ss
elif kk.startswith("rmse"):
sum_err[kk] += ee * ee * ss
else:
raise RuntimeError("unknown error type")
njzjz-bot
left a comment
There was a problem hiding this comment.
I reviewed the current head after the follow-up fixes. The previously reported atomic-label, legacy TensorFlow spin dispatch, graph-route aparam, multi-chunk coverage, and repeated-header issues are addressed. The focused LMDB/DPA4/common regressions and the PyTorch stress/property/weighted-average regressions pass locally (16 tests total).
One remaining blocking detail-output issue is inline: atomic tensor results from systems or LMDB groups with different selected-atom counts are appended into one ragged file. CI for this head is still in progress.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
iProzd
left a comment
There was a problem hiding this comment.
Thanks for addressing grouped tensor detail paths and zero-selected-atom handling; both prior runtime blockers are resolved on 01fe8fd. One blocking issue remains: the new three-frame stress fixture still affects the inherited DPTest methods, which run all frames but compare the detail output against a single validation batch. The prior CI failure—576 force rows versus 192 in three tests—is therefore unchanged; the follow-up commit did not touch this fixture or helper. Please fix this case and add or adjust the focused regression before merging.
njzjz-bot
left a comment
There was a problem hiding this comment.
Approve — I reviewed the current head (b014db8) against the previously requested changes and the fresh diff.
The memory-bound chunked-testing design holds up: ModelTester.run walks a system via data.iter_test(chunk_atoms=...) (DP_TEST_CHUNK_ATOMS), LmdbTestData now groups by atom count with index arithmetic and decodes only retained frames (_select_frames/_read_frames), and merge_weighted_errors correctly recovers MAE/RMSE from partial (weighted) results — I verified the math on chunked vs single-chunk outputs in the stress regression.
All previously requested changes are resolved in the head:
atom_dipole/atom_polarizabilitykey normalization (d2756cd) and the legacy TF spin dispatch (has_spin or get_ntypes_spin() != 0) with theSpinEnerTesterregression — confirmed inmodel_test/__init__.py:build_tester.aparamflattened to the flat node axis on the output-stat graph route (base_atomic_model._get_forward_wrapper_func) withtest_output_stat_graph_wrapper_flattens_aparam.- Multi-chunk coverage (
DP_TEST_CHUNK_ATOMS=192) and one-header-per-file behavior in the stress regression. sel_natoms == 0guarded inTensorTester(returns{}before any division,rmse_sqrtn/rmse_nonly whensel_natoms).LmdbTestDataNlocView.iter_testhonors_frame_indices, and per-frame detail files are disambiguated bydetail_group+frame_offset(Verilon:_detail_output_path).- iProzd's blocking issue fixed via
_run_dp_testdefaulting tonumb_test=1, whiletest_stressexplicitly runs all three frames. - CodeQL findings gone: no unused
logglobals and no runtimeDeepPotimport in the newly splitmodel_testmodules.
I rebuilt the head locally, imported the refactored modules, and ran the new lightweight regressions (test_dp_test_ener_split.py, the four new test_lmdb_data.py cases, and the DPA4 output-stat graph test) — all pass, and ruff is clean on the changed Python.
Two non-blocking notes the author may want to address separately (not blocking approval):
deepmd/entrypoints/test.py:140siz_coll = []is now dead after the refactor; harmless but removable.DosTester/PropertyTestertightenatom_*labels frommust=Falsetomust=True; this surfaces a clear error instead of a lateKeyError, which is an improvement, just worth a line in the PR description.
Coding agent: opencode
opencode version: 1.18.9
Model: ustc/deepseek-v4-flash
Reasoning effort: max
njzjz
left a comment
There was a problem hiding this comment.
Approve
I reviewed the current head (b014db81c) in full, including the earlier review threads and the fresh diff, and I am approving.
Core design is sound
- Memory-bound chunked testing:
ModelTester.runwalks each system viadata.iter_test(chunk_atoms=DP_TEST_CHUNK_ATOMS); the weighted aggregation accumulatessum(e_i * w_i)/sum(e_i^2 * w_i)before rooting, so a merged (error, weight) partial is reusable and chunked MAE/RMSE equals the one-shot result exactly. - LMDB:
_select_frames/_read_framesdo frame selection by index arithmetic from metadata and decode only retained frames, grouping by nloc while honoring availability subgroups — this is what fixes the 38.8M-frame/54GB all-into-memory pathology with the GPU idle. - Output-bias stats: graph-native atomic models now route through a carry-all
NeighborGraph(build_neighbor_graph) instead of the fixed-capacity dense neighbor list (get_sel-sized), removing the unbounded-allocation failure mode.
Blocking issues from earlier rounds are all resolved
- Canonical
atom_dipole/atom_polarizabilitylabel keys (wereatomic_*→ KeyError) — fixed. - Legacy TF spin dispatch restored (
build_testerselectsSpinEnerTesteronhas_spinorget_ntypes_spin() != 0) with a regression — fixed. aparamflattened to(nframes*nloc, nda)beforeforward_common_atomic_graphon the graph route — fixed.- Repeated mid-file detail headers suppressed (header written only on a new/empty file) with a single-header-per-file multi-chunk assertion — fixed.
- Ragged tensor detail file across differing
sel_natomsgroups and zero-selected-atom handling — fixed. - Three-frame stress fixture breaking the inherited
DPTestmethods (576 vs 192 force rows) — fixed via_run_dp_testdefaulting tonumb_test=1.
Verification
- Full CI is green at head: Test Python (all 12 shards), Test C++ (all 4), CUDA/rocm builds, CodeQL, and codecov all pass.
Remaining thread comments are non-blocking (dead siz_coll list; a must=True atomic-label improvement note; CodeQL unused-log lint). Nothing blocks merge.
iProzd
left a comment
There was a problem hiding this comment.
Re-reviewed the current head after the test follow-up. Setting the inherited dp test checks to one frame now aligns the generated detail files with the direct single-batch model result, resolving the three-frame 576-versus-192 row mismatch. The earlier grouped tensor detail-path and zero-selected-atom issues remain fixed, the obsolete size collection was removed cleanly, and I found no new blocker.
Summary
dp testin atom-bounded chunks, with lazy on-demand decoding for LMDB inputsdeepmd.infer.model_testNeighborGraphfor graph-native models and the existing dense neighbor list for other modelsWhy
Large LMDB datasets were fully decoded before frame selection, so even a small test run could require loading tens of millions of frames. Separately, graph-native atomic models were routed through a fixed-capacity dense neighbor list while calibrating output bias; because those models do not declare a finite neighbor capacity, that path could allocate tens of gigabytes and fail before fine-tuning began.
This change bounds LMDB decoding and model-evaluation memory while preserving the existing evaluation and output-statistics contracts. Ordinary
DeepmdDatainputs still materialize one test system before chunked evaluation.Additional correctness fixes
Checks
ruff check .ruff format --check .git diff --checkThe TensorFlow atomic dipole/polar integration cases could not be run locally because this checkout has no built TensorFlow backend; they are left to CI.
Summary by CodeRabbit