[pull] master from deepmodeling:master - #310
Merged
Merged
Conversation
…5903) ## Summary - Add shared `resolve_auto_graph_builder(device, nf)` for **inference / DeepEval only** (training keeps `resolve_neighbor_graph_method` from #5913). - Resolve `neighbor_graph_method="auto"` **at eval call time** with the batch frame count: CUDA prefers `nv`; `vesin` only when `nf == 1` and importable; otherwise `dense`. Matches `_select_neighbor_builder`. - Multi-frame `auto_batch_size` / `dp test` batches therefore stay off vesin's per-frame Python loop. - Parametrize vesin (and nv) vs dense energy/force parity over `nf in {1, 4}`. This is **not** a model-level / training default flip — those already shipped in #5912 / #5913. The remaining change is the inference auto ladder: re-introducing vesin only under the `nf == 1` gate that review asked for on #5912. ## Validation - `ruff check` / `ruff format` on touched files - `pytest` resolver ladder + DeepEval resolution + vesin parity for `nf=1` and `nf=4` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Improved automatic neighbor-graph builder selection during inference. - CUDA prioritizes optimized GPU processing, with Vesin used for eligible single-frame cases and dense processing as a fallback. - CPU inference uses Vesin for single-frame cases when available; otherwise, it uses dense processing. - Training continues to use dense processing for automatic selection. - **Bug Fixes** - Ensured automatic and unspecified builder settings produce consistent energy and force results. - **Tests** - Expanded coverage for single- and multi-frame backend selection and CPU/CUDA fallback behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
## Summary - align the pt_expt DPA4 parameter set and precision behavior with the PT reference - keep compiled training and evaluation graphs mode-specific - make native-spin DPA4 fine-tuning from a spin-free checkpoint function-preserving and trainable ## Details The pt_expt wrapper treated the bare NumPy weights in `FrameExpand` and `FrameContract` as buffers, so they never reached the optimizer. It also used NumPy-only operations in `ReducedEquivariantRMSNorm`, coupled inference AMP to the training switch, and reused a train-mode compiled graph during evaluation. This change promotes the missing trainable arrays, makes the norm tensor-safe, separates training and inference precision policies, and caches compiled lowers independently for train and eval mode. Native-spin fine-tuning had two related state-transfer problems. Per-type spin gates derived from `use_spin` were persisted in checkpoints, allowing an all-zero gate from a spin-free pretraining to override the fine-tune configuration. In addition, dormant randomly initialized spin routes became active at full amplitude when a magnetic type was introduced. Configuration-derived gates are now non-persistent and archived copies are ignored when loading older checkpoints; output-bias calibration receives the native moment; and native-spin models no longer fabricate a virtual-atom scale. DPA4 spin routes now initialize at the zero function. `prepare_finetune` resets dormant routes only when activating a spin-free checkpoint, preserves already-trained matching routes, and rejects reassignment to a different magnetic-type set. The environment-spin gate moves after the quadratic form so it retains a nonzero gradient at zero. Versioned migration preserves the function of existing checkpoints, and the reset capability is forwarded through atomic-model compositions such as ZBL bridging. This PR contains no DPA4C-specific or LMDB changes. ## Checks - SeZM descriptor, Triton-dispatch, model, compile, native-spin, and LoRA tests: 119 passed and 27 skipped in the full run; the two Warp-dependent tests that the sandbox could not compile were rerun outside it and passed with all 24 subtests - DPA4 dpmodel, PT/pt_expt gradient parity, pt_expt descriptor/model, native-spin, fine-tune, and compiled dynamic-shape tests: 130 passed - full `source/tests/pt_expt/test_training.py`: 50 passed - pre-commit hooks and `git diff --check` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Native-spin models now accept per-atom spin data during prediction and compiled training. * Added independent training and inference controls for AMP and TF32 precision. * Compilation now caches separate training and evaluation graphs. * Native-spin fine-tuning validates pretrained model compatibility and magnetic element settings. * **Bug Fixes** * Improved migration and loading of older DPA4 and SeZM checkpoints. * Corrected restoration of spin configuration and dormant parameters. * Improved compatibility with legacy exported model state. * Fixed spin embedding initialization and gradient behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
2 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.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )