deepmodeling / deepmodeling/deepmd-kit
refactor(train): centralize stable gradient clipping and non-finite guards
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Current status
The core implementation is present on `master`:
- the canonical helper lives in `deepmd/pt_expt/train/gradient.py`;
- both PT and pt_expt import the same stable clipping and non-finite guard implementation;
- large finite gradients use an overflow-safe scaled reduction;
- foreach operations are used on the ordinary tensor path;
- the non-finite condition is accumulated on device without a per-step host synchronization;
- regular and EMA checkpoint publication through `save_checkpoint()` is guarded;
- sharded gradients use an explicit distributed-native path instead of a shard-local stable norm.
The remaining gap is not test coverage or the validation calculation itself. It is one unguarded checkpoint-writing path.
## Remaining problem
The training loop runs full validation after an optimizer step and before the periodic checkpoint boundary. When full validation selects a top-K candidate, pt_expt calls:
- `_save_full_validation_checkpoint()`, or
- `_save_full_validation_ema_checkpoint()`.
Both methods write through `_save_checkpoint_to_path()` directly and bypass `NonFiniteGradGuard.raise_if_nonfinite()`. Therefore a non-finite gradient recorded since the previous checkpoint boundary can still be persisted as a validation-best checkpoint.
Full validation evaluates in inference mode, but the relevant side effect here is checkpoint publication from inside the training loop.
## Proposed design
Centralize the checkpoint safety gate at the lowest common publication boundary used by every training checkpoint namespace:
1. expose one trainer method that validates and resets the accumulated non-finite state immediately before serialization;
2. call it from regular, EMA, validation-best, and EMA-validation-best checkpoint entry points;
3. perform the check before any checkpoint bytes or metadata are written;
4. keep the guard outside model compilation and synchronize with the host only when a checkpoint is actually about to be published.
The gate must run once per logical checkpoint boundary. A regular checkpoint that materializes both live and EMA weights should validate once before either file is written; it must not reset between the two files.
## Acceptance criteria
- No regular, EMA, validation-best, or EMA-validation-best checkpoint is written after a non-finite gradient norm in the current checkpoint interval.
- A successful boundary check resets the accumulated state exactly once.
- A failed boundary check raises before serialization or latest/top-K metadata publication.
- The normal optimizer-step path remains free of host synchronization.
- Existing stable clipping behavior for PT, pt_expt, DDP, and the explicit sharded path remains unchanged.
Refs #5755.
Contributor guide
Assessment
This issue has not been assessed yet.