deepmodeling / deepmodeling/deepmd-kit

tf2: align backend support with argcheck schema

Open
#5,757 1 comment 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

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

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.