pt_expt: align backend support with argcheck schema

## Summary

While comparing the unified training schema in `deepmd/utils/argcheck.py` with the `pt_expt` backend implementation, I found several options that are exposed as configurable (or documented as PyTorch-supported) but are not implemented in `pt_expt` yet.

This issue tracks the gaps so we can either implement them in `pt_expt` or make the schema/docs distinguish legacy `pt` from `pt_expt` more clearly.

## Gaps

### 1. `model.type: dpa4` advertises `fitting_net.type: property`, but `pt_expt` rejects it

Schema/docs:

- `deepmd/utils/argcheck.py` registers the DPA4/SeZM model (`model_args_plugin.register("dpa4", alias=["DPA4", "SeZM", "sezm"])`).
- Its docs say the DPA4/SeZM scaffold has a selectable energy or invariant-property fitting.
- The DPA4 model schema includes both `dpa4_ener` and `property` in the `fitting_net.type` variant.

Implementation:

- `deepmd/pt_expt/model/get_model.py` `get_sezm_model()` only accepts `fitting_net.type` in `("dpa4_ener", "sezm_ener")` and raises otherwise.

Impact:

- A schema-valid `pt_expt` input with `model.type=dpa4` and `fitting_net.type=property` fails during model construction.

### 2. `fitting_net.type: population` / `loss.type: population` are exposed, but `pt_expt` has no implementation

Schema/docs:

- `deepmd/utils/argcheck.py` registers `fitting_args_plugin.register("population", doc=doc_only_pt_supported)`.
- `deepmd/utils/argcheck.py` registers `loss_args_plugin.register("population")`.

Implementation:

- `deepmd/pt_expt/fitting/` registers `ener`, `dpa4_ener`/`sezm_ener`, `dos`, `property`, `polar`, `dipole`, and `invar`, but no `population` fitting.
- `deepmd/pt_expt/train/training.py` loss dispatch handles `ener`, `dos`, `ener_spin`, `tensor`, and `property`, but no `population`.

Impact:

- `population` is presented as PyTorch-supported in the common schema, but `pt_expt` users cannot train it.

### 3. `loss.type: dens` is exposed for DPA4/SeZM, but `pt_expt` does not dispatch it

Schema/docs:

- `deepmd/utils/argcheck.py` registers `loss_args_plugin.register("dens")`.
- The loss variant docs say `dens` can be used with `ener` fitting for DPA4/SeZM.

Implementation:

- `deepmd/pt_expt/loss/` exposes `DOSLoss`, `EnergyLoss`, `EnergySpinLoss`, `PropertyLoss`, and `TensorLoss`, but no `DensLoss`.
- `deepmd/pt_expt/train/training.py` loss dispatch has no `dens` branch.

Impact:

- `loss.type=dens` passes schema-level expectations but fails in `pt_expt` training.

### 4. `learning_rate.type: cosine/wsd` are exposed, but `pt_expt` hard-codes `LearningRateExp`

Schema/docs:

- `deepmd/utils/argcheck.py` registers `learning_rate.type` variants `exp`, `cosine`, and `wsd`.

Implementation:

- `deepmd/pt_expt/train/training.py` constructs `self.lr_schedule = LearningRateExp(**lr_params)` directly, without dispatching on `learning_rate.type`.

Impact:

- `cosine`/`wsd` learning-rate configs are either rejected by unexpected parameters or silently treated as the wrong schedule.

### 5. Optimizer schema exposes `LKF`, `AdaMuon`, and `HybridMuon`, but `pt_expt` only supports `Adam`/`AdamW`

Schema/docs:

- `deepmd/utils/argcheck.py` registers optimizer variants `Adam`, `AdamW`, `LKF`, `AdaMuon`, and `HybridMuon`.

Implementation:

- `deepmd/pt_expt/train/training.py` only dispatches `Adam` and `AdamW`; all other optimizer names raise `ValueError("Unsupported optimizer type: ...")`.

Impact:

- PyTorch-only optimizer options in the schema are not available through `pt_expt`.

### 6. `training.zero_stage` appears schema-exposed, but `pt_expt` uses plain DDP

Schema/docs:

- `deepmd/utils/argcheck.py` documents `training.zero_stage` for DDP/ZeRO/FSDP modes.

Implementation observation:

- `deepmd/pt_expt/train/training.py` wraps distributed training with `torch.nn.parallel.DistributedDataParallel` and directly constructs Adam/AdamW optimizers; I did not find ZeRO/FSDP wrapping in the `pt_expt` path.

Impact:

- If `zero_stage` is meant to apply to `pt_expt`, it appears to be accepted but not honored. If it is legacy-`pt` only, the docs/schema should say so.

## Expected resolution

Either:

1. implement these options in `pt_expt`, or
2. update `argcheck.py`/docs to distinguish `pt` vs `pt_expt` support so users do not get schema-valid configurations that fail at backend runtime.

Filed by OpenClaw 2026.6.11 (e085fa1), model: custom-chat-jinzhezeng-group/gpt-5.5.

## Implementation checklist

### Model, loss, and optimizer support

- [ ] Support `fitting_net.type: property` for DPA4/SeZM
- [ ] Implement `fitting_net.type: population` and `loss.type: population`
- [ ] Implement `loss.type: dens` for DPA4/SeZM
- [x] Support `learning_rate.type: cosine` and `learning_rate.type: wsd` (#5776)
- [ ] Support the schema-exposed `LKF`, `AdaMuon`, and `HybridMuon` optimizers (#5815 covers normalized optimizer wiring and `HybridMuon`)
- [ ] Implement or explicitly reject/document `training.zero_stage` (#5822)

### Training support

- [ ] Derive training steps from `training.numb_epoch` and its aliases (#5814)
- [ ] Support multi-task `training.num_epoch_dict`
- [ ] Consume the schema-shaped `optimizer` block instead of backend-only training fields (#5815)
- [ ] Support `training.save_dir` (#5817)
- [ ] Support `training.ckpt_keep_ratio`
- [ ] Support `training.enable_ema`, `training.ema_decay`, and `training.ema_ckpt_keep` (#5818)
- [ ] Support `training.disp_avg` (#5819)
- [ ] Support `training.profiling` and `training.enable_profiler` (#5820)
- [ ] Support `training.tensorboard*` options (#5821)
- [ ] Support `training.mixed_precision`
- [x] Apply `training.change_bias_after_training` after the final step (#5745)
- [ ] Support `training.training_data.min_pair_dist`

### Training stability and distributed data

- [ ] Centralize stable gradient clipping and non-finite guards (#5816)
- [ ] Shard and reshuffle LMDB batches across ranks and epochs (#5823)

### Graph-native DPA1 follow-ups

- [ ] Decouple DPA1 graph normalization and statistics from dense `sel` (#5824)
- [ ] Make graph-native preprocessing and export independent of `get_sel()` (#5825)