Model Training Example with MACE

[ys-teh]

ALCHEMI Toolkit Pull Request

Description

This PR adds an advanced training example for a charged MACE model and the supporting code modifications needed to train it with available ALCHEMI tools.

Note: This can only be merged after #126 is merged.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Performance improvement
• Documentation update
• Refactoring (no functional changes)
• CI/CD or infrastructure change

Related Issues

Changes Made

• Adds a MACE training example examples/advanced/10_mace_training.py (along with config examples/advanced/10_vanilla_mace.yaml) using nvalchemi model training pipeline.
• Adds examples/advanced/_mace_training_helpers.py with additional training utilities including stress unit conversion, training loss logging, validation, parameter counting, and gradient clipping hook.
• Adds examples/advanced/_mace_models.py with builders for vanilla MACE model, including cuEquivariance config support.
• Adds MACE training user guide docs/userguide/mace_training_example.md.

Testing

• Unit tests pass locally (make pytest)
• Linting passes (make lint)
• New tests added for new functionality meets coverage expectations?

Run training

Checklist

• I have read and understand the Contributing Guidelines
• I have updated the CHANGELOG.md
• I have performed a self-review of my code
• I have added docstrings to new functions/classes
• I have updated the documentation (if applicable)

Additional Notes

Below are benchmarking results comparing ALCHEMI MACE training against the official MACE CLI on the same R2SCAN test set:

Model	Train time	Energy MAE	Forces MAE	Stress MAE	Stress MAE
Official MACE CLI	2 hr	31.1 meV/atom	155.7 meV/A	4.8 meV/A^3	0.769 GPa
ALCHEMI MACE	80 min	25.5 meV/atom	144.8 meV/A	4.39 meV/A^3	0.704 GPa

• ALCHEMI result used the latest trained EMA checkpoint . Official result is from epoch-49_swa.pt, labeled “Stage two model”.
• The ALCHEMI run used 68k training steps, which is nearly equivalent to the 50 epochs used for the MACE CLI run.
• Both cases were trained on a H100 GPU and 8 CPUs.

Below are NVT and NVE stability results for the trained MACE model on a 324-atom MgVF4 3x3x3 MatPES-r2SCAN test structure. These runs completed all 20,000 steps without numerical instability or force/temperature warnings.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR adds an end-to-end MACE model training example for the ALCHEMI toolkit, including a Hydra entrypoint (10_mace_training.py), a default config (10_vanilla_mace.yaml), and two helper modules for model building and training utilities. The example covers distributed training, EMA, gradient clipping, two-stage loss/LR scheduling, and validation with checkpointing.

• _mace_training_helpers.py provides the two-stage cosine LR schedule, dtype/stress transforms, gradient clipping hook, and a TrainingMetricsLogger that reduces metrics across ranks before logging.
• _mace_models.py wraps ScaleShiftMACE with cuEquivariance support and attaches a checkpointable model spec for restart.
• 10_vanilla_mace.yaml ships with MatPES r2SCAN–derived dataset metadata (E0s, avg neighbors, scale/shift) and fully documented hyperparameters.

Important Files Changed

Filename	Overview
examples/advanced/10_mace_training.py	New Hydra entrypoint orchestrating the full MACE training lifecycle: data loading, model construction, loss, hooks, validation, and checkpointing. Logic is well-structured; a few edge cases around distributed teardown and sampler assignment are minor.
examples/advanced/_mace_training_helpers.py	Training utilities including LR schedule, dtype cast, stress unit conversion, gradient clipping, and metrics logging. The JsonLinesLogger is created on all ranks but only written by rank 0, creating unnecessary file handles on non-zero ranks.
examples/advanced/_mace_models.py	MACE model builders wrapping ScaleShiftMACE with cuEquivariance support and checkpoint spec. Module-level torch.serialization.add_safe_globals([slice]) is a process-wide side effect on import, noted in prior review.
examples/advanced/10_vanilla_mace.yaml	Default Hydra config for the MACE training example. Placeholder dataset paths are clearly documented; all parameters are well-commented.
CHANGELOG.md	One-line entry for the new MACE training example, correctly placed under Added in the unreleased section.

_{Reviews (9): Last reviewed commit: "update training descriptions and diagram" | Re-trigger Greptile}

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.