fix(train): scope statistics cache lifecycle - #5883
Merged
Merged
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
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 (12)
🚧 Files skipped from review as they are similar to previous changes (12)
📝 WalkthroughWalkthroughThe PR introduces shared ChangesUnified statistics cache
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
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}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
deepmd/pt_expt/train/training.py (1)
1535-1535: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider documenting the two-phase
share_paramscall.The flow is correct (verified): chief recomputes/merges shared stats with
resume=Falseand ties params, values are broadcast to peers, thenresume=Trueis called unconditionally on every rank (including chief, redundantly) to (re-)establish the intra-process parameter tying using the now-synced values — since object-identity sharing can't be broadcast across processes, only tensor values can. This is subtle enough that a short inline comment explaining whyshare_paramsruns twice with differentresumevalues would meaningfully help future readers.Also applies to: 1547-1571
🤖 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` at line 1535, Add a concise inline comment around the two-phase share_params flow near synchronize_model_state, documenting that the chief merges and broadcasts shared-stat values with resume=False, then every rank calls resume=True to restore intra-process parameter identity after cross-process synchronization. Keep the existing call order and behavior unchanged.deepmd/pt/train/training.py (1)
1007-1035: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd
strict=Trueto the multi-taskzip()formodel_key_prob_map.Static analysis flags
zip(self.model_keys, self.model_prob)(Line 1019) withoutstrict=. Althoughself.model_probis always derived fromself.model_keysand should match in length, the sibling JAX implementation (deepmd/jax/train/trainer.py,_share_model_params) already guards the equivalent call withzip(self.model_keys, model_prob, strict=True). Addingstrict=Truehere costs nothing and fails fast instead of silently truncating/misaligning task↔probability pairs if this invariant is ever broken by a future refactor.🔧 Proposed fix
share_kwargs = { - "model_key_prob_map": dict(zip(self.model_keys, self.model_prob)), + "model_key_prob_map": dict( + zip(self.model_keys, self.model_prob, strict=True) + ), "data_stat_protect": _data_stat_protect[0], }🤖 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 1007 - 1035, Update the model_key_prob_map construction in the multi-task shared-parameter block to call zip(self.model_keys, self.model_prob, strict=True), matching the invariant enforced by the sibling implementation while preserving the existing dictionary construction.Source: Linters/SAST tools
🤖 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/train/training.py`:
- Line 1535: Add a concise inline comment around the two-phase share_params flow
near synchronize_model_state, documenting that the chief merges and broadcasts
shared-stat values with resume=False, then every rank calls resume=True to
restore intra-process parameter identity after cross-process synchronization.
Keep the existing call order and behavior unchanged.
In `@deepmd/pt/train/training.py`:
- Around line 1007-1035: Update the model_key_prob_map construction in the
multi-task shared-parameter block to call zip(self.model_keys, self.model_prob,
strict=True), matching the invariant enforced by the sibling implementation
while preserving the existing dictionary construction.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: c7b8fe8a-bd3e-4a95-b78a-9606e19e13b6
📒 Files selected for processing (35)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/descriptor.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/repflows.pydeepmd/dpmodel/descriptor/repformers.pydeepmd/dpmodel/descriptor/se_e2_a.pydeepmd/dpmodel/descriptor/se_r.pydeepmd/dpmodel/descriptor/se_t.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/train/data.pydeepmd/dpmodel/utils/env_mat_stat.pydeepmd/dpmodel/utils/stat.pydeepmd/jax/common.pydeepmd/jax/train/trainer.pydeepmd/pt/entrypoints/main.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/task/fitting.pydeepmd/pt/train/training.pydeepmd/pt/utils/stat.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/train/training.pydeepmd/tf2/entrypoints/train.pydeepmd/tf2/train/trainer.pydeepmd/utils/argcheck.pydeepmd/utils/stat_file.pysource/tests/common/dpmodel/test_train_data.pysource/tests/common/stat_file.pysource/tests/common/test_stat_file.pysource/tests/jax/test_model_factory.pysource/tests/jax/test_training.pysource/tests/pt/test_stat_file_mode.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_training.pysource/tests/tf2/test_training.py
💤 Files with no reviewable changes (1)
- source/tests/pt_expt/test_entrypoint.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5883 +/- ##
==========================================
- Coverage 79.06% 78.86% -0.21%
==========================================
Files 1067 1068 +1
Lines 123805 124023 +218
Branches 4510 4515 +5
==========================================
- Hits 97891 97810 -81
- Misses 24304 24605 +301
+ Partials 1610 1608 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/train/training.py`:
- Around line 1018-1021: Update the model_key_prob_map construction in the
share_kwargs block to call zip(self.model_keys, self.model_prob) with
strict=True, preserving the existing dictionary conversion while enforcing
equal-length inputs.
In `@deepmd/pt/utils/stat.py`:
- Around line 114-125: Update _restore_from_file to return (None, None) whenever
cached_keys does not include every requested key, allowing compute_output_stats
to recompute statistics for newly added outputs; retain the current dictionaries
only when all keys are cached.
- Line 116: Replace the mutable default value for keys in the relevant function
with a safe immutable or None-based default, initializing ["energy"] inside the
function when no keys are provided. Preserve the existing behavior for callers
that supply keys and ensure Ruff B006 passes.
🪄 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: 40a39df4-d3e1-41c3-a1c0-0a004587f8a2
📒 Files selected for processing (35)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/descriptor.pydeepmd/dpmodel/descriptor/dpa1.pydeepmd/dpmodel/descriptor/repflows.pydeepmd/dpmodel/descriptor/repformers.pydeepmd/dpmodel/descriptor/se_e2_a.pydeepmd/dpmodel/descriptor/se_r.pydeepmd/dpmodel/descriptor/se_t.pydeepmd/dpmodel/descriptor/se_t_tebd.pydeepmd/dpmodel/fitting/general_fitting.pydeepmd/dpmodel/train/data.pydeepmd/dpmodel/utils/env_mat_stat.pydeepmd/dpmodel/utils/stat.pydeepmd/jax/common.pydeepmd/jax/train/trainer.pydeepmd/pt/entrypoints/main.pydeepmd/pt/model/atomic_model/sezm_atomic_model.pydeepmd/pt/model/task/fitting.pydeepmd/pt/train/training.pydeepmd/pt/utils/stat.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/train/training.pydeepmd/tf2/entrypoints/train.pydeepmd/tf2/train/trainer.pydeepmd/utils/argcheck.pydeepmd/utils/stat_file.pysource/tests/common/dpmodel/test_train_data.pysource/tests/common/stat_file.pysource/tests/common/test_stat_file.pysource/tests/jax/test_model_factory.pysource/tests/jax/test_training.pysource/tests/pt/test_stat_file_mode.pysource/tests/pt_expt/test_entrypoint.pysource/tests/pt_expt/test_training.pysource/tests/tf2/test_training.py
💤 Files with no reviewable changes (1)
- source/tests/pt_expt/test_entrypoint.py
🚧 Files skipped from review as they are similar to previous changes (27)
- deepmd/dpmodel/descriptor/se_r.py
- deepmd/dpmodel/train/data.py
- deepmd/dpmodel/descriptor/repformers.py
- source/tests/tf2/test_training.py
- source/tests/common/stat_file.py
- deepmd/dpmodel/atomic_model/base_atomic_model.py
- deepmd/dpmodel/descriptor/se_t.py
- source/tests/common/dpmodel/test_train_data.py
- deepmd/dpmodel/descriptor/se_e2_a.py
- deepmd/dpmodel/utils/env_mat_stat.py
- deepmd/dpmodel/descriptor/descriptor.py
- source/tests/pt_expt/test_training.py
- deepmd/pt/model/task/fitting.py
- source/tests/jax/test_model_factory.py
- deepmd/jax/common.py
- deepmd/dpmodel/descriptor/dpa1.py
- deepmd/jax/train/trainer.py
- deepmd/tf2/entrypoints/train.py
- deepmd/pt/entrypoints/main.py
- source/tests/jax/test_training.py
- deepmd/dpmodel/utils/stat.py
- deepmd/utils/stat_file.py
- deepmd/dpmodel/fitting/general_fitting.py
- source/tests/common/test_stat_file.py
- deepmd/pt_expt/train/training.py
- deepmd/pt_expt/entrypoints/main.py
- source/tests/pt/test_stat_file_mode.py
njzjz-bot
approved these changes
Jul 29, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Reviewed the statistics-cache lifecycle, transactional output-pair handling, legacy cache compatibility, and backend synchronization paths. I found no actionable issues.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
njzjz
approved these changes
Jul 29, 2026
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
readandupdatemodesRoot cause
Statistics files were opened through
DPH5Path, whose process-widelru_cacheretained liveh5py.Fileobjects for the entire training process. Update-mode cache hits could therefore keep a writable HDF5 handle open and modify file metadata even when no statistics needed to be recomputed. Cache completeness and distributed initialization were also implemented independently by each backend, which allowed partial output pairs and divergent rank-local statistics.Behavior
Existing HDF5 caches in
updatemode are opened read-only and promoted to a writer only after a cache miss. A complete cache hit does not sample training data, acquire a writer handle, or change the file bytes. Every scoped handle is closed when statistics initialization finishes.Output bias and standard-deviation datasets are treated as complete pairs. Replacements use a transient transaction marker, remove stale requested pairs, write only newly computed complete pairs, flush them, and then remove the marker. A successful cache retains the legacy HDF5 layout; an interrupted transaction is recomputed in
updatemode and rejected in strictreadmode.This change intentionally does not modify Paddle, legacy TF1, or HDF5 training-data access.
Validation
git diff --checkpassedPT-expt compile-specific training tests were not completed locally because the Python process aborted inside a generated TorchInductor CPU kernel. The corresponding non-compile paths passed.
Summary by CodeRabbit
stat_file_mode: read/update) using a sharedStatFileSpecabstraction across training backends.