[OutisLi]

Summary

Fixes #5816. Supersedes #6061 with a fresh implementation on current master.

• Run an all-rank non-finite-gradient preflight before any checkpoint state changes or serialization. Agree on the accumulated flag before gradient diagnostics as well, since sharded diagnostics can themselves be collective.
• Full validation proposes a top-K candidate without changing state, checks only when a checkpoint will actually be saved, then records the selected state before serialization. Restart metadata therefore includes the checkpoint being written.
• Share the PT regular/final/zero-step publication path and check once for each live-plus-EMA save. PT-expt includes every rank in checkpoint preflight, including ordinary DDP and the post-bias final save.

The optimizer-step path and gradient-clipping algorithm are unchanged. Existing serialization/reconciliation I/O-error behavior is retained; this fix does not introduce a general checkpoint rollback or file-recovery protocol.

Validation

• Focused guard, checkpoint, validator, and real PT training-loop regressions: 63 passed, 15 subtests passed.
• Additional existing PT/PT-expt full-validation, EMA, zero-step, retention, and restart training regressions: 24 passed, 2 subtests passed. These overlap some focused coverage.
• Repository-wide Ruff check/format, changed-file isort, and git diff --check: passed.
• Added a bounded two-process CPU Gloo regression with real DDP and ZeRO-1 checkpoint collection, covering both backends and regular/live-best/EMA-best boundaries. It injects a failure only on rank 1 and includes finite-save controls. Not locally validated: this environment blocks Gloo sockets, and elevated execution failed before Python started. The test remains enabled on Gloo-capable CI; it was deselected for the passing local runs above.
• Hosted CI on 04f93d9: all 12 Test Python shards passed. Shard 6 ran all 23 cases in test_checkpoint_guard.py successfully, including the two-process Gloo/DDP/ZeRO-1 regression; none of those 23 cases were skipped. CodeQL, pre-commit, and documentation checks also passed.
• Local testing did not include GPUs or the full suite. CUDA test jobs were skipped, so no GPU validation is claimed.

Summary by CodeRabbit

• Bug Fixes
  • Checkpoint saves now check for non-finite gradients consistently across distributed training ranks, including during periodic, validation, and final saves.
  • If any rank has detected a non-finite gradient, training raises an error before saving a checkpoint or updating best-checkpoint records.
  • Finite-gradient saves continue to work for both sharded and unsharded training, including zero-step and post-bias-adjustment saves.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing. Review the Copilot code review billing settings.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 6895285d-e60c-48d0-b3a6-653094268bbd

📥 Commits

Reviewing files that changed from the base of the PR and between 04f93d9 and 91364f6.

📒 Files selected for processing: 1

• source/tests/pt_expt/test_lmdb_training.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 6 remain after this review.

📝 Walkthrough

The training paths check accumulated gradient finiteness before checkpoint serialization. Distributed ranks agree on detected failures. Full validation prepares selected checkpoints across ranks before rank 0 publishes proposed top-K records.

Changes: Distributed checkpoint safety

Layer / File(s)	Summary
Synchronized gradient guard deepmd/pt_expt/train/gradient.py, source/tests/pt_expt/test_train_gradient.py	The guard can synchronize divergence flags across ranks and resets its accumulated state on each check. A test covers detection of a remote rank’s failure.
Checkpoint preflight and save paths deepmd/pt/train/training.py, deepmd/pt_expt/train/training.py, source/tests/pt/test_training.py, source/tests/pt_expt/test_checkpoint_guard.py, source/tests/pt_expt/test_lmdb_training.py	Trainer checkpoint paths check gradient finiteness before saving. Distributed checkpoint paths coordinate ranks, with tests covering regular, EMA, best-checkpoint, and sharded saves. LMDB training tests register cleanup for temporary directories and trainer data systems.
Validation checkpoint proposal and commit deepmd/pt_expt/train/validation.py, source/tests/pt/test_validation.py	Full validation separates proposing top-K records from publishing them. For a selected checkpoint, all ranks prepare the checkpoint before rank 0 commits the proposed records.

Priority: ⬇️ Low

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant Rank0
  participant FullValidator
  participant Trainer
  participant StateStore
  Rank0->>FullValidator: Run validation and propose top-K records
  FullValidator->>Trainer: Prepare selected checkpoint on all ranks
  FullValidator->>StateStore: Commit proposed records on rank 0

Loading

Merge Risk: ⚪ Minimal · up to 91364

The changed checkpoint paths guard serialization and metadata publication correctly, with no concrete merge-blocking risk remaining.

Pre-merge checks: 3 · 2

Failed checks: 2 warnings

Check name	Status	Explanation	Resolution
Out of Scope Changes check		The changes in source/tests/pt_expt/test_lmdb_training.py only change temporary-directory and data-system cleanup. They do not implement or test checkpoint preflight, full-validation publication, non-…	Remove the unrelated cleanup changes from source/tests/pt_expt/test_lmdb_training.py, or provide a direct issue-related requirement that makes those changes necessary.
Docstring Coverage		Docstring coverage is 49.09% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 55 functions across 9 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 3 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title "fix(train): preflight checkpoint publication on every rank" matches the main change. The PR adds a synchronized non-finite-gradient check before checkpoint serialization and publication, an…
Linked Issues check		Issue #5816 is open and directly linked. The reviewed changes add checkpoint preflight for PT and pt_expt regular, EMA, validation-best, and EMA-validation-best saves. The preflight agrees the accumul…

Full details: Title check

Explanation

The title "fix(train): preflight checkpoint publication on every rank" matches the main change. The PR adds a synchronized non-finite-gradient check before checkpoint serialization and publication, and it makes that preflight collective across all ranks in the PT and pt_expt trainers. The objectives state the same goal, so the title is concise and specific. It uses the conventional commit prefix and names the affected area.

Full details: Linked Issues check

Explanation

Issue #5816 is open and directly linked. The reviewed changes add checkpoint preflight for PT and pt_expt regular, EMA, validation-best, and EMA-validation-best saves. The preflight agrees the accumulated non-finite flag across ranks before diagnostics and checkpoint collection. Full validation proposes state before the check and publishes top-K state only after a successful check. The guard resets after each check, and the optimizer-step and stable clipping paths remain unchanged. The added tests cover failed publication, unchanged metadata, recovery, and distributed DDP/ZeRO-1 paths.

Full details: Out of Scope Changes check

Explanation

The changes in source/tests/pt_expt/test_lmdb_training.py only change temporary-directory and data-system cleanup. They do not implement or test checkpoint preflight, full-validation publication, non-finite-gradient handling, or clipping for issue #5816. The remaining reviewed production and test changes support the linked issue.

Resolution

Remove the unrelated cleanup changes from source/tests/pt_expt/test_lmdb_training.py, or provide a direct issue-related requirement that makes those changes necessary.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[njzjz-bot]

Agent: dot

Reviewed exact head 04f93d9d8710d5e9517b7bebcc806db071d8d0ca against 4f78c2d0de347e4a40920c105a04a0506118f7af, including all eight changed files, the shared trainer loop, and the superseded #6061 review history. I found no additional actionable regression in this bounded review.

The replacement addresses the earlier preflight/serialization ordering concern structurally: even a rank without a local guard update joins the flag agreement before diagnostics; selected top-K state remains provisional until preflight succeeds and is then present during serialization. PT regular/final/zero-step publication converges on one helper, and ordinary DDP nonchief ranks now join PT-expt preflight before returning without writing. Each regular live-plus-EMA boundary checks once. Not-due, unselected, and save_best=False validation paths do not consume the accumulated guard. Reset on inspection, including rejection, is retained behavior, not a new success-only reset contract.

Independent verification: 25 exact-source host-control scenarios passed with blocking two-thread collective doubles (5 guard, 12 validator, 8 regular live-plus-EMA save cases), plus trainer wiring/participation checks and syntax inspection. These check control flow and top-K reload bookkeeping; they are not native Torch serialization, real Gloo/ZeRO, FSDP, GPU, or full-training execution. The added two-process test source does exercise real trainer writers, rank-1-only historical failure, and finite controls, but its hosted execution remains pending while current-head jobs are queued. Author-reported local test totals are separate evidence.

No general I/O rollback/recovery requirement is being reintroduced into this deliberately narrower replacement. This is a scoped review record, not full distributed-runtime qualification.

[codecov]

Codecov Report

❌ Patch coverage is 88.67925% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.84%. Comparing base (4f78c2d) to head (91364f6).
⚠️ Report is 2 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/pt/train/training.py	89.28%	3 Missing ⚠️
deepmd/pt_expt/train/training.py	57.14%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6089      +/-   ##
==========================================
- Coverage   78.08%   77.84%   -0.24%     
==========================================
  Files        1171     1171              
  Lines      141657   141655       -2     
  Branches     5072     5078       +6     
==========================================
- Hits       110615   110278     -337     
- Misses      29102    29439     +337     
+ Partials     1940     1938       -2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[njzjz-bot]

Agent: dot

Qualification follow-up to review 5574354987, at unchanged head 04f93d9d8710d5e9517b7bebcc806db071d8d0ca: hosted CPU CI is now complete, with all 12 Python shards and all 4 C++ jobs passing.

Python shard 6 reports all 23 cases in source/tests/pt_expt/test_checkpoint_guard.py passed, with no skips in that module. I checked the parametrized case count and verified that the tested merge e68787c1dd5b95d9057f123956272854e3f6b516 contains this head and base 4f78c2d; its checkpoint-test file matches the head. This includes the real two-process Gloo regression covering PT/PT-expt, ordinary DDP and ZeRO-1, rank-1-only historical failure, and finite-save controls at regular/live-best/EMA-best boundaries. The earlier pending CPU distributed-checkpoint qualification is therefore established by hosted execution, rather than the host-control doubles in my source review.

The CUDA workflow is aggregate green but both actual CUDA test jobs were skipped. This does not establish FSDP or GPU execution. No source re-review is being repeated.

[wanghan-iapcm]

LGTM. Agreeing on the non-finite flag across ranks before any checkpoint write, and proposing the validation top-K update without side effects until the save is confirmed, closes the unguarded save paths.

[njzjz-bot]

Agent: dot

Incremental review only: 04f93d9 → 91364f6. No additional actionable issue found in this one-file test cleanup change; the earlier checkpoint review is not being repeated.

The temporary-directory cleanup is registered first, so unittest's LIFO cleanups close all seven trainer/probe creation sites before removing the LMDB files. Registration immediately after construction also covers failures in get_batch/run/assertions, and directory cleanup now runs after a failed setUp. All training assertions remain intact. I checked Trainer._close_data_systems and the LMDB adapter/iterator/reader close paths: they release the prefetch pipeline before the reader, with idempotent iterator/reader closure, so cleanup after trainer.run's own closure is safe. Removing ignore_errors makes a remaining resource leak visible instead of concealing it.

Verification: exact-source syntax inspection and three stdlib cleanup-order/error-path scenarios passed; these are not real LMDB/Torch training execution. Current-head CI has successful Python shards plus others still running, with CUDA Python skipped; no GPU or full-suite pass is claimed. Review-only: no code, branches, auto-merge, or merge changed.