jax: align training/model support with argcheck schema

[njzjz-bot]

Summary

While comparing the unified input schema in deepmd/utils/argcheck.py with the JAX backend implementation, I found several schema-exposed options that are not implemented or not wired up in deepmd/jax yet.

This issue tracks the gaps so we can either implement them for JAX or make the schema/docs explicitly say that these options are unavailable in the JAX backend.

Gaps

1. Non-energy fittings/models are registered, but JAX training only accepts loss.type: ener

Schema/backend registrations:

• deepmd/utils/argcheck.py exposes fitting types including ener, dos, property, polar, and dipole.
• The JAX backend registers matching fitting/model classes for dos, property, dipole, and polar under deepmd/jax/fitting/fitting.py and deepmd/jax/model/*.py.

Implementation:

• deepmd/jax/train/trainer.py _build_losses() only accepts loss_type == "ener"; any other loss raises RuntimeError("unknown loss type ...").

Impact:

• JAX can construct non-energy model/fitting objects, but JAX training cannot train the corresponding DOS/property/tensor tasks because the loss dispatch rejects non-energy losses.

2. model.spin is accepted by the common schema, but JAX explicitly rejects spin models

Schema/docs:

• deepmd/utils/argcheck.py exposes the top-level model.spin block.

Implementation:

• deepmd/jax/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 JAX model construction time.

3. Learning-rate schema exposes exp, cosine, and wsd, but JAX only implements exp

Schema/docs:

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

Implementation:

• deepmd/jax/train/trainer.py _get_lr_and_coef() dispatches only type == "exp" to LearningRateExp; all other types raise RuntimeError("unknown learning_rate type ...").

Impact:

• learning_rate.type=cosine or learning_rate.type=wsd is schema-exposed but unavailable in JAX training.

4. The optimizer schema is not wired into JAX training

Schema/docs:

• deepmd/utils/argcheck.py exposes an optimizer block and optimizer variants.

Implementation:

• deepmd/jax/train/trainer.py builds an Optax chain directly with optax.scale_by_adam() and _scale_by_global_learning_rate(); it does not read config["optimizer"] or dispatch optimizer variants/parameters.

Impact:

• Optimizer options accepted by the common input schema have no effect in JAX training. At minimum this should be documented as unsupported; ideally JAX should consume the optimizer config it can support.

5. Several model-level schema variants are not present in the JAX backend registry

Schema/docs:

• deepmd/utils/argcheck.py exposes model variants such as frozen, pairtab, and linear_ener in the common model schema.

Implementation observation:

• JAX model registrations are limited to the task models (ener, dos, dipole, polar, property) plus zbl; I did not find JAX BaseModel.register(...) implementations for frozen, top-level pairtab, or linear_ener.
• model.type: dpa4 / dpa4_ener are marked PyTorch-only in the schema, so those are not counted here as JAX implementation gaps.

Impact:

• Users can discover these model variants from the unified schema, but JAX cannot construct them unless there is a separate compatibility path.

Expected resolution

Either:

1. implement/wire these options in the JAX backend, or
2. update argcheck.py/docs to mark them as unsupported for JAX so schema-valid configurations do not fail later in backend-specific code.

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

Implementation checklist

Model, loss, and optimizer support

• Train non-energy DOS, property, dipole, and polar models with matching losses
• Implement JAX spin models or reject them during backend-aware validation
• Support learning_rate.type: cosine and learning_rate.type: wsd (fix(train): dispatch registered learning rate schedules #5776)
• Consume the optimizer schema and support the applicable optimizer options (Add HybridMuon optimizer support to JAX and TF2 for DPA4 training #5771 tracks DPA4 HybridMuon)
• Implement or explicitly reject/document frozen, pairtab, and linear_ener model variants

Training support

• Derive training steps from training.numb_epoch and its aliases
• Support multi-task training.num_epoch_dict
• Honor training.stat_file when computing or loading statistics
• Implement training.profiling and training.enable_profiler
• Implement training.tensorboard* output
• Apply training.mixed_precision to training
• Apply training.change_bias_after_training after the final step (fix(train): share post-training bias adjustment #5745)

[njzjz-bot]

Addendum: training option coverage

I re-checked the training_data_args / validation_data_args / mixed_precision_args / training_args path against the JAX backend on current master (f73de32). Additional training-option gaps:

• training.numb_epoch / aliases and multi-task training.num_epoch_dict: the schema allows epoch-based training, but the shared trainer config still falls back to int(training_params["numb_steps"]), and the JAX path constructs/runs the trainer without an epoch -> step conversion. I did not find the legacy-TF-style conversion using total_numb_batch, nor the multi-task num_epoch_dict -> model probabilities / total steps derivation.
• training.stat_file: TrainingTaskConfig carries stat_file, but deepmd/jax/entrypoints/train.py's task factory returns (train_data, valid_data, None), and JAX compute_or_load_stat(...) is called without a stat_file_path. So a configured stat_file appears to be ignored in the JAX training path.
• training.profiling, training.enable_profiler, and training.tensorboard*: deepmd/jax/train/trainer.py reads these fields into attributes, but I did not find profiler or TensorBoard writer hooks using them during training.
• training.mixed_precision: the trainer reads self.mixed_prec = tr_data.get("mixed_precision", None), but I did not find it applied to the model, train step, optimizer, or loss computation.
• training.change_bias_after_training: the trainer reads the flag, but I only found bias changes in fine-tuning setup; I did not find a post-training branch analogous to the documented “after the last training step” behavior.

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