fix(pt_expt): reuse the stored min_nbor_dist and batch the neighbor statistics - #5956
fix(pt_expt): reuse the stored min_nbor_dist and batch the neighbor statistics#5956yckbz wants to merge 5 commits into
Conversation
dp convert-backend stores it there, so compress recomputed it from the training data on every run. Add --recompute-min-nbor-dist to force a recompute, and log where the value comes from.
Sending a whole set to the device at once needs hundreds of GiB. Use AutoBatchSize, as the pt, pd, jax and tf backends already do.
|
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 skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe compression command adds a PyTorch Exportable model option to recompute ChangesCompression recomputation flow
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant CompressionCLI
participant enable_compression
participant ExportableModel
participant TrainingData
participant NeighborStat
CompressionCLI->>enable_compression: pass recompute_min_nbor_dist
enable_compression->>ExportableModel: read saved min_nbor_dist
alt recomputation requested
enable_compression->>TrainingData: load training script data
TrainingData->>NeighborStat: provide frames
NeighborStat-->>enable_compression: computed minimum neighbor distance
else saved value available
ExportableModel-->>enable_compression: saved minimum neighbor distance
end
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt_expt/test_compress_min_nbor_dist.py (1)
66-70: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftAdd execution-path tests for forced recomputation.
This test only confirms parser state. Add focused tests that verify
enable_compressionignores a saved value, assigns the value returned byUpdateSel.get_min_nbor_dist, and rejects a missingtraining_script.As per coding guidelines, use pytest for single test cases instead of the full test suite.
🤖 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_expt/test_compress_min_nbor_dist.py` around lines 66 - 70, Extend the compression tests around enable_compression to cover forced recomputation: verify the saved minimum-neighbor-distance value is ignored, the result from UpdateSel.get_min_nbor_dist is assigned, and a missing training_script is rejected. Use focused pytest test cases with mocked dependencies rather than invoking the full test suite.Source: Coding guidelines
🤖 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/pt_expt/test_compress_min_nbor_dist.py`:
- Around line 66-70: Extend the compression tests around enable_compression to
cover forced recomputation: verify the saved minimum-neighbor-distance value is
ignored, the result from UpdateSel.get_min_nbor_dist is assigned, and a missing
training_script is rejected. Use focused pytest test cases with mocked
dependencies rather than invoking the full test suite.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 3dae69f2-f343-4c91-bb44-52c1212ea168
📒 Files selected for processing (5)
deepmd/main.pydeepmd/pt_expt/entrypoints/compress.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/utils/neighbor_stat.pysource/tests/pt_expt/test_compress_min_nbor_dist.py
There was a problem hiding this comment.
Pull request overview
This PR improves the dp --pt-expt compress workflow by avoiding unnecessary recomputation of min_nbor_dist when it is already stored in the model (including under @variables), and by reducing peak device memory use when computing neighbor statistics via frame batching.
Changes:
- Teach
pt_expt compressto reusemin_nbor_distfrom the serialized model buffer, the top-levelmin_nbor_distkey, or@variables, with a new--recompute-min-nbor-distoverride flag. - Batch
NeighborStat.iteratorevaluations viaAutoBatchSize.execute_allto avoid loading entire datasets onto device at once. - Add pytest coverage for the
min_nbor_distread precedence and CLI flag default behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
source/tests/pt_expt/test_compress_min_nbor_dist.py |
Adds tests for stored min_nbor_dist precedence and the new CLI flag default. |
deepmd/pt_expt/utils/neighbor_stat.py |
Uses AutoBatchSize to batch neighbor-stat execution over frames to reduce device memory pressure. |
deepmd/pt_expt/entrypoints/main.py |
Wires the new recompute_min_nbor_dist flag through to the compress entrypoint. |
deepmd/pt_expt/entrypoints/compress.py |
Adds @variables lookup for min_nbor_dist and supports forced recomputation via a new parameter. |
deepmd/main.py |
Exposes --recompute-min-nbor-dist on the compress CLI parser (scoped in help to pt-expt). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Ce Liu <lc6866@outlook.com>
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
Suppressed comments (2)
deepmd/pt_expt/entrypoints/compress.py:43
_read_saved_min_nbor_distassumesmodel_dict["@variables"]is a dict by doing(model_dict.get("@variables") or {}).get(...). If@variablesis present but not a mapping (corrupt/older metadata, or other producers), this will raiseAttributeErrorduring compress. Other code paths in the repo defensively guard@variableswithisinstance(..., dict)(e.g. pt_expt DPA4 normalization).
return float(min_nbor_dist), "the model file"
min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist")
if min_nbor_dist is not None:
return float(min_nbor_dist), "the model file (@variables)"
return None, ""
deepmd/main.py:677
--recompute-min-nbor-distis added to the top-leveldp compressparser, so it will be accepted for all backends. However, only the PyTorch Exportable backend actually reads/usesrecompute_min_nbor_dist; other backends ignore it (e.g.deepmd/pt/entrypoints/main.pyanddeepmd/jax/entrypoints/main.pydon't pass it through). This can mislead users because the CLI will accept the flag but it will have no effect unless the backend is PyTorch Exportable.
Consider validating at dispatch time (or in each backend entrypoint) that this flag is only allowed with the PyTorch Exportable backend, and error out otherwise.
parser_compress.add_argument(
"--recompute-min-nbor-dist",
action="store_true",
help="(Supported backend: PyTorch Exportable) Ignore the minimal neighbor "
"distance saved in the model and recompute it from the training data. "
"Requires -t,--training-script",
)
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.
Suppressed comments (2)
deepmd/pt_expt/entrypoints/compress.py:41
model_dict.get("@variables")is assumed to be a dict, but if a malformed/older model file stores a truthy non-mapping (e.g. a list/string),(model_dict.get("@variables") or {}).get(...)will raiseAttributeError. Since this data comes from disk, it’s safer to guard the type and treat non-dict values as “not present” (or raise a clearerValueError).
min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist")
if min_nbor_dist is not None:
return float(min_nbor_dist), "the model file (@variables)"
return None, ""
deepmd/pt_expt/utils/neighbor_stat.py:82
- This change routes neighbor-stat computation through
AutoBatchSize.execute_all, but there’s no unit-level regression test to ensure batching is actually invoked (and keeps working) for the pt_expt backend. Consider adding a lightweight test (e.g. monkeypatchAutoBatchSize.execute_allto assert it’s called) so future refactors don’t accidentally revert to whole-set execution.
minrr2, max_nnei = self.auto_batch_size.execute_all(
self._execute,
data_set_data["coord"].shape[0],
data_set.get_natoms(),
data_set_data["coord"],
dp --pt-expt compressrecomputes the minimal neighbor distance from thetraining data on every run even when the model already carries it, and the
recomputation sends a whole set to the device at once.
Reading
min_nbor_distfrom@variablesdp convert-backendstores the value under@variablesinmodel.json— thelocation
deepmd/pt/utils/serialization.pywrites and the PyTorch and Paddlebackends read back.
enable_compressionlooked only atmodel.get_min_nbor_dist()and a top-levelmin_nbor_distkey, so it neverfound the value and fell back to a full pass over the training data. It now
checks
@variablestoo, and logs the source:--recompute-min-nbor-distforces a recompute, for a model compressed againsta data set other than the one it was trained on.
Batching the neighbor statistics
NeighborStat.iteratorpassed a whole set to_execute. The intermediatetensor is
[nframes, nloc, nall, 3], andnallis 27·nlocunder periodicboundaries: a 470-frame, 280-atom set allocates ~24 GiB per intermediate with
several live at once, and larger sets do not fit. It now goes through
AutoBatchSize, as thept,pd,jaxandtfbackends already do.Verification
On a 6-system, 1257-frame carbon nanotube data set (A800-80G),
min_nbor_distis
1.240096678690whether recomputed with batching (~3 s) or read from@variables(statistics skipped). The serialized compressed model, tabulateddata included, is identical in both cases and matches what the current code
produces.
Summary by CodeRabbit
New Features
--recompute-min-nbor-distoption for compression workflows, allowing supported PyTorch Exportable models to recalculate minimum neighbor distances from training data.Bug Fixes