fix(data): preserve LMDB label availability across batches - #5839
Conversation
|
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 (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughLMDB batching and validation now group frames by atom count and scalar ChangesLMDB availability-aware batching
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant LmdbDataReader
participant SameNlocBatchSampler
participant collate_lmdb_frames
participant FullValidator
LmdbDataReader->>SameNlocBatchSampler: group frames by nloc and find signature
SameNlocBatchSampler->>collate_lmdb_frames: emit homogeneous batches
collate_lmdb_frames->>FullValidator: provide scalar-consistent find flags
FullValidator->>FullValidator: evaluate grouped validation frames
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.
Actionable comments posted: 4
🤖 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/utils/lmdb_data.py`:
- Around line 1195-1210: Update the sampling flow around
group_indices_by_find_signature and _expand_indices_by_blocks to allocate each
block’s global target across all (nloc, signature) groups before expansion,
using remainder-preserving integer allocation so per-group rounding preserves
the exact block total. Pass each group’s allocated target into expansion, and
update the estimator logic around _block_total_actual to use the same allocation
rather than independently recomputing targets per signature.
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Around line 202-213: Update the test around DataRequirementItem registration
to realize or capture the lazy sampler iterator before calling
ds.add_data_requirements(). Assert that registration replaces the previous
iterator, then obtain subsequent batches and verify they use the newly
registered requirement signatures.
In `@source/tests/pt/test_lmdb_dataloader.py`:
- Around line 388-423: Extend the test around EnergyStdLoss to cover the force
branch using start_pref_f and limit_pref_f. Invoke the loss with a zero-force
model and force/find_force labels, asserting zero loss when find_force is 0.0
and positive loss when it is 1.0, alongside the existing energy assertions.
- Around line 407-408: Update the zero_energy_model closure in the loop so it
binds the current batch value at definition time, avoiding a late-bound
loop-variable closure flagged by Ruff B023 while preserving the existing energy
tensor behavior. Run ruff check . to verify the lint issue 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: 2bd6db4e-0f72-43dd-8c28-134427b8cdbe
📒 Files selected for processing (7)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt/train/validation.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt_expt/utils/lmdb_dataset.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt/test_validation.pysource/tests/pt_expt/test_lmdb_training.py
Partition LMDB frames by atom count and scalar find-flag signature before training, statistics, auto-probability, and distributed batching. Allocate auto-probability expansion targets exactly across homogeneous groups so repeated rounding cannot over-sample a block. Evaluate full LMDB validation in the same homogeneous groups so default-filled optional labels do not bias metrics. Cover standard PyTorch energy and force masking, iterator resets, distributed lengths, statistics loaders, pt_expt batches, and partial-label validation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
2fa9dae to
a5d7da6
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_lmdb_data.py (1)
752-753: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
batch.countto explicitly count occurrences.Currently,
index in batchevaluates to a boolean. This correctly counts the total frames here because the mock reader's batch size is explicitly set to 1. However, if the batch size were ever increased to pack multiple identical indices, this expression would only count the number of batches containing the index rather than the total occurrences. Usingbatch.count(index)is more robust and semantically clearer.♻️ Proposed refactor
batches = list(sampler) - counts = [sum(index in batch for batch in batches) for index in (0, 1)] + counts = [sum(batch.count(index) for batch in batches) for index in (0, 1)]🤖 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/common/dpmodel/test_lmdb_data.py` around lines 752 - 753, Update the counts comprehension in the sampler test to use batch.count(index) instead of membership testing, so it sums all occurrences of each index across batches rather than only counting batches that contain 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.
Nitpick comments:
In `@source/tests/common/dpmodel/test_lmdb_data.py`:
- Around line 752-753: Update the counts comprehension in the sampler test to
use batch.count(index) instead of membership testing, so it sums all occurrences
of each index across batches rather than only counting batches that contain it.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2bf6ab22-4ac4-4374-82bd-0d46f30282c7
📒 Files selected for processing (8)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt/train/validation.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt_expt/utils/lmdb_dataset.pysource/tests/common/dpmodel/test_lmdb_data.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt/test_validation.pysource/tests/pt_expt/test_lmdb_training.py
🚧 Files skipped from review as they are similar to previous changes (7)
- deepmd/pt_expt/utils/lmdb_dataset.py
- deepmd/pt/train/validation.py
- deepmd/pt/utils/lmdb_dataset.py
- source/tests/pt_expt/test_lmdb_training.py
- source/tests/pt/test_validation.py
- source/tests/pt/test_lmdb_dataloader.py
- deepmd/dpmodel/utils/lmdb_data.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5839 +/- ##
==========================================
- Coverage 79.03% 78.79% -0.25%
==========================================
Files 1055 1055
Lines 122233 122357 +124
Branches 4401 4401
==========================================
- Hits 96607 96411 -196
- Misses 24061 24380 +319
- Partials 1565 1566 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The core fix for #5636 is correct and well-tested — grouping by (nloc, availability-signature) and splitting validation metrics genuinely resolve the mixed-availability loss/metric contamination, and the new tests fail on pre-fix code. One blocking concern though: the grouping key-set and the collate guard disagree, which turns a previously-silent case into a hard training crash for exactly the merged multi-system LMDBs this PR targets. Detail inline. Plus a matching test-coverage gap.
Group training and validation frames by every availability field exposed to collation, including unrequested raw labels. Add coverage for optional virials and combined atom-count/label variation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both findings resolved; clearing my earlier request for changes.
The grouping/collation mismatch is genuinely closed. _availability_signature_keys now unions the registered requirements, the optional model inputs, and every non-structural field actually present in the frame, and because signatures are (find_key, bool) tuples a frame carrying an extra field gains an extra entry and therefore lands in its own group - so the case I raised (two systems merged into one LMDB, differing only by an unrequested virial) is partitioned before collate_lmdb_frames can reject it. test_unrequested_labels_form_homogeneous_batches reproduces exactly that setup and would hit the ValueError on the previous code.
The validation coverage gap is closed too: the new fixture varies both dimensions at once, and the test asserts against what actually reaches _evaluate_system, pinning all four (nloc, find_energy, find_force) groups rather than inferring them.
One consequence worth being aware of rather than acting on: grouping now splits on any differing non-structural field, so a dataset whose frames carry heterogeneous incidental labels will produce more partial batches, which feeds num_steps for epoch-based runs. That is the correct trade - it is precisely the condition collation rejects - but it is a behaviour change in batch counts.
3 participants
Fixes #5636
Summary
find_*signatureSimply stacking
find_*values would not be safe here. Current PyTorch and dpmodel losses reduce a whole batch before applying scalar availability flags, and display/validation paths convert those flags to booleans. A vector flag would broadcast incorrectly, produce vector-valued losses, or fail boolean conversion. Homogeneous batching is the smaller compatible fix.Why existing tests missed this
The direct collator test used
find_energy=1in every frame and did not assert the flag value. All LMDB factories stored energy and force in every record, so same-nloc batching was also label-homogeneous by accident. Reader consistency tests inspect individual frames before collation, and loss tests construct homogeneous scalar flags without connecting them to LMDB batches. Auto-probability tests also used one availability signature, so they could not expose independent per-signature rounding that changed a block target of 3 into 4 sampled frames.Full-validation coverage only varied atom count.
LmdbTestDatatherefore always observed globally present labels, and no test combined real labels with default-filled frames or checked order-independent optional-label metrics.Validation
find_*flag is zeroruff format .ruff check .Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
find_*settings.Performance / Behavior
Tests