perf(pt_expt): use scalable graph builders during training - #5913
Conversation
Resolve the carry-all graph backend once per training run so eager, compiled, and validation paths consistently use NV on supported CUDA environments with a dense fallback. Centralize graph-builder dispatch, validate explicit backend choices, and keep neighbor-capacity estimation compatible with nvalchemiops 0.3 and 0.4. Strengthen graph parity coverage by preserving edge multiplicity, exercising pair exclusions, and testing runtime and compiled backend routing.
Consume the normalized optimizer configuration for Adam, AdamW, and HybridMuon, including runtime parameter names required for correct HybridMuon routing. Apply the shared compiler compatibility patches and Inductor training options so compiled PT-expt execution follows the same runtime contract as the PT backend.
|
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:
📝 WalkthroughWalkthroughPT experimental neighbor-graph construction now supports configurable ChangesNeighbor-graph dispatch
Training runtime and optimizer updates
Estimated code review effort: 4 (Complex) | ~50 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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)
deepmd/pt_expt/utils/nv_graph_builder.py (1)
298-303: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueAdjust the capacity-comment baseline wording.
estimate_max_neighborsacceptsatomic_densityin the supported nvalchemiops 0.3+ path, but the comment describes the estimator default as0.2; change this to the actual effective baseline rather than to another arbitrary value like0.35.🤖 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_expt/utils/nv_graph_builder.py` around lines 298 - 303, Update the comment above initial_capacity to state the estimator’s actual effective atomic-density baseline used by estimate_max_neighbors, while retaining the explanation of the 0.25 density margin and deprecated safety_factor context; do not change the calculation.
🤖 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 `@deepmd/pt_expt/utils/nv_graph_builder.py`:
- Around line 298-303: Update the comment above initial_capacity to state the
estimator’s actual effective atomic-density baseline used by
estimate_max_neighbors, while retaining the explanation of the 0.25 density
margin and deprecated safety_factor context; do not change the calculation.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 1802ecc1-b696-48e1-9fca-b5e617badeea
📒 Files selected for processing (9)
deepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/graph_builder.pydeepmd/pt_expt/utils/nv_graph_builder.pydeepmd/utils/argcheck.pysource/tests/common/dpmodel/test_neighbor_graph_builder.pysource/tests/pt_expt/model/test_graph_builder_dispatch.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/utils/test_nv_graph_builder.py
💤 Files with no reviewable changes (1)
- source/tests/common/dpmodel/test_neighbor_graph_builder.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5913 +/- ##
==========================================
- Coverage 79.37% 79.12% -0.25%
==========================================
Files 1070 1071 +1
Lines 124791 124834 +43
Branches 4531 4531
==========================================
- Hits 99050 98778 -272
- Misses 24119 24435 +316
+ Partials 1622 1621 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt/test_hybrid_muon.py (1)
295-297: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winVerify the no-decay Adam route, not only Adam state creation.
Adam and AdamW both initialize
exp_avgand omitmomentum_buffer, so these assertions would pass even ifbwere incorrectly routed to AdamW. Also assert thatmodel.bis inoptimizer._routing[0]["adam_no_decay"]and not inadam_decay, or verify the update behavior with nonzero weight decay.🤖 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/test_hybrid_muon.py` around lines 295 - 297, Strengthen the assertions for model.b in the hybrid optimizer test to verify the no-decay Adam route, not merely Adam-compatible state fields. Assert that model.b appears in optimizer._routing[0]["adam_no_decay"] and does not appear in optimizer._routing[0]["adam_decay"], preserving the existing state assertions.
🤖 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/test_hybrid_muon.py`:
- Around line 295-297: Strengthen the assertions for model.b in the hybrid
optimizer test to verify the no-decay Adam route, not merely Adam-compatible
state fields. Assert that model.b appears in
optimizer._routing[0]["adam_no_decay"] and does not appear in
optimizer._routing[0]["adam_decay"], preserving the existing state assertions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: de9f0e93-4a02-47ab-be4f-9fa2ff5e664a
📒 Files selected for processing (2)
deepmd/pt/optimizer/hybrid_muon.pysource/tests/pt/test_hybrid_muon.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for the inline compile-safety gap.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt_expt/train/training.py (1)
104-107: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMove compile patches out of module import.
apply_global_compile_patches()runs on everydeepmd.pt_expt.train.trainingimport at line 106, thenenable_compile/check_compile_torch_version()are evaluated later in__init__. This makes non-compiled runs install process-global Dynamo/Inductor options and can break imports on Torch versions that are unsupported only fortorch.compile; move the call into theenable_compilebranch after the version check, or make the patch helper safe for every supported non-compile runtime.🤖 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_expt/train/training.py` around lines 104 - 107, Remove the module-level call to apply_global_compile_patches() in training.py, and invoke it only within the enable_compile path in the relevant initializer after check_compile_torch_version() succeeds. Ensure non-compiled runs do not install process-global compiler workarounds or trigger compile-version compatibility logic.
🤖 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.
Outside diff comments:
In `@deepmd/pt_expt/train/training.py`:
- Around line 104-107: Remove the module-level call to
apply_global_compile_patches() in training.py, and invoke it only within the
enable_compile path in the relevant initializer after
check_compile_torch_version() succeeds. Ensure non-compiled runs do not install
process-global compiler workarounds or trigger compile-version compatibility
logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 83a6b360-7bb9-4701-820b-0091554788a6
📒 Files selected for processing (1)
deepmd/pt_expt/train/training.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Reviewed the graph-builder dispatch, the optimizer-config move and the compile-contract alignment. Three findings below; the dispatch plumbing itself checks out -- the attribute is installed after every model-replacing step (finetune/init_model/change-bias) and before _compile_model and FullValidator, and getattr(..., "dense") preserves prior behaviour on CPU.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
One follow-up question on the compile version gate.
Validate optimizer support before reading variant-specific settings, prefer the NV graph builder on CUDA with an actionable fallback warning, and apply compiler workarounds only at explicit compile boundaries. Keep compile-only tests aligned with the production PyTorch version gate and strengthen optimizer routing coverage.
|
Regarding the CodeRabbit capacity-comment nit: this is not applicable. Upstream Upstream source: https://github.com/NVIDIA/nvalchemi-toolkit-ops/blob/main/nvalchemiops/neighbors/neighbor_utils.py |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pt_expt/test_training.py (1)
289-315: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winAdd 60-second timeouts to the changed training tests.
As per coding guidelines, this training test file requires a maximum 60-second timeout.
REQUIRES_SUPPORTED_COMPILEonly skips unsupported runtimes; it does not bound execution. Add@pytest.mark.timeout(60)to the changed tests/classes, following the existing pattern at Line 374.Also applies to: 317-340, 462-480, 545-545, 607-607, 1024-1025, 1119-1120, 1459-1460, 1504-1504, 1715-1715
🤖 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_training.py` around lines 289 - 315, Add `@pytest.mark.timeout`(60) to each changed training test identified in the comment, including test_neighbor_graph_method_defaults_to_auto, test_trainer_installs_resolved_graph_method, and test_explicit_graph_method_rejects_ineligible_model, following the existing decorator pattern near line 374. Apply the same timeout to the additional changed tests or enclosing classes listed in the comment, without altering their test behavior.Source: Coding guidelines
🧹 Nitpick comments (1)
source/tests/pt_expt/compile_utils.py (1)
6-8: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the compile-version gate in one implementation.
This helper imports the utility gate, while production training retains a separate identical implementation at
deepmd/pt_expt/train/training.py:129-136. Move the production check to the shared utility (or import it there) so tests cannot silently skip or run under a different version policy after a future update.🤖 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/compile_utils.py` around lines 6 - 8, Consolidate the duplicate compile-version gate by removing the production implementation in the training flow around the version check and reusing check_compile_torch_version from deepmd.pt.utils.compile_compat. Ensure production training and compile_utils tests invoke the same shared gate and preserve the existing version-policy 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.
Outside diff comments:
In `@source/tests/pt_expt/test_training.py`:
- Around line 289-315: Add `@pytest.mark.timeout`(60) to each changed training
test identified in the comment, including
test_neighbor_graph_method_defaults_to_auto,
test_trainer_installs_resolved_graph_method, and
test_explicit_graph_method_rejects_ineligible_model, following the existing
decorator pattern near line 374. Apply the same timeout to the additional
changed tests or enclosing classes listed in the comment, without altering their
test behavior.
---
Nitpick comments:
In `@source/tests/pt_expt/compile_utils.py`:
- Around line 6-8: Consolidate the duplicate compile-version gate by removing
the production implementation in the training flow around the version check and
reusing check_compile_torch_version from deepmd.pt.utils.compile_compat. Ensure
production training and compile_utils tests invoke the same shared gate and
preserve the existing version-policy behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 37011fcd-a349-400e-8dbb-e674388873f6
📒 Files selected for processing (9)
deepmd/pt/utils/compile_compat.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/graph_builder.pysource/tests/pt/test_hybrid_muon.pysource/tests/pt_expt/compile_utils.pysource/tests/pt_expt/model/test_graph_builder_dispatch.pysource/tests/pt_expt/test_multitask.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/test_training_ddp.py
🚧 Files skipped from review as they are similar to previous changes (3)
- source/tests/pt/test_hybrid_muon.py
- deepmd/pt_expt/utils/graph_builder.py
- source/tests/pt_expt/model/test_graph_builder_dispatch.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Thanks -- the four points from the previous round check out against HEAD 534a52e75, with one exception that the CI on this very commit now settles.
Addressed. The optimizer dispatch guard is in the right place: opt_type is validated against the three supported optimizers at training.py#L1589-L1592, before adam_beta1 is read at L1600, so LKF raises ValueError rather than KeyError. test_unsupported_optimizer_has_clear_error genuinely fails pre-fix (the old code raised KeyError, not ValueError). test_supported_optimizers_construct covers both new branches through the real update_deepmd_input + normalize path, and optimizer_hybrid_muon() supplies every one of the nine keys the HybridMuon branch reads, so a renamed kwarg or missing key fails construction. Both tests ran and passed in the CPU Python jobs on this commit. I accept your reasoning on not adding .get() fallbacks for a config that skipped update_deepmd_input().
The compile gate. Your answer resolves my question -- the gate is deliberate and inherited on purpose. Noting the consequence rather than contesting it, in the inline comment below.
The auto default -- this is the one I would like resolved before merge. I flagged the CI-coverage gap as the actionable half of that thread, and the run you pointed to has since finished: Test Python on CUDA failed on 534a52e75, in the NV path this PR newly makes the default. Details inline. (The three Test C++ failures are the unrelated paddle-inference tarball extraction problem, not this PR.)
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pt_expt/utils/test_nv_graph_builder.py (1)
28-40: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winPreserve frame identity in the multi-frame test oracle.
_edge_multisetcompares only local(src, dst, edge_vec)data, while the batched test now aggregates every frame into one multiset. A graph with incorrect frame assignments can therefore compare equal when those local tuples coincide. Include the stored frame/batch index in each key, or retain per-frame comparisons.Also applies to: 78-82
🤖 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/utils/test_nv_graph_builder.py` around lines 28 - 40, Update the _edge_multiset test oracle to include each edge’s stored frame or batch index in its comparison key, alongside src, dst, and edge_vec. Obtain the index from the graph’s existing frame/batch metadata and preserve the sorted multiset behavior so edges from different frames cannot compare equal when their local data matches.
🤖 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_expt/utils/nv_graph_builder.py`:
- Around line 272-273: Update the graph-building flow around the atype and coord
preparation to convert atype with torch.as_tensor before reading its shape.
Derive nf and nloc from the normalized tensor, then reuse that same tensor for
virtual-atom and pair-exclusion handling instead of reconverting or accessing
shape on the original Any input.
---
Outside diff comments:
In `@source/tests/pt_expt/utils/test_nv_graph_builder.py`:
- Around line 28-40: Update the _edge_multiset test oracle to include each
edge’s stored frame or batch index in its comparison key, alongside src, dst,
and edge_vec. Obtain the index from the graph’s existing frame/batch metadata
and preserve the sorted multiset behavior so edges from different frames cannot
compare equal when their local data matches.
🪄 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 Plus
Run ID: 972248d1-3e7a-4799-a02e-a861e27de2f0
📒 Files selected for processing (2)
deepmd/pt_expt/utils/nv_graph_builder.pysource/tests/pt_expt/utils/test_nv_graph_builder.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The NV multi-frame regression is fixed, and the fix is better than what I suggested.
nv_graph_builder.py now derives the shape from atype:
nf, nloc = atype.shape[:2]
coord = coord.reshape(nf, nloc, 3)atype is (nf, nloc) in both layouts, so it is unambiguous where coord.ndim was not. I had proposed deriving nf from box or asserting the two agree; keying off atype is cleaner and also fixes nloc, which the old code inferred from a coordinate array that had already collapsed. nv_search_matrix's own nf = coord.shape[0] is fine now, since the caller guarantees the reshape its docstring assumes.
Reusing the existing multi-frame test rather than adding a CUDA case was the right economy, and the new layout crosses the cell that was missing: rng.random((3, 5, 3)).reshape(3, -1) is the flattened training form with three frames, which is what _call_common_graph passes. On the old code that collapses to nf=1 and the three-frame box fails to reshape, so it fails before the fix and passes after.
Approving on the code. Worth noting for whoever merges: the Test CUDA label has come off, so Test Python on CUDA is currently skipping and neither this fix nor its test has actually executed -- Pass testing on CUDA reports green because it aggregates over skipped jobs. The merge queue does run test_cuda.yml, so the finetune case that originally surfaced this will be exercised there before merge; I mention it only because the CUDA run is the only thing that can see this code path, and it would be easy to read the current green as coverage.
Conflict in deepmd/pt_expt/utils/serialization.py: upstream deepmodeling#5913 relocated the graph edge-dtype helper to deepmd/pt_expt/model/graph_lower.py as a public graph_edge_dtype(), while this branch had rewritten the same helper in place to answer from the atomic-model capability instead of reaching through a single .descriptor. Resolution keeps BOTH: upstream's location (serialization.py imports it from graph_lower) with this branch's implementation. The reach-in version is wrong for compositions -- a LinearEnergyAtomicModel has no .descriptor, so a bridged/ZBL model would silently report float64 regardless of its children -- which is exactly what the capability aggregation fixes. Dropped the now-unused torch import from graph_lower.py.
Summary
Checks
Not run per request; the changes were tested separately on another machine.
Summary by CodeRabbit
training.neighbor_graph_method(auto,dense,nv) with consistent runtime selection for graph-eligible energy models.Adam,AdamW, andHybridMuon).training.neighbor_graph_methodbehavior and resolution timing.