feat(dpa4): multiple updates for DPA4/SeZM#5503
Open
OutisLi wants to merge 7 commits into
Open
Conversation
feat(sezm): add activation checkpoint option feat(sezm): add cross node tensor product feat(sezm): add message node tensor product feat(sezm): seperate node lmax with edge lmax feat(sezm): add custom kernel for lmax=5-10 feat(sezm): add so3 grid projection refactor: change default values feat(sezm): tf32 infer
Contributor
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
This PR expands the PyTorch backend’s neighbor-list strategy support (adding an NVIDIA Toolkit-Ops O(N) option) and updates SeZM/DPA4 configuration, docs, and tests to reflect new grid features and inference-time env controls.
Changes:
- Add
NvNeighborListstrategy (Toolkit-Ops /nvalchemi-toolkit-ops) plus dedicated unit tests and broadennlist_backenddispatch/equivalence tests to cover multiple strategies. - Extend SeZM/DPA4 descriptor argument schema (e.g.,
extra_node_l,kmax,ffn_so3_grid, grid branch controls, and validation-timetf32_infer) and wire env-var translation in the PT trainer. - Refactor/replace several SeZM grid/nonlinearity components and update example inputs + DPA4 docs accordingly; relax certain numeric tolerances in tests.
Reviewed changes
Copilot reviewed 47 out of 47 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| source/tests/pt/model/test_sezm_model.py | Adjust tolerances for compiled/eager comparisons (GPU). |
| source/tests/pt/model/test_nv_nlist.py | New unit tests for NvNeighborList edge topology/geometry equivalence and error behavior. |
| source/tests/pt/model/test_nlist_backend.py | Generalize backend strategy tests (vesin/nv) and improve dispatch coverage. |
| source/tests/pt/model/test_descriptor_sezm_s2_equivariance.py | Removes S2 equivariance tests (file deleted). |
| source/tests/pt/model/test_descriptor_sezm.py | Add tests for new zonal gather and special Wigner blocks; adjust Wigner tests and remove one deserialize regression test. |
| examples/water/dpa4/lora_ft.json | Update example descriptor/optimizer settings to newer compact DPA4 config. |
| examples/water/dpa4/input_multitask.json | Update multitask example config to newer compact DPA4 config. |
| examples/water/dpa4/input_dens.json | Update dens example config to newer compact DPA4 config. |
| examples/water/dpa4/input.json | Update baseline example config and add validating inference flags. |
| examples/water/dpa4/input-zbl.json | Update ZBL example config to newer compact DPA4 config. |
| examples/water/dpa4/input-spin.json | Update spin example config to newer compact DPA4 config. |
| examples/water/dpa4/README.md | Align README text with updated “compact/DPA4-Neo-style” examples. |
| doc/model/dpa4.md | Major DPA4/SeZM documentation expansion and updates, including env var behavior. |
| deepmd/utils/argcheck.py | Add get_argument, extend DPA4/SeZM schema and validation options, adjust defaults/variants. |
| deepmd/pt/utils/nv_nlist.py | New NVIDIA Toolkit-Ops neighbor-list strategy implementation. |
| deepmd/pt/train/training.py | Translate validating.* options into env vars before model construction. |
| deepmd/pt/model/model/init.py | Default descriptor.type to dpa4 in SeZM scaffolds. |
| deepmd/pt/model/descriptor/sezm_nn/triton/kernels_generic.py | Remove old Triton kernels (file deleted). |
| deepmd/pt/model/descriptor/sezm_nn/triton/kernels_edge_geometry_rbf.py | Remove old Triton kernels (file deleted). |
| deepmd/pt/model/descriptor/sezm_nn/triton/dispatch.py | Remove old Triton dispatch (file deleted). |
| deepmd/pt/model/descriptor/sezm_nn/triton/constants.py | Remove old Triton constants/flags (file deleted). |
| deepmd/pt/model/descriptor/sezm_nn/triton/autograd.py | Remove old Triton autograd wrapper/API (file deleted). |
| deepmd/pt/model/descriptor/sezm_nn/triton/init.py | Switch Triton exports to new compile-composable rotation API. |
| deepmd/pt/model/descriptor/sezm_nn/so2.py | Extend SO(2) conv with grid nets/branches + inference-time Triton rotation toggle. |
| deepmd/pt/model/descriptor/sezm_nn/projection.py | New projection utilities (S2 + SO3 projectors and grid resolution helpers). |
| deepmd/pt/model/descriptor/sezm_nn/indexing.py | Add build_gie_zonal_index helper for zonal coupling gather. |
| deepmd/pt/model/descriptor/sezm_nn/ffn.py | Refactor FFN nonlinearity onto new grid-net abstractions and add SO3-grid option. |
| deepmd/pt/model/descriptor/sezm_nn/embedding.py | Use new zonal index helper and allow passing precomputed zonal coupling. |
| deepmd/pt/model/descriptor/sezm_nn/edge_cache.py | Remove geometry/RBF Triton path and simplify edge cache construction. |
| deepmd/pt/model/descriptor/sezm_nn/cute/init.py | New CuTe accelerated SO(2) rotation package exports. |
| deepmd/pt/model/descriptor/sezm_nn/block.py | Add node/message grid branches, node-lmax decoupling, and inference checkpointing hooks. |
| deepmd/pt/model/descriptor/sezm_nn/activation.py | Remove grid-projector/grid-activation implementations (moved/refactored). |
| deepmd/pt/model/descriptor/sezm_nn/init.py | Re-export new grid/projector utilities and indexing helpers. |
| deepmd/pt/model/atomic_model/sezm_atomic_model.py | Use node_l_schedule if present for dens-head reconstruction. |
| deepmd/pt/infer/deep_eval.py | Add nv strategy handling, refactor strategy eval path, and broaden backend validation. |
| deepmd/pt/entrypoints/freeze_pt2.py | Patch Inductor config (triton.max_tiles) during AOTI packaging. |
| backend/find_pytorch.py | Add optional dependency for Toolkit-Ops neighbor list on Python >= 3.11. |
Comments suppressed due to low confidence (4)
deepmd/utils/argcheck.py:1
- The
extra_checkcan raiseTypeErrorfor non-bool, non-sequence inputs (e.g.,len(1)), and it also accepts any length-3 sequence without verifying element types arebool. Make the check robust by first verifyingisinstance(x, list)(orSequence) and then validatinglen(x) == 3andall(isinstance(v, bool) for v in x).
deepmd/pt/utils/nv_nlist.py:1 NvNeighborListis a GPU Toolkit-Ops path, butbuild()does not enforce CUDA tensors. If called with CPU tensors (or in a CPU-only environment), this will likely fail insidenvalchemiopswith a less actionable error. Consider adding an explicit check likecoord.device.type == 'cuda'(and similarly foratype/box) and raising a clearValueErrorwhen not on CUDA.
deepmd/pt/utils/nv_nlist.py:1- This builds a flattened
validmask of sizetotal_atoms * max_neighborsand then materializesedge_idxvianonzero, which can be very memory/latency heavy for large systems (the primary target for this strategy). Consider constructing the flattened edge indices without forming the full 2D mask (e.g., usingrepeat_interleave(num_neighbors)to builddstindices andarange(sum(num_neighbors))-based offsets to buildslotindices), so the intermediate scales with the number of real edges rather than the dense capacity.
source/tests/pt/model/test_descriptor_sezm.py:1 - A deserialize/backward-compatibility regression test was removed in this area, and the PR also deletes
test_descriptor_sezm_s2_equivariance.pyentirely. Since this PR introduces substantial new grid/projection behavior and new serialization-relevant config keys (e.g.,extra_node_l, grid branches, SO3-grid options), it would be good to retain equivalent coverage: add/restore tests that (1) validate descriptors deserialize correctly when older config fields are missing and (2) assert key equivariance/invariance properties for the new grid-net paths.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
|
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:
📝 WalkthroughWalkthroughAdds SeZM grid/projector/net stack, new Triton/CuTe rotation kernels, NV Toolkit-Ops neighbor-list backend, AOT/Inductor inference lowering with Triton tiling patch, descriptor/config/docs updates, trainer env scoping, and many aligned tests. ChangesSeZM grid, descriptor, and config surface
Descriptor and grid implementation
Rotation kernels, Triton/CuTe, and triton package surface
Neighbor-list backend, model compile/export, and runtime
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 11
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
examples/water/dpa4/input.json (1)
1-124:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUse
input_torch.jsonfor this PyTorch example.This file is a PyTorch training config (
model.type: "DPA4"). Keeping it asinput.jsonconflicts with the repo rule for TensorFlowinput.jsonusage and creates command ambiguity for users/docs.As per coding guidelines,
**/input.json: TensorFlow backend training configuration should useinput.jsonformat and be invoked withdp train input.json.🤖 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 `@examples/water/dpa4/input.json` around lines 1 - 124, The config is a PyTorch training config (model.type: "DPA4") but is named as the TensorFlow-style input.json, causing ambiguity; rename this file to input_torch.json (and update any references) so PyTorch examples use input_torch.json, ensure any documentation or scripts that reference this example (e.g., training invocation) point to input_torch.json instead of input.json, and verify consumers expect a PyTorch config when parsing model.type == "DPA4".Source: Coding guidelines
🧹 Nitpick comments (1)
source/tests/pt/model/test_descriptor_sezm_triton.py (1)
76-177: ⚡ Quick winThe new dense rotation path still has no coverage.
Every fixture builds
coeff_index = build_m_major_index(lmax, 1, ...), sorotate_to_local/rotate_backalways auto-select the block-diagonal kernels. The dense kernels, inverse-index path, and dense registered backwards introduced inso2_rotation.pycan regress without any test failing. Please add at least one non-mmax==1case or direct*_denseassertions.🤖 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/model/test_descriptor_sezm_triton.py` around lines 76 - 177, Tests always build coeff_index via build_m_major_index(lmax, 1, ...) so rotate_to_local and rotate_back always pick the block-diagonal (mmax==1) path; add a test case that constructs a coeff_index with mmax>1 (or directly call the *_dense variants) so the dense/inverse-index code paths and their registered backwards in so2_rotation.py are exercised; update test_rotate_to_local_matches_reference and test_rotate_back_matches_reference (or add a new subTest) to call build_m_major_index(lmax, mmax>1, ...) or to invoke rotate_to_local_dense / rotate_back_dense equivalents and assert outputs and gradients match the *_reference functions (and still apply the mask checks for dense behavior).
🤖 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/pt/infer/deep_eval.py`:
- Around line 309-343: The nvalchemi ('nv') backend selection in the
_setup_nlist_backend logic (symbols: nlist_backend, builder, NvNeighborList,
VesinNeighborList, is_nv_available, is_vesin_torch_available) must be gated on
the actual inference device (DEVICE from deepmd/pt/utils/env.py) rather than
host CUDA availability; update both the explicit "nv" branch and the "auto"
branch to check DEVICE.type == "cuda" (or the device of the input tensors)
before choosing NvNeighborList and, if DEVICE.type != "cuda" but nlist_backend
== "nv" was requested, raise a clear ValueError explaining that nv requires
CUDA-device inference; keep the existing checks for vesin (using
is_vesin_torch_available) but ensure vesin selection also verifies the
input/DEVICE compatibility if needed.
In `@deepmd/pt/model/descriptor/sezm_nn/__init__.py`:
- Around line 45-51: The module removed package-level re-exports (notably
SwiGLUS2Activation) which breaks downstream imports; restore a deprecated
re-export shim in __init__.py that re-exports the original symbol name (e.g.,
SwiGLUS2Activation) while importing/forwarding to the new implementation
location or wrapper and emit a DeprecationWarning when the shim is used; apply
the same pattern for the other removed symbols referenced in the file (the
exports around the later ranges) so callers keep compatibility for one release
before fully removing them.
In `@deepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.py`:
- Around line 833-839: The CuTe fast-path guard (_cute_usable) must also verify
kernel channel-width constraints: reject CuTe when channel count C would produce
invalid launch geometry or tails; update _cute_usable to compute the channel
dimension C (from the relevant tensor shape used by so2_rotation) and return
False unless C >= _TN and C % _TN == 0 (i.e., divisible by the kernel tile size
_TN), in addition to the existing checks (SEZM_CUTE_AVAILABLE, is_cuda, dtype);
apply the same additional gating logic to the other CuTe guard sites referenced
(the checks around lines 875-879 and 912-916) so they all fall back to the eager
path when channel widths are unsupported.
In `@deepmd/pt/model/descriptor/sezm_nn/so2.py`:
- Around line 1327-1343: The S2GridNet for the edge-local S2 branch
(instantiated as node_wise_grid_product) is using the default "packed"
coefficient layout while x_local and x_dst_local are reduced m-major tensors;
change the S2GridNet constructor to pass layout="m_major" so the projector
matches the m-major ordering; also update the other analogous S2GridNet
instantiation used for the edge-local branch later in the file (the second
S2GridNet for the same local branch) to use layout="m_major" as well.
In `@deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py`:
- Around line 1168-1191: The dispatch helper _block_layout_lmax currently only
checks shapes and may misidentify non-canonical reduced layouts as the m-major
mmax=1 block layout; update _block_layout_lmax to also validate that coeff_index
contains the canonical ordering produced by build_m_major_index(lmax, 1) (i.e.,
compare coeff_index values to the expected index tensor using value-safe
operations) and return -1 if it does not match so the code falls back to the
dense path or raises; apply the same value-validation wherever you currently use
this shape-only detection (the rotate_to_local / rotate_back block dispatch and
the explicit *_block helpers) so block kernels are only used for truly canonical
layouts.
- Line 649: The loop variable named `l` used in the `for ... in
tl.static_range(0, LMAX + 1):` constructs triggers Ruff E741; rename the
lone-letter variable to a non-ambiguous identifier (e.g. `ell` or `l_idx`) in
every occurrence in this module (all loops using `tl.static_range(0, LMAX + 1)`
including the other three occurrences flagged) and update any inner references
accordingly so the code compiles and passes `ruff check .`.
In `@deepmd/pt/model/descriptor/sezm.py`:
- Around line 1491-1524: The concatenation in _build_gie_zonal_coupling mixes
frames because mp_coupling comes from edge_cache.Dt_full (which was built with
training-time random-γ) while extra_coupling is recomputed from raw edge_vec
without that γ; fix by reusing the same augmented quaternion/frame used for
Dt_full when computing extra_coupling. Concretely, obtain the edge
quaternion/frame that matches Dt_full (e.g. edge_cache.edge_quat or
edge_cache.random_gamma-applied edge_vec) instead of recomputing from raw
edge_vec, then call gie_zonal_wigner_calc.forward_zonal with that quaternion so
extra_coupling and mp_coupling share the same local frame before concatenation;
update references in _build_gie_zonal_coupling, mp_coupling, extra_coupling,
build_edge_quaternion and gie_zonal_wigner_calc.forward_zonal accordingly.
- Line 624: The list comprehension assigning self.ebed_dims uses the ambiguous
one-letter variable name "l" which triggers Ruff E741; update the comprehension
in the expression self.ebed_dims = [get_so3_dim_of_lmax(l) for l in
self.l_schedule] to use a clear, unambiguous variable name (e.g., "ell" or
"l_val") and update the call to get_so3_dim_of_lmax accordingly so it becomes
self.ebed_dims = [get_so3_dim_of_lmax(ell) for ell in self.l_schedule], ensuring
only the variable name is changed and no other behavior is altered.
- Around line 619-623: The kmax validation currently compares to the raw
constructor arg int(lmax) instead of the resolved schedule max; update the check
so self.kmax is validated against the resolved effective lmax (self.lmax) after
l_schedule is applied — i.e., ensure l_schedule resolution runs before
validating self.kmax and replace the comparison using int(lmax) with a
comparison to self.lmax (or int(self.lmax) if needed) in the initialization code
around self.kmax and self.lmax.
- Around line 846-848: The zip calls over the schedules should enforce equal
lengths: update the loop zip(self.l_schedule, self.node_l_schedule,
self.m_schedule) inside the SEZM descriptor and the zip(self.m_schedule,
self.l_schedule) used in the any(m > l for m, l in zip(...)) check to include
strict=True (e.g., zip(..., strict=True)) so mismatched schedule lengths raise
immediately; locate these by referencing the attributes self.l_schedule,
self.node_l_schedule, self.m_schedule and the any(...) check and add the
strict=True argument to each zip call.
In `@deepmd/pt/train/training.py`:
- Around line 184-192: The code uses os.environ.setdefault for
DP_COMPILE_INFER/DP_TF32_INFER (based on validating_params) which makes these
flags process-global and sticky; change this to temporarily set the env vars
only while constructing the SeZMModel instance and then restore the previous
environment values immediately after construction so other Trainer instances are
not affected—specifically, around the SeZMModel creation in training.py capture
os.environ.get for DP_COMPILE_INFER and DP_TF32_INFER, set/override them
according to validating_params, instantiate SeZMModel, and finally restore the
saved values (or delete if previously absent).
---
Outside diff comments:
In `@examples/water/dpa4/input.json`:
- Around line 1-124: The config is a PyTorch training config (model.type:
"DPA4") but is named as the TensorFlow-style input.json, causing ambiguity;
rename this file to input_torch.json (and update any references) so PyTorch
examples use input_torch.json, ensure any documentation or scripts that
reference this example (e.g., training invocation) point to input_torch.json
instead of input.json, and verify consumers expect a PyTorch config when parsing
model.type == "DPA4".
---
Nitpick comments:
In `@source/tests/pt/model/test_descriptor_sezm_triton.py`:
- Around line 76-177: Tests always build coeff_index via
build_m_major_index(lmax, 1, ...) so rotate_to_local and rotate_back always pick
the block-diagonal (mmax==1) path; add a test case that constructs a coeff_index
with mmax>1 (or directly call the *_dense variants) so the dense/inverse-index
code paths and their registered backwards in so2_rotation.py are exercised;
update test_rotate_to_local_matches_reference and
test_rotate_back_matches_reference (or add a new subTest) to call
build_m_major_index(lmax, mmax>1, ...) or to invoke rotate_to_local_dense /
rotate_back_dense equivalents and assert outputs and gradients match the
*_reference functions (and still apply the mask checks for dense behavior).
🪄 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
Run ID: 4fc847a9-e9f4-477b-a673-749e1a3a07e5
📒 Files selected for processing (47)
backend/find_pytorch.pydeepmd/pt/entrypoints/freeze_pt2.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/__init__.pydeepmd/pt/model/descriptor/sezm_nn/activation.pydeepmd/pt/model/descriptor/sezm_nn/block.pydeepmd/pt/model/descriptor/sezm_nn/cute/__init__.pydeepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.pydeepmd/pt/model/descriptor/sezm_nn/edge_cache.pydeepmd/pt/model/descriptor/sezm_nn/embedding.pydeepmd/pt/model/descriptor/sezm_nn/ffn.pydeepmd/pt/model/descriptor/sezm_nn/grid_net.pydeepmd/pt/model/descriptor/sezm_nn/indexing.pydeepmd/pt/model/descriptor/sezm_nn/projection.pydeepmd/pt/model/descriptor/sezm_nn/so2.pydeepmd/pt/model/descriptor/sezm_nn/triton/__init__.pydeepmd/pt/model/descriptor/sezm_nn/triton/autograd.pydeepmd/pt/model/descriptor/sezm_nn/triton/constants.pydeepmd/pt/model/descriptor/sezm_nn/triton/custom_ops.pydeepmd/pt/model/descriptor/sezm_nn/triton/dispatch.pydeepmd/pt/model/descriptor/sezm_nn/triton/kernels_edge_geometry_rbf.pydeepmd/pt/model/descriptor/sezm_nn/triton/kernels_generic.pydeepmd/pt/model/descriptor/sezm_nn/triton/kernels_small.pydeepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.pydeepmd/pt/model/descriptor/sezm_nn/wignerd.pydeepmd/pt/model/model/__init__.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/train/training.pydeepmd/pt/utils/nv_nlist.pydeepmd/utils/argcheck.pydoc/model/dpa4.mdexamples/water/dpa4/README.mdexamples/water/dpa4/input-spin.jsonexamples/water/dpa4/input-zbl.jsonexamples/water/dpa4/input.jsonexamples/water/dpa4/input_dens.jsonexamples/water/dpa4/input_multitask.jsonexamples/water/dpa4/lora_ft.jsonsource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_descriptor_sezm_grid_projection.pysource/tests/pt/model/test_descriptor_sezm_s2_equivariance.pysource/tests/pt/model/test_descriptor_sezm_triton.pysource/tests/pt/model/test_nlist_backend.pysource/tests/pt/model/test_nv_nlist.pysource/tests/pt/model/test_sezm_model.py
💤 Files with no reviewable changes (8)
- deepmd/pt/model/descriptor/sezm_nn/triton/constants.py
- deepmd/pt/model/descriptor/sezm_nn/triton/kernels_small.py
- deepmd/pt/model/descriptor/sezm_nn/triton/dispatch.py
- deepmd/pt/model/descriptor/sezm_nn/triton/kernels_generic.py
- deepmd/pt/model/descriptor/sezm_nn/triton/kernels_edge_geometry_rbf.py
- source/tests/pt/model/test_descriptor_sezm_s2_equivariance.py
- deepmd/pt/model/descriptor/sezm_nn/triton/autograd.py
- deepmd/pt/model/descriptor/sezm_nn/triton/custom_ops.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5503 +/- ##
==========================================
+ Coverage 81.41% 82.00% +0.58%
==========================================
Files 871 716 -155
Lines 96952 72256 -24696
Branches 4242 979 -3263
==========================================
- Hits 78938 59251 -19687
+ Misses 16711 12428 -4283
+ Partials 1303 577 -726 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/pt/model/descriptor/sezm_nn/triton/so2_rotation.py`:
- Around line 1192-1201: The check in _canonical_block_index_values currently
returns True for FakeTensor (fake_mode) and meta tensors, allowing
_block_layout_lmax to pick the block-diagonal/triton path during torch.compile
tracing even when the real coeff_index is noncanonical; change the logic so
fake/meta tensors do not unconditionally return True — instead detect fake/meta
and conservatively return False (or defer decision) so the non-block codepath is
used during compile-time tracing; update references in
_canonical_block_index_values, _block_layout_lmax and any callsites relying on
coeff_index.fake_mode or coeff_index.device.type == "meta" to ensure the block
kernels are only chosen when coeff_index is explicitly canonical (e.g., by
actually comparing to build_m_major_index for real tensors).
In `@deepmd/pt/model/descriptor/sezm.py`:
- Line 1687: The generator expression uses a single-letter variable name `l`
which triggers Ruff E741; update the expression in the condition that reads "if
any(m > l for m, l in zip(self.m_schedule, self.l_schedule, strict=True))" to
use a clearer binding (e.g. `l_val` or `sched_l`) instead of `l`, and update the
comparison accordingly so the check uses `m > l_val` while keeping the same
zip(self.m_schedule, self.l_schedule, strict=True) semantics.
🪄 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
Run ID: 7213fdd3-1a94-4027-8ff2-372f1789b514
📒 Files selected for processing (11)
deepmd/pt/infer/deep_eval.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/block.pydeepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.pydeepmd/pt/model/descriptor/sezm_nn/edge_cache.pydeepmd/pt/model/descriptor/sezm_nn/so2.pydeepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.pydeepmd/pt/train/training.pysource/tests/pt/model/test_descriptor_sezm_grid_projection.pysource/tests/pt/model/test_descriptor_sezm_triton.pysource/tests/pt/model/test_nv_nlist.py
🚧 Files skipped from review as they are similar to previous changes (5)
- source/tests/pt/model/test_nv_nlist.py
- source/tests/pt/model/test_descriptor_sezm_grid_projection.py
- deepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.py
- deepmd/pt/model/descriptor/sezm_nn/so2.py
- deepmd/pt/model/descriptor/sezm_nn/block.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
deepmd/pt/model/descriptor/sezm.py (1)
1687-1687:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRename the zipped
lbinding before Ruff blocks CI.
any(m > l for m, l in ...)still trips Ruff E741 here, so this file will continue to fail the requiredruff check .step.Suggested fix
- if any(m > l for m, l in zip(self.m_schedule, self.l_schedule, strict=True)): + if any( + m > degree + for m, degree in zip(self.m_schedule, self.l_schedule, strict=True) + ):As per coding guidelines,
**/*.py: Install linter and runruff check .before committing changes or the CI will fail; Ruff also reports E741 on this line.🤖 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/pt/model/descriptor/sezm.py` at line 1687, The generator expression uses a single-letter variable name "l" which triggers Ruff E741; update the zip unpacking in the conditional that references self.m_schedule and self.l_schedule (the expression inside the if using zip(..., strict=True)) to use a different identifier such as l_val or l_sched (e.g., any(m > l_val for m, l_val in zip(self.m_schedule, self.l_schedule, strict=True))) so the binding is no longer the letter "l" and the linter error is resolved.Sources: Coding guidelines, Linters/SAST tools
deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py (1)
1164-1183:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRestore value-level validation for the block API.
_block_layout_lmax()now accepts anycoeff_indexwith the right length, butrotate_to_local_block()/rotate_back_block()rebuild the canonicalmmax=1layout and ignore the passed values. A same-length permutation therefore returns incorrect coefficients and gradients instead of raising.Suggested fix
def _block_layout_lmax(coeff_index: Tensor, dim_full: int) -> int: @@ if lmax < 1 or numel != 3 * lmax + 1: return -1 + expected = build_m_major_index(lmax, 1, device=coeff_index.device) + if not torch.equal(coeff_index, expected): + return -1 return lmaxIf fake/meta tracing needs the relaxed path, keep that as a narrowly scoped exception rather than the default for real tensors.
🤖 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/pt/model/descriptor/sezm_nn/triton/so2_rotation.py` around lines 1164 - 1183, The current _block_layout_lmax() only checks shape and length, allowing any same-length coeff_index through; change it so after computing lmax and numel it also validates the actual values of coeff_index against the canonical m-major mmax=1 ordering and return -1 if they differ. Implement by constructing the expected index sequence for the computed lmax (same logic used by rotate_to_local_block()/rotate_back_block()) and comparing with coeff_index using a concrete equality test (e.g., torch.equal or elementwise comparison), but preserve the relaxed path for non-concrete/trace/meta shapes by detecting those cases (e.g., meta tensors or when coeff_index has no real data) and only perform the value check for real tensors; keep this logic in _block_layout_lmax so rotate_to_local_block and rotate_back_block can continue to rely on it.
🤖 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 `@source/tests/pt/model/test_nv_nlist.py`:
- Line 177: The test imports only unittest but calls unittest.mock.patch.object
(causing AttributeError); update the imports so the mock module is available
(e.g., add "import unittest.mock" or "from unittest import mock") and then use
the mock namespace when patching NV_CELL_LIST_THRESHOLD in the test (the patch
call around NV_CELL_LIST_THRESHOLD should reference mock.patch.object or
unittest.mock depending on which import you add); ensure the change is applied
where the test uses NV_CELL_LIST_THRESHOLD and the context manager remains the
same.
---
Duplicate comments:
In `@deepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.py`:
- Around line 1164-1183: The current _block_layout_lmax() only checks shape and
length, allowing any same-length coeff_index through; change it so after
computing lmax and numel it also validates the actual values of coeff_index
against the canonical m-major mmax=1 ordering and return -1 if they differ.
Implement by constructing the expected index sequence for the computed lmax
(same logic used by rotate_to_local_block()/rotate_back_block()) and comparing
with coeff_index using a concrete equality test (e.g., torch.equal or
elementwise comparison), but preserve the relaxed path for
non-concrete/trace/meta shapes by detecting those cases (e.g., meta tensors or
when coeff_index has no real data) and only perform the value check for real
tensors; keep this logic in _block_layout_lmax so rotate_to_local_block and
rotate_back_block can continue to rely on it.
In `@deepmd/pt/model/descriptor/sezm.py`:
- Line 1687: The generator expression uses a single-letter variable name "l"
which triggers Ruff E741; update the zip unpacking in the conditional that
references self.m_schedule and self.l_schedule (the expression inside the if
using zip(..., strict=True)) to use a different identifier such as l_val or
l_sched (e.g., any(m > l_val for m, l_val in zip(self.m_schedule,
self.l_schedule, strict=True))) so the binding is no longer the letter "l" and
the linter error is resolved.
🪄 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
Run ID: fe323a2a-a138-4cf0-a506-7409bb56c4d5
📒 Files selected for processing (12)
deepmd/pt/infer/deep_eval.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/block.pydeepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.pydeepmd/pt/model/descriptor/sezm_nn/edge_cache.pydeepmd/pt/model/descriptor/sezm_nn/so2.pydeepmd/pt/model/descriptor/sezm_nn/triton/__init__.pydeepmd/pt/model/descriptor/sezm_nn/triton/so2_rotation.pydeepmd/pt/train/training.pysource/tests/pt/model/test_descriptor_sezm_grid_projection.pysource/tests/pt/model/test_descriptor_sezm_triton.pysource/tests/pt/model/test_nv_nlist.py
🚧 Files skipped from review as they are similar to previous changes (6)
- deepmd/pt/train/training.py
- deepmd/pt/infer/deep_eval.py
- deepmd/pt/model/descriptor/sezm_nn/edge_cache.py
- deepmd/pt/model/descriptor/sezm_nn/cute/so2_rotation.py
- source/tests/pt/model/test_descriptor_sezm_grid_projection.py
- deepmd/pt/model/descriptor/sezm_nn/block.py
Contributor
|
I reviewed the new nlist/backend and SeZM grid paths and found a few issues that should be fixed before merge:
— OpenClaw 2026.5.28 (model: custom-chat-jinzhezeng-group/gpt-5.5) |
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
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 by CodeRabbit
New Features
Performance
Bug Fixes
Documentation
Tests