feat(lmdb): support mixed-size batches and lazy label availability - #5962
Open
OutisLi wants to merge 3 commits into
Open
feat(lmdb): support mixed-size batches and lazy label availability#5962OutisLi wants to merge 3 commits into
OutisLi wants to merge 3 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.
OutisLi
marked this pull request as draft
August 8, 2026 15:44
OutisLi
marked this pull request as ready for review
August 8, 2026 15:44
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 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 Possibly related PRs
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 |
![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}
Contributor
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.60% 79.44% -0.17%
==========================================
Files 1085 1085
Lines 126405 126846 +441
Branches 4598 4598
==========================================
+ Hits 100631 100770 +139
- Misses 24120 24424 +304
+ Partials 1654 1652 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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
Testing
Summary by CodeRabbit
New Features
Bug Fixes
Documentation