deepmodeling / deepmodeling/deepmd-kit

feat(tf2): align optimizer dispatch with the common schema

Open
#5,878 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 optimizer dispatch supports Adam and AdamW, while the common schema also exposes LKF, AdaMuon, and HybridMuon. Unsupported choices currently fail only in the TF2 trainer.

DPA4 HybridMuon support is tracked in #5771; this issue covers the broader backend-aware optimizer contract, including explicit support/restriction for the remaining schema variants.

## Scope

- Consume and validate the normalized top-level `optimizer` block.
- Implement TF2-native optimizers that are intended to be supported.
- Reject intentionally unsupported variants during backend-aware validation.
- Wire optimizer-specific parameters and preserve state through checkpoints.
- Do not silently fall back to Adam or ignore normalized fields.
- Coordinate the shared dispatch/factory boundary with #5771.

## Acceptance criteria

- Every schema-exposed optimizer has an explicit TF2 support status.
- Supported optimizer parameters affect deterministic one-step update tests.
- Unsupported variants fail before the training loop with an actionable message.
- Adam and AdamW remain behaviorally compatible.
- Checkpoint/resume restores complete optimizer state and reproduces the next update.
- HybridMuon integration can land through #5771 without duplicating configuration parsing.

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

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.