[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 consistently detect non-finite gradients across distributed training ranks, including during validation and final saves.
  • When a non-finite gradient is detected, training raises an error before saving or updating best-checkpoint records, preventing incomplete checkpoint updates.
  • Finite-gradient checkpoint saves continue to work across sharded and unsharded training.

[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

• Configuration used: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 67c53b02-d6d3-4939-bf0e-f72ea76a603e

📥 Commits

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

📒 Files selected for processing: 1

• source/tests/pt/test_validation.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	Trainer checkpoint paths check gradient finiteness before saving. Distributed checkpoint paths coordinate ranks, with tests covering regular, EMA, best-checkpoint, and sharded saves.
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: ➖ Normal

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 04f93

Distributed checkpoint preflight runs before rank-specific writes, and failed validation preparation leaves top-K metadata unpublished. Existing tests cover that failure path; no concrete merge-blocking risk remains.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 44.68% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 47 functions across 8 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: checkpoint publication is preflighted on every rank.
Linked Issues check		Issue #5816 is active and directly linked. The changes route PT and pt_expt regular, EMA, validation-best, and EMA-validation-best publication through finite-gradient preflight. The preflight agrees t…
Out of Scope Changes check		The changes stay within issue #5816. The production changes centralize checkpoint preflight, defer full-validation state publication, and include all required ranks in collective checkpoint paths. The…

Full details: Linked Issues check

Explanation

Issue #5816 is active and directly linked. The changes route PT and pt_expt regular, EMA, validation-best, and EMA-validation-best publication through finite-gradient preflight. The preflight agrees the accumulated flag across ranks, runs before checkpoint collection or serialization, and resets only after a successful boundary check. Full validation proposes state before the check and commits it after the check. The normal optimizer-step path and clipping algorithm remain unchanged. Tests cover failed publication, unchanged metadata, recovery, restart state, and distributed paths.

Full details: Out of Scope Changes check

Explanation

The changes stay within issue #5816. The production changes centralize checkpoint preflight, defer full-validation state publication, and include all required ranks in collective checkpoint paths. The tests support these requirements. No unrelated optimizer, clipping, rollback, or file-recovery feature is added.

• Fix all pre-merge checks with AI

✨ 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.78%. Comparing base (4f78c2d) to head (04f93d9).

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.78%   -0.31%     
==========================================
  Files        1171     1171              
  Lines      141657   141649       -8     
  Branches     5072     5072              
==========================================
- Hits       110615   110180     -435     
- Misses      29102    29531     +429     
+ 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.