feat: add optional MACE cuEquivariance training

[njzjz]

Summary

• add an optional enable_cueq MACE model parameter
• pass MACE's CuEquivarianceConfig into ScaleShiftMACE when requested
• fall back to eager MACE if cuEquivariance modules cannot be scripted
• guard dense_edge_index fake registration when running source Python against an older installed OP library

Benchmark

Installed cuEquivariance packages:

• cuequivariance==0.10.0
• cuequivariance-torch==0.10.0
• cuequivariance-ops-torch-cu12==0.10.0

Command shape:

srun --gres=gpu:1 dp --pt train input.json

Environment:

• GPU: NVIDIA GeForce RTX 5090
• PyTorch: 2.10.0+cu128
• DeePMD-kit: 3.1.3
• MACE: 0.3.15
• model: 0.203M parameters, 192 atoms/frame, 100 training steps

Results:

Mode	Total real time	DeePMD average training time
e3nn baseline	34.59 s	0.1128 s/batch
cuEquivariance	60.03 s	0.1012 s/batch

The stable training throughput improves by about 10%, but total wall time is worse for this short job because the first cuEquivariance batch takes about 30 s for initialization/compilation.

Notes

With MACE 0.3.15 and cuEquivariance 0.10.0, the cuEquivariance MACE submodel does not script cleanly due to the conv-fusion wrapper (SegmentedPolynomial.original_forward). This PR keeps training usable by falling back to eager mode when enable_cueq is true. This does not imply frozen TorchScript inference support.

Tests

• python -m ruff check deepmd_gnn/mace.py deepmd_gnn/argcheck.py deepmd_gnn/edge.py
• default MACE construction smoke test
• srun --gres=gpu:1 dp --pt train input.json for baseline and enable_cueq=true

Known local test limitations:

• tests/test_pt_expt.py does not collect with local deepmd-kit==3.1.3 because _needs_with_comm_artifact is unavailable.
• tests/test_mace_comm.py has one dtype assertion failure unrelated to the cuEquivariance path.

[codecov]

Codecov Report

❌ Patch coverage is 42.50000% with 69 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.52%. Comparing base (d431f51) to head (f40820b).

Files with missing lines	Patch %	Lines
deepmd_gnn/mace.py	35.51%	69 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #129      +/-   ##
==========================================
- Coverage   80.97%   78.52%   -2.46%     
==========================================
  Files          12       12              
  Lines        1472     1574     +102     
==========================================
+ Hits         1192     1236      +44     
- Misses        280      338      +58     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.