deepmodeling / deepmodeling/deepmd-kit

feat(tf2): implement training.mixed_precision

Open
#5,881 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

## Summary

TF2 explicitly raises `NotImplementedError` when the schema-exposed `training.mixed_precision` block is set.

## Scope

- Translate the normalized configuration into a TF2 mixed-precision policy.
- Apply compute, variable, output, and reduction dtypes consistently to model and loss execution.
- Add dynamic or fixed loss scaling where required.
- Detect non-finite gradients before optimizer/checkpoint updates.
- Preserve policy and optimizer state across checkpoint/resume.
- Reject unsupported dtype/device combinations before training.

## Acceptance criteria

- Supported mixed-precision configurations change the actual compute policy.
- Energy, force, loss, and gradient results remain within documented tolerances of full precision.
- Loss scaling and non-finite handling are covered by tests.
- Checkpoint/resume preserves the policy and reproduces the next update.
- Unsupported configurations fail with actionable messages.
- Full-precision TF2 behavior remains unchanged.

Refs #5757.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

Contributor guide

Open the contributing guide

Research direction

Start by locating the TF2 path that currently raises NotImplementedError for the schema-exposed training.mixed_precision block, then trace how normalized configuration reaches model, loss, optimizer, and checkpoint handling. Use the acceptance criteria as the completion checklist: supported policies affect computation, loss scaling and non-finite gradients are tested, checkpoint/resume is reproducible, unsupported combinations fail clearly, and full-precision behavior is unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, tensorflow
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.