feat(lmdb): support mixed-size batches and lazy label availability - #5962
feat(lmdb): support mixed-size batches and lazy label availability#5962OutisLi wants to merge 7 commits into
Conversation
LMDB batches previously required every frame to have the same atom count, which can leave sparse size groups under-filled and give their frames a disproportionate optimizer weight. Add `batch_size: "mix:N"` so frames of different sizes can share one atom-budgeted batch. Use two layouts according to the model contract. Eligible graph models consume one concatenated node axis with per-frame `n_node` counts; other models retain rectangular batches whose shorter frames are padded with `atype = -1`. Keep native-spin models on the rectangular public path because their output translation is spin-specific. Exclude phantom rows from neighbor graphs and model evaluation, scatter per-atom outputs back at the public boundary, and make losses and validation weight only real atoms. Consolidate LMDB sampling and decoding around an explicit batch layout so serial and worker-process decoding preserve the same field shapes and frame order. Cover the ragged training path with the existing DPA1 graph lower, alongside padding, compaction, loss-reduction, sampler, and decoder regressions.
Large LMDBs must not pay O(frame count) random I/O or Python-object allocation before the active training contract is known. Read metadata through sequential readahead, keep frame tables in compact NumPy arrays, and choose readahead according to the access pattern of each reader. Defer availability resolution until requirements are registered. Probe only optional tracked fields; uniform datasets start without a full scan, while detected mixed datasets build one compact cached signature index through a sequential reader with bounded progress logging. Mandatory fields fail at decode, default-backed inputs remain available per frame, and derived fields are computed from normalized structure data. Apply the contract consistently to statistics, samplers, full validation, and both PT training paths. Declare only active loss labels, preserve explicit values beside defaults, and gate force-derived losses by the availability of their force target. Keep filtered frame and system indices, mixed-nloc packing, and validation views in one index domain, and retain the block-allocation and batch-layout fixes found while consolidating the data path.
|
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:
📝 WalkthroughWalkthroughThis PR adds mixed-atom-count LMDB batching, ragged graph execution, phantom-node filtering, source-policy metadata, and pooled per-atom loss reductions across DPModel and PyTorch paths. ChangesMixed-NLOC and ragged execution
Estimated code review effort: 5 (Critical) | ~120 minutes 🚥 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 |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pt/train/training.py (1)
292-328: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winNon-distributed LMDB sampler drops the configured seed.
The distributed branch passes
seed=_training_params.get("seed")toDistributedLmdbBatchSamplerat Line 312. The non-distributed branch does not pass aseedtoLmdbBatchSamplerat Line 316.LmdbBatchSampleraccepts an optionalseedand falls back to OS entropy when it isNone, so single-process training with a configuredtraining.seedproduces a different batch shuffle on every run, while multi-rank training stays reproducible. Pass the same seed to both samplers.🐛 Proposed fix to pass the configured seed to the non-distributed sampler
else: _inner_sampler = LmdbBatchSampler( _data._reader, shuffle=True, + seed=_training_params.get("seed"), block_targets=_block_targets, )🤖 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/train/training.py` around lines 292 - 328, Pass the configured training seed from _training_params.get("seed") to the non-distributed LmdbBatchSampler construction, matching the existing DistributedLmdbBatchSampler branch while preserving the current sampler options.deepmd/pt_expt/train/wrapper.py (1)
223-238: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick winRemove the unsafe ragged-spin dispatch risk.
_forward_without_losspassesinput_dictdirectly tomodel.forward_ragged()whenn_nodeis set.forward_raggeddoes not acceptspinor**kwargs, butinput_dictcan includespinfor spin-capable models. Add a local guard or dispatch path so a ragged input for a spin model fails explicitly instead of asTypeError: forward_ragged() got an unexpected keyword argument 'spin'.🤖 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/wrapper.py` around lines 223 - 238, Update _forward_without_loss so ragged inputs with spin are detected before calling model.forward_ragged; explicitly reject this combination with a clear supported error, while preserving the existing forward_ragged dispatch for ragged inputs without spin and the regular model call for non-ragged inputs.
🧹 Nitpick comments (4)
deepmd/pt_expt/train/training.py (1)
1704-1708: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPass the normalized per-task data maps instead of the raw constructor arguments.
_configure_batch_layoutre-implements the dict-or-bare dispatch that_as_task_mapalready performed at lines 1593-1602.self.training_data_by_taskandself.validation_data_by_taskare keyed byself.model_keysand are available at this point. Using them removes the second normalization path and lets_configure_batch_layoutdrop itsisinstance(data_map, dict)branch.♻️ Proposed change
- self._configure_batch_layout(training_data, validation_data) + self._configure_batch_layout( + self.training_data_by_task, self.validation_data_by_task + )🤖 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 1704 - 1708, Update the call to _configure_batch_layout in the training initialization flow to pass self.training_data_by_task and self.validation_data_by_task instead of the raw training_data and validation_data constructor arguments. Then simplify _configure_batch_layout to consume these normalized per-task maps directly and remove its redundant isinstance(data_map, dict) dispatch.deepmd/utils/data.py (1)
1226-1232: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the allowed values from the alias.
The literal set duplicates
DataRequirementSourcePolicy. A future value added to the alias would pass type checking but fail at runtime. Usetyping.get_argsso one declaration governs both.♻️ Proposed refactor
- if source_policy not in {"tracked", "default", "derived"}: + if source_policy not in get_args(DataRequirementSourcePolicy): raise ValueError( "source_policy must be 'tracked', 'default', or 'derived', " f"got {source_policy!r}" )Add the import next to
Literal:from typing import ( Literal, get_args, )🤖 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/utils/data.py` around lines 1226 - 1232, Update the validation in the function containing source_policy to derive allowed values with typing.get_args(DataRequirementSourcePolicy) instead of duplicating the literal set, and import get_args alongside Literal. Preserve the existing ValueError and message behavior for invalid values.source/tests/pt/test_loss_default_pf.py (1)
232-236: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDerive the leading dimension of
drdqfrom the label batch.The test hard-codes a batch dimension of 1.
self.label_with_pref["force"]carries the batch size of the loaded water system. The mismatch is harmless today, becausefind_force = 0.0suppresses the generalized-force branch beforedrdqis used. If that gating changes, the test would fail for a shape reason instead of the reason it checks.♻️ Proposed refactor
label["drdq"] = torch.ones( - (1, self.nloc * 3 * numb_generalized_coord), + (label["force"].shape[0], self.nloc * 3 * numb_generalized_coord), dtype=label["force"].dtype, device=label["force"].device, )🤖 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_loss_default_pf.py` around lines 232 - 236, Update the drdq initialization in the relevant test setup to derive its leading dimension from self.label_with_pref["force"] rather than hard-coding 1, while preserving the existing coordinate dimension, dtype, and device.source/tests/pt_expt/utils/test_nv_matrix_decode.py (1)
171-182: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winRounded floats make set equality brittle.
_graph_edgesrounds eachedge_veccomponent to 8 decimals and puts the result in a set. Two builders that produce values differing by less than 1e-8 still land in different sets when a value sits on a rounding boundary, for example 1.234567895. Set membership gives no tolerance.Key the comparison on the integer endpoints only, then compare the matched vectors with
np.testing.assert_allclose.♻️ Proposed refactor
-def _graph_edges(graph) -> set: - """Edges as (src, dst, rounded edge_vec), so two builders can be compared.""" - keep = graph.edge_mask - return { - (int(s), int(d), *(round(float(x), 8) for x in v)) - for s, d, v in zip( - graph.edge_index[0][keep], - graph.edge_index[1][keep], - graph.edge_vec[keep], - strict=True, - ) - } +def _graph_edges(graph) -> dict: + """Edge vectors keyed by (src, dst), so two builders can be compared.""" + keep = graph.edge_mask + return { + (int(s), int(d)): np.asarray(v, dtype=np.float64) + for s, d, v in zip( + graph.edge_index[0][keep], + graph.edge_index[1][keep], + graph.edge_vec[keep], + strict=True, + ) + } + + +def _assert_same_edges(actual, expected) -> None: + assert set(actual) == set(expected) + for key, vector in expected.items(): + np.testing.assert_allclose(actual[key], vector, atol=1e-10)Then call
_assert_same_edges(_graph_edges(nv), _graph_edges(dense))at line 222.🤖 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_matrix_decode.py` around lines 171 - 182, Update _graph_edges to key edges by integer (src, dst) endpoints while retaining each edge vector without rounding, and add or update _assert_same_edges to match endpoint keys and compare corresponding vectors with np.testing.assert_allclose. Replace the comparison near the indicated call site with _assert_same_edges(_graph_edges(nv), _graph_edges(dense)), preserving detection of missing or extra endpoint pairs.
🤖 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/model/make_model.py`:
- Around line 892-910: Update the graph Hessian path around
_cal_hessian_ext_graph and _WrapperForwardEnergyGraph to use the same compacted
node set as the forward branch. Pass the compacted coordinates, atom types,
batch metadata, and parameters—or apply the equivalent valid-node mask when
rebuilding each frame—so phantom padded nodes are excluded from Hessian graph
construction.
In `@deepmd/pt_expt/utils/graph_builder.py`:
- Around line 11-13: Remove the module-scope PHANTOM_ATOM_TYPE import from
graph_builder.py and relocate the constant to a lightweight dependency-free
module, then update graph_builder references to use that module. Ensure the
graph-building path no longer imports deepmd.dpmodel.utils.lmdb_data or its
lmdb/msgpack dependencies.
In `@deepmd/pt/model/model/sezm_model.py`:
- Around line 3277-3279: Update _make_inter_potential_edge_mask to accept and
use the already-computed real_atom mask from core_compute instead of
re-sanitizing descriptor_atype. Pass real_atom at the core_compute call site,
preserve the existing atom-exclusion logic, and ensure phantom atoms remain
excluded even when edge builders include them.
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Line 851: Rename the keyword-only compile parameter in _run to enable_compile,
then update its call site and the assignment to
config["training"]["enable_compile"] to use the new name while preserving the
existing behavior.
---
Outside diff comments:
In `@deepmd/pt_expt/train/wrapper.py`:
- Around line 223-238: Update _forward_without_loss so ragged inputs with spin
are detected before calling model.forward_ragged; explicitly reject this
combination with a clear supported error, while preserving the existing
forward_ragged dispatch for ragged inputs without spin and the regular model
call for non-ragged inputs.
In `@deepmd/pt/train/training.py`:
- Around line 292-328: Pass the configured training seed from
_training_params.get("seed") to the non-distributed LmdbBatchSampler
construction, matching the existing DistributedLmdbBatchSampler branch while
preserving the current sampler options.
---
Nitpick comments:
In `@deepmd/pt_expt/train/training.py`:
- Around line 1704-1708: Update the call to _configure_batch_layout in the
training initialization flow to pass self.training_data_by_task and
self.validation_data_by_task instead of the raw training_data and
validation_data constructor arguments. Then simplify _configure_batch_layout to
consume these normalized per-task maps directly and remove its redundant
isinstance(data_map, dict) dispatch.
In `@deepmd/utils/data.py`:
- Around line 1226-1232: Update the validation in the function containing
source_policy to derive allowed values with
typing.get_args(DataRequirementSourcePolicy) instead of duplicating the literal
set, and import get_args alongside Literal. Preserve the existing ValueError and
message behavior for invalid values.
In `@source/tests/pt_expt/utils/test_nv_matrix_decode.py`:
- Around line 171-182: Update _graph_edges to key edges by integer (src, dst)
endpoints while retaining each edge vector without rounding, and add or update
_assert_same_edges to match endpoint keys and compare corresponding vectors with
np.testing.assert_allclose. Replace the comparison near the indicated call site
with _assert_same_edges(_graph_edges(nv), _graph_edges(dense)), preserving
detection of missing or extra endpoint pairs.
In `@source/tests/pt/test_loss_default_pf.py`:
- Around line 232-236: Update the drdq initialization in the relevant test setup
to derive its leading dimension from self.label_with_pref["force"] rather than
hard-coding 1, while preserving the existing coordinate dimension, dtype, and
device.
🪄 Autofix
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: 70419c7e-1272-4870-b2b2-4070566e7cf3
📒 Files selected for processing (53)
deepmd/dpmodel/loss/dos.pydeepmd/dpmodel/loss/ener.pydeepmd/dpmodel/loss/ener_spin.pydeepmd/dpmodel/loss/reduction.pydeepmd/dpmodel/loss/tensor.pydeepmd/dpmodel/model/make_model.pydeepmd/dpmodel/utils/__init__.pydeepmd/dpmodel/utils/batch.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/dpmodel/utils/neighbor_graph/__init__.pydeepmd/dpmodel/utils/neighbor_graph/ase_builder.pydeepmd/dpmodel/utils/neighbor_graph/from_ijs.pydeepmd/dpmodel/utils/neighbor_graph/graph.pydeepmd/pt/loss/dens.pydeepmd/pt/loss/dos.pydeepmd/pt/loss/ener.pydeepmd/pt/loss/tensor.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/model/sezm_native_spin_model.pydeepmd/pt/train/training.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt/utils/nv_nlist.pydeepmd/pt_expt/model/ener_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/train/validation.pydeepmd/pt_expt/train/wrapper.pydeepmd/pt_expt/utils/edge_schema.pydeepmd/pt_expt/utils/graph_builder.pydeepmd/pt_expt/utils/lmdb_dataset.pydeepmd/pt_expt/utils/nv_graph_builder.pydeepmd/pt_expt/utils/vesin_graph_builder.pydeepmd/pt_expt/utils/vesin_neighbor_list.pydeepmd/utils/argcheck.pydeepmd/utils/data.pydeepmd/utils/data_system.pydoc/data/system.mddoc/train/training-advanced.mdsource/tests/common/dpmodel/test_from_ijs.pysource/tests/common/dpmodel/test_graph_ragged.pysource/tests/common/dpmodel/test_lmdb_data.pysource/tests/common/dpmodel/test_loss_ener.pysource/tests/common/dpmodel/test_loss_padding.pysource/tests/common/dpmodel/test_loss_reduction.pysource/tests/consistent/test_lmdb_data.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt/test_loss_default_pf.pysource/tests/pt/test_loss_padding.pysource/tests/pt_expt/model/test_dpa4_native_spin.pysource/tests/pt_expt/test_lmdb_training.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/utils/test_nv_matrix_decode.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5962 +/- ##
==========================================
- Coverage 79.63% 79.49% -0.14%
==========================================
Files 1085 1085
Lines 126577 127139 +562
Branches 4592 4592
==========================================
+ Hits 100798 101068 +270
- Misses 24127 24419 +292
Partials 1652 1652 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Mixed-size batching is worth having and much of the machinery here is careful. Things I checked rather than assumed, and found correct: the frame-offset arithmetic in the from_ijs ragged rewrite, and that all four call sites were updated to pass full((nf,), nloc); compact_nodes / expand_node_values being an exact round trip, since keep_index holds distinct positions so the segment_sum scatter inverts the gather; the claim that every builder already refuses an edge touching a phantom, which holds in ase_builder, nv_graph_builder and vesin_graph_builder alike; and the sezm_model rebinding, which actually fixes a pre-existing case where atom_excl was applied to raw possibly-negative atype. One reported problem I want to explicitly withdraw before it wastes your time: I looked hard at whether the padded generalized-force path was mis-normalised by phantom rows, and it is not -- _finalize_atomic_ret sets mask unconditionally, so that path takes the masked branch.
Four things I do want changed, and two of them are the same bug.
The shared root cause: ragged is switched on for models that never asked for it. _configure_batch_layout enables the ragged layout for every LMDB data system whenever the model is non-spin, graph-lower and has forward_ragged -- with no reference to the batch-size rule and no reference to whether frames actually differ in size. So every pt_expt LMDB run on a dpa2/dpa3-class model changes layout, including runs with perfectly uniform frames that never opted into mixed-size batching. That is what turns two narrow new-feature concerns into regressions on existing configurations. Gating it on the reader actually needing it fixes both inline comments below at once.
The other regression is the LMDB availability rewrite, and it is the one I would fix first. This is not anchorable inline because the file is too large for GitHub to serve a patch, so here it is with permalinks.
PR #5839 (e37cc0095), which fixed #5636, deliberately did two things: an exact per-frame availability signature used to partition batches, and a hard ValueError in collate_lmdb_frames if the find_* flags disagreed inside a batch. This PR removes both. The scan becomes _probe_uniform_availability over _AVAILABILITY_PROBE_FRAMES = 256 evenly spaced frames, and when the probe reports uniform, availability_groups returns a single group -- so no partitioning happens for the rest of the dataset at all. The guard becomes _batch_find_flags, which on disagreement silently sets the flag to 0 for the whole batch. No exception, and no log line on either path.
Composed, those two changes reproduce #5636 exactly: on a dataset larger than 256 frames where the frames missing a label fall between probe points, the label is silently discarded for every frame in the batch. I am not neutral about this one -- I approved #5839 specifically because it added that guard, and the regression test that reassured me, test_unrequested_labels_form_homogeneous_batches, uses 8 frames, so _evenly_spaced returns every index and the probe is exact for it. It will keep passing while the property it was written to protect is gone. That is worse than the original bug, because the test now certifies it.
The PR's own comment concedes the probe is "a finding and not a proof" and even names caching the signature in __metadata__ as the right fix. I would take that: keep exactness and make it cheap. If you want to keep the probe, then collate_lmdb_frames must keep a defensive check -- silently zeroing a label is precisely the failure mode #5839 existed to prevent.
Two smaller notes I am not blocking on. masked_pair_mean and the coverage gap are inline. Also, deepmd/pt/loss/dens.py still ignores intensive_ener_virial in its energy term while the adjacent rmse_e uses inv_natoms**2, so the displayed metric and the optimised quantity disagree -- that predates this PR (it is not yours to fix here), but you were rewriting those exact lines, so it is worth knowing.
CI is green, which given the above is itself informative: none of these have a test that can see them.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
The mixed-size implementation is substantial and well covered, but two compatibility blockers remain: ragged layout is enabled for existing non-mixing LMDB runs, and the rectangular graph Hessian path does not use the compacted real-node set used by the forward path.
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
|
@wanghan-iapcm Thank you for the careful review. On LMDB label availability, I agree with the factual concern: a bounded probe cannot prove uniformity, and when it misses a rare signature the affected batch can lose otherwise valid supervision. I am intentionally not restoring either the eager O(N) startup scan or the batch-time ValueError in this PR. The former defeats lazy startup for large legacy LMDB datasets, while the latter makes a missed probe turn into a shuffle-dependent training abort. The current batch-level AND reduction is conservative: it may disable a disputed label for that batch, but it never treats a default-filled value as real supervision. Exactness needs to be established when the LMDB is generated, by recording a schema-versioned per-frame availability signature in metadata and using that as the reader source of truth. That work is tracked in #5954. I have made both the legacy tradeoff and the long-term exact solution explicit in the PR body. The remaining requested changes are addressed in 38cc7b9 and replied to individually in their threads. |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pt_expt/model/make_model.py (1)
355-417: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winPass
spinandcharge_spinthrough the graph Hessian path.Lines 923-942 pass
spin_flatandcharge_spinto normal graph evaluation. Lines 385-403 and 976-989 omit both inputs when constructing the Hessian wrapper.A native-spin or charge-conditioned model can then compute a Hessian for a different energy function. It can also fail when the model requires the omitted conditioning.
Add
spinandcharge_spinparameters to_cal_hessian_ext_graphand_WrapperForwardEnergyGraph. Compactspinwithnode_index. Pass the per-framecharge_spinvalue unchanged. Add a conditioned Hessian parity test.Also applies to: 976-989
🤖 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/model/make_model.py` around lines 355 - 417, Update _cal_hessian_ext_graph and _WrapperForwardEnergyGraph to accept and forward spin and charge_spin, matching the normal graph evaluation path. In _cal_hessian_ext_graph, compact each frame’s spin using node_index and pass the corresponding per-frame charge_spin unchanged when constructing the wrapper; update all callers, including the additional Hessian path around the referenced wrapper construction. Add a parity test covering conditioned/native-spin Hessian evaluation.source/tests/pt_expt/test_lmdb_training.py (1)
777-950: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winAdd a 60-second pytest timeout to
TestRaggedTrainingBatches.
source/tests/pt_expt/test_lmdb_training.pyimportspytestelsewhere,pyproject.tomlinstallspytest-timeout, and existing PT training tests use@pytest.mark.timeout(60). Add that class-level marker to this training-test class.🤖 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_lmdb_training.py` around lines 777 - 950, Add the class-level pytest timeout marker to TestRaggedTrainingBatches, using the existing project convention for a 60-second timeout. Keep the marker applied to the entire class so all its training tests are covered.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.
Inline comments:
In `@deepmd/dpmodel/loss/ener.py`:
- Around line 303-313: In the ragged-batch energy path, update the
coefficient-weighted atomic energy reduction to use the existing n_node-derived
frame IDs and segment_sum, producing one value per frame with shape (_nf, 1)
before comparison with energy_hat. Apply this in the enable_atom_ener_coeff flow
near the is_ragged handling, while preserving the existing non-weighted
reduction behavior.
- Around line 303-313: Update the ragged-frame normalization in the is_ragged
branch of the loss calculation to avoid dividing by zero when included_n_node is
zero. Compute a safe denominator and make inv zero for fully excluded frames,
while preserving the existing reciprocal normalization for frames with included
nodes.
---
Outside diff comments:
In `@deepmd/pt_expt/model/make_model.py`:
- Around line 355-417: Update _cal_hessian_ext_graph and
_WrapperForwardEnergyGraph to accept and forward spin and charge_spin, matching
the normal graph evaluation path. In _cal_hessian_ext_graph, compact each
frame’s spin using node_index and pass the corresponding per-frame charge_spin
unchanged when constructing the wrapper; update all callers, including the
additional Hessian path around the referenced wrapper construction. Add a parity
test covering conditioned/native-spin Hessian evaluation.
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Around line 777-950: Add the class-level pytest timeout marker to
TestRaggedTrainingBatches, using the existing project convention for a 60-second
timeout. Keep the marker applied to the entire class so all its training tests
are covered.
🪄 Autofix
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: 24db1f31-08c7-4f17-bd52-0868462927a6
📒 Files selected for processing (8)
deepmd/dpmodel/loss/ener.pydeepmd/dpmodel/loss/reduction.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pysource/tests/common/dpmodel/test_loss_padding.pysource/tests/pt_expt/model/test_dpa2_graph_lower.pysource/tests/pt_expt/test_lmdb_training.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/dpmodel/loss/reduction.py
- deepmd/pt_expt/train/training.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
The previous two blockers are fixed on this head: non-mixed batch rules remain rectangular, and the padded graph Hessian is compacted to real nodes before being scattered back. The current implementation still has correctness blockers in optional-label partitioning, atomic-label alias handling, loss-aware layout selection, ragged coefficient-weighted energy reduction, and Hessian dispatch/conditioning. The inline comments contain concrete reproductions and the required 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
iProzd
left a comment
iProzd
left a comment
There was a problem hiding this comment.
The backend/layout work now looks coherent, and the current head addresses the earlier ragged-model, loss-contract, mask, and Hessian issues. One blocking correctness issue remains: the bounded legacy-LMDB availability probe discussed in #5962 (comment) can still miss a rare signature and silently demote a valid label for the entire mixed batch. I reproduced this on dc4c99a with 512 frames: frame 1 omitted energy, the 256-frame probe returned one 512-frame group, and decoding frames 0 and 1 changed find_energy from 1 for frame 0 alone to 0 for the batch. This makes the training objective depend on batch composition without an error. Please make availability exact for tracked optional labels (metadata when available, otherwise an exact scan), or fail loudly when mixed availability reaches collation. I do not think a probabilistic correctness boundary is safe for training.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
All six findings from my earlier review are addressed at dc4c99a.
- The ragged layout is now genuinely opt-in:
use_ragged_batchesends withragged and self.mixed_nloc, andmixed_nlocisself._mix_rule is not None, so only amix:Nrule can select the flat node axis. Existing dpa2/dpa3 LMDB runs keep the rectangular layout, which also removes the root cause of the generalized-force abort. - The ragged loss branch consumes the model mask whenever it is present.
frame_id_from_n_nodeplussegment_sum(mask, frame_id)give the included counts,n_nodenow only sets the frame boundaries, and every per-atom term goes throughmasked_atom_meanon both layouts, soatom_exclude_typesaffects numerator and denominator alike. - The generalized-force guard fires only under
is_ragged, which is the representation it actually cannot serve. masked_pair_meanis now a documented third reduction category with the quadratic-growth rationale, rather than an unexplained holdout.- The layout-agreement test is replaced by one that can fail:
n_node=[2,3], a real padding slot, a model-excluded physical atom, force, atomic-energy, virial and atom-prefactor terms enabled, and a 1e-14 comparison of the total loss and every metric across the two layouts. - The reduction tests were re-derived algebraically through a
frame_weightsparameter rather than re-baselined, so no golden value was updated and no tolerance was widened.
Two things I am accepting as documented decisions rather than defects. First, pooling retires the bit-identical guarantee from #5738/#5783 for existing mixed_type NPY data whose frames hold different real atom counts; the Behavioral changes section states this and uniform-count batches are unchanged. Second, the bounded label-availability probe may discard valid supervision for a batch on a missed rare signature, with exact metadata tracked in #5954.
Worth keeping in mind on merge: the CUDA jobs report skipping, so a PR that reworks graph-model batching and the pt_expt training loop has no GPU coverage in CI.
There was a problem hiding this comment.
🧹 Nitpick comments (2)
deepmd/pt_expt/model/native_spin_model.py (1)
106-133: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocument the parameters of
forward_ragged.Every other public method in this class carries a full numpydoc
Parametersblock, including the siblingforwardandforward_lower_graph_exportable.forward_raggedis a public entry whose positional order (spinat index 4, aftern_node) differs fromforward(spinat index 3). A reader cannot see the flat shapes or the argument order from the single summary line.Add a
ParametersandReturnsblock stating thatcoordis(N, 3)frame-major overn_node,atypeis(N,),spinis(N, 3), and that per-atom outputs keep the flat axis whilen_nodeis returned alongside them.🤖 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/model/native_spin_model.py` around lines 106 - 133, Expand the docstring of NativeSpinModel.forward_ragged with numpydoc Parameters and Returns sections. Document the positional arguments, including spin after n_node, their flat frame-major shapes—coord (N, 3), atype (N,), and spin (N, 3)—plus the optional inputs, and state that per-atom outputs retain the flat axis while n_node is returned alongside them.source/tests/pt_expt/loss/test_ener_spin.py (1)
345-354: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a ragged case without
mask.
model_raggedalways carriesmask, so the test only covers the ragged branch that derivesincluded_n_nodefromsegment_sum.EnergySpinLoss.callalso has a branch for a ragged batch with nomask, which readsincluded_n_nodestraight fromn_nodeand skips themask_magintersection. A ragged producer reaches that branch whenever it emits nomask;_translate_eager_callcopiesmaskonly when the key is present in the model result.Add a second ragged dict without the
maskkey and assert the same parity. It is a two-line addition that locks the branch the graph path can actually take.♻️ Proposed addition
rectangular_loss, rectangular_more = loss_fn(1.0, 3, model_rect, label_rect) ragged_loss, ragged_more = loss_fn(1.0, 5, model_ragged, label_ragged) torch.testing.assert_close(ragged_loss, rectangular_loss) assert ragged_more.keys() == rectangular_more.keys() for key in ragged_more: torch.testing.assert_close(ragged_more[key], rectangular_more[key]) + + # A ragged producer that emits no ``mask`` must reduce identically. + model_no_mask = {k: v for k, v in model_ragged.items() if k != "mask"} + no_mask_loss, no_mask_more = loss_fn(1.0, 5, model_no_mask, label_ragged) + torch.testing.assert_close(no_mask_loss, rectangular_loss) + for key in no_mask_more: + torch.testing.assert_close(no_mask_more[key], rectangular_more[key])🤖 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/loss/test_ener_spin.py` around lines 345 - 354, Add a second ragged test case alongside model_ragged without the "mask" key, while retaining the other relevant fields, and assert that its loss matches the existing expected parity result. This should exercise EnergySpinLoss.call’s no-mask ragged branch using n_node directly.
🤖 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/model/native_spin_model.py`:
- Around line 106-133: Expand the docstring of NativeSpinModel.forward_ragged
with numpydoc Parameters and Returns sections. Document the positional
arguments, including spin after n_node, their flat frame-major shapes—coord (N,
3), atype (N,), and spin (N, 3)—plus the optional inputs, and state that
per-atom outputs retain the flat axis while n_node is returned alongside them.
In `@source/tests/pt_expt/loss/test_ener_spin.py`:
- Around line 345-354: Add a second ragged test case alongside model_ragged
without the "mask" key, while retaining the other relevant fields, and assert
that its loss matches the existing expected parity result. This should exercise
EnergySpinLoss.call’s no-mask ragged branch using n_node directly.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: e96d2f0a-a9b7-4c4c-8f79-c02bc6cb312f
📒 Files selected for processing (22)
deepmd/dpmodel/loss/ener_spin.pydeepmd/dpmodel/model/native_spin_model.pydeepmd/dpmodel/model/spin_model.pydeepmd/pt/model/model/spin_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/model/native_spin_model.pydeepmd/pt_expt/model/spin_ener_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/train/validation.pydeepmd/pt_expt/train/wrapper.pydeepmd/pt_expt/utils/graph_builder.pydeepmd/pt_expt/utils/vesin_graph_builder.pydeepmd/utils/argcheck.pysource/tests/common/dpmodel/test_loss_padding.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt_expt/loss/test_ener_spin.pysource/tests/pt_expt/model/test_dpa2_graph_lower.pysource/tests/pt_expt/model/test_dpa4_native_spin.pysource/tests/pt_expt/model/test_get_model_dpa4.pysource/tests/pt_expt/model/test_spin_ener_model.pysource/tests/pt_expt/test_lmdb_training.pysource/tests/pt_expt/test_training.py
🚧 Files skipped from review as they are similar to previous changes (10)
- deepmd/pt_expt/utils/vesin_graph_builder.py
- source/tests/pt_expt/test_training.py
- deepmd/pt_expt/train/validation.py
- deepmd/utils/argcheck.py
- deepmd/pt_expt/train/wrapper.py
- source/tests/pt_expt/model/test_dpa2_graph_lower.py
- deepmd/pt_expt/utils/graph_builder.py
- source/tests/pt/model/test_sezm_model.py
- source/tests/common/dpmodel/test_loss_padding.py
- source/tests/pt_expt/model/test_get_model_dpa4.py
Summary
mix:NLMDB batches containing frames with different atom counts, using padded rectangular batches for dense models and a flat ragged node axis for eligible graph modelsBehavioral changes
mixed_typeNPY datasets whose frames have different real atom counts: those frames are weighted by their real label counts rather than equally. Uniform-atom-count batches are unchanged. Hessian pair terms remain normalized per frame so their quadratic component count does not make large structures dominate a batch.find_*flags at collation. A missed rare signature may discard valid supervision for the affected batch, but default-filled values are never treated as real labels. Recording exact availability metadata when generating LMDB datasets is tracked in Record exact label availability when building LMDB datasets #5954.Testing
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Closes #5965