deepmodeling / deepmodeling/deepmd-kit
tf2: align backend support with argcheck schema
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
## Summary
While comparing `deepmd/utils/argcheck.py` with the `deepmd/tf2` backend, I found several common-schema options that are not implemented or not fully wired up in TF2.
Some options are documented as `Supported Backend: TensorFlow`; if those are intended to mean only the legacy `deepmd/tf` backend and not `deepmd/tf2`, the docs/schema should make that distinction explicit.
## Gaps
### 1. Legacy TensorFlow-only descriptors are schema-exposed, but TF2 does not register them
Schema/docs:
- `deepmd/utils/argcheck.py` exposes TensorFlow-supported descriptor variants including `loc_frame`, `se_a_tpe`/`se_a_ebd`, `se_a_ebd_v2`/`se_a_tpe_v2`, and `se_a_mask`.
Implementation observation:
- `deepmd/tf2/descriptor/` registers the common dpmodel-style descriptors (`se_e2_a`, `se_e2_r`, `se_e3`, `se_e3_tebd`, `se_atten`, `se_atten_v2`, `dpa2`, `dpa3`, `hybrid`), but not `loc_frame`, `se_a_tpe`/`se_a_ebd`, `se_a_ebd_v2`/`se_a_tpe_v2`, or `se_a_mask`.
Impact:
- A user reading `Supported Backend: TensorFlow` may reasonably expect these to work with the TF2 backend, but TF2 cannot construct them.
### 2. TF-only model features are schema-exposed, but TF2 does not implement them
Schema/docs:
- `deepmd/utils/argcheck.py` exposes TF-only model-level features such as `type_embedding`, `modifier`, `compress`, and hybrid model `pairwise_dprc`.
- It also exposes model variants such as `frozen`, `pairtab`, and `linear_ener` in the common model schema.
Implementation:
- `deepmd/tf2/model/model.py` explicitly rejects model-level `type_embedding` with a `ValueError` and does not implement the legacy TF model-level path.
- TF2 model registrations are limited to task models (`ener`, `dos`, `dipole`, `polar`, `property`) plus `zbl`; I did not find TF2 `BaseModel.register(...)` implementations for `frozen`, top-level `pairtab`, `linear_ener`, or `pairwise_dprc`.
Impact:
- Schema-visible TensorFlow/common model options are unavailable in TF2 or fail during model construction.
### 3. `model.spin` is accepted by the common schema, but TF2 explicitly rejects spin models
Schema/docs:
- `deepmd/utils/argcheck.py` exposes the top-level `model.spin` block.
Implementation:
- `deepmd/tf2/model/model.py` raises `NotImplementedError("Spin model is not implemented yet.")` when `model.type == "standard"` and `spin` is present.
Impact:
- Schema-valid spin model configs fail at TF2 model construction time.
### 4. Learning-rate schema exposes `exp`, `cosine`, and `wsd`, but TF2 hard-codes `LearningRateExp`
Schema/docs:
- `deepmd/utils/argcheck.py` registers `learning_rate.type` variants `exp`, `cosine`, and `wsd`.
Implementation:
- `deepmd/tf2/train/trainer.py` directly constructs `self.lr_schedule = LearningRateExp(**lr_params)` without dispatching on `learning_rate.type`.
Impact:
- `learning_rate.type=cosine` or `learning_rate.type=wsd` is schema-exposed but unavailable in TF2 training.
### 5. Some schema loss types are not implemented in TF2
Schema/docs:
- `deepmd/utils/argcheck.py` exposes loss variants including `ener`, `dens`, `ener_spin`, `dos`, `population`, `property`, and `tensor`.
Implementation:
- `deepmd/tf2/train/trainer.py` `get_loss()` only supports `ener`, `dos`, `tensor`, and `property`.
- There is no TF2 dispatch for `dens`, `ener_spin`, or `population`.
Impact:
- These loss types are visible in the common schema but fail in TF2 training. If they are intentionally limited to other backends, the schema/docs should spell out the backend restriction.
### 6. Optimizer support is narrower than the common schema suggests
Schema/docs:
- `deepmd/utils/argcheck.py` exposes optimizer variants `Adam`, `AdamW`, `LKF`, `AdaMuon`, and `HybridMuon`.
Implementation:
- `deepmd/tf2/train/trainer.py` `_build_optimizer()` supports `Adam` and `AdamW` only; other optimizer names raise `ValueError("Unsupported optimizer type for tf2: ...")`.
Impact:
- The common optimizer schema is broader than TF2 support. If only `Adam`/`AdamW` are intended for TF2, the docs/schema should make that explicit.
## Expected resolution
Either:
1. implement/wire these options in the TF2 backend, or
2. update `argcheck.py`/docs to clearly distinguish legacy `tf` vs `tf2` support and mark unsupported TF2 options before users hit backend runtime errors.
Filed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5.
## Implementation checklist
### Model, loss, and optimizer support
- [ ] Implement or explicitly reject/document legacy TensorFlow-only descriptors in TF2 (#5874)
- [ ] Implement or explicitly reject/document TF-only model features and model variants in TF2 (#5875)
- [ ] Implement TF2 spin models or reject them during backend-aware validation (#5876)
- [x] Support `learning_rate.type: cosine` and `learning_rate.type: wsd` (#5776)
- [ ] Implement or explicitly reject/document `dens`, `ener_spin`, and `population` losses (#5877)
- [ ] Support or explicitly restrict schema-exposed optimizers beyond `Adam` and `AdamW` (#5878; #5771 tracks DPA4 `HybridMuon`)
### Training support
- [ ] Derive training steps from `training.numb_epoch` and its aliases (#5879)
- [ ] Support multi-task `training.num_epoch_dict` (#5879)
- [ ] Implement `training.profiling` and `training.enable_profiler` (#5880)
- [ ] Implement `training.mixed_precision` (#5881)
- [x] Apply `training.change_bias_after_training` after the final step (#5745)
Contributor guide
Assessment
This issue has not been assessed yet.