fix(pt-expt): group LMDB statistics by atom count#5828
Conversation
Expose fixed-nloc statistical partitions without changing the LMDB's logical dataset identity, preventing mixed-size batches from being concatenated during model-stat initialization. Add regression coverage for statistics packing and end-to-end mixed-nloc training.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthrough
ChangesLMDB statistical batching
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Trainer
participant model_stat
participant LmdbDataSystem
participant LMDB
Trainer->>model_stat: make_stat_input(data)
model_stat->>LmdbDataSystem: get_stat_nsystems()
model_stat->>LmdbDataSystem: get_stat_numb_batches(sys_idx)
loop statistical systems and batches
model_stat->>LmdbDataSystem: get_stat_batch(sys_idx)
LmdbDataSystem->>LMDB: load nloc-group frames
LMDB-->>LmdbDataSystem: frame data
LmdbDataSystem-->>model_stat: collated batch
end
model_stat-->>Trainer: statistical input
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: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/utils/model_stat.py`:
- Line 46: Rename the unused loop variable jj to _jj in both loops in
deepmd/utils/model_stat.py at lines 46-46 and 84-84, preserving the loop
behavior while satisfying Ruff B007.
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Around line 300-301: Update test_mixed_nloc_statistics_and_training with the
project’s pytest timeout decorator configured for 60 seconds, and add the pytest
import if it is not already present. Preserve the test’s existing body and
behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 227d9aed-6591-4c45-9079-5381272dd296
📒 Files selected for processing (3)
deepmd/pt_expt/utils/lmdb_dataset.pydeepmd/utils/model_stat.pysource/tests/pt_expt/test_lmdb_training.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5828 +/- ##
==========================================
- Coverage 78.57% 78.33% -0.24%
==========================================
Files 1050 1050
Lines 120549 120677 +128
Branches 4351 4354 +3
==========================================
- Hits 94718 94530 -188
- Misses 24267 24587 +320
+ Partials 1564 1560 -4 ☔ 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.
The core fix is correct and well-tested: grouping frames by atom count into fixed-nloc partitions is sound, the min(nbatches, ceil(nframes/bs)) cap and sequential offset-with-wraparound visit every frame once per pass, the _collate_indices extraction is behavior-preserving, and the getattr-based shims fall back byte-identically to the legacy get_nsystems()/get_batch(sys_idx=) path for non-LMDB datasets (TF/jax/pd). test_stat_input_partitions_mixed_nloc_batches genuinely fails pre-fix (a single-system LMDB would concatenate ragged mixed-nloc shapes). LGTM.
Two non-blocking follow-ups:
-
Batch-count test coverage. The new test uses 4 frames with batch_size=2 — evenly divisible, so it cannot distinguish the correct ceil-division in get_stat_numb_batches from a floor-division regression, and it never exercises the offset-wraparound reset or the undersampling case (nbatches smaller than available). This is exactly the class of bug that shipped in #5488 (floor-division under-reporting batch counts in the sibling pt LMDB dataloader), caught only when reviewers added non-divisible-boundary tests. Please add a case with a non-divisible frame count (e.g. 5 frames at batch_size=2) and an undersampling case to lock the counting contract.
-
Stale Parameters docstring. In deepmd/utils/model_stat.py the make_stat_input summary now mentions that data sources may expose dedicated statistical-system methods, but the Parameters blocks of both make_stat_input and collect_batches still document the contract as only "must support get_nsystems() and get_batch(sys_idx=)". Since the implementation now probes get_stat_nsystems/get_stat_numb_batches/get_stat_batch and they change the effective per-shape system count, please name those three optional hooks in the Parameters section so the documented contract matches the code.
Summary
nlocstatistical partitions for PT-experimental LMDB data while preserving one logical training datasetnlocstatistics packing and end-to-end PT-experimental trainingTest plan
python -m pytest source/tests/pt_expt/test_lmdb_training.py source/tests/consistent/test_make_stat_input.py -vnloctests)Summary by CodeRabbit