.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/advanced/10_mace_training.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_advanced_10_mace_training.py: MACE Training with ALCHEMI Training Utilities ============================================= This example walks through a complete model-training lifecycle on the ALCHEMI Toolkit, using a baseline ScaleShiftMACE model trained on the MatPES r2SCAN dataset as the reference workflow. The training script run configuration is loaded with `Hydra `__ from :download:`10_vanilla_mace.yaml <10_vanilla_mace.yaml>`. At a high level, the ALCHEMI training workflow has the following structure: .. code-block:: text [Graph Data] -> [Model Architecture] -> [Supervised Objective] -> [Runtime Hooks] -> [TrainingStrategy] **Data** — MatPES r2SCAN 2025.2 structures, obtained from `MatPES `__, are read from ALCHEMI-compatible Zarr splits. Each sample contains graph inputs (positions, atom types, periodic boundary metadata) and supervised labels (energy, forces, stress). **Model** — A 9.06M-parameter ScaleShiftMACE model from `ACEsuit `__ is wrapped with :class:`~nvalchemi.models.mace.MACEWrapper` so it can be used by :class:`~nvalchemi.training.TrainingStrategy`. NVIDIA cuEquivariance kernels are enabled by default in the Hydra config (`model.cueq.enabled: true`). **Loss** — Energies, forces, and stresses are fit with a weighted sum of Huber losses. :class:`~nvalchemi.training.PiecewiseWeight` schedules are used to change the loss-term weights at a configured optimizer step for the second training stage. **Runtime** — Distributed wrapping, EMA, neighbor-list rebuild, gradient clipping, metrics logging, and checkpointing are attached through runtime hooks rather than being implemented directly in the core trainin loop. Validation is configured separately using :class:`~nvalchemi.training.ValidationConfig` on :class:`~nvalchemi.training.TrainingStrategy`. Validation runs automatically during :meth:`~nvalchemi.training.TrainingStrategy.run`. Dataset-derived metadata (`E0s`, `avg_num_neighbors`, `atomic_inter_shift` / `atomic_inter_scale`), must be precomputed and set in ``cfg.model`` before training. The default YAML includes values computed from the MatPES r2SCAN training split. .. GENERATED FROM PYTHON SOURCE LINES 58-60 .. code-block:: Python :dedent: 1 .. rst-class:: sphx-glr-script-out .. code-block:: none /home/kelvin/Repos/nvalchemi-toolkit/.venv/lib/python3.13/site-packages/torch/jit/_script.py:1488: DeprecationWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`. warnings.warn( GPU information: NVIDIA GB10, 12, 1, 128, 48, 140, 1700 .. GENERATED FROM PYTHON SOURCE LINES 141-246 Loading train and validation data --------------------------------- The source data is the `MatPES r2SCAN 2025.2 release `__ on Hugging Face. This example expects separate train, validation, and test Zarr stores. Load each MatPES split file, for example the train, validation, and test JSON/JSONL files, and convert each split independently to an ALCHEMI Zarr store with :class:`~nvalchemi.data.datapipes.AtomicDataZarrWriter`. A minimal converter for this workflow should map each pymatgen/MSON ``structure`` dictionary to atomic numbers, Cartesian positions, cell, and PBC tensors; write ``energy`` as a system label, ``forces`` as an atom label, and convert Voigt-6 ``stress`` to a ``3 x 3`` system tensor. .. code-block:: python import periodictable as pt import torch from nvalchemi.data import AtomicData from nvalchemi.data.atomic_data import voigt_to_matrix from nvalchemi.data.datapipes import AtomicDataZarrWriter def atomic_numbers_from_element_symbols(sites): return torch.as_tensor( [ int(pt.elements.symbol(site["species"][0]["element"]).number) for site in sites ], dtype=torch.int32, ) writer = AtomicDataZarrWriter("r2scan-2025.2-train.zarr") chunk_size = 8192 chunk = [] initialized = False for record in jsonl_records: structure = record["structure"] chunk.append( AtomicData( atomic_numbers=atomic_numbers_from_element_symbols(structure["sites"]), positions=torch.as_tensor([site["xyz"] for site in structure["sites"]]), cell=torch.as_tensor(structure["lattice"]["matrix"]).reshape(1, 3, 3), pbc=torch.as_tensor(structure["lattice"].get("pbc", [True] * 3)).reshape(1, 3), energy=torch.as_tensor([[record["energy"]]]), forces=torch.as_tensor(record["forces"]), stress=voigt_to_matrix(torch.as_tensor(record["stress"])).reshape(1, 3, 3), ) ) if len(chunk) >= chunk_size: writer.append(chunk) if initialized else writer.write(chunk) initialized = True chunk.clear() if chunk: writer.append(chunk) if initialized else writer.write(chunk) This pipeline reads those Zarr splits with :class:`~nvalchemi.data.datapipes.AtomicDataZarrReader`. :class:`~nvalchemi.data.datapipes.InMemoryDataset` materializes each split once as a :class:`~nvalchemi.data.Batch` on the target device, and :class:`~nvalchemi.data.datapipes.DataLoader` selects shuffled or sequential batches from that in-memory batch. The default configuration uses a per-process training batch size of 256 and a validation batch size of 512. Given that the structure sizes in this dataset range from 1 atom to 240 atoms, :class:`~nvalchemi.dynamics.sampler.SizeAwareSampler` can also be used as an alternative to cap the atom count per batch when memory is tight. .. code-block:: python from pathlib import Path import torch from nvalchemi.data.datapipes import AtomicDataZarrReader, DataLoader, InMemoryDataset device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_dataset = InMemoryDataset( reader=AtomicDataZarrReader(Path("/path/to/r2scan-2025.2-train.zarr")), device=device, skip_validation=True, ) train_batches = DataLoader( train_dataset, batch_size=256, shuffle=True, ) val_dataset = InMemoryDataset( reader=AtomicDataZarrReader(Path("/path/to/r2scan-2025.2-valid.zarr")), device=device, skip_validation=True, ) val_batches = DataLoader( val_dataset, batch_size=512, shuffle=False, ) The runnable script wraps this pattern in ``_loader(...)`` so Hydra can supply paths, batch sizes, and optional stress scaling transforms. .. GENERATED FROM PYTHON SOURCE LINES 246-248 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 299-328 Building the MACE model ----------------------- The default configuration trains ScaleShiftMACE to predict energy, force, and stress. Any model object passed to :class:`~nvalchemi.training.TrainingStrategy` must follow :class:`~nvalchemi.models.base.BaseModelMixin`. :class:`~nvalchemi.models.mace.MACEWrapper` handles input adaptation, neighbor-list metadata, and routes model outputs for MACE model variants. Before building the model, populate the Hydra config with dataset-derived metadata: ``E0s`` (from structure-energy regression or isolated-atom DFT), ``avg_num_neighbors``, and the ScaleShiftMACE pair ``atomic_inter_shift`` / ``atomic_inter_scale``. The default YAML includes values precomputed from the training split. .. code-block:: python import torch from mace.modules import ScaleShiftMACE from nvalchemi.models.mace import MACEWrapper mace_model = ScaleShiftMACE(...) model = MACEWrapper(mace_model.to(device=device, dtype=torch.float32)) model.model_config.active_outputs = {"energy", "forces", "stress"} The runnable script reads architecture hyperparameters from Hydra and builds the wrapped model through ``_build_model(cfg, device)``, which calls :func:`~examples.advanced._mace_models.build_training_mace_model` to set ``active_outputs`` and attach a checkpointable model spec. .. GENERATED FROM PYTHON SOURCE LINES 328-330 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 358-403 Defining the loss ----------------- The default configuration fits energies, forces, and stresses. The loss is a weighted sum of Huber terms composed with ``+`` and ``*`` into a :class:`~nvalchemi.training.ComposedLossFunction`. Stage-one weights hold until ``stage_two_start``, then switch (for example ``1/10/100`` to ``10/1/10`` at step 54,400 of 68,000). .. code-block:: python from nvalchemi.training import ( ComposedLossFunction, EnergyHuberLoss, ForceHuberLoss, PiecewiseWeight, StressHuberLoss, ) stage_two_start = 54_400 loss_fn: ComposedLossFunction = ( PiecewiseWeight( boundaries=(stage_two_start,), values=(1.0, 10.0), per_epoch=False, ) * EnergyHuberLoss(per_atom=True, delta=0.01) + PiecewiseWeight( boundaries=(stage_two_start,), values=(10.0, 1.0), per_epoch=False, ) * ForceHuberLoss(normalize_by_atom_count=False, delta=0.01) + PiecewiseWeight( boundaries=(stage_two_start,), values=(100.0, 10.0), per_epoch=False, ) * StressHuberLoss(delta=0.01) ) loss_fn.normalize_weights = False The runnable script builds the same composition from ``cfg.training.loss`` through ``_build_mace_huber_loss(cfg.training.loss)``. .. GENERATED FROM PYTHON SOURCE LINES 403-405 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 482-513 Configuring the optimizer and scheduler --------------------------------------- Schedulers are attached through :class:`~nvalchemi.training.OptimizerConfig`. The runnable example uses :class:`~examples.advanced._mace_training_helpers.TwoStageCosineConstantLR` — cosine annealing for stage one, then a constant stage-two learning rate; any ``torch.optim.lr_scheduler.LRScheduler`` subclass can be passed via ``scheduler_cls`` and ``scheduler_kwargs``. .. code-block:: python import torch from _mace_training_helpers import TwoStageCosineConstantLR from nvalchemi.training import OptimizerConfig optimizer_config = OptimizerConfig( optimizer_cls=torch.optim.AdamW, optimizer_kwargs={ "lr": 5.0e-3, }, scheduler_cls=TwoStageCosineConstantLR, scheduler_kwargs={ "first_stage_steps": 54_400, "second_stage_lr": 1.0e-3, "eta_min": 1.0e-3, }, ) Hydra supplies learning-rate and schedule values; ``_optimizer(cfg)`` maps them onto :class:`~nvalchemi.training.OptimizerConfig`. .. GENERATED FROM PYTHON SOURCE LINES 513-515 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 543-589 Adding runtime hooks -------------------- Hooks extend the core training loop without embedding that logic in the loop itself. For example, :class:`~nvalchemi.training.DDPHook` wraps the model in DDP at the :class:`~nvalchemi.training.TrainingStage` ``SETUP`` stage when ``training.distributed.enabled`` is true. :class:`~nvalchemi.training.EMAHook` maintains shadow weights for validation at ``AFTER_OPTIMIZER_STEP``, and :class:`~nvalchemi.hooks.NeighborListHook` rebuilds the interaction graph at ``BEFORE_FORWARD`` before every forward pass. .. code-block:: python from pathlib import Path from _mace_training_helpers import ( GradientClipHook, TrainingMetricsLogger, ) from nvalchemi.hooks import NeighborListHook from nvalchemi.training import ( CheckpointHook, DDPHook, EMAHook, TrainingStage, ) hooks = [ DDPHook(backend="nccl", sampler_kwargs={"seed": 42}), EMAHook(model_key="main", decay=0.995), GradientClipHook(max_norm=2.0), NeighborListHook( model.model_config.neighbor_config, max_neighbors=256, method="batch_naive_tile", stage=TrainingStage.BEFORE_FORWARD, ), TrainingMetricsLogger(every=100), CheckpointHook( checkpoint_dir=Path("outputs/checkpoints"), step_interval=10_000, ), ] ``GradientClipHook`` and ``TrainingMetricsLogger`` are implemented in this example's helper module. The other hooks shown above are public ALCHEMI training APIs. The runnable script assembles the full hook list from Hydra through ``_hooks(cfg, model)``. .. GENERATED FROM PYTHON SOURCE LINES 589-591 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 673-700 Configuring validation ---------------------- Validation is configured with :class:`~nvalchemi.training.ValidationConfig`. The configuration specifies the validation data, validation function, loss function, evaluation cadence, and whether to use EMA weights. During :class:`~nvalchemi.training.TrainingStrategy.run`, the strategy evaluates validation at this cadence and once more at the end of training. The latest validation summary is stored on ``strategy.last_validation``. In multi-GPU runs, each rank evaluates a disjoint validation shard through a DistributedSampler. The runnable script builds this configuration with ``_build_validation_config(...)`` after the validation loader and loss function have been constructed. .. code-block:: python from nvalchemi.training import ValidationConfig, default_training_fn validation_config = ValidationConfig( validation_data=val_batches, validation_fn=default_training_fn, loss_fn=loss_fn, every_n_steps=1000, grad_mode="auto", use_ema="auto", name="validation", ) .. GENERATED FROM PYTHON SOURCE LINES 700-702 .. code-block:: Python :dedent: 1 .. GENERATED FROM PYTHON SOURCE LINES 739-797 Running TrainingStrategy ------------------------ The final step is to assemble the objects created above and hand them to :class:`~nvalchemi.training.TrainingStrategy`, which runs the training loop. On each step, it calls the training function, steps the optimizer and scheduler, invokes hooks at their registered stages, runs validation when configured, and tracks checkpointable training state. :class:`~nvalchemi.distributed.DistributedManager` provides distributed runtime information such as rank, local rank, world size, and device placement. The same code path works for single-GPU and multi-GPU launches. In distributed runs, :class:`~nvalchemi.training.DDPHook` uses the distributed manager to wrap the model and coordinate rank-specific behavior. .. code-block:: python from nvalchemi.distributed import DistributedManager from nvalchemi.training import TrainingStrategy, default_training_fn DistributedManager.initialize() manager = DistributedManager() device = torch.device(manager.device) strategy = TrainingStrategy( models=model, optimizer_configs=optimizer_config, num_steps=68_000, training_fn=default_training_fn, loss_fn=loss_fn, devices=[device], distributed_manager=manager, hooks=hooks, validation_config=validation_config, ) strategy.run(train_loader) Run the Hydra entrypoint on one or more GPUs: Single GPU: .. code-block:: bash uv run --extra cu12 --extra mace python examples/advanced/10_mace_training.py Multi-GPU: .. code-block:: bash uv run --extra cu12 --extra mace torchrun --standalone --nproc_per_node=8 \ examples/advanced/10_mace_training.py \ --config-name=10_vanilla_mace distributed.enabled=true The commands above use ``--extra cu12``; change ``uv run --extra cuXX`` to match your NVIDIA driver / CUDA toolkit availability (for example, ``cu13`` on CUDA 13.x). Note that ``training.batch_size`` is the per-process batch size. The global batch size is therefore ``training.batch_size * nproc_per_node``. .. GENERATED FROM PYTHON SOURCE LINES 797-799 .. code-block:: Python :dedent: 1 .. rst-class:: sphx-glr-script-out .. code-block:: none Skipping Hydra training during docs build. Run with: uv run --extra cu12 --extra mace python examples/advanced/10_mace_training.py (use --extra cu13 instead depending on your CUDA availability) .. GENERATED FROM PYTHON SOURCE LINES 896-914 Validation curves and reference results --------------------------------------- The figure below shows validation Huber losses from a full default-config run on 1× H100 GPU, which took about 80 minutes of wall time. Actual wall time may differ depending on system configuration, hardware, and software stack. The sharp transition near step 54,400 marks the stage-two loss-weight schedule configured by ``training.loss.stage_two.start_step``. .. image:: ../_static/vanilla_mace_validation_metrics_260701.png :align: center :width: 70% With this default config (68,000 optimizer steps, about 50 epochs on the MatPES r2SCAN train set), the trained model reaches held-out test MAEs of energy 25.5 meV/atom, forces 145 meV/Å, and stress 0.703 GPa. These values are comparable to the MatPES r2SCAN benchmarks reported in `the MatPES paper `__ and to training with the `MACE CLI `__. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 1.149 seconds) .. _sphx_glr_download_examples_advanced_10_mace_training.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 10_mace_training.ipynb <10_mace_training.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 10_mace_training.py <10_mace_training.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 10_mace_training.zip <10_mace_training.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_