pt_expt: align backend support with argcheck schema

[njzjz-bot]

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
• Support learning_rate.type: cosine and learning_rate.type: wsd (fix(train): dispatch registered learning rate schedules #5776)
• Support the schema-exposed LKF, AdaMuon, and HybridMuon optimizers (feat(pt_expt): consume optimizer config and support HybridMuon #5815 covers normalized optimizer wiring and HybridMuon)
• Implement or explicitly reject/document training.zero_stage (feat(pt_expt): support zero_stage through distribution strategies #5822)

Training support

• Derive training steps from training.numb_epoch and its aliases (feat(pt_expt): support epoch-based training horizons #5814)
• Support multi-task training.num_epoch_dict
• Consume the schema-shaped optimizer block instead of backend-only training fields (feat(pt_expt): consume optimizer config and support HybridMuon #5815)
• Support training.save_dir (feat(pt_expt): honor training.save_dir through a checkpoint manager #5817)
• Support training.ckpt_keep_ratio
• Support training.enable_ema, training.ema_decay, and training.ema_ckpt_keep (feat(pt_expt): add compile-aware EMA lifecycle and checkpoint retention #5818)
• Support training.disp_avg (feat(train): support interval metric aggregation for disp_avg #5819)
• Support training.profiling and training.enable_profiler (feat(pt_expt): add torch profiler lifecycle and trace export #5820)
• Support training.tensorboard* options (feat(pt_expt): add rank-aware TensorBoard observation #5821)
• Support training.mixed_precision
• Apply training.change_bias_after_training after the final step (fix(train): share post-training bias adjustment #5745)
• Support training.training_data.min_pair_dist

Training stability and distributed data

• Centralize stable gradient clipping and non-finite guards (refactor(train): centralize stable gradient clipping and non-finite guards #5816)
• Shard and reshuffle LMDB batches across ranks and epochs (fix(pt_expt): shard and reshuffle LMDB batches across ranks and epochs #5823)

Graph-native DPA1 follow-ups

• Decouple DPA1 graph normalization and statistics from dense sel (feat(pt_expt): decouple DPA1 graph normalization and statistics from dense sel #5824)
• Make graph-native preprocessing and export independent of get_sel() (refactor(pt_expt): make graph-native preprocessing and export independent of get_sel #5825)

[njzjz-bot]

Addendum: training option coverage

I re-checked the training_*_args exposed by deepmd/utils/argcheck.py against pt_expt on current master (f73de32). Additional training-schema gaps not called out in the initial issue:

• training.numb_epoch / aliases and multi-task training.num_epoch_dict: the schema allows epoch-based step derivation, but deepmd/pt_expt/train/training.py still requires training_params["numb_steps"] directly in DPTrainer.__init__; I did not find epoch -> numb_steps conversion or the multi-task num_epoch_dict -> model probability / total steps derivation in the pt_expt path.
• Optimizer config shape mismatch: the common schema exposes an optimizer block (optimizer.type, optimizer-specific parameters), but pt_expt currently reads training.opt_type and training.weight_decay instead. So the initial note that pt_expt only dispatches Adam/AdamW is true but incomplete: even the schema-shaped optimizer: {type: AdamW, ...} block is not wired to this backend.
• training.save_dir, training.ckpt_keep_ratio, training.enable_ema, training.ema_decay, training.ema_ckpt_keep, training.disp_avg, training.profiling, training.enable_profiler, training.tensorboard*, and training.mixed_precision: these are schema-visible training options, but I did not find corresponding pt_expt handling beyond the older save_ckpt / save_freq / max_ckpt_keep / display/timing fields.
• training.change_bias_after_training: pt_expt has a bias-adjustment helper used for fine-tuning / pretrained initialization paths, but DPTrainer.run() has no post-training branch that applies change_bias_after_training after the last step.
• training.training_data.min_pair_dist: the schema marks this as PyTorch-supported. The legacy pt trainer adds the data requirement and filters batches by min_pair_dist, but I did not find equivalent data-requirement wiring or filtering in the pt_expt trainer.

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