deepmd.pt_expt.train.gradient

Numerical safeguards around the gradient of a training step.

Two hazards are handled here, both of which have to be addressed without slowing the step down.

A gradient that is large but still representable can overflow the sum of squares of a naive L2 reduction, so a finite gradient would be reported as infinite and clipped as if training had diverged. Scaling the gradients by their largest magnitude before the reduction removes that failure mode.

A gradient that is genuinely non-finite must stop the run before it reaches a checkpoint, but reading that condition back on every step would serialize the host against the device. The condition is therefore accumulated on device and read only at a checkpoint boundary.

Classes

NonFiniteGradGuard

Detect non-finite gradient norms without a per-step host synchronization.

Functions

clip_grad_norm_(→ torch.Tensor)	Clip gradients in place so their global L2 norm does not exceed max_norm.
raise_nonfinite_gradient_norm(→ None)	Raise a RuntimeError reporting the current non-finite gradients.

Module Contents

deepmd.pt_expt.train.gradient.clip_grad_norm_(parameters: collections.abc.Iterable[torch.nn.Parameter], max_norm: float, stable: bool = True) → torch.Tensor

Clip gradients in place so their global L2 norm does not exceed max_norm.

The norm is computed and applied on device and returned without a host synchronization. Non-finite gradients are not reported here; see NonFiniteGradGuard.

Parameters:

parametersIterable[torch.nn.Parameter]

Parameters whose gradients are clipped in place.

max_normfloat

Maximum allowed global L2 norm of the gradients.

stablebool, optional

Norm reduction strategy. True scales the gradients by their largest magnitude before the float64 reduction, keeping the norm finite for an arbitrarily large but finite gradient of any dtype. False uses the native reduction, which propagates DTensor sharding under FSDP2 but is not overflow-safe.

Returns:

torch.Tensor

The global gradient norm before clipping, or a CPU zero when no gradient is present.

class deepmd.pt_expt.train.gradient.NonFiniteGradGuard

Detect non-finite gradient norms without a per-step host synchronization.

update() accumulates the non-finite condition on device; the result is read back only by raise_if_nonfinite(). Calling the check before each checkpoint keeps a diverged interval from being written while leaving the training step free of host reads.

_nonfinite: torch.Tensor | None = None

update(total_norm: torch.Tensor) → None

Accumulate whether total_norm is non-finite.

Parameters:

total_normtorch.Tensor

The gradient norm returned by clip_grad_norm_().

raise_if_nonfinite(named_parameters: collections.abc.Callable[[], collections.abc.Iterable[tuple[str, torch.nn.Parameter]]]) → None

Raise if any norm accumulated since the previous call was non-finite.

On failure the current gradient state is reported via raise_nonfinite_gradient_norm().

Parameters:

named_parametersCallable[[], Iterable[tuple[str, torch.nn.Parameter]]]

Accessor for the model’s named parameters, consulted only when raising.

Raises:

RuntimeError

If a non-finite gradient norm was recorded.

deepmd.pt_expt.train.gradient.raise_nonfinite_gradient_norm(named_parameters: collections.abc.Iterable[tuple[str, torch.nn.Parameter]]) → None

Raise a RuntimeError reporting the current non-finite gradients.

Parameters whose current gradient is non-finite are listed by name and shape. When every current individual gradient is finite, the accumulated guard flag came from an earlier step in the checkpoint interval or from the norm reduction.

Parameters:

named_parametersIterable[tuple[str, torch.nn.Parameter]]

The model’s named parameters.

Raises:

RuntimeError

Always; this is the divergence-reporting path.