deepmodeling / deepmodeling/deepmd-kit
feat(tf2): implement training.mixed_precision
- 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
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