Skip to content

fix(pt_expt): reuse the stored min_nbor_dist and batch the neighbor statistics - #5956

Open
yckbz wants to merge 18 commits into
deepmodeling:masterfrom
yckbz:fix-pt-expt-min-nbor-dist
Open

yckbz wants to merge 18 commits into
deepmodeling:masterfrom
yckbz:fix-pt-expt-min-nbor-dist

Conversation

@yckbz

@yckbz yckbz commented Aug 4, 2026 •

Copy link
Copy Markdown

dp --pt-expt compress recomputes the minimal neighbor distance from the
training 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_dist from @variables

dp convert-backend stores the value under @variables in model.json — the
location deepmd/pt/utils/serialization.py writes and the PyTorch and Paddle
backends read back. enable_compression looked only at
model.get_min_nbor_dist() and a top-level min_nbor_dist key, so it never
found the value and fell back to a full pass over the training data. It now
checks @variables too, and logs the source:

DEEPMD INFO  Minimal neighbor distance read from the model file (@variables): 1.240097

--recompute-min-nbor-dist forces a recompute, for a model compressed against
a data set other than the one it was trained on.

Batching the neighbor statistics

NeighborStat.iterator passed a whole set to _execute. The intermediate
tensor is [nframes, nloc, nall, 3], and nall is 27·nloc under periodic
boundaries: 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 the pt, pd, jax and tf backends already do.

Verification

On a 6-system, 1257-frame carbon nanotube data set (A800-80G), min_nbor_dist
is 1.240096678690 whether recomputed with batching (~3 s) or read from
@variables (statistics skipped). The serialized compressed model, tabulated
data included, is identical in both cases and matches what the current code
produces.

Summary by CodeRabbit

  • New Features

    • Added a --recompute-min-nbor-dist option for compression workflows, enabling supported PyTorch Exportable models to recalculate minimum neighbor distances from training data.
    • Added validation requiring a training script when recomputation is enabled.
    • Preserved computed neighbor-distance information in graph-lowered compressed models.
    • Added memory-bounded batching for neighbor-statistics processing.
  • Bug Fixes

    • Improved detection and preservation of saved minimum neighbor-distance values across supported model formats.
    • Preserved existing validation when required distance information is unavailable.
    • Improved compression behavior across CPU and GPU environments.

yckbz added 2 commits August 4, 2026 12:33
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.
Copilot AI lite review requested due to automatic review settings August 4, 2026 04:37
@dosubot dosubot Bot added the bug label Aug 4, 2026
@github-actions github-actions Bot added the Python label Aug 4, 2026
@coderabbitai

coderabbitai Bot commented Aug 4, 2026 •

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: b9ced153-6582-42f2-ae6e-a1f1c3355172

📥 Commits

Reviewing files that changed from the base of the PR and between 8cfd46e and fcd067c.

📒 Files selected for processing (8)
  • deepmd/main.py
  • deepmd/pt_expt/entrypoints/compress.py
  • deepmd/pt_expt/entrypoints/main.py
  • deepmd/pt_expt/model/spin_model.py
  • deepmd/pt_expt/utils/auto_batch_size.py
  • deepmd/pt_expt/utils/neighbor_stat.py
  • source/tests/pt_expt/model/test_model_compression.py
  • source/tests/pt_expt/test_compress_min_nbor_dist.py
🚧 Files skipped from review as they are similar to previous changes (8)
  • deepmd/main.py
  • deepmd/pt_expt/entrypoints/main.py
  • deepmd/pt_expt/utils/auto_batch_size.py
  • source/tests/pt_expt/test_compress_min_nbor_dist.py
  • source/tests/pt_expt/model/test_model_compression.py
  • deepmd/pt_expt/model/spin_model.py
  • deepmd/pt_expt/utils/neighbor_stat.py
  • deepmd/pt_expt/entrypoints/compress.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

📝 Walkthrough

The compression command adds a PyTorch Exportable model option to recompute min_nbor_dist from training data. Saved-value lookup covers multiple model locations. Neighbor statistics use device-aware automatic batching. Graph-lower exports preserve the computed value.

Changes: Compression recomputation flow

Layer / File(s) Summary
Batch neighbor statistics
deepmd/pt_expt/utils/auto_batch_size.py, deepmd/pt_expt/utils/neighbor_stat.py, source/tests/pt_expt/test_compress_min_nbor_dist.py
NeighborStat uses a device-aware AutoBatchSize controller to process frame batches. Tests verify CUDA-only batch growth.
Distance selection and recomputation
deepmd/pt_expt/entrypoints/compress.py, deepmd/pt_expt/model/spin_model.py
Compression reads min_nbor_dist from supported locations or recomputes it when requested. Recompute mode requires a training script. SpinModel delegates the property to its backbone. Graph-lower exports include the value.
CLI wiring and validation coverage
deepmd/main.py, deepmd/pt_expt/entrypoints/main.py, source/tests/pt_expt/test_compress_min_nbor_dist.py, source/tests/pt_expt/model/test_model_compression.py
The CLI exposes and forwards the new flag. Tests cover parsing, lookup precedence, missing metadata, export preservation, and model propagation.

Estimated code review effort: 3 (Moderate) | ~20 minutes

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
  enable_compression->>ExportableModel: write min_nbor_dist to graph-lower export
Loading

Suggested reviewers: njzjz-bot

Merge Risk: ⚪ Minimal · up to fcd06

The PR changes how stored neighbor-distance data is reused and how statistics are batched; no actionable merge-blocking risk remains beyond normal checks and review.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 70.83% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 24 functions across 8 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the two main changes: reusing stored min_nbor_dist values and batching neighbor-statistics computation.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
source/tests/pt_expt/test_compress_min_nbor_dist.py (1)

66-70: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift

Add execution-path tests for forced recomputation.

This test only confirms parser state. Add focused tests that verify enable_compression ignores a saved value, assigns the value returned by UpdateSel.get_min_nbor_dist, and rejects a missing training_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

📥 Commits

Reviewing files that changed from the base of the PR and between 6330a2f and 110e60b.

📒 Files selected for processing (5)
  • deepmd/main.py
  • deepmd/pt_expt/entrypoints/compress.py
  • deepmd/pt_expt/entrypoints/main.py
  • deepmd/pt_expt/utils/neighbor_stat.py
  • source/tests/pt_expt/test_compress_min_nbor_dist.py

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 compress to reuse min_nbor_dist from the serialized model buffer, the top-level min_nbor_dist key, or @variables, with a new --recompute-min-nbor-dist override flag.
  • Batch NeighborStat.iterator evaluations via AutoBatchSize.execute_all to avoid loading entire datasets onto device at once.
  • Add pytest coverage for the min_nbor_dist read 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.

Comment thread deepmd/pt_expt/entrypoints/compress.py
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Ce Liu <lc6866@outlook.com>
Copilot AI review requested due to automatic review settings August 4, 2026 04:44

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_dist assumes model_dict["@variables"] is a dict by doing (model_dict.get("@variables") or {}).get(...). If @variables is present but not a mapping (corrupt/older metadata, or other producers), this will raise AttributeError during compress. Other code paths in the repo defensively guard @variables with isinstance(..., 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-dist is added to the top-level dp compress parser, so it will be accepted for all backends. However, only the PyTorch Exportable backend actually reads/uses recompute_min_nbor_dist; other backends ignore it (e.g. deepmd/pt/entrypoints/main.py and deepmd/jax/entrypoints/main.py don'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",
    )

Copilot AI review requested due to automatic review settings August 4, 2026 04:51

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Copilot AI review requested due to automatic review settings August 5, 2026 11:34

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 raise AttributeError. Since this data comes from disk, it’s safer to guard the type and treat non-dict values as “not present” (or raise a clearer ValueError).
    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. monkeypatch AutoBatchSize.execute_all to 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"],

@iProzd iProzd added the P0 Blocks the DPA4/DPA4C release. label Aug 10, 2026
@njzjz njzjz added this to the v3.2.0 milestone Aug 10, 2026

@njzjz-bot njzjz-bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The saved-distance lookup and frame batching address the main performance problem, but three integration gaps remain: graph-lower outputs discard the recovered metadata, virtual-spin wrappers do not apply it to the backbone that performs compression, and explicit CPU execution on CUDA hosts uses a GPU-style expanding batch policy that can again exhaust host memory. The inline comments describe the affected paths and fixes.

Codex quota is about to reset, so I am using the remaining token budget to complete a concentrated review pass over the outstanding PRs.

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

Comment thread deepmd/pt_expt/utils/neighbor_stat.py
Comment thread deepmd/pt_expt/entrypoints/compress.py Outdated
Comment thread deepmd/pt_expt/entrypoints/compress.py Outdated
yckbz and others added 4 commits August 11, 2026 10:50
DEVICE is CPU whenever DEVICE=cpu is set, even on a CUDA host, where the
CUDA-OOM handler cannot recover from a host OOM. Base the growth policy
on the selected device instead of CUDA availability.
model.serialize() does not carry the runtime buffer, so the graph branch
dropped the value it had just recovered and the next run fell back to -t.
SpinModel.__getattr__ only delegates reads, so the assignment landed on
the wrapper while get_min_nbor_dist and enable_compression kept reading
the backbone.
@OutisLi

OutisLi commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator

Should wait for the merge of DPA4C, which refactors this.

@iProzd iProzd removed the P0 Blocks the DPA4/DPA4C release. label Aug 18, 2026
@yckbz

yckbz commented Aug 21, 2026

Copy link
Copy Markdown
Author

Should wait for the merge of DPA4C, which refactors this.

Merged master now that DPA4C has landed. Only compress.py conflicted:
#5972's compression_needs_min_nbor_dist() short-circuit is now an elif
after the saved-value lookup, so behaviour without --recompute-min-nbor-dist
is unchanged.

@yckbz
yckbz marked this pull request as ready for review August 21, 2026 08:54
@coderabbitai

coderabbitai Bot commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

@njzjz-bot njzjz-bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NEEDS HUMAN REVIEW

Re-reviewed the current head after the previous requested changes. The three earlier blocking issues are addressed in this revision: pt_expt batching now follows the selected DEVICE, graph-lower compression preserves min_nbor_dist in the exported archive, and SpinModel delegates min_nbor_dist assignment to its backbone. I reviewed all 8 changed files and the added tests and did not find another high-confidence functional blocker.

I cannot approve this head yet because relevant CI has not completed: the pull-request GitHub Actions runs are in action_required state with no jobs started, and the Read the Docs status is still pending. pre-commit.ci and CodeRabbit are successful. Please let the required Actions/RTD checks run to completion before approval.

Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: fcd067c
Trigger: scheduled review-request monitoring

Comment thread deepmd/pt_expt/utils/auto_batch_size.py Fixed
Comment thread source/tests/pt_expt/test_compress_min_nbor_dist.py Fixed
@codecov

codecov Bot commented Aug 28, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 72.50000% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.60%. Comparing base (025f4f6) to head (bd1c734).

Files with missing lines Patch % Lines
deepmd/pt_expt/entrypoints/compress.py 64.00% 9 Missing ⚠️
deepmd/pt_expt/utils/neighbor_stat.py 33.33% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #5956      +/-   ##
==========================================
- Coverage   77.85%   77.60%   -0.25%     
==========================================
  Files        1170     1171       +1     
  Lines      140914   140945      +31     
  Branches     5056     5056              
==========================================
- Hits       109705   109378     -327     
- Misses      29325    29684     +359     
+ Partials     1884     1883       -1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

yckbz and others added 3 commits August 29, 2026 09:23
Importing the value froze the device at import time, which defeats the
point of a policy meant to follow the selected device.
It was reached both by `import` and by `import from`.
@yckbz
yckbz requested a review from njzjz-bot August 29, 2026 13:24

@njzjz-bot njzjz-bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NEEDS HUMAN REVIEW: I re-reviewed the current head after the update. The current changes address the previously relevant device-selection concern by resolving the selected PyTorch device at call time, and I did not find another high-confidence functional, API/backward-compatibility, edge-case, concurrency/resource/security, or test-design defect in the reviewed diff that warrants an inline blocker.

I cannot approve this exact head yet because the repository's relevant GitHub Actions runs for this SHA have completed with action_required rather than actually running/passing (including Test Python, Test CUDA, Test C++, Build C/C++, CodeQL, and packaging). The available external checks (pre-commit.ci and Read the Docs) are successful, but they do not exercise the changed compression/neighbor-statistics paths sufficiently to replace the project test suite. Please allow/run the workflows and re-request review after the relevant checks pass.

Agent: ChatGPT
Model: GPT-5.6 Sol
GitHub account: njzjz-bot
Reviewed head: 2c69cff
Trigger: scheduled review-request monitoring

@njzjz-bot njzjz-bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed the current head. The previously reported integration gaps are addressed: graph-lowered archives retain min_nbor_dist, the spin wrapper delegates the property to the backbone that compression reads, and pt_expt batching follows the selected env.DEVICE rather than CUDA availability. The stored-value lookup order and recompute path are coherent, and the neighbor-stat batching matches the memory-bounded design used by the other backends. I did not find a blocking issue. Upstream currently exposes only the successful triage check on this head; the PR includes focused regression validation.

Agent: ChatGPT
Model: GPT-5.6 Sol

@OutisLi OutisLi left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Keep unused compression arguments separate from measured distance metadata

In deepmd/pt_expt/entrypoints/compress.py:85-96, --recompute-min-nbor-dist discards the saved distance before checking compression_needs_min_nbor_dist(). For DPA4C, that check then skips statistics and substitutes 0.0, which is subsequently persisted as min_nbor_dist in the output archive.

I reproduced this through the actual CLI entry point with a DPA4C .pte, a saved distance of 0.75 angstrom, and a valid training script whose dataset has a measured minimum distance of 1.25 angstrom. Ordinary compression preserves 0.75; adding --recompute-min-nbor-dist skips statistics and writes 0.0 instead.

Please keep DPA4C's statistics-free compression behavior. This is not a request to scan data that DPA4C does not need. Its tabulation domain is fixed independently of the minimum distance, so this does not directly change its energy or force predictions. The defect is treating an unused argument's placeholder as a measured statistic and overwriting valid metadata.

Please structure the decision around whether compression consumes this statistic:

  • If it does, reuse the saved value by default; recompute from the supplied training data when explicitly requested, retaining the required training-script validation.
  • If it does not, skip statistics even when the flag is supplied, and clearly warn that the recompute option is inapplicable to this model. Preserve an existing real distance unchanged; leave it absent when no measurement is available.
  • If the common compression interface still requires a numeric argument, a local 0.0 placeholder is sufficient for a descriptor that ignores it. Do not assign or serialize that placeholder as measured min_nbor_dist. No broad compression-interface redesign is needed.

Please cover both saved-value and missing-value cases for DPA4C, with and without the flag, checking the output metadata and that no dataset scan occurs. The flag help/docstring should also reflect that it applies only to models whose compression needs the statistic.

@wanghan-iapcm wanghan-iapcm left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed at bd1c734.

Verification:

  • Ran source/tests/pt_expt/test_compress_min_nbor_dist.py and source/tests/pt_expt/model/test_model_compression.py on CPU with the custom op library: 18 passed on this head. With the base versions of the six source files restored, all 11 new tests fail; test_compress_min_nbor_dist_from_variables fails with the original ValueError (no min_nbor_dist without -t), so it does expose the @variables bug.
  • The neighbor-statistics batching gives identical results to the whole-set path in a local check, but no test covers it (see inline).
  • Also ran test_spin_ener_model.py (10 passed), test_entrypoint.py (10 passed) and test_dpa4c_cpu.py (2 passed, 31 skipped on CPU).

Findings are inline: the recompute flag and the placeholder 0.0 being persisted as a measured distance for models that do not use it (same issue as the open change request), and two test gaps.

Comment thread deepmd/pt_expt/entrypoints/compress.py Outdated
Comment thread deepmd/pt_expt/utils/neighbor_stat.py
Comment thread source/tests/pt_expt/test_compress_min_nbor_dist.py
--recompute-min-nbor-dist dropped the saved distance before checking
whether compression consumes it, so a model that skips the statistics
wrote the 0.0 placeholder over a valid one. Decide on that first, and
record only a measured distance.
@yckbz

yckbz commented Oct 15, 2026

Copy link
Copy Markdown
Author

[P2] Keep unused compression arguments separate from measured distance metadata

In deepmd/pt_expt/entrypoints/compress.py:85-96, --recompute-min-nbor-dist discards the saved distance before checking compression_needs_min_nbor_dist(). For DPA4C, that check then skips statistics and substitutes 0.0, which is subsequently persisted as min_nbor_dist in the output archive.

I reproduced this through the actual CLI entry point with a DPA4C .pte, a saved distance of 0.75 angstrom, and a valid training script whose dataset has a measured minimum distance of 1.25 angstrom. Ordinary compression preserves 0.75; adding --recompute-min-nbor-dist skips statistics and writes 0.0 instead.

Please keep DPA4C's statistics-free compression behavior. This is not a request to scan data that DPA4C does not need. Its tabulation domain is fixed independently of the minimum distance, so this does not directly change its energy or force predictions. The defect is treating an unused argument's placeholder as a measured statistic and overwriting valid metadata.

Please structure the decision around whether compression consumes this statistic:

  • If it does, reuse the saved value by default; recompute from the supplied training data when explicitly requested, retaining the required training-script validation.
  • If it does not, skip statistics even when the flag is supplied, and clearly warn that the recompute option is inapplicable to this model. Preserve an existing real distance unchanged; leave it absent when no measurement is available.
  • If the common compression interface still requires a numeric argument, a local 0.0 placeholder is sufficient for a descriptor that ignores it. Do not assign or serialize that placeholder as measured min_nbor_dist. No broad compression-interface redesign is needed.

Please cover both saved-value and missing-value cases for DPA4C, with and without the flag, checking the output metadata and that no dataset scan occurs. The flag help/docstring should also reflect that it applies only to models whose compression needs the statistic.

Fixed in 3b3f493. You were right that the flag discarded the saved distance
before the compression_needs_min_nbor_dist() check, so the skip branch
handed back the ABI placeholder as if it were a measurement.

The decision is now made on whether compression consumes the statistic:

  • Consumed: reuse the saved value by default, recompute from the supplied
    training data when explicitly requested, with the training-script validation
    kept on both routes into measuring.
  • Not consumed: skip the statistics whether or not the flag is given, and
    warn that it is inapplicable. An existing distance is carried through
    untouched; with no measurement the field is left out of the archive rather
    than filled with a placeholder.

enable_compression still reads the value off the model, so a descriptor that
ignores it is handed 0.0 — but only for that call. The model is restored
right after, and only a measured distance is serialized, in both the
graph-lower and the fallback export path.

Tests cover DPA4C with and without a saved value, each with and without the
flag, asserting the archive metadata, that the model is left with no
placeholder, and that no training-data scan happens (scans == []). The
paths where the statistic is consumed are covered too, so the reuse,
recompute and missing-script cases are pinned. The flag's --help and the
docstring now state that it only applies to models whose compression needs
the statistic.

Nothing in source/tests imported the pt_expt NeighborStat and the
consistent suite skipped that backend, so the batching had no test.
Pin the split itself, since comparing results alone passes unbatched
too. Also run the consistent suite against pt-expt, and recompute
with nothing saved.
@yckbz

yckbz commented Oct 15, 2026

Copy link
Copy Markdown
Author

Reviewed at bd1c734.

Verification:

  • Ran source/tests/pt_expt/test_compress_min_nbor_dist.py and source/tests/pt_expt/model/test_model_compression.py on CPU with the custom op library: 18 passed on this head. With the base versions of the six source files restored, all 11 new tests fail; test_compress_min_nbor_dist_from_variables fails with the original ValueError (no min_nbor_dist without -t), so it does expose the @variables bug.
  • The neighbor-statistics batching gives identical results to the whole-set path in a local check, but no test covers it (see inline).
  • Also ran test_spin_ener_model.py (10 passed), test_entrypoint.py (10 passed) and test_dpa4c_cpu.py (2 passed, 31 skipped on CPU).

Findings are inline: the recompute flag and the placeholder 0.0 being persisted as a measured distance for models that do not use it (same issue as the open change request), and two test gaps.

Thanks for the thorough verification — both findings were right.

The placeholder (compress.py). Same defect @OutisLi raised, fixed in
3b3f493, which landed after the bd1c734 you reviewed. The decision is now
made on compression_needs_min_nbor_dist() first, exactly the shape you
suggested: when the statistic is not consumed the flag is ignored with a
warning, a real saved value is carried through unchanged, and the local 0.0
goes to enable_compression only — the model is restored right after, and
only a measured distance is serialized, in both export branches. Against your
table:

stored flag bd1c734 now
0.75 off 0.75 0.75
0.75 on 0.0 0.75 (+ warning)
none off 0.0 key absent
none on — key absent (+ warning)

The flag's --help and the docstring now say it only applies to models whose
compression uses the statistic.

The untested batching. You were right that nothing covered it — I had only
checked it by hand. Added source/tests/pt_expt/utils/test_neighbor_stat.py
in 88bf6a8, following your suggestion: DP_INFER_BATCH_SIZE is set before
the NeighborStat is built, _execute is wrapped to record the frames per
call, and the assertions pin the split rather than just the result —

whole_calls   == [4, 4]      # one call per set
batched_calls == [1] * 8     # one call per frame

plus min_nbor_dist / max_nbor_size equality between the two runs on a
multi-frame, multi-set system. To confirm it is not decorative, I restored the
base neighbor_stat.py and reran: the split assertion fails
(assert [4, 4] == [1, 1, ...]) while the results-only test still passes —
which is the point you made about equality checks.

I also added test_neighbor_stat_pt_expt to the consistent suite, since
pt-expt is a valid backend key and that suite was skipping it.

The recompute branches. Also addressed in 3b3f493: enable_compression
is now exercised with recompute_min_nbor_dist=True for both a model that
consumes the statistic and one that does not, with and without a stored value,
asserting the written metadata, that the model keeps no placeholder, and that
no training-data scan happens (scans == []). The ValueError without a
training script is covered on both routes into measuring.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: Todo

Development

Successfully merging this pull request may close these issues.

8 participants