deepmodeling / deepmodeling/deepmd-kit

refactor(train): centralize stable gradient clipping and non-finite guards

Open
#5,816 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.