Repository navigation
fix(pt_expt): fold vacuum reference on CPU tracing device (#6076) - #6079
Closed
OutisLi-Bot wants to merge 1 commit into
Closed
OutisLi-Bot wants to merge 1 commit into
OutisLi-Bot wants to merge 1 commit into
Conversation
Freezing with vacuum_ref ran fold_vacuum_reference before moving the
model to the CPU tracing device, so vacuum_descriptor could OOM on CUDA.
Moving to CPU first avoids that allocation on the GPU.
fold_vacuum_reference replaces bias_atom_e (and may set vacuum_table)
through numpy assignment. dpmodel_setattr previously always placed those
buffers on env.DEVICE, so a fold after Module.to("cpu") pulled tensors
back to CUDA and broke FakeTensor tracing. Place new buffers on the
module's current device instead.
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation |
|---|---|---|
| Description Check |  |
Check skipped - CodeRabbit’s high-level summary is enabled. |
| Title check | |
The title clearly and concisely describes the primary change: folding the vacuum reference on the CPU tracing device. |
| Linked Issues check | |
For directly linked issue #6076, _trace_and_export_impl now calls model.to("cpu") before model.fold_vacuum_reference(). This prevents vacuum-reference evaluation on CUDA during export. `module… |
| Out of Scope Changes check | |
The changed implementation and tests stay within issue #6076. The serialization change controls the device used for vacuum-reference folding. The common helper fixes the related NumPy-to-buffer device… |
| Docstring Coverage | |
Docstring coverage is 81.82% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 11 functions across 4 files. |
✨ Finishing Touches: 💡 1
-
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
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 help to get the list of available commands.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| def __init__(self) -> None: | ||
| self.stat = np.zeros(3, dtype=np.float64) | ||
|
|
||
| def call(self, x: np.ndarray) -> np.ndarray: |
OutisLi
requested review from
OutisLi and
wanghan-iapcm
and removed request for
OutisLi and
wanghan-iapcm
October 15, 2026 09:16
Collaborator
|
Superseded by #6080, which retains this PR's buffer-device fix and additionally removes DPA4C's dense vacuum-reference allocation and preserves CUDA-target deployment kernels while folding on CPU. Closing this PR in favor of the combined fix; validation details and limitations are recorded in #6080. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #6079 +/- ##
==========================================
- Coverage 77.85% 77.60% -0.25%
==========================================
Files 1170 1170
Lines 140914 140925 +11
Branches 5056 5056
==========================================
- Hits 109705 109364 -341
- Misses 29325 29679 +354
+ Partials 1884 1882 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
yckbz
pushed a commit
to yckbz/deepmd-kit
that referenced
this pull request
Oct 19, 2026
…els (deepmodeling#6080) ## Summary Fixes deepmodeling#6076 and supersedes deepmodeling#6079, retaining its module-local NumPy buffer-device fix and CPU-before-fold export ordering. This addresses both excessive vacuum-reference allocation and a CUDA-target DPA4 export regression introduced by moving the eager fold to CPU. ## Changes - Evaluate DPA4/DPA4C isolated references using one node per type and two masked guard edges, instead of sizing a dense neighbor list from `get_nnei()`. DPA4C's `[999999]` graph sentinel must not become almost one million padded neighbors per type. - Introduce an explicit `supports_graph_vacuum()` capability. DPA4/DPA4C opt in; legacy DPA1/DPA2 reference semantics and explicit graph-disable routing remain unchanged. Preserve type, charge and native-spin conditioning. - During export, temporarily reconstruct only descriptors requiring references on the CPU with reference kernels. Fold under `torch.no_grad()`, retaining bias/reference tables on the original fittings, then restore deployment descriptors and their optimized bindings before fake tracing. Runtime graph-disable state and environment settings are restored, including failure paths. - Retain deepmodeling#6079's device-selection fix for NumPy buffer assignments, including newly materialized nonpersistent reference tables. Simply moving a CUDA-bound float32 DPA4 model to CPU does not unbind its CUDA-only operators. Conversely, disabling deployment kernels throughout tracing would silently lose acceleration. The scoped reference-descriptor context avoids both problems without reconstructing the whole folded model or losing `vacuum_table` through serialization. ## Regression coverage - Allocation guards prohibit querying dense capacity or calling the dense adapter for graph references, including sentinel values 999999 and 1999999. - Float32/float64 DPA4/DPA4C, charge/native-spin combinations, empty/masked edges, and nonzero learned conditioning. - Nonzero-mean DPA1 legacy reference energies and disabled graph lowering. - Export order, CUDA-only operator dispatch without a CPU implementation, preservation of the optimized operator in fake tracing, fparam/aparam reference tables, restoration after failures, and no-reference fast exit. ## Validation - **40 passed**: new graph-reference tests, new export-reference tests, and common wrapper tests, using current Python source and a freshly rebuilt CPU PyTorch extension (MPI disabled because this host lacks MPI headers; TensorFlow build disabled). - **44 passed**: graph-reference tests plus existing dpmodel vacuum and capability coverage. - **36 passed**: broader export/kernel-policy/vacuum suite, including actual CPU `.pt2`/`.pte` compilation, fparam `.pt2`, and archive reserialization/refreezing. This run used the existing installed CPU native library with current Python source. - Additional independent checks covered mixed DPA4+DPA4C linear and native-spin models, nonpersistent tables and restoration of original descriptor objects. - Repository-wide Ruff check/format and `git diff --check` passed. ### Memory sanity check An untrained default-size DPA4C descriptor with the bundled 89-element OMat24 type map, float64 and charge/native-spin conditioning, retained its actual `get_nnei() == 999999`. Vacuum evaluation produced finite `(89, 207)` output in approximately 0.01–0.03 s, with about 34 MiB incremental process peak RSS. This is a CPU sanity check, not a benchmark of the reporter's pretrained checkpoint or a directly comparable speedup claim. ### Not verified - Real CUDA kernel execution and CUDA-target AOT compilation on GPU hardware. - Full large pretrained-checkpoint GPU/host peak-memory benchmarking. The CPU dispatch regression uses an operator without a CPU implementation and checks fake-trace preservation; it is not presented as a GPU integration test. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Vacuum-reference evaluation now supports graph-based DPA4 and DPA4C models, including when dense neighbor capacity is unavailable. Descriptors without graph-based vacuum support continue to use dense evaluation. * **Bug Fixes** * Improved consistency of vacuum-reference folding during model export, including preserving model settings and keeping folding on the CPU. * Corrected buffer assignment so arrays follow the model’s device when possible, rather than always using the default device. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: OutisLi_Bot <198997215+OutisLi-Bot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #6076: freezing a
vacuum_ref=TrueDPA4C (and other) checkpoint withpt_exptOOMs on CUDA duringfold_vacuum_reference(), because export folded while the model still resided on the GPU. Swapping to fold-after-to("cpu")avoided the OOM but exposed a FakeTensor CPU/CUDA mismatch: numpy buffer reassignment throughdpmodel_setattralways placed new buffers onenv.DEVICE.Design
serialization.py): move the model to the CPU tracing device beforefold_vacuum_reference(). The vacuum descriptor evaluates an isolated atom of every type through the descriptor; doing that on CUDA can allocate several GiB even when export already holds most of the GPU.common.py):_module_buffer_device()places a numpy → buffer assignment on the module's current device (existing same-named buffer, else any parameter/buffer, elseenv.DEVICE), matching themin_nbor_distsetter. Foldedbias_atom_e/vacuum_tabletherefore stay on CPU with the rest of the model afterModule.to("cpu"), so a secondto("cpu")hack is unnecessary..to("cpu")keeps CPU residency even whenenv.DEVICEis patched tometa, plus a fold residency check on a small DPA4 vacuum-ref model.Files
deepmd/pt_expt/utils/serialization.pyto("cpu")before folddeepmd/pt_expt/common.py_module_buffer_devicefor numpy buffer setattrsource/tests/pt_expt/utils/test_common.pysource/tests/pt_expt/model/test_dpa4_vacuum_ref.pyValidation
Ran on box (CPU-only)
python -m pytest source/tests/pt_expt/utils/test_common.py— 7 passed (incl.test_dpmodel_setattr_keeps_module_device)python -m pytest source/tests/pt_expt/model/test_dpa4_vacuum_ref.py::test_fold_vacuum_reference_keeps_cpu_residency source/tests/pt_expt/model/test_dpa4_vacuum_ref.py::test_fold_reproduces_the_referenced_model— 2 passedNot verified here
vacuum_refcheckpoint (dp --pt-expt freeze … --lower-kind graph) and numerical equivalence of the resulting.pt2(issue environment; this box is CPU-only)Related
Fixes #6076
Summary by CodeRabbit