Add HybridMuon optimizer support to JAX and TF2 for DPA4 training

[OutisLi]

Summary

DPA4/SeZM training is consistently recommended to use DeePMD-kit's custom HybridMuon optimizer, and all current DPA4 training examples use it. However, the new JAX and TF2 DPA4 training paths cannot select this optimizer:

• JAX constructs a fixed Optax Adam chain and does not consume the optimizer configuration.
• TF2 only implements Adam and AdamW; selecting HybridMuon raises an unsupported-optimizer error.

With DPA4 training being added to JAX and TF2 in #5748 and #5749, these backends can construct/train the model but cannot use its recommended optimizer.

Related umbrella issues: #5756 and #5757. Those track general backend/schema gaps; this issue tracks the concrete DPA4-specific HybridMuon implementation and parity requirements across JAX and TF2.

Why this matters for DPA4

HybridMuon is not merely another optimizer exposed by the common schema. DPA4's parameter layout and naming conventions are designed around its routing rules:

• eligible matrix parameters use Muon updates;
• adam_* parameters use the Adam path without weight decay;
• adamw_* and other excluded matrix parameters use the AdamW-style path;
• biases and other non-matrix parameters use the Adam path;
• higher-rank DPA4 tensors use the configured muon_mode (the recommended/default mode is slice).

The DPA4 examples for regular energy training, multitask training, property fitting, LoRA fine-tuning, spin, density loss, and native ZBL all configure HybridMuon.

Falling back to Adam changes the training algorithm and may change convergence speed, stability, sample efficiency, and final accuracy. It also makes cross-backend DPA4 benchmarks difficult to interpret because optimizer differences are mixed with backend differences.

Current behavior

A common configuration such as:

"optimizer": {
  "type": "HybridMuon"
}

has backend-dependent behavior:

• PyTorch: constructs HybridMuonOptimizer.
• JAX: optimizer configuration is not dispatched; training uses the hard-coded Adam transformation.
• TF2: _build_optimizer() rejects the optimizer because only Adam/AdamW are supported.

Requested implementation

Please add backend-native HybridMuon support to both JAX and TF2 while preserving the PyTorch implementation's effective semantics:

1. Consume and validate the common optimizer.type: HybridMuon configuration.
2. Implement consistent parameter routing for Muon, Adam, and AdamW-style paths, including DPA4's adam_ / adamw_ naming contract and effective-rank rules.
3. Support the shared HybridMuon options, especially muon_mode, momentum, Adam betas, weight decay, learning-rate adjustment, Newton-Schulz/Gram configuration, and Magma damping where applicable.
4. Save and restore all optimizer state correctly, including Muon momentum, Adam moments, step counters, and any damping state.
5. Handle distributed training and shared/multitask parameters without duplicating optimizer state or updating a shared parameter twice.
6. Add cross-backend unit tests for route classification and one-step updates on representative 1D, 2D, and higher-rank DPA4 parameters.
7. Add an end-to-end DPA4 training/checkpoint-resume smoke test for each backend using the same HybridMuon configuration.
8. Document any backend-specific performance limitations, but do not silently fall back to Adam.

Acceptance criteria

• The same DPA4 input using optimizer.type: HybridMuon is accepted by PT, JAX, and TF2.
• JAX and TF2 demonstrably execute Muon updates for eligible matrix parameters and Adam/AdamW-style updates for the intended auxiliary parameters.
• Checkpoint/resume preserves optimizer state and produces the same next update as uninterrupted training.
• DPA4 examples and backend benchmarks can use a common optimizer configuration, so backend comparisons do not conflate optimizer choice with backend behavior.