Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,6 @@
+outputs/
+checkpoints/
+.trainer_work/
+__pycache__/
+demo/__pycache__/
+*.log

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 The Atomic Architects
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,255 @@
+---
+license: mit
+tasks:
+  - materials-simulation
+  - energy-prediction
+  - force-prediction
+  - molecular-dynamics
+frameworks:
+  - pytorch
+language:
+  - en
+tags:
+  - OneScience
+  - EquiformerV3
+  - materials-science
+  - machine-learning-potential
+  - graph-neural-network
+  - equivariant-neural-network
+  - training
+  - inference
+datasets:
+  - OneScience-Group/oc20
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Equiformer V3</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Equiformer V3 is an SE(3)-equivariant graph attention potential model for three-dimensional atomic systems. It predicts the energy, atomic forces, and stress of atomic structures.
+
+# Model Description
+
+Equiformer V3 is based on an equivariant graph neural network architecture. Building on Equiformer V2, it improves equivariant normalization, smooth cutoff attention, and SwiGLU-S² activation for materials potential energy surface modeling. This release includes OneScience-adapted model code, pretrained weights, single-point inference, and an OC20 training example.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :---: |
+| Single-point energy, forces, and stress | Predict energy, forces, and stress for ASE-readable structures |
+| Structure relaxation and formation energy | Use the model and elemental reference energies for structure optimization and formation-energy calculations |
+| Elastic tensor | Calculate the elastic properties of periodic materials |
+| Phonons | Calculate the phonon properties of periodic materials |
+| OC20 S2EF training | Train on energy and force labels using preprocessed data from `OneScience-Group/oc20` |
+| Custom data migration | Convert custom data to the ASE-LMDB format required by the configuration, then update the training paths |
+
+# Usage
+
+## 1. Using OneCode
+
+Try intelligent, one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware requirements**
+
+- An NVIDIA GPU or Hygon DCU is recommended for inference and training.
+- A CPU can be used to inspect configurations and prepare data but is not recommended for production training.
+- DCU users must load a DTK environment compatible with the current PyTorch build.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Group/Equiformer_v3 --local-dir ./Equiformer_v3
+cd Equiformer_v3
+```
+
+### Install the Runtime Environment
+
+**DCU environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[matchem-dcu] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+**GPU environment**
+
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[matchem-gpu] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+### Training Dataset
+
+Training data is not bundled with this model repository. The OC20 training data is available from the Hugging Face dataset **OneScience-Group/oc20** and corresponds to the Open Catalyst 2020 (OC20) S2EF task. See the [OC20 dataset paper](https://doi.org/10.1021/acscatal.0c04525) and the [FAIR Chemistry OC20 documentation](https://fair-chem.github.io/oc20/) for the original data, task definitions, and citation information. When using the data, follow the licenses and terms on the OneScience dataset page, from the Open Catalyst Project, and for the original data.
+
+```bash
+hf download --dataset OneScience-Group/oc20 --local-dir ./data
+```
+
+Convert the downloaded data to the ASE-LMDB format expected by the training configuration. The current configuration reads:
+
+```text
+data/oc20_finetune/
+├── train/                 # ASE-LMDB training data
+└── val/                   # ASE-LMDB validation data
+```
+
+If the downloaded OC20 data consists of `.extxyz` files or other ASE-readable source files, convert it with the bundled UMA data-processing script:
+
+```bash
+python scripts/create_uma_finetune_dataset.py \
+  --train-dir ./data/oc20/s2ef_200k_uncompressed \
+  --val-dir ./data/oc20/s2ef_val_id_uncompressed \
+  --uma-task oc20 \
+  --regression-tasks ef \
+  --output-dir ./data/oc20_finetune \
+  --num-workers 8
+```
+
+The script generates `train/` and `val/` ASE-LMDB shards together with a UMA data configuration file. Input files must contain energy and force labels provided through an ASE calculator. Files that fail conversion are recorded in a `.failed` log corresponding to each output shard. `create_finetune_dataset.py` is the low-level ASE data conversion implementation, while `create_uma_finetune_dataset.py` is the complete OC20/UMA data-processing entry point. Both scripts come from the UMA repository and are provided in their original form.
+
+Under the OneScience MatChem data root, the corresponding paths are:
+
+```text
+${ONESCIENCE_DATASETS_DIR}/matchem/oc20/uma_oc20_finetune/
+├── train/
+└── val/
+```
+
+If the data is stored elsewhere, update the `train` and `val` fields in `demo/configs/oc20_scratch_8dcu.yaml` and `demo/configs/oc20_scratch_8dcu_smoke.yaml`, or set `ONESCIENCE_DATASETS_DIR`. The Equiformer V3 training YAML must continue to use its own `train`, `val`, and `transforms.element_references.energy.file` fields; the UMA fine-tuning YAML generated by the UMA script does not directly replace the Equiformer V3 training configuration. After changing the training data, refit the elemental reference-energy coefficients:
+
+```bash
+python fit_element_references.py \
+  --config demo/configs/oc20_scratch_8dcu.yaml \
+  --output demo/reference_data/oc20_subset_energy_element_references.npz
+```
+
+### Model Weights
+
+This repository includes the following Equiformer V3 weights:
+
+```text
+weight/
+├── Jd.pt
+├── mptrj_gradient.pt
+├── omat24_direct.pt
+├── omat24_gradient.pt
+└── omat24-mptrj-salex_gradient.pt
+```
+
+| Weight | Training task or domain | Intended use |
+| --- | --- | --- |
+| `mptrj_gradient.pt` | MPtrj gradient | Inference or evaluation on MPtrj material structures |
+| `omat24_direct.pt` | OMat24 direct + DeNS | Inference or evaluation with the OMat24 direct model |
+| `omat24_gradient.pt` | OMat24 gradient | Inference or evaluation with the OMat24 gradient model |
+| `omat24-mptrj-salex_gradient.pt` | OMat24 + MPtrj + sAlex gradient | General materials-structure inference; used by the inference scripts by default |
+
+`Jd.pt` contains the Wigner rotation basis and is loaded automatically from `weight/Jd.pt` by the entry-point scripts.
+
+### Fine-Tuning
+
+Run the OC20 smoke-test configuration:
+
+```bash
+export ONESCIENCE_DATASETS_DIR=/path/to/onescience_datasets
+bash demo/run.sh --config configs/oc20_scratch_8dcu_smoke.yaml
+```
+
+The smoke-test configuration uses eight training samples, eight validation samples, and a reduced one-layer model. It performs a single update to verify the data pipeline, distributed communication, forward and backward passes, optimizer, and checkpoint-saving workflow.
+
+Run the full OC20 configuration:
+
+```bash
+bash demo/run.sh --config configs/oc20_scratch_8dcu.yaml
+```
+
+This configuration uses one node, eight DCUs, and 12 epochs. FP32 is the currently validated stable path. Controlled experiments with the same full model in FP16/BF16 encountered a DCU kernel VMFault, so the release configuration uses `amp: false`.
+
+### Inference
+
+Single-point energy, forces, and stress:
+
+```bash
+python single_point.py --device cuda --output outputs/single_point.json
+```
+
+Formation energy:
+
+```bash
+python formation_energy.py --device cuda --output outputs/formation_energy.json
+```
+
+Elastic tensor:
+
+```bash
+python elastic.py --relax --device cuda --output outputs/elastic.json
+```
+
+Phonons:
+
+```bash
+python phonons.py \
+  --supercell 3 3 3 \
+  --bandpath GXWKGL \
+  --device cuda \
+  --output-dir outputs/phonons
+```
+
+Custom structures or weights:
+
+```bash
+python single_point.py --input structure.cif
+python elastic.py --input POSCAR --output outputs/elastic.json
+python phonons.py --input structure.cif --supercell 2 2 2
+python single_point.py --checkpoint weight/omat24_gradient.pt
+```
+
+`--input` supports CIF, POSCAR, XYZ, trajectory, and other formats readable by ASE. All four inference scripts use `weight/omat24-mptrj-salex_gradient.pt` by default.
+
+## Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+---
+
+## Citation and License
+
+- The Equiformer V3 architecture code is based on upstream software released under the MIT License. See `LICENSE` in the repository root.
+- The OC20 dataset is licensed under CC BY 4.0 as described on its official page. The repository's MIT License does not cover the dataset license.
+- If you use Equiformer V3 or OC20 in research, cite the Equiformer V3 paper, the OC20 dataset paper, and the source of the dataset actually used.
+
+```bibtex
+@article{equiformer_v3,
+  title={EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers},
+  author={Yi-Lun Liao and Alexander J. Hoffman and Sabrina C. Shen and Alexandre Duval and Sam Walton Norwood and Tess Smidt},
+  journal={arXiv preprint arXiv:2604.09130},
+  year={2026}
+}
+
+@article{oc20,
+  title={Open Catalyst 2020 (OC20) Dataset and Community Challenges},
+  author={Lowik Chanussot and Abhishek Das and Siddharth Goyal and others},
+  journal={ACS Catalysis},
+  year={2021},
+  doi={10.1021/acscatal.0c04525}
+}
+```

demo/_parse_config.py

@@ -0,0 +1,113 @@
+#!/usr/bin/env python3
+"""Extract Equiformer V3 demo metadata and its fine-tuning configuration."""
+
+from __future__ import annotations
+
+import shlex
+import sys
+
+import yaml
+
+
+META_KEYS = {"name", "description", "launch", "slurm", "nccl"}
+
+
+def _assignment(name: str, value) -> None:
+    print(f"{name}={shlex.quote(str(value))}")
+
+
+def _positive_int(value, field: str) -> int:
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as error:
+        raise ValueError(f"{field} must be a positive integer") from error
+    if parsed < 1:
+        raise ValueError(f"{field} must be a positive integer")
+    return parsed
+
+
+def _config(path: str) -> dict:
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle) or {}
+
+
+def _launch(cfg: dict) -> dict:
+    return cfg.get("launch", {}) or {}
+
+
+def _slurm(cfg: dict) -> dict:
+    return cfg.get("slurm", {}) or {}
+
+
+def _print_launch(cfg: dict) -> None:
+    launch = _launch(cfg)
+    mode = launch.get("mode", "local")
+    if mode not in {"local", "submit"}:
+        raise ValueError("launch.mode must be 'local' or 'submit'")
+    _assignment("RUN_MODE", mode)
+    _assignment("NODES", _positive_int(launch.get("num_nodes", 1), "launch.num_nodes"))
+    _assignment(
+        "GPUS_PER_NODE",
+        _positive_int(launch.get("num_gpus", 1), "launch.num_gpus"),
+    )
+    _assignment(
+        "OMP_NUM_THREADS",
+        _positive_int(launch.get("omp_num_threads", 1), "launch.omp_num_threads"),
+    )
+
+
+def _print_slurm(cfg: dict) -> None:
+    slurm = _slurm(cfg)
+    _assignment("PARTITION", slurm.get("partition", "hx1hdexclu12"))
+    _assignment("TIME", slurm.get("time", "24:00:00"))
+    _assignment(
+        "CPUS_PER_TASK",
+        _positive_int(slurm.get("cpus_per_task", 16), "slurm.cpus_per_task"),
+    )
+    _assignment("NODELIST", slurm.get("nodelist", ""))
+
+
+def _print_env(cfg: dict) -> None:
+    nccl = cfg.get("nccl", {}) or {}
+    if nccl.get("socket_ifname"):
+        _assignment("export NCCL_SOCKET_IFNAME", nccl["socket_ifname"])
+    if nccl.get("ib_hca"):
+        _assignment("export NCCL_IB_HCA", nccl["ib_hca"])
+    if nccl.get("proto"):
+        _assignment("export NCCL_PROTO", nccl["proto"])
+    print("export HSA_FORCE_FINE_GRAIN_PCIE=1")
+
+
+def _print_finetune_config(cfg: dict) -> None:
+    finetune = {key: value for key, value in cfg.items() if key not in META_KEYS}
+    yaml.safe_dump(
+        finetune,
+        sys.stdout,
+        sort_keys=False,
+        default_flow_style=False,
+        allow_unicode=True,
+    )
+
+
+def main() -> None:
+    if len(sys.argv) != 3:
+        raise SystemExit(
+            "usage: _parse_config.py <config.yaml> <name|launch|slurm|env|finetune-config>"
+        )
+    cfg = _config(sys.argv[1])
+    action = sys.argv[2]
+    actions = {
+        "name": lambda: print(cfg.get("name", "equiformer_v3_finetune")),
+        "launch": lambda: _print_launch(cfg),
+        "slurm": lambda: _print_slurm(cfg),
+        "env": lambda: _print_env(cfg),
+        "finetune-config": lambda: _print_finetune_config(cfg),
+    }
+    try:
+        actions[action]()
+    except KeyError as error:
+        raise SystemExit(f"unknown action: {action}") from error
+
+
+if __name__ == "__main__":
+    main()

demo/configs/oc20_scratch_8dcu.yaml

@@ -0,0 +1,103 @@
+# Complete FP32 OC20 scratch training on the bundled ASE-LMDB subset with eight DCUs.
+# The official model scale and effective global batch of 64 are retained.
+# Per-DCU batch 1 plus accumulation 8 retains global batch 64. Energy element
+# references are fitted offline, following the official OC20 configuration.
+# FP32 is intentional: controlled 20-step FP16/BF16 runs raised DCU kernel
+# VMFaults, while the otherwise equivalent FP32 run completed and checkpointed.
+name: equiformer_v3_oc20_scratch_8dcu
+description: "Complete Equiformer V3 OC20 scratch training on eight DCUs in FP32"
+
+launch: {mode: local, num_nodes: 1, num_gpus: 8, omp_num_threads: 8}
+slurm: {partition: hx1hdnormal01, time: "24:00:00", cpus_per_task: 64}
+nccl: {}
+
+mode: train_from_scratch
+train: "${ONESCIENCE_DATASETS_DIR}/matchem/oc20/uma_oc20_finetune/train"
+val: "${ONESCIENCE_DATASETS_DIR}/matchem/oc20/uma_oc20_finetune/val"
+output: "checkpoints/equiformer_v3_oc20_scratch_8dcu.pt"
+device: cuda
+backend: nccl
+seed: 0
+amp: false
+load_balancing: atoms
+load_balancing_on_error: raise
+
+epochs: 12
+batch_size: 1
+eval_batch_size: 1
+workers: 0
+grad_accumulation_steps: 8
+log_every_n_steps: 10
+log_every_n_validation_batches: 100
+clip_grad_norm: 100
+ema_decay: 0.999
+
+transforms:
+  element_references:
+    energy:
+      file: "${ONESCIENCE_EQUIFORMER_V3_DIR}/demo/reference_data/oc20_subset_energy_element_references.npz"
+  normalizers:
+    energy: {mean: 0.0, rmsd: 24.901469505465872}
+    forces: {mean: 0.0, rmsd: 0.5111534595489502}
+
+losses:
+  energy: {fn: mae, coefficient: 4}
+  forces: {fn: l2mae, coefficient: 100, free_atoms_only: true}
+
+optimizer:
+  name: AdamW
+  lr: 2.0e-4
+  weight_decay: 1.0e-3
+  betas: [0.9, 0.999]
+  eps: 1.0e-8
+
+scheduler:
+  name: cosine
+  warmup_factor: 0.2
+  warmup_epochs: 0.1
+  lr_min_factor: 0.01
+
+model:
+  name: equiformer_v3
+  use_pbc: true
+  use_pbc_single: false
+  otf_graph: true
+  regress_forces: true
+  regress_stress: false
+  direct_prediction: true
+  max_neighbors: 20
+  max_radius: 12.0
+  num_radial_basis: 128
+  max_num_elements: 128
+  num_layers: 8
+  num_channels: 128
+  attn_hidden_channels: 64
+  num_heads: 8
+  attn_alpha_channels: 64
+  attn_value_channels: 16
+  ffn_hidden_channels: 512
+  norm_type: merge_layer_norm
+  lmax: 6
+  mmax: 2
+  attn_grid_resolution_list: [20, 8]
+  ffn_grid_resolution_list: [20, 20]
+  edge_channels: 128
+  use_atom_edge_embedding: true
+  use_envelope: true
+  attn_activation: sep-merge_gates2_swiglu
+  use_attn_renorm: true
+  use_add_merge: false
+  use_rad_l_parametrization: true
+  softcap: null
+  ffn_activation: sep-merge_gates2_swiglu
+  use_grid_mlp: true
+  use_gate_force_head: true
+  alpha_drop: 0.0
+  attn_mask_rate: 0.0
+  attn_weights_drop: 0.1
+  value_drop: 0.0
+  drop_path_rate: 0.05
+  proj_drop: 0.0
+  ffn_drop: 0.0
+  gradient_checkpointing_block_list: [1, 1, 1, 1, 1, 1, 1, 1]
+  enforce_max_neighbors_strictly: true

demo/configs/oc20_scratch_8dcu_smoke.yaml

@@ -0,0 +1,91 @@
+# Eight-DCU FP32 DDP smoke test. This tiny model is not a scientific run.
+# FP32 matches the precision path that passed the controlled 20-step full-model test.
+name: equiformer_v3_oc20_scratch_8dcu_smoke
+description: "One-update OC20 scratch FP32 DDP smoke test"
+
+launch: {mode: submit, num_nodes: 1, num_gpus: 8, omp_num_threads: 8}
+slurm: {partition: hx1hdnormal01, time: "00:20:00", cpus_per_task: 64}
+nccl: {}
+
+mode: train_from_scratch
+train: "${ONESCIENCE_DATASETS_DIR}/matchem/oc20/uma_oc20_finetune/train"
+val: "${ONESCIENCE_DATASETS_DIR}/matchem/oc20/uma_oc20_finetune/val"
+output: "checkpoints/equiformer_v3_oc20_scratch_8dcu_smoke.pt"
+device: cuda
+backend: nccl
+seed: 17
+amp: false
+load_balancing: atoms
+load_balancing_on_error: raise
+epochs: 1
+max_steps: 1
+batch_size: 1
+eval_batch_size: 1
+workers: 0
+max_train_samples: 8
+max_val_samples: 8
+grad_accumulation_steps: 1
+log_every_n_steps: 10
+log_every_n_validation_batches: 100
+clip_grad_norm: 100
+ema_decay: 0.9
+
+transforms:
+  element_references:
+    energy:
+      file: "${ONESCIENCE_EQUIFORMER_V3_DIR}/demo/reference_data/oc20_subset_energy_element_references.npz"
+  normalizers:
+    energy: {mean: 0.0, rmsd: 24.901469505465872}
+    forces: {mean: 0.0, rmsd: 0.5111534595489502}
+
+losses:
+  energy: {fn: mae, coefficient: 4}
+  forces: {fn: l2mae, coefficient: 100, free_atoms_only: true}
+
+optimizer: {name: AdamW, lr: 2.0e-4, weight_decay: 1.0e-3, betas: [0.9, 0.999], eps: 1.0e-8}
+scheduler: {name: cosine, warmup_factor: 0.2, warmup_epochs: 0.1, lr_min_factor: 0.01}
+
+model:
+  name: equiformer_v3
+  use_pbc: true
+  use_pbc_single: false
+  otf_graph: true
+  regress_forces: true
+  regress_stress: false
+  direct_prediction: true
+  max_neighbors: 20
+  max_radius: 6.0
+  num_radial_basis: 16
+  max_num_elements: 128
+  num_layers: 1
+  num_channels: 16
+  attn_hidden_channels: 8
+  num_heads: 2
+  attn_alpha_channels: 8
+  attn_value_channels: 4
+  ffn_hidden_channels: 32
+  norm_type: merge_layer_norm
+  lmax: 2
+  mmax: 2
+  attn_grid_resolution_list: [8, 4]
+  ffn_grid_resolution_list: [8, 8]
+  edge_channels: 16
+  use_atom_edge_embedding: true
+  use_envelope: true
+  attn_activation: sep-merge_gates2_swiglu
+  use_attn_renorm: true
+  use_add_merge: false
+  use_rad_l_parametrization: true
+  softcap: null
+  ffn_activation: sep-merge_gates2_swiglu
+  use_grid_mlp: true
+  use_gate_force_head: true
+  alpha_drop: 0.0
+  attn_mask_rate: 0.0
+  attn_weights_drop: 0.0
+  value_drop: 0.0
+  drop_path_rate: 0.0
+  proj_drop: 0.0
+  ffn_drop: 0.0
+  gradient_checkpointing_block_list: [0]
+  enforce_max_neighbors_strictly: true

demo/reference_data/oc20_subset_energy_element_references.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1ed53d158dfd54bf6f8a1f74c1e56d36e974651a9f87c7ca70fd9a5af94e03e
+size 740

demo/run.sh

@@ -0,0 +1,160 @@
+#!/bin/bash
+# Run Equiformer V3 after the central MatChem environment has been loaded.
+set -euo pipefail
+
+DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+EQV3_DIR="$(cd "$DEMO_DIR/.." && pwd)"
+REPO_ROOT="$EQV3_DIR"
+export ONESCIENCE_EQUIFORMER_V3_DIR="$EQV3_DIR"
+export ONESCIENCE_EQUIFORMER_V3_JD_PATH="${ONESCIENCE_EQUIFORMER_V3_JD_PATH:-$EQV3_DIR/weight/Jd.pt}"
+PARSER="$DEMO_DIR/_parse_config.py"
+CONFIG=""
+SUBMIT=false
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --config) CONFIG="$2"; shift 2 ;;
+        --config=*) CONFIG="${1#*=}"; shift ;;
+        --submit) SUBMIT=true; shift ;;
+        -h|--help)
+            echo "Usage: bash demo/run.sh --config configs/<name>.yaml"
+            echo "Set launch.mode: submit in YAML to submit a Slurm job."
+            exit 0
+            ;;
+        *) echo "Unknown argument: $1" >&2; exit 2 ;;
+    esac
+done
+
+[[ -n "$CONFIG" ]] || { echo "Please specify --config configs/<name>.yaml" >&2; exit 2; }
+[[ "$CONFIG" = /* ]] || CONFIG="$DEMO_DIR/$CONFIG"
+[[ -f "$CONFIG" ]] || { echo "Config not found: $CONFIG" >&2; exit 2; }
+
+missing_env=()
+for variable in CONDA_PREFIX ONESCIENCE_DATASETS_DIR; do
+    [[ -n "${!variable:-}" ]] || missing_env+=("$variable")
+done
+if (( ${#missing_env[@]} )); then
+    echo "The OneScience MatChem environment is not loaded (missing: ${missing_env[*]})." >&2
+    echo "Activate the conda environment and set ONESCIENCE_DATASETS_DIR first." >&2
+    exit 2
+fi
+export MATCHEM_CONDA_NAME="${MATCHEM_CONDA_NAME:-$(basename "$CONDA_PREFIX")}"
+
+NAME="$(python3 "$PARSER" "$CONFIG" name)"
+eval "$(python3 "$PARSER" "$CONFIG" launch)"
+eval "$(python3 "$PARSER" "$CONFIG" slurm)"
+ENV_EXPORTS="$(python3 "$PARSER" "$CONFIG" env)"
+
+if [[ "$RUN_MODE" == "submit" ]]; then
+    SUBMIT=true
+fi
+
+if ! $SUBMIT; then
+    AUTO_SUBMIT_REASON=""
+    if (( NODES > 1 )); then
+        AUTO_SUBMIT_REASON="the config requests $NODES nodes"
+    else
+        AVAILABLE_GPUS="$(python3 -c 'import torch; print(torch.cuda.device_count() if torch.cuda.is_available() else 0)' 2>/dev/null || echo 0)"
+        if (( AVAILABLE_GPUS < GPUS_PER_NODE )); then
+            AUTO_SUBMIT_REASON="the config requests $GPUS_PER_NODE DCUs but only $AVAILABLE_GPUS are visible"
+        fi
+    fi
+    if [[ -n "$AUTO_SUBMIT_REASON" ]]; then
+        if ! command -v sbatch >/dev/null 2>&1; then
+            echo "Cannot run locally: $AUTO_SUBMIT_REASON, and sbatch is unavailable." >&2
+            exit 2
+        fi
+        echo "Local resources are insufficient: $AUTO_SUBMIT_REASON. Submitting to Slurm."
+        SUBMIT=true
+    fi
+fi
+
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+OUTPUT_ROOT="${EQUIFORMER_V3_OUTPUT_ROOT:-$EQV3_DIR/outputs}"
+OUTPUT_DIR="$OUTPUT_ROOT/${NAME}_${TIMESTAMP}"
+mkdir -p "$OUTPUT_DIR/checkpoints"
+cp "$CONFIG" "$OUTPUT_DIR/config.yaml"
+FINETUNE_CONFIG="$OUTPUT_DIR/finetune.yaml"
+python3 "$PARSER" "$CONFIG" finetune-config > "$FINETUNE_CONFIG"
+
+if $SUBMIT; then
+    SLURM_SCRIPT="$OUTPUT_DIR/submit.sh"
+    cat > "$SLURM_SCRIPT" <<EOF
+#!/bin/bash
+#SBATCH --job-name=$NAME
+#SBATCH --partition=$PARTITION
+#SBATCH --nodes=$NODES
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=$CPUS_PER_TASK
+#SBATCH --gres=dcu:$GPUS_PER_NODE
+#SBATCH --time=$TIME
+#SBATCH --output=$OUTPUT_DIR/slurm_%j.out
+#SBATCH --error=$OUTPUT_DIR/slurm_%j.err
+EOF
+
+    if [[ -n "$NODELIST" ]]; then
+        printf '#SBATCH --nodelist=%s\n' "$NODELIST" >> "$SLURM_SCRIPT"
+    fi
+
+    cat >> "$SLURM_SCRIPT" <<EOF
+
+set -euo pipefail
+export PYTHONPATH="$REPO_ROOT:\${PYTHONPATH:-}"
+export ONESCIENCE_EQUIFORMER_V3_DIR="$EQV3_DIR"
+export ONESCIENCE_EQUIFORMER_V3_JD_PATH="$ONESCIENCE_EQUIFORMER_V3_JD_PATH"
+# A Slurm allocation starts a new shell, so activate the same conda environment.
+source "$(dirname "$CONDA_EXE")/../etc/profile.d/conda.sh"
+conda activate "$MATCHEM_CONDA_NAME"
+export OMP_NUM_THREADS="$OMP_NUM_THREADS"
+# A login/compute shell may carry a single-device selection into sbatch.
+# Let Slurm expose all devices requested by this YAML.
+if (( $GPUS_PER_NODE > 1 )); then
+    unset HIP_VISIBLE_DEVICES
+fi
+$ENV_EXPORTS
+cd "$OUTPUT_DIR"
+EOF
+
+    if (( NODES > 1 )); then
+        cat >> "$SLURM_SCRIPT" <<EOF
+
+nodes=(\$(scontrol show hostnames "\$SLURM_JOB_NODELIST"))
+export MASTER_ADDR="\${nodes[0]}"
+export MASTER_PORT=\$((20000 + SLURM_JOB_ID % 20000))
+echo "Equiformer V3 multi-node DDP: nodes=\$SLURM_NNODES, devices/node=$GPUS_PER_NODE"
+echo "MASTER_ADDR=\$MASTER_ADDR MASTER_PORT=\$MASTER_PORT"
+
+srun --nodes="\$SLURM_NNODES" --ntasks="\$SLURM_NNODES" --ntasks-per-node=1 \\
+  bash -c 'exec torchrun \\
+    --nnodes=$NODES \\
+    --node_rank="\${SLURM_NODEID}" \\
+    --nproc_per_node=$GPUS_PER_NODE \\
+    --rdzv_id="\${SLURM_JOB_ID}" \\
+    --rdzv_backend=c10d \\
+    --rdzv_endpoint="\${MASTER_ADDR}:\${MASTER_PORT}" \\
+    "$EQV3_DIR/finetune.py" --config "$FINETUNE_CONFIG"'
+EOF
+    elif (( GPUS_PER_NODE > 1 )); then
+        cat >> "$SLURM_SCRIPT" <<EOF
+
+torchrun --standalone --nproc_per_node=$GPUS_PER_NODE \\
+  "$EQV3_DIR/finetune.py" --config "$FINETUNE_CONFIG"
+EOF
+    else
+        echo "python \"$EQV3_DIR/finetune.py\" --config \"$FINETUNE_CONFIG\"" >> "$SLURM_SCRIPT"
+    fi
+
+    chmod u+x "$SLURM_SCRIPT"
+    echo "Submitting Equiformer V3 job: $SLURM_SCRIPT"
+    sbatch "$SLURM_SCRIPT"
+    exit 0
+fi
+
+eval "$ENV_EXPORTS"
+export PYTHONPATH="$REPO_ROOT:${PYTHONPATH:-}"
+cd "$OUTPUT_DIR"
+if (( GPUS_PER_NODE > 1 )); then
+    exec torchrun --standalone --nproc_per_node="$GPUS_PER_NODE" \
+        "$EQV3_DIR/finetune.py" --config "$FINETUNE_CONFIG"
+fi
+exec python "$EQV3_DIR/finetune.py" --config "$FINETUNE_CONFIG"

elastic.py

@@ -0,0 +1,104 @@
+"""Fit an elastic tensor from Equiformer V3 stress predictions."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+from ase.build import bulk
+from ase.io import read, write
+
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3Calculator,
+    relax_structure,
+    run_elastic_workflow,
+    write_workflow_result,
+)
+
+
+def default_checkpoint() -> str:
+    return str(
+        Path(__file__).resolve().parent
+        / "weight"
+        / "omat24-mptrj-salex_gradient.pt"
+    )
+
+
+def load_structure(path: str | None):
+    if path:
+        return read(path)
+    return bulk("Cu")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint")
+    parser.add_argument(
+        "--input",
+        help="periodic structure readable by ASE; defaults to built-in bulk Cu",
+    )
+    parser.add_argument(
+        "--relax",
+        action="store_true",
+        help="relax the cell and atomic positions before applying strains",
+    )
+    parser.add_argument(
+        "--normal-strains", type=float, nargs="+", default=(-0.01, 0.01)
+    )
+    parser.add_argument(
+        "--shear-strains", type=float, nargs="+", default=(-0.02, 0.02)
+    )
+    parser.add_argument(
+        "--relax-positions",
+        action="store_true",
+        help="relax atomic positions at fixed cell for every deformation",
+    )
+    parser.add_argument("--relax-fmax", type=float, default=0.02)
+    parser.add_argument("--relax-steps", type=int, default=100)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--output", default="outputs/elastic.json")
+    args = parser.parse_args()
+
+    checkpoint = args.checkpoint or default_checkpoint()
+    calculator = EquiformerV3Calculator.from_checkpoint(
+        checkpoint, device=args.device
+    )
+    structure = load_structure(args.input)
+    initial_relaxation = None
+    if args.relax:
+        structure, initial_relaxation = relax_structure(
+            structure,
+            calculator,
+            relax_cell=True,
+            fmax=args.relax_fmax,
+            steps=args.relax_steps,
+        )
+        relaxed_output = Path(args.output).with_name("elastic_relaxed.cif")
+        relaxed_output.parent.mkdir(parents=True, exist_ok=True)
+        write(relaxed_output, structure)
+    result = run_elastic_workflow(
+        structure,
+        calculator,
+        normal_strains=args.normal_strains,
+        shear_strains=args.shear_strains,
+        relax_positions=args.relax_positions,
+        relax_fmax=args.relax_fmax,
+        relax_steps=args.relax_steps,
+    )
+    result["checkpoint"] = str(checkpoint)
+    result["initial_relaxation"] = initial_relaxation
+    output = write_workflow_result(result, args.output)
+
+    print("formula:", result["formula"])
+    print("bulk modulus, Hill (GPa):", result["bulk_modulus_gpa"]["hill"])
+    print("result:", output)
+
+
+if __name__ == "__main__":
+    main()

evaluate.py

@@ -0,0 +1,262 @@
+"""Evaluate an Equiformer V3 checkpoint on an independent ASE database."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+import torch
+
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3CheckpointTransforms,
+    load_equiformer_v3_checkpoint,
+)
+
+from finetune import _loader
+
+
+@dataclass
+class ErrorAccumulator:
+    """Accumulate MAE/RMSE inputs without retaining the full dataset."""
+
+    absolute_sum: float = 0.0
+    squared_sum: float = 0.0
+    count: int = 0
+
+    def update(self, error: torch.Tensor) -> None:
+        error = error.detach().reshape(-1).double()
+        self.absolute_sum += float(error.abs().sum())
+        self.squared_sum += float(error.square().sum())
+        self.count += error.numel()
+
+    def result(self) -> dict[str, float]:
+        if self.count == 0:
+            raise ValueError("cannot compute metrics for an empty tensor")
+        return {
+            "mae": self.absolute_sum / self.count,
+            "rmse": (self.squared_sum / self.count) ** 0.5,
+        }
+
+
+@dataclass
+class MeanAccumulator:
+    total: float = 0.0
+    count: int = 0
+
+    def update(self, values: torch.Tensor) -> None:
+        values = values.detach().reshape(-1).double()
+        self.total += float(values.sum())
+        self.count += values.numel()
+
+    def result(self) -> float:
+        if self.count == 0:
+            raise ValueError("cannot compute a mean for an empty tensor")
+        return self.total / self.count
+
+
+def _force_cosine(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    prediction = prediction.detach().float().reshape(-1, 3)
+    target = target.detach().float().reshape(-1, 3)
+    if prediction.shape != target.shape or prediction.shape[0] == 0:
+        raise ValueError(
+            "cannot compute force cosine similarity for an empty/mismatched tensor"
+        )
+    return torch.cosine_similarity(prediction, target, dim=1).detach().cpu()
+
+
+def _force_magnitude_error(
+    prediction: torch.Tensor, target: torch.Tensor
+) -> torch.Tensor:
+    prediction = prediction.detach().float().reshape(-1, 3)
+    target = target.detach().float().reshape(-1, 3)
+    if prediction.shape != target.shape or prediction.shape[0] == 0:
+        raise ValueError(
+            "cannot compute force magnitude error for an empty/mismatched tensor"
+        )
+    return (
+        torch.linalg.vector_norm(prediction, dim=1)
+        - torch.linalg.vector_norm(target, dim=1)
+    ).abs().cpu()
+
+
+def _energy_force_success(
+    energy_error: torch.Tensor,
+    force_error: torch.Tensor,
+    selected_natoms: torch.Tensor,
+) -> torch.Tensor:
+    """Match FairChem's OC20 energy/force threshold metric per structure."""
+
+    energy_error = energy_error.detach().reshape(-1).abs()
+    force_error = force_error.detach().float().reshape(-1, 3).abs()
+    selected_natoms = selected_natoms.detach().reshape(-1).long()
+    if energy_error.numel() != selected_natoms.numel():
+        raise ValueError("energy count and per-structure atom counts differ")
+    if int(selected_natoms.sum()) != force_error.shape[0]:
+        raise ValueError("force count and per-structure atom counts differ")
+    if bool((selected_natoms == 0).any()):
+        raise ValueError("a structure contains no selected atoms for force evaluation")
+
+    successes = []
+    for structure_energy_error, structure_force_error in zip(
+        energy_error, torch.split(force_error, selected_natoms.tolist())
+    ):
+        successes.append(
+            (structure_energy_error < 0.02)
+            & (structure_force_error.max() < 0.03)
+        )
+    return torch.stack(successes).detach().cpu()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--data", required=True, help="Independent ASE DB/ASE-LMDB")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--workers", type=int, default=0)
+    parser.add_argument("--max-samples", type=int)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument(
+        "--free-atoms-only",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Evaluate forces only on atoms not constrained by FixAtoms (default: true).",
+    )
+    parser.add_argument(
+        "--include-stress",
+        action="store_true",
+        help="Also report stress when both checkpoint and dataset provide it.",
+    )
+    parser.add_argument(
+        "--include-oc20-threshold",
+        action="store_true",
+        help="Report FairChem's OC20 energy/force threshold success rate.",
+    )
+    parser.add_argument("--output")
+    args = parser.parse_args()
+
+    device = torch.device(args.device)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        raise RuntimeError("CUDA/DCU was requested but is unavailable")
+
+    model = load_equiformer_v3_checkpoint(args.checkpoint).to(device)
+    model.eval()
+    transforms = EquiformerV3CheckpointTransforms.from_checkpoint(
+        args.checkpoint
+    ).to(device)
+    loader = _loader(
+        args.data,
+        args.batch_size,
+        args.workers,
+        max_samples=args.max_samples,
+        train=False,
+        seed=args.seed,
+    )
+
+    energy_errors = ErrorAccumulator()
+    energy_per_atom_errors = ErrorAccumulator()
+    force_errors = ErrorAccumulator()
+    force_cosines = MeanAccumulator()
+    force_magnitude_errors = MeanAccumulator()
+    energy_force_successes = MeanAccumulator()
+    stress_errors = ErrorAccumulator()
+    with torch.enable_grad():
+        for batch in loader:
+            batch = batch.to(device)
+            prediction = model(batch)
+            pred_energy = transforms.denormalize_prediction(
+                "energy", prediction["energy"], batch
+            )
+            target_energy = batch.energy.reshape_as(pred_energy)
+            energy_error = pred_energy - target_energy
+            energy_errors.update(energy_error)
+            natoms = batch.natoms.to(pred_energy).reshape(
+                (-1,) + (1,) * (pred_energy.ndim - 1)
+            )
+            energy_per_atom_errors.update(energy_error / natoms)
+
+            pred_forces = transforms.denormalize_prediction(
+                "forces", prediction["forces"], batch
+            )
+            target_forces = batch.forces.reshape_as(pred_forces)
+            selected_natoms = batch.natoms
+            if args.free_atoms_only and hasattr(batch, "fixed"):
+                free_mask = batch.fixed.reshape(-1) == 0
+                selected_natoms = torch.stack(
+                    [
+                        structure_mask.sum()
+                        for structure_mask in torch.split(
+                            free_mask, batch.natoms.tolist()
+                        )
+                    ]
+                )
+                pred_forces = pred_forces.reshape(-1, 3)[free_mask]
+                target_forces = target_forces.reshape(-1, 3)[free_mask]
+            if not pred_forces.numel():
+                selection = "free atoms" if args.free_atoms_only else "atoms"
+                raise ValueError(f"an evaluation batch contains no selected {selection}")
+            force_error = pred_forces - target_forces
+            force_errors.update(force_error)
+            force_cosines.update(_force_cosine(pred_forces, target_forces))
+            force_magnitude_errors.update(
+                _force_magnitude_error(pred_forces, target_forces)
+            )
+            if args.include_oc20_threshold:
+                energy_force_successes.update(
+                    _energy_force_success(
+                        energy_error, force_error, selected_natoms
+                    )
+                )
+            if (
+                args.include_stress
+                and "stress" in prediction
+                and hasattr(batch, "stress")
+            ):
+                pred_stress = transforms.denormalize_prediction(
+                    "stress", prediction["stress"], batch
+                )
+                target_stress = batch.stress.reshape_as(pred_stress)
+                stress_errors.update(pred_stress - target_stress)
+
+    result = {
+        "checkpoint": str(Path(args.checkpoint).expanduser()),
+        "data": str(Path(args.data).expanduser()),
+        "samples": len(loader.dataset),
+        "force_atoms": force_cosines.count,
+        "free_atoms_only": args.free_atoms_only,
+        "energy_total_eV": energy_errors.result(),
+        "energy_per_atom_eV": energy_per_atom_errors.result(),
+    }
+    if force_errors.count == 0:
+        selection = "free atoms" if args.free_atoms_only else "atoms"
+        raise ValueError(f"the evaluation dataset contains no selected {selection}")
+    result["forces_eV_per_A"] = force_errors.result()
+    result["forces_cosine_similarity"] = {"mean": force_cosines.result()}
+    result["forces_magnitude_error_eV_per_A"] = {
+        "mean": force_magnitude_errors.result()
+    }
+    if args.include_oc20_threshold:
+        result["energy_forces_within_threshold"] = {
+            "fraction": energy_force_successes.result(),
+            "energy_threshold_eV": 0.02,
+            "force_threshold_eV_per_A": 0.03,
+        }
+    if stress_errors.count:
+        result["stress_eV_per_A3"] = stress_errors.result()
+    print(json.dumps(result, indent=2, sort_keys=True))
+    if args.output:
+        output = Path(args.output).expanduser()
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
+
+
+if __name__ == "__main__":
+    main()

finetune.py

@@ -0,0 +1,392 @@
+"""Compatibility entrypoint for Equiformer V3 checkpoint fine-tuning.
+
+New training configurations are handled by :mod:`train`, which supports
+scratch training, checkpoint initialization, and full-state resume. The
+loader remains here for compatibility with ``evaluate.py`` and external code.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Subset
+from torch.utils.data.distributed import DistributedSampler
+
+from onescience.datapipes.materials.custom_stack import data_list_collater
+from onescience.datapipes.materials.custom_stack.storage.ase_datasets import AseDBDataset
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3CheckpointTransforms,
+    load_equiformer_v3_checkpoint,
+)
+from onescience.utils.uma.normalization.element_references import (
+    fit_linear_references,
+)
+
+
+@dataclass(frozen=True)
+class DistributedContext:
+    """Runtime information for a normal Python process or a torchrun worker."""
+
+    rank: int = 0
+    world_size: int = 1
+    local_rank: int = 0
+
+    @property
+    def enabled(self) -> bool:
+        return self.world_size > 1
+
+    @property
+    def is_main(self) -> bool:
+        return self.rank == 0
+
+
+def _init_distributed(device_name: str, backend: str) -> DistributedContext:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size == 1:
+        if device_name.startswith("cuda") and torch.cuda.is_available():
+            torch.cuda.set_device(0)
+        return DistributedContext()
+
+    if not torch.distributed.is_available():
+        raise RuntimeError("torch.distributed is required for multi-device fine-tuning.")
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    if device_name.startswith("cuda"):
+        if not torch.cuda.is_available():
+            raise RuntimeError(
+                "torchrun requested multiple CUDA/DCU devices, but CUDA is unavailable."
+            )
+        torch.cuda.set_device(local_rank)
+    torch.distributed.init_process_group(backend=backend, rank=rank, world_size=world_size)
+    return DistributedContext(rank=rank, world_size=world_size, local_rank=local_rank)
+
+
+def _close_distributed(context: DistributedContext) -> None:
+    if context.enabled and torch.distributed.is_initialized():
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
+def _loader(
+    path: str | list[str],
+    batch_size: int,
+    workers: int,
+    max_samples: int | None = None,
+    context: DistributedContext | None = None,
+    train: bool = False,
+    seed: int = 0,
+) -> DataLoader:
+    dataset = AseDBDataset(
+        {
+            "src": path,
+            "a2g_args": {
+                "r_edges": False,
+                "r_energy": True,
+                "r_forces": True,
+                "r_stress": True,
+            },
+        }
+    )
+    if max_samples is not None:
+        sample_count = min(max_samples, len(dataset))
+        generator = torch.Generator().manual_seed(seed)
+        indices = torch.randperm(len(dataset), generator=generator)[:sample_count].tolist()
+        dataset = Subset(dataset, indices)
+    context = context or DistributedContext()
+    sampler = None
+    if context.enabled:
+        sampler = DistributedSampler(
+            dataset,
+            num_replicas=context.world_size,
+            rank=context.rank,
+            shuffle=train,
+            drop_last=False,
+        )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=sampler is None and train,
+        sampler=sampler,
+        num_workers=workers,
+        collate_fn=lambda items: data_list_collater(items, otf_graph=True),
+    )
+
+
+def _loss(
+    pred,
+    batch,
+    energy_weight: float,
+    force_weight: float,
+    stress_weight: float,
+    transforms: EquiformerV3CheckpointTransforms,
+):
+    losses = {}
+    if energy_weight:
+        energy_target = transforms.normalize_target(
+            "energy", batch.energy, pred["energy"], batch
+        )
+        energy_error = pred["energy"] - energy_target
+        natoms_shape = (-1,) + (1,) * (energy_error.ndim - 1)
+        natoms = batch.natoms.to(energy_error).reshape(natoms_shape)
+        losses["energy"] = (energy_error / natoms).square().mean()
+    if force_weight:
+        force_target = transforms.normalize_target(
+            "forces", batch.forces, pred["forces"], batch
+        )
+        losses["forces"] = (pred["forces"] - force_target).square().mean()
+    if stress_weight:
+        if not hasattr(batch, "stress"):
+            raise ValueError("stress_weight is nonzero, but the batch has no stress labels")
+        if "stress" not in pred:
+            raise ValueError("stress_weight is nonzero, but the model returned no stress")
+        stress_target = transforms.normalize_target(
+            "stress", batch.stress, pred["stress"], batch
+        )
+        losses["stress"] = (pred["stress"] - stress_target).square().mean()
+    total = energy_weight * losses.get("energy", 0.0)
+    total = total + force_weight * losses.get("forces", 0.0)
+    total = total + stress_weight * losses.get("stress", 0.0)
+    return total, {key: float(value.detach()) for key, value in losses.items()}
+
+
+def _run_epoch(
+    model,
+    loader,
+    device,
+    optimizer,
+    weights,
+    transforms: EquiformerV3CheckpointTransforms,
+    context: DistributedContext,
+):
+    training = optimizer is not None
+    model.train(training)
+    total = 0.0
+    batches = 0
+    metric_names = tuple(
+        name
+        for name, weight in zip(("energy", "forces", "stress"), weights)
+        if weight
+    )
+    metrics = {name: 0.0 for name in metric_names}
+    for batch in loader:
+        batch = batch.to(device)
+        if training:
+            optimizer.zero_grad(set_to_none=True)
+        prediction = model(batch)
+        loss, batch_metrics = _loss(prediction, batch, *weights, transforms)
+        if training:
+            loss.backward()
+            optimizer.step()
+        total += float(loss.detach())
+        batches += 1
+        for key, value in batch_metrics.items():
+            metrics[key] = metrics.get(key, 0.0) + value
+    if batches == 0:
+        raise RuntimeError("The dataset contains no samples.")
+    values = torch.tensor(
+        [total, *metrics.values(), float(batches)],
+        dtype=torch.float64,
+        device=device,
+    )
+    if context.enabled:
+        torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.SUM)
+    global_batches = values[-1].item()
+    return {
+        "loss": values[0].item() / global_batches,
+        **{
+            key: values[index].item() / global_batches
+            for index, key in enumerate(metrics, start=1)
+        },
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", help="YAML configuration path")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--train", help="ASE DB or ASE-LMDB training path")
+    parser.add_argument("--val", help="ASE DB or ASE-LMDB validation path")
+    parser.add_argument("--output")
+    parser.add_argument("--device")
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--workers", type=int)
+    parser.add_argument("--lr", type=float)
+    parser.add_argument("--energy-weight", type=float)
+    parser.add_argument("--force-weight", type=float)
+    parser.add_argument("--stress-weight", type=float)
+    parser.add_argument("--max-train-samples", type=int)
+    parser.add_argument("--max-val-samples", type=int)
+    parser.add_argument("--backend", help="torch.distributed backend for torchrun")
+    parser.add_argument("--seed", type=int)
+    parser.add_argument(
+        "--fit-element-references",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+        help="fit energy element references on the training data",
+    )
+    args = parser.parse_args()
+
+    if not args.config:
+        parser.error("--config is required; use a YAML file from demo/configs")
+    config_path = args.config
+    with Path(config_path).expanduser().open() as handle:
+        config = yaml.safe_load(handle) or {}
+    for key, value in config.items():
+        if getattr(args, key.replace("-", "_"), None) is None:
+            setattr(args, key.replace("-", "_"), value)
+    for key in ("checkpoint", "train", "val", "output"):
+        value = getattr(args, key)
+        if value is not None:
+            if isinstance(value, list):
+                value = [
+                    os.path.expandvars(os.path.expanduser(str(item)))
+                    for item in value
+                ]
+            else:
+                value = os.path.expandvars(os.path.expanduser(str(value)))
+            setattr(args, key, value)
+    required = ("checkpoint", "train", "val", "output")
+    missing = [key for key in required if not getattr(args, key)]
+    if missing:
+        parser.error("missing required config fields: " + ", ".join(missing))
+    args.backend = args.backend or "nccl"
+    args.seed = 0 if args.seed is None else args.seed
+    args.fit_element_references = bool(args.fit_element_references)
+    if not any((args.energy_weight, args.force_weight, args.stress_weight)):
+        parser.error(
+            "at least one of energy_weight, force_weight, or stress_weight "
+            "must be nonzero"
+        )
+
+    if args.device.startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false.")
+    context = _init_distributed(args.device, args.backend)
+    try:
+        if args.device.startswith("cuda"):
+            device = torch.device(f"cuda:{context.local_rank}")
+        else:
+            device = torch.device(args.device)
+        torch.manual_seed(args.seed + context.rank)
+
+        model = load_equiformer_v3_checkpoint(args.checkpoint).to(device)
+        transforms = EquiformerV3CheckpointTransforms.from_checkpoint(
+            args.checkpoint
+        )
+        if args.fit_element_references:
+            reference_dataset = _loader(
+                args.train, args.batch_size, args.workers
+            ).dataset
+            fitted_references = fit_linear_references(
+                targets=["energy"],
+                dataset=reference_dataset,
+                batch_size=args.batch_size,
+                num_workers=args.workers,
+                log_metrics=False,
+                shuffle=False,
+            )
+            transforms.elementrefs["energy"] = fitted_references["energy"]
+            if context.is_main:
+                print("fitted energy element references from training data", flush=True)
+        transforms = transforms.to(device)
+        if context.enabled:
+            model = DistributedDataParallel(
+                model,
+                device_ids=[context.local_rank] if device.type == "cuda" else None,
+                output_device=context.local_rank if device.type == "cuda" else None,
+            )
+        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
+        train_loader = _loader(
+            args.train,
+            args.batch_size,
+            args.workers,
+            args.max_train_samples,
+            context=context,
+            train=True,
+            seed=args.seed,
+        )
+        val_loader = _loader(
+            args.val,
+            args.batch_size,
+            args.workers,
+            args.max_val_samples,
+            context=context,
+            train=False,
+            seed=args.seed + 1,
+        )
+        weights = (args.energy_weight, args.force_weight, args.stress_weight)
+
+        history = []
+        for epoch in range(args.epochs):
+            if isinstance(train_loader.sampler, DistributedSampler):
+                train_loader.sampler.set_epoch(epoch)
+            train_metrics = _run_epoch(
+                model, train_loader, device, optimizer, weights, transforms, context
+            )
+            # Force/stress outputs are gradients of the energy, so validation also
+            # needs autograd even though model parameters are not updated.
+            val_metrics = _run_epoch(
+                model, val_loader, device, None, weights, transforms, context
+            )
+            record = {"epoch": epoch, "train": train_metrics, "val": val_metrics}
+            if context.is_main:
+                history.append(record)
+                print(json.dumps(record, sort_keys=True), flush=True)
+
+        if context.is_main:
+            output = Path(args.output)
+            output.parent.mkdir(parents=True, exist_ok=True)
+            source = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+            base_model = model.module if context.enabled else model
+            checkpoint = {
+                "config": source["config"],
+                "normalizers": source.get("normalizers", {}),
+                "state_dict": {
+                    key: value.detach().cpu()
+                    for key, value in base_model.state_dict().items()
+                },
+                "metadata": dict(source.get("metadata", {})),
+            }
+            checkpoint["elementrefs"] = {
+                name: {
+                    key: value.detach().cpu()
+                    for key, value in elementref.state_dict().items()
+                }
+                for name, elementref in transforms.elementrefs.items()
+            }
+            checkpoint["metadata"].update(
+                {
+                    "onescience_equiformer_v3_history": history,
+                    "source_checkpoint": args.checkpoint,
+                    "world_size": context.world_size,
+                    "loss_space": "checkpoint_normalized",
+                    "element_references": (
+                        "fitted_from_training_data"
+                        if args.fit_element_references
+                        else "source_checkpoint"
+                    ),
+                }
+            )
+            del source
+            torch.save(checkpoint, output)
+            print(f"saved checkpoint: {output}", flush=True)
+    finally:
+        _close_distributed(context)
+
+
+if __name__ == "__main__":
+    from train import main as training_main
+
+    training_main()

fit_element_references.py

@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+"""Fit OC20 energy element references before starting model training."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+from onescience.utils.uma.normalization.element_references import (
+    fit_linear_references,
+)
+
+from train import _expand_path, _loader
+
+
+def fit_references(
+    config_path: str | Path,
+    output_path: str | Path,
+    batch_size: int = 32,
+    num_batches: int | None = None,
+) -> Path:
+    """Fit and save legacy ``coeff`` references consumed by the OC20 YAML."""
+
+    if batch_size < 2:
+        raise ValueError("reference fitting batch_size must be at least 2")
+    if num_batches is not None and num_batches < 1:
+        raise ValueError("num_batches must be positive")
+
+    config_path = Path(config_path).expanduser().resolve()
+    with config_path.open(encoding="utf-8") as handle:
+        config = yaml.safe_load(handle) or {}
+
+    train_path = _expand_path(config.get("train"))
+    if not train_path:
+        raise ValueError(f"training data is not configured in {config_path}")
+
+    loader = _loader(
+        train_path,
+        batch_size=batch_size,
+        workers=int(config.get("workers", 0)),
+        max_samples=config.get("max_train_samples"),
+        train=False,
+        seed=int(config.get("seed", 0)),
+        max_atoms=config.get("max_atoms"),
+        load_balancing=False,
+    )
+    references = fit_linear_references(
+        targets=["energy"],
+        dataset=loader.dataset,
+        batch_size=batch_size,
+        num_batches=num_batches,
+        num_workers=int(config.get("workers", 0)),
+        shuffle=False,
+    )
+
+    output_path = Path(output_path).expanduser().resolve()
+    if output_path.suffix != ".npz":
+        raise ValueError("output path must use the .npz extension")
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    coefficients = (
+        references["energy"].element_references.detach().cpu().numpy()
+    )
+    np.savez(output_path, coeff=coefficients)
+    print(f"Saved energy element references to {output_path}")
+    return output_path
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", required=True, help="OC20 training YAML")
+    parser.add_argument("--output", required=True, help="output .npz path")
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-batches", type=int)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    fit_references(
+        args.config,
+        args.output,
+        batch_size=args.batch_size,
+        num_batches=args.num_batches,
+    )
+
+
+if __name__ == "__main__":
+    main()

formation_energy.py

@@ -0,0 +1,132 @@
+"""Calculate an uncorrected formation energy with Equiformer V3."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+from ase.build import bulk, molecule
+from ase.io import read
+
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3Calculator,
+    calculate_element_reference_energies,
+    calculate_formation_energy,
+    load_element_reference_energies,
+    relax_structure,
+    write_workflow_result,
+)
+
+
+def default_checkpoint() -> str:
+    return str(
+        Path(__file__).resolve().parent
+        / "weight"
+        / "omat24-mptrj-salex_gradient.pt"
+    )
+
+
+def load_compound(path: str | None):
+    if path:
+        return read(path)
+    return bulk("MgO", "rocksalt", a=4.21)
+
+
+def model_reference_structures():
+    oxygen = molecule("O2")
+    oxygen.center(vacuum=8.0)
+    oxygen.pbc = True
+    return {
+        "Mg": bulk("Mg", "hcp", a=3.21, c=5.21),
+        "O": oxygen,
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint")
+    parser.add_argument(
+        "--input",
+        help="compound structure readable by ASE; defaults to built-in MgO",
+    )
+    parser.add_argument(
+        "--reference-energies",
+        help="JSON/YAML mapping of element to trusted reference energy in eV/atom",
+    )
+    parser.add_argument(
+        "--relax",
+        action="store_true",
+        help="relax the compound and model-evaluated reference phases",
+    )
+    parser.add_argument("--fmax", type=float, default=0.03)
+    parser.add_argument("--steps", type=int, default=100)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--output", default="outputs/formation_energy.json")
+    args = parser.parse_args()
+
+    checkpoint = args.checkpoint or default_checkpoint()
+    calculator = EquiformerV3Calculator.from_checkpoint(
+        checkpoint, device=args.device
+    )
+    compound = load_compound(args.input)
+    compound_relaxation = None
+    if args.relax:
+        compound, compound_relaxation = relax_structure(
+            compound,
+            calculator,
+            relax_cell=bool(compound.pbc.all()),
+            fmax=args.fmax,
+            steps=args.steps,
+        )
+
+    reference_relaxations = {}
+    if args.reference_energies:
+        reference_energies = load_element_reference_energies(
+            args.reference_energies
+        )
+        reference_source = str(Path(args.reference_energies))
+    else:
+        reference_structures = model_reference_structures()
+        if args.relax:
+            relaxed_references = {}
+            for element, atoms in reference_structures.items():
+                relaxed, metadata = relax_structure(
+                    atoms,
+                    calculator,
+                    relax_cell=element != "O",
+                    fmax=args.fmax,
+                    steps=args.steps,
+                )
+                relaxed_references[element] = relaxed
+                reference_relaxations[element] = metadata
+            reference_structures = relaxed_references
+        reference_energies = calculate_element_reference_energies(
+            reference_structures, calculator
+        )
+        reference_source = "Equiformer V3 evaluation of Mg(hcp) and O2"
+
+    result = calculate_formation_energy(
+        compound,
+        calculator,
+        reference_energies,
+    )
+    result["reference_source"] = reference_source
+    result["checkpoint"] = str(checkpoint)
+    result["compound_relaxation"] = compound_relaxation
+    result["reference_relaxations"] = reference_relaxations
+    output = write_workflow_result(result, args.output)
+
+    print("formula:", result["formula"])
+    print("formation energy (eV/atom):", result["formation_energy_ev_per_atom"])
+    print("corrections applied:", result["corrections_applied"])
+    print("result:", output)
+
+
+if __name__ == "__main__":
+    main()

model/__init__.py

@@ -0,0 +1,14 @@
+"""Equiformer V3 atomistic models adapted to OneScience."""
+
+from .equiformer_v3 import EquiformerV3_OC
+from .equiformer_v3_dens import EquiformerV3DeNS_OC
+
+EquiformerV3 = EquiformerV3_OC
+EquiformerV3DeNS = EquiformerV3DeNS_OC
+
+__all__ = [
+    "EquiformerV3",
+    "EquiformerV3DeNS",
+    "EquiformerV3DeNS_OC",
+    "EquiformerV3_OC",
+]

model/equiformer_v3.py

@@ -0,0 +1,771 @@
+"""Equiformer V3 model adapted to the OneScience FairChem runtime."""
+
+import math
+
+import torch
+from torch.utils.checkpoint import checkpoint
+
+from onescience.modules.func_utils.uma_graph.mixin import GraphModelMixin
+from onescience.modules.layer.equiformer_v3.edge_rot_mat import init_edge_rot_mat
+from onescience.modules.layer.equiformer_v3.envelope import PolynomialEnvelope
+from onescience.modules.layer.equiformer_v3.input_block import EdgeDegreeEmbedding
+from onescience.modules.layer.equiformer_v3.layer_norm import (
+    EquivariantLayerNorm,
+    EquivariantMergeLayerNorm,
+    EquivariantSeparableLayerNorm,
+    RMSNorm,
+    get_normalization_layer,
+)
+from onescience.modules.layer.equiformer_v3.output_block import (
+    EquivariantGraphAttentionStressHead,
+    FeedForwardNetworkStressHead,
+    ScalarFeedForwardNetwork,
+)
+from onescience.modules.layer.equiformer_v3.radial_function import (
+    GaussianSmearing,
+    RadialFunction,
+)
+from onescience.modules.layer.equiformer_v3.so3 import (
+    SO3Rotation,
+    SO3Linear,
+)
+from onescience.modules.layer.equiformer_v3.transformer_block import (
+    EquivariantGraphAttention,
+    FeedForwardNetwork,
+    TransBlockV3,
+)
+from onescience.modules.layer.equiformer_v3.wigner import configure_jd_path
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.common.utils import conditional_grad
+
+
+# Statistics of IS2RE 100K
+_AVG_NUM_NODES = 77.81317
+_AVG_DEGREE = 23.395238876342773    # IS2RE: 100k, max_radius = 5, max_neighbors = 100
+
+_NORM_SCALE_NODES = math.sqrt(_AVG_NUM_NODES)   # 8.82117735906041
+_NORM_SCALE_DEGREE = math.sqrt(_AVG_DEGREE)     # 4.836862503353054
+
+
+@registry.register_model("equiformer_v3")
+class EquiformerV3_OC(torch.nn.Module, GraphModelMixin):
+    """
+    Args:
+        use_pbc (bool):         Use periodic boundary conditions
+        use_pbc_single (bool, optional):    Process batch PBC graphs one at a time.
+                                            This argument is set to True when training on MPTraj.
+        otf_graph (bool):       Compute graph On The Fly (OTF)
+
+        regress_forces (bool):  Compute forces
+        regress_stress (bool):  Compute stress
+        direct_prediction (bool):   Whether to use direct methods to predict forces and stress
+
+        max_neighbors (int):    Maximum number of neighbors per atom
+        max_radius (float):     Maximum distance between nieghboring atoms in Angstroms
+        num_radial_basis (int): Number of radial basis functions
+        max_num_elements (int): Maximum atomic number
+
+        num_layers (int):           Number of layers in the GNN
+        num_channels (int):         Number of channels in node embeddings
+        attn_hidden_channels (int): Number of hidden channels in equivariant graph attention
+        num_heads (int):            Number of attention heads
+        attn_alpha_channels (int):  Number of channels for alpha vector in each attention head
+        attn_value_channels (int):  Number of channels for value vector in each attention head
+        ffn_hidden_channels (int):  Number of hidden channels in feedforward network
+        norm_type (str):            Type of normalization layer
+                                    (['sep_layer_norm', 'merge_layer_norm',
+                                    'merge_layer_norm_attn_rms_norm', 'merge_rms_norm'])
+
+        lmax (int):                 Maximum degrees (l)
+        mmax (int):                 Maximum order (m)
+        attn_grid_resolution_list (list:int):
+                                    Grid resolution list in class `SO3Grid` in attention
+        ffn_grid_resolution_list (list:int):
+                                    Grid resolution list in class `SO3Grid` in feedforward network
+
+        edge_channels (int):                Number of channels for edge-wise invariant features
+        use_atom_edge_embedding (bool):     Whether to use atomic embedding along with relative distance for edge scalar features
+        use_envelope (bool):        Whether to apply an envelope function to attention
+
+        attn_activation (str):      Type of activation function in equivariant graph attention
+        use_attn_renorm (bool):     Whether to re-normalize attention weights
+        use_add_merge (bool):       Default: False
+                                    If True, use addition to merge the source/target node features instead of concat,
+                                    which can save 2x compute when rotating with Wigner-D matrices.
+        use_rad_l_parametrization (bool):
+                                    Default: True
+                                    If True, all the m components within the same type-L vector will share the same
+                                    weight from the radial function.
+        softcap (float):            Default: None
+                                    If not None, use soft capping to limit the range of attention logits to
+                                    [- `softcap`, + `softcap`].
+        attn_eps (float):           Default: 1e-16
+                                    Epsilon value used in the softmax operation of attention
+        ffn_activation (str):       Type of activation function for feedforward network
+        use_grid_mlp (bool):        If `True`, use projecting to grids and performing MLPs for FFNs.
+
+        use_gate_force_head (bool): If `True`, use `GateActivation` in the equivariant attention of the force prediction head.
+
+        alpha_drop (float):         Dropout rate for the hidden features in non-linear MLP attention
+        attn_mask_rate (float):     Mask rate for neighbors considered in attention
+        attn_weights_drop (float):  Dropout rate for attention weights
+        value_drop (float):         Dropout rate for the hidden features in non-linear value vectors
+        drop_path_rate (float):     Drop path rate
+        proj_drop (float):          Dropout rate for outputs of attention and FFN in Transformer blocks
+        ffn_drop (float):           Dropout rate for the hidden features in FFN
+        use_head_reg (bool):        Whether to apply regularization to output head (dummy argument for backend compatibility)
+
+        gradient_checkpointing_block_list (list):
+                                    A list indicating which block we apply gradient/activation checkpointing to save memory.
+
+        avg_num_nodes (float):   Normalization factor for sum aggregation over nodes
+        avg_degree (float):      Normalization factor for sum aggregation over edges
+
+        enforce_max_neighbors_strictly (bool):      When edges are subselected based on the `max_neighbors` arg, arbitrarily select amongst equidistant / degenerate edges to have exactly the correct number.
+    """
+    def __init__(
+        self,
+
+        use_pbc=True,
+        use_pbc_single=False,
+        otf_graph=True,
+
+        regress_forces=True,
+        regress_stress=False,
+        direct_prediction=True,
+
+        max_neighbors=20,
+        max_radius=12.0,
+        num_radial_basis=600,
+        max_num_elements=128,
+
+        num_layers=12,
+        num_channels=128,
+        attn_hidden_channels=64,
+        num_heads=8,
+        attn_alpha_channels=32,
+        attn_value_channels=16,
+        ffn_hidden_channels=128,
+        norm_type='merge_layer_norm',
+
+        lmax=6,
+        mmax=2,
+        attn_grid_resolution_list=[20, 8],
+        ffn_grid_resolution_list=[20, 20],
+
+        edge_channels=128,
+        use_atom_edge_embedding=True,
+        use_envelope=True,
+
+        attn_activation='sep-merge_gates2_swiglu',
+        use_attn_renorm=True,
+        use_add_merge=False,
+        use_rad_l_parametrization=True,
+        softcap=None,
+        attn_eps=1e-16,
+        ffn_activation='sep-merge_gates2_swiglu',
+        use_grid_mlp=True,
+
+        use_gate_force_head=True,
+
+        alpha_drop=0.0,
+        attn_mask_rate=0.0,
+        attn_weights_drop=0.1,
+        value_drop=0.0,
+        drop_path_rate=0.05,
+        proj_drop=0.0,
+        ffn_drop=0.0,
+        use_head_reg=False,
+
+        gradient_checkpointing_block_list=None,
+
+        avg_num_nodes=_AVG_NUM_NODES,
+        avg_degree=_AVG_DEGREE,
+
+        enforce_max_neighbors_strictly=True,
+        jd_path=None,
+    ):
+        super().__init__()
+
+        configure_jd_path(jd_path)
+
+        self.use_pbc = use_pbc
+        self.use_pbc_single = use_pbc_single
+        self.otf_graph = otf_graph
+
+        self.regress_forces = regress_forces
+        self.regress_stress = regress_stress
+        self.direct_prediction = direct_prediction
+        self.direct_forces = direct_prediction
+
+        self.max_neighbors = max_neighbors
+        self.max_radius = max_radius
+        self.cutoff = max_radius
+        self.num_radial_basis = num_radial_basis
+        self.max_num_elements = max_num_elements
+
+        self.num_layers = num_layers
+        self.num_channels = num_channels
+        self.attn_hidden_channels = attn_hidden_channels
+        self.num_heads = num_heads
+        self.attn_alpha_channels = attn_alpha_channels
+        self.attn_value_channels = attn_value_channels
+        self.ffn_hidden_channels = ffn_hidden_channels
+        self.norm_type = norm_type
+
+        self.lmax = lmax
+        self.mmax = mmax
+        self.attn_grid_resolution_list = attn_grid_resolution_list
+        self.ffn_grid_resolution_list = ffn_grid_resolution_list
+
+        self.edge_channels = edge_channels
+        self.use_atom_edge_embedding = use_atom_edge_embedding
+        self.use_envelope = use_envelope
+
+        self.attn_activation = attn_activation
+        self.use_attn_renorm = use_attn_renorm
+        self.use_add_merge = use_add_merge
+        self.use_rad_l_parametrization = use_rad_l_parametrization
+        self.softcap = softcap
+        self.attn_eps = attn_eps
+        self.ffn_activation = ffn_activation
+        self.use_grid_mlp = use_grid_mlp
+
+        self.use_gate_force_head = use_gate_force_head
+
+        self.alpha_drop = alpha_drop
+        self.attn_mask_rate = attn_mask_rate
+        self.attn_weights_drop = attn_weights_drop
+        self.value_drop = value_drop
+        self.drop_path_rate = drop_path_rate
+        self.proj_drop = proj_drop
+        self.ffn_drop = ffn_drop
+        self.use_head_reg = use_head_reg
+
+        self.gradient_checkpointing_block_list = gradient_checkpointing_block_list
+        if self.gradient_checkpointing_block_list is not None:
+            assert len(self.gradient_checkpointing_block_list) == self.num_layers
+        else:
+            self.gradient_checkpointing_block_list = [0] * self.num_layers
+
+        self.avg_num_nodes = avg_num_nodes
+        self.avg_degree = avg_degree
+
+        self.enforce_max_neighbors_strictly = enforce_max_neighbors_strictly
+
+        # Atom-type embedding
+        self.sphere_embedding = torch.nn.Embedding(self.max_num_elements, self.num_channels)
+
+        # Radial basis function
+        self.distance_expansion = GaussianSmearing(
+            0.0,
+            self.cutoff,
+            self.num_radial_basis,
+            2.0,
+        )
+        edge_input_channels = int(self.distance_expansion.num_output)
+
+        # The sizes of radial functions (input channels and 2 hidden channels)
+        self.edge_channels_list = [edge_input_channels] + [self.edge_channels] * 2
+
+        # Envelope function
+        self.envelope_func = PolynomialEnvelope(
+            cutoff=self.cutoff,
+            exponent=5
+        ) if self.use_envelope else None
+
+        # Computing Wigner-D matrices
+        self.so3_rotation = SO3Rotation(self.lmax, self.mmax, use_rotation_mask=(not self.direct_prediction))
+
+        # Edge-degree embedding
+        self.edge_degree_embedding = EdgeDegreeEmbedding(
+            num_channels=self.num_channels,
+            lmax=self.lmax,
+            mmax=self.mmax,
+            so3_rotation=self.so3_rotation,
+            max_num_elements=self.max_num_elements,
+            edge_channels_list=self.edge_channels_list,
+            use_atom_edge_embedding=self.use_atom_edge_embedding,
+            rescale_factor=self.avg_degree
+        )
+
+        # Transformer block
+        self.blocks = torch.nn.ModuleList()
+        for i in range(self.num_layers):
+            if self.gradient_checkpointing_block_list[i] == 1:
+                attn_activation = self.attn_activation.replace('_mem', '')
+                ffn_activation  = self.ffn_activation.replace('_mem', '')
+            else:
+                attn_activation = self.attn_activation
+                ffn_activation  = self.ffn_activation
+            block_config_dict = dict(
+                num_in_channels=self.num_channels,
+                attn_hidden_channels=self.attn_hidden_channels,
+                num_heads=self.num_heads,
+                attn_alpha_channels=self.attn_alpha_channels,
+                attn_value_channels=self.attn_value_channels,
+                ffn_hidden_channels=self.ffn_hidden_channels,
+                num_out_channels=self.num_channels,
+                lmax=self.lmax,
+                mmax=self.mmax,
+                so3_rotation=self.so3_rotation,
+                attn_grid_resolution_list=self.attn_grid_resolution_list,
+                ffn_grid_resolution_list=self.ffn_grid_resolution_list,
+                max_num_elements=self.max_num_elements,
+                edge_channels_list=self.edge_channels_list,
+                use_atom_edge_embedding=self.use_atom_edge_embedding,
+                attn_activation=attn_activation,
+                use_attn_renorm=self.use_attn_renorm,
+                use_add_merge=self.use_add_merge,
+                use_rad_l_parametrization=self.use_rad_l_parametrization,
+                softcap=self.softcap,
+                attn_eps=self.attn_eps,
+                ffn_activation=ffn_activation,
+                use_grid_mlp=self.use_grid_mlp,
+                norm_type=self.norm_type,
+                alpha_drop=self.alpha_drop,
+                attn_mask_rate=self.attn_mask_rate,
+                attn_weights_drop=attn_weights_drop,
+                value_drop=self.value_drop,
+                drop_path_rate=self.drop_path_rate,
+                proj_drop=self.proj_drop,
+                ffn_drop=self.ffn_drop
+            )
+            block_class = TransBlockV3
+            self.blocks.append(block_class(**block_config_dict))
+
+        # Output blocks for energy and forces (and optionally stress)
+        self.norm = get_normalization_layer(
+            self.norm_type,
+            lmax=self.lmax,
+            num_channels=self.num_channels
+        )
+        self.energy_block = ScalarFeedForwardNetwork(
+            num_in_channels=self.num_channels,
+            num_hidden_channels=self.ffn_hidden_channels,
+            num_out_channels=1,
+            dropout=0.0
+        )
+        if self.direct_prediction:
+            if self.regress_forces:
+                self.force_block = EquivariantGraphAttention(
+                    num_in_channels=self.num_channels,
+                    num_hidden_channels=self.attn_hidden_channels,
+                    num_heads=self.num_heads,
+                    attn_alpha_channels=self.attn_alpha_channels,
+                    attn_value_channels=self.attn_value_channels,
+                    num_out_channels=1,
+                    lmax=self.lmax,
+                    mmax=self.mmax,
+                    so3_rotation=self.so3_rotation,
+                    grid_resolution_list=self.attn_grid_resolution_list,
+                    max_num_elements=self.max_num_elements,
+                    edge_channels_list=self.edge_channels_list,
+                    use_atom_edge_embedding=self.use_atom_edge_embedding,
+                    activation=('sep_s2' if not self.use_gate_force_head else 'gate'),
+                    use_attn_renorm=self.use_attn_renorm,
+                    use_add_merge=self.use_add_merge,
+                    use_rad_l_parametrization=self.use_rad_l_parametrization,
+                    softcap=self.softcap,
+                    eps=self.attn_eps,
+                    alpha_drop=0.0,
+                    attn_mask_rate=0.0,
+                    attn_weights_drop=0.0,
+                    value_drop=0.0
+                )
+                if 'rms_norm' in norm_type:
+                    if self.force_block.alpha_norm is not None:
+                        del self.force_block.alpha_norm
+                        self.force_block.alpha_norm = RMSNorm(self.attn_alpha_channels)
+            if self.regress_stress:
+                self.stress_block = FeedForwardNetworkStressHead(
+                    num_in_channels=self.num_channels,
+                    num_hidden_channels=self.ffn_hidden_channels,
+                    num_out_channels=1,
+                    lmax=self.lmax,
+                    mmax=self.mmax,
+                    grid_resolution_list=self.ffn_grid_resolution_list,
+                    activation='gate',
+                    use_grid_mlp=False,
+                    dropout=0.0,
+                )
+
+        self.apply(self._init_weights)
+
+
+    def _forward_edge(
+        self,
+        edge_distance,
+        edge_distance_vec
+    ):
+        # Compute 3x3 rotation matrix per edge
+        edge_rot_mat = self._init_edge_rot_mat(edge_distance_vec)
+
+        # Compute Wigner-D matrices
+        self.so3_rotation.set_wigner(edge_rot_mat)
+
+        # Envelope function
+        edge_envelope_weight = self.envelope_func(edge_distance) if self.envelope_func is not None else None
+
+        # Radial basis function
+        edge_distance = self.distance_expansion(edge_distance)
+
+        return edge_distance, edge_envelope_weight
+
+
+    def _forward_embedding(
+        self,
+        atomic_numbers,
+        edge_distance,
+        edge_index,
+        edge_envelope_weight
+    ):
+        num_atoms = len(atomic_numbers)
+
+        # Initialize node embedding
+        x = torch.zeros(
+            (
+                num_atoms,
+                ((self.lmax + 1) ** 2),
+                self.num_channels
+            ),
+            device=self.device,
+            dtype=self.dtype
+        )
+
+        # Atom-type embedding
+        atom_embedding = self.sphere_embedding(atomic_numbers)
+        x[:, 0, :] = atom_embedding
+
+        # Edge-degree embedding
+        edge_degree = self.edge_degree_embedding(
+            atomic_numbers,
+            edge_distance,
+            edge_index,
+            edge_envelope_weight
+        )
+        x = x + edge_degree
+
+        return x
+
+
+    def _forward_blocks(
+        self,
+        x,
+        source_atomic_numbers,
+        target_atomic_numbers,
+        edge_distance,
+        edge_index,
+        edge_envelope_weight,
+        batch
+    ):
+        # Transformer blocks
+        for i in range(self.num_layers):
+            if self.gradient_checkpointing_block_list[i] == 0:
+                x = self.blocks[i](
+                    x,
+                    source_atomic_numbers,
+                    target_atomic_numbers,
+                    edge_distance,
+                    edge_index,
+                    edge_envelope_weight,
+                    batch,     # for GraphDropPath
+                )
+            elif self.gradient_checkpointing_block_list[i] == 1:
+                x = checkpoint(
+                    self.blocks[i],
+                    x,
+                    source_atomic_numbers,
+                    target_atomic_numbers,
+                    edge_distance,
+                    edge_index,
+                    edge_envelope_weight,
+                    batch,     # for GraphDropPath
+                    use_reentrant=False
+                )
+            else:
+                raise ValueError
+
+        # Final layer norm
+        x = self.norm(x)
+        x_scalar = x.narrow(1, 0, 1)
+        x_scalar = x_scalar.view(x_scalar.shape[0], self.num_channels)
+        return x_scalar, x
+
+
+    def _forward_direct(self, data):
+        self.batch_size = len(data.natoms)
+        self.dtype = data.pos.dtype
+        self.device = data.pos.device
+
+        (
+            edge_index,
+            edge_distance,
+            edge_distance_vec,
+            cell_offsets,
+            _,  # cell offset distances
+            neighbors,
+        ) = self.generate_graph(
+            data,
+            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
+            use_pbc_single=self.use_pbc_single
+        )
+
+        atomic_numbers = data.atomic_numbers.long()
+        source_atomic_numbers = atomic_numbers[edge_index[0]]
+        target_atomic_numbers = atomic_numbers[edge_index[1]]
+
+        edge_distance, edge_envelope_weight = self._forward_edge(edge_distance, edge_distance_vec)
+        x = self._forward_embedding(atomic_numbers, edge_distance, edge_index, edge_envelope_weight)
+        x_scalar, x = self._forward_blocks(
+            x,
+            source_atomic_numbers,
+            target_atomic_numbers,
+            edge_distance,
+            edge_index,
+            edge_envelope_weight,
+            data.batch
+        )
+
+        outputs = {}
+
+        # Energy prediction
+        node_energy = self.energy_block(x_scalar)
+        energy = torch.zeros(self.batch_size, device=node_energy.device, dtype=node_energy.dtype)
+        energy.index_add_(0, data.batch, node_energy.view(-1))
+        energy = energy / self.avg_num_nodes
+        outputs['energy'] = energy
+
+        # Force prediction
+        if self.regress_forces:
+            forces = self.force_block(
+                x,
+                source_atomic_numbers,
+                target_atomic_numbers,
+                edge_distance,
+                edge_index,
+                edge_envelope_weight
+            )
+            forces = forces.narrow(1, 1, 3)
+            forces = forces.view(-1, 3)
+            outputs['forces'] = forces
+
+        # Stress Prediction
+        if self.regress_stress:
+            stress = self.stress_block(
+                x,
+                batch_size=self.batch_size,
+                batch=data.batch
+            )
+            outputs['stress'] = stress
+
+        return outputs
+
+
+    @conditional_grad(torch.enable_grad())
+    def _forward_gradient(self, data):
+        """
+            1.  We have additional `@conditional_grad` as the decorator since the decorator might not be compatible with
+                `torch.compile()` in direct methods.
+        """
+        self.batch_size = len(data.natoms)
+        self.dtype = data.pos.dtype
+        self.device = data.pos.device
+
+        """
+            For gradient methods
+        """
+        displacement = None
+        orig_cell = None
+        if self.regress_stress and self.regress_forces:
+            displacement = torch.zeros(
+                (3, 3),
+                dtype=self.dtype,
+                device=self.device,
+            )
+            displacement = displacement.view(-1, 3, 3).expand(self.batch_size, 3, 3)
+            displacement.requires_grad = True
+            symmetric_displacement = 0.5 * (
+                displacement + displacement.transpose(-1, -2)
+            )
+
+            data.pos.requires_grad = True
+            data.pos = data.pos + torch.bmm(
+                data.pos.unsqueeze(-2),
+                torch.index_select(symmetric_displacement, 0, data.batch)
+            ).squeeze(-2)
+
+            orig_cell = data.cell
+            data.cell = data.cell + torch.bmm(
+                data.cell, symmetric_displacement
+            )
+        elif not self.regress_stress and self.regress_forces:
+            data.pos.requires_grad = True
+
+        (
+            edge_index,
+            edge_distance,
+            edge_distance_vec,
+            cell_offsets,
+            _,  # cell offset distances
+            neighbors,
+        ) = self.generate_graph(
+            data,
+            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
+            use_pbc_single=self.use_pbc_single
+        )
+
+        atomic_numbers = data.atomic_numbers.long()
+        source_atomic_numbers = atomic_numbers[edge_index[0]]
+        target_atomic_numbers = atomic_numbers[edge_index[1]]
+
+        edge_distance, edge_envelope_weight = self._forward_edge(edge_distance, edge_distance_vec)
+        x = self._forward_embedding(atomic_numbers, edge_distance, edge_index, edge_envelope_weight)
+        x_scalar, x = self._forward_blocks(
+            x,
+            source_atomic_numbers,
+            target_atomic_numbers,
+            edge_distance,
+            edge_index,
+            edge_envelope_weight,
+            data.batch
+        )
+
+        outputs = {}
+
+        # Energy prediction
+        node_energy = self.energy_block(x_scalar)
+        energy = torch.zeros(self.batch_size, device=node_energy.device, dtype=node_energy.dtype)
+        energy.index_add_(0, data.batch, node_energy.view(-1))
+        energy = energy / self.avg_num_nodes
+        outputs['energy'] = energy
+
+        if self.regress_stress and self.regress_forces:
+            # Stress and forces prediction
+            grads = torch.autograd.grad(
+                [energy.sum()],
+                [data.pos, displacement],
+                create_graph=self.training,
+            )
+            forces = torch.neg(grads[0])
+            virial = grads[1].view(-1, 3, 3)
+            volume = torch.det(data.cell).abs().unsqueeze(-1)
+            stress = virial / volume.view(-1, 1, 1)
+            virial = torch.neg(virial)
+            outputs['forces'] = forces
+            outputs['stress'] = stress.view(-1, 9)
+            data.cell = orig_cell
+
+        elif not self.regress_stress and self.regress_forces:
+            forces = (
+                -1
+                * torch.autograd.grad(
+                    energy.sum(), data.pos, create_graph=self.training
+                )[0]
+            )
+            outputs['forces'] = forces
+
+        return outputs
+
+
+    def forward(self, data):
+        if self.direct_prediction:
+            outputs = self._forward_direct(data)
+        else:
+            outputs = self._forward_gradient(data)
+        return outputs
+
+
+    # Initialize the edge rotation matrics
+    def _init_edge_rot_mat(self, edge_distance_vec):
+        return init_edge_rot_mat(edge_distance_vec, use_rotation_mask=(not self.direct_prediction))
+
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+
+    def _init_weights(self, m):
+        if (isinstance(m, torch.nn.Linear)
+            or isinstance(m, SO3Linear)
+        ):
+            if m.bias is not None:
+                torch.nn.init.constant_(m.bias, 0)
+        elif isinstance(m, torch.nn.LayerNorm):
+            torch.nn.init.constant_(m.bias, 0)
+            torch.nn.init.constant_(m.weight, 1.0)
+        elif (isinstance(m, RadialFunction)):
+            m.apply(self._uniform_init_linear_weights)
+
+
+    def _uniform_init_linear_weights(self, m):
+        if isinstance(m, torch.nn.Linear):
+            if m.bias is not None:
+                torch.nn.init.constant_(m.bias, 0)
+            std = 1 / math.sqrt(m.in_features)
+            torch.nn.init.uniform_(m.weight, -std, std)
+
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        no_wd_list = []
+        named_parameters_list = [name for name, _ in self.named_parameters()]
+        for module_name, module in self.named_modules():
+            if (isinstance(module, torch.nn.Embedding)
+                or isinstance(module, torch.nn.Linear)
+                or isinstance(module, SO3Linear)
+                or isinstance(module, torch.nn.LayerNorm)
+                or isinstance(module, RMSNorm)
+                or isinstance(module, EquivariantLayerNorm)
+                or isinstance(module, EquivariantSeparableLayerNorm)
+                or isinstance(module, EquivariantMergeLayerNorm)
+            ):
+                for parameter_name, _ in module.named_parameters():
+                    if (isinstance(module, torch.nn.Linear)
+                        or isinstance(module, SO3Linear)
+                    ):
+                        if 'weight' in parameter_name:
+                            continue
+                    global_parameter_name = module_name + '.' + parameter_name
+                    assert global_parameter_name in named_parameters_list
+                    no_wd_list.append(global_parameter_name)
+        return set(no_wd_list)
+
+
+    @torch._dynamo.disable
+    def generate_graph(
+        self,
+        data,
+        cutoff=None,
+        max_neighbors=None,
+        use_pbc=None,
+        otf_graph=None,
+        enforce_max_neighbors_strictly=None,
+        use_pbc_single=False,
+    ):
+        graph_data = super().generate_graph(
+            data,
+            cutoff=cutoff,
+            max_neighbors=max_neighbors,
+            use_pbc=use_pbc,
+            otf_graph=otf_graph,
+            enforce_max_neighbors_strictly=enforce_max_neighbors_strictly,
+            use_pbc_single=use_pbc_single,
+        )
+
+        edge_index   = graph_data.edge_index
+        edge_dist    = graph_data.edge_distance
+        distance_vec = graph_data.edge_distance_vec
+        cell_offsets = graph_data.cell_offsets
+        cell_offset_distances = graph_data.offset_distances
+        neighbors    = graph_data.neighbors
+
+        return (
+            edge_index,
+            edge_dist,
+            distance_vec,
+            cell_offsets,
+            cell_offset_distances,
+            neighbors,
+        )

model/equiformer_v3_dens.py

@@ -0,0 +1,502 @@
+import math
+
+import e3nn
+import torch
+
+from onescience.modules.layer.equiformer_v3.layer_norm import (
+    RMSNorm,
+    get_normalization_layer,
+)
+from onescience.modules.layer.equiformer_v3.so3 import SO3Linear
+from onescience.modules.layer.equiformer_v3.transformer_block import (
+    EquivariantGraphAttention,
+    FeedForwardNetwork,
+)
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.common.utils import conditional_grad
+
+from .equiformer_v3 import EquiformerV3_OC
+
+
+# Statistics of IS2RE 100K
+_AVG_NUM_NODES = 77.81317
+_AVG_DEGREE = 23.395238876342773    # IS2RE: 100k, max_radius = 5, max_neighbors = 100
+
+
+@registry.register_model("equiformer_v3_dens")
+class EquiformerV3DeNS_OC(EquiformerV3_OC):
+    """
+    Args:
+        use_pbc (bool):         Use periodic boundary conditions
+        use_pbc_single (bool, optional):    Process batch PBC graphs one at a time.
+                                            This argument is set to True when training on MPTraj.
+        otf_graph (bool):       Compute graph On The Fly (OTF)
+
+        regress_forces (bool):  Compute forces
+        regress_stress (bool):  Compute stress
+        direct_prediction (bool):   Whether to use direct methods to predict forces and stress
+
+        max_neighbors (int):    Maximum number of neighbors per atom
+        max_radius (float):     Maximum distance between nieghboring atoms in Angstroms
+        num_radial_basis (int): Number of radial basis functions
+        max_num_elements (int): Maximum atomic number
+
+        num_layers (int):           Number of layers in the GNN
+        num_channels (int):         Number of channels in node embeddings
+        attn_hidden_channels (int): Number of hidden channels in equivariant graph attention
+        num_heads (int):            Number of attention heads
+        attn_alpha_channels (int):  Number of channels for alpha vector in each attention head
+        attn_value_channels (int):  Number of channels for value vector in each attention head
+        ffn_hidden_channels (int):  Number of hidden channels in feedforward network
+        norm_type (str):            Type of normalization layer
+                                    (['sep_layer_norm', 'sep_layer_norm_liger', 'merge_layer_norm',
+                                    'merge_layer_norm_attn_rms_norm', 'merge_rms_norm'])
+
+        lmax (int):                 Maximum degrees (l)
+        mmax (int):                 Maximum order (m)
+        attn_grid_resolution_list (list:int):
+                                    Grid resolution list in class `SO3Grid` in attention
+        ffn_grid_resolution_list (list:int):
+                                    Grid resolution list in class `SO3Grid` in feedforward network
+
+        edge_channels (int):                Number of channels for edge-wise invariant features
+        use_atom_edge_embedding (bool):     Whether to use atomic embedding along with relative distance for edge scalar features
+        use_envelope (bool):        Whether to apply an envelope function to attention
+
+        attn_activation (str):      Type of activation function in equivariant graph attention
+        use_attn_renorm (bool):     Whether to re-normalize attention weights
+        use_add_merge (bool):       Default: False
+                                    If True, use addition to merge the source/target node features instead of concat,
+                                    which can save 2x compute when rotating with Wigner-D matrices.
+        use_rad_l_parametrization (bool):
+                                    Default: True
+                                    If True, all the m components within the same type-L vector will share the same
+                                    weight from the radial function.
+        softcap (float):            Default: None
+                                    If not None, use soft capping to limit the range of attention logits to
+                                    [- `softcap`, + `softcap`].
+        attn_eps (float):           Default: 1e-16
+                                    Epsilon value used in the softmax operation of attention
+        ffn_activation (str):       Type of activation function for feedforward network
+        use_grid_mlp (bool):        If `True`, use projecting to grids and performing MLPs for FFNs.
+
+        use_gate_force_head (bool): If `True`, use `GateActivation` in the equivariant attention of the force prediction head.
+
+        alpha_drop (float):         Dropout rate for the hidden features in non-linear MLP attention
+        attn_mask_rate (float):     Mask rate for neighbors considered in attention
+        attn_weights_drop (float):  Dropout rate for attention weights
+        value_drop (float):         Dropout rate for the hidden features in non-linear value vectors
+        drop_path_rate (float):     Drop path rate
+        proj_drop (float):          Dropout rate for outputs of attention and FFN in Transformer blocks
+        ffn_drop (float):           Dropout rate for the hidden features in FFN
+        use_head_reg (bool):        Whether to apply regularization to output head (dummy argument for backend compatibility)
+
+        gradient_checkpointing_block_list (list):
+                                    A list indicating which block we apply gradient/activation checkpointing to save memory.
+
+        avg_num_nodes (float):   Normalization factor for sum aggregation over nodes
+        avg_degree (float):      Normalization factor for sum aggregation over edges
+
+        enforce_max_neighbors_strictly (bool):      When edges are subselected based on the `max_neighbors` arg, arbitrarily select amongst equidistant / degenerate edges to have exactly the correct number.
+    """
+    def __init__(
+        self,
+
+        use_pbc=True,
+        use_pbc_single=False,
+        otf_graph=True,
+
+        regress_forces=True,
+        regress_stress=False,
+        direct_prediction=True,
+
+        max_neighbors=20,
+        max_radius=12.0,
+        num_radial_basis=600,
+        max_num_elements=128,
+
+        num_layers=12,
+        num_channels=128,
+        attn_hidden_channels=64,
+        num_heads=8,
+        attn_alpha_channels=32,
+        attn_value_channels=16,
+        ffn_hidden_channels=128,
+        norm_type='merge_layer_norm',
+
+        lmax=6,
+        mmax=2,
+        attn_grid_resolution_list=[20, 8],
+        ffn_grid_resolution_list=[20, 20],
+
+        edge_channels=128,
+        use_atom_edge_embedding=True,
+        use_envelope=True,
+
+        attn_activation='sep-merge_gates2_swiglu',
+        use_attn_renorm=True,
+        use_add_merge=False,
+        use_rad_l_parametrization=True,
+        softcap=None,
+        attn_eps=1e-16,
+        ffn_activation='sep-merge_gates2_swiglu',
+        use_grid_mlp=True,
+
+        use_gate_force_head=True,
+
+        alpha_drop=0.0,
+        attn_mask_rate=0.0,
+        attn_weights_drop=0.1,
+        value_drop=0.0,
+        drop_path_rate=0.05,
+        proj_drop=0.0,
+        ffn_drop=0.0,
+        use_head_reg=False,
+
+        gradient_checkpointing_block_list=None,
+
+        avg_num_nodes=_AVG_NUM_NODES,
+        avg_degree=_AVG_DEGREE,
+
+        enforce_max_neighbors_strictly=True,
+        jd_path=None,
+    ):
+        super().__init__(
+            use_pbc=use_pbc,
+            use_pbc_single=use_pbc_single,
+            otf_graph=otf_graph,
+            regress_forces=regress_forces,
+            regress_stress=regress_stress,
+            direct_prediction=direct_prediction,
+            max_neighbors=max_neighbors,
+            max_radius=max_radius,
+            num_radial_basis=num_radial_basis,
+            max_num_elements=max_num_elements,
+            num_layers=num_layers,
+            num_channels=num_channels,
+            attn_hidden_channels=attn_hidden_channels,
+            num_heads=num_heads,
+            attn_alpha_channels=attn_alpha_channels,
+            attn_value_channels=attn_value_channels,
+            ffn_hidden_channels=ffn_hidden_channels,
+            norm_type=norm_type,
+            lmax=lmax,
+            mmax=mmax,
+            attn_grid_resolution_list=attn_grid_resolution_list,
+            ffn_grid_resolution_list=ffn_grid_resolution_list,
+            edge_channels=edge_channels,
+            use_atom_edge_embedding=use_atom_edge_embedding,
+            use_envelope=use_envelope,
+            attn_activation=attn_activation,
+            use_attn_renorm=use_attn_renorm,
+            use_add_merge=use_add_merge,
+            use_rad_l_parametrization=use_rad_l_parametrization,
+            softcap=softcap,
+            attn_eps=attn_eps,
+            ffn_activation=ffn_activation,
+            use_grid_mlp=use_grid_mlp,
+            use_gate_force_head=use_gate_force_head,
+            alpha_drop=alpha_drop,
+            attn_mask_rate=attn_mask_rate,
+            attn_weights_drop=attn_weights_drop,
+            value_drop=value_drop,
+            drop_path_rate=drop_path_rate,
+            proj_drop=proj_drop,
+            ffn_drop=ffn_drop,
+            use_head_reg=use_head_reg,
+            gradient_checkpointing_block_list=gradient_checkpointing_block_list,
+            avg_num_nodes=avg_num_nodes,
+            avg_degree=avg_degree,
+            enforce_max_neighbors_strictly=enforce_max_neighbors_strictly,
+            jd_path=jd_path,
+        )
+
+        # Force encoding
+        self.irreps_sh = e3nn.o3.Irreps.spherical_harmonics(lmax=self.lmax, p=1)
+        self.force_embedding = SO3Linear(
+            in_features=1,
+            out_features=self.num_channels,
+            lmax=self.lmax,
+            bias=True
+        )
+
+        if self.regress_forces:
+            self.dens_block = EquivariantGraphAttention(
+                num_in_channels=self.num_channels,
+                num_hidden_channels=self.attn_hidden_channels,
+                num_heads=self.num_heads,
+                attn_alpha_channels=self.attn_alpha_channels,
+                attn_value_channels=self.attn_value_channels,
+                num_out_channels=1,
+                lmax=self.lmax,
+                mmax=self.mmax,
+                so3_rotation=self.so3_rotation,
+                grid_resolution_list=self.attn_grid_resolution_list,
+                max_num_elements=self.max_num_elements,
+                edge_channels_list=self.edge_channels_list,
+                use_atom_edge_embedding=self.use_atom_edge_embedding,
+                activation=('sep_s2' if not self.use_gate_force_head else 'gate'),
+                use_attn_renorm=self.use_attn_renorm,
+                use_add_merge=self.use_add_merge,
+                use_rad_l_parametrization=self.use_rad_l_parametrization,
+                softcap=self.softcap,
+                eps=self.attn_eps,
+                alpha_drop=0.0,
+                attn_mask_rate=0.0,
+                attn_weights_drop=0.0,
+                value_drop=0.0
+            )
+            if 'rms_norm' in norm_type:
+                if self.dens_block.alpha_norm is not None:
+                    del self.dens_block.alpha_norm
+                    self.dens_block.alpha_norm = RMSNorm(self.attn_alpha_channels)
+
+        self.apply(self._init_weights)
+
+
+    def _forward_direct(self, data):
+        self.batch_size = len(data.natoms)
+        self.dtype = data.pos.dtype
+        self.device = data.pos.device
+
+        (
+            edge_index,
+            edge_distance,
+            edge_distance_vec,
+            cell_offsets,
+            _,  # cell offset distances
+            neighbors,
+        ) = self.generate_graph(
+            data,
+            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
+            use_pbc_single=self.use_pbc_single
+        )
+
+        atomic_numbers = data.atomic_numbers.long()
+        source_atomic_numbers = atomic_numbers[edge_index[0]]
+        target_atomic_numbers = atomic_numbers[edge_index[1]]
+
+        edge_distance, edge_envelope_weight = self._forward_edge(edge_distance, edge_distance_vec)
+        x = self._forward_embedding(atomic_numbers, edge_distance, edge_index, edge_envelope_weight)
+        force_embedding, noise_mask_tensor, dens_batch_mask_tensor, dens_mask_tensor = self._forward_dens_force_encoding(data)
+        x = x + force_embedding
+        x_scalar, x = self._forward_blocks(
+            x,
+            source_atomic_numbers,
+            target_atomic_numbers,
+            edge_distance,
+            edge_index,
+            edge_envelope_weight,
+            data.batch
+        )
+
+        outputs = {}
+
+        # Energy prediction
+        node_energy = self.energy_block(x_scalar)
+        energy = torch.zeros(self.batch_size, device=node_energy.device, dtype=node_energy.dtype)
+        energy.index_add_(0, data.batch, node_energy.view(-1))
+        energy = energy / self.avg_num_nodes
+        outputs['energy'] = energy
+
+        # Force prediction
+        if self.regress_forces:
+            forces = self.force_block(
+                x,
+                source_atomic_numbers,
+                target_atomic_numbers,
+                edge_distance,
+                edge_index,
+                edge_envelope_weight
+            )
+            forces = forces.narrow(1, 1, 3)
+            forces = forces.view(-1, 3)
+
+            # for DeNS
+            denoising_pos_vec = self.dens_block(
+                x,
+                source_atomic_numbers,
+                target_atomic_numbers,
+                edge_distance,
+                edge_index,
+                edge_envelope_weight
+            )
+            denoising_pos_vec = denoising_pos_vec.narrow(1, 1, 3)
+            denoising_pos_vec = denoising_pos_vec.view(-1, 3)
+
+            outputs['forces'] = forces * (~noise_mask_tensor) + denoising_pos_vec * noise_mask_tensor
+
+        # Stress Prediction
+        if self.regress_stress:
+            stress = self.stress_block(
+                x,
+                batch_size=self.batch_size,
+                batch=data.batch
+            )
+            outputs['stress'] = stress * (~dens_batch_mask_tensor) # not predict stress during DeNS
+
+        return outputs
+
+
+    @conditional_grad(torch.enable_grad())
+    def _forward_gradient(self, data):
+        """
+            1.  We have additional `@conditional_grad` as the decorator since the decorator might not be compatible with
+                `torch.compile()` in direct methods.
+        """
+        self.batch_size = len(data.natoms)
+        self.dtype = data.pos.dtype
+        self.device = data.pos.device
+
+        """
+            For gradient methods
+        """
+        displacement = None
+        orig_cell = None
+        if self.regress_stress and self.regress_forces:
+            displacement = torch.zeros(
+                (3, 3),
+                dtype=self.dtype,
+                device=self.device,
+            )
+            displacement = displacement.view(-1, 3, 3).expand(self.batch_size, 3, 3)
+            displacement.requires_grad = True
+            symmetric_displacement = 0.5 * (
+                displacement + displacement.transpose(-1, -2)
+            )
+
+            data.pos.requires_grad = True
+            data.pos = data.pos + torch.bmm(
+                data.pos.unsqueeze(-2),
+                torch.index_select(symmetric_displacement, 0, data.batch)
+            ).squeeze(-2)
+
+            orig_cell = data.cell
+            data.cell = data.cell + torch.bmm(
+                data.cell, symmetric_displacement
+            )
+        elif not self.regress_stress and self.regress_forces:
+            data.pos.requires_grad = True
+
+        (
+            edge_index,
+            edge_distance,
+            edge_distance_vec,
+            cell_offsets,
+            _,  # cell offset distances
+            neighbors,
+        ) = self.generate_graph(
+            data,
+            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
+            use_pbc_single=self.use_pbc_single
+        )
+
+        atomic_numbers = data.atomic_numbers.long()
+        source_atomic_numbers = atomic_numbers[edge_index[0]]
+        target_atomic_numbers = atomic_numbers[edge_index[1]]
+
+        edge_distance, edge_envelope_weight = self._forward_edge(edge_distance, edge_distance_vec)
+        x = self._forward_embedding(atomic_numbers, edge_distance, edge_index, edge_envelope_weight)
+        force_embedding, noise_mask_tensor, dens_batch_mask_tensor, dens_mask_tensor = self._forward_dens_force_encoding(data)
+        x = x + force_embedding
+        x_scalar, x = self._forward_blocks(
+            x,
+            source_atomic_numbers,
+            target_atomic_numbers,
+            edge_distance,
+            edge_index,
+            edge_envelope_weight,
+            data.batch
+        )
+
+        outputs = {}
+
+        # Energy prediction
+        node_energy = self.energy_block(x_scalar)
+        energy = torch.zeros(self.batch_size, device=node_energy.device, dtype=node_energy.dtype)
+        energy.index_add_(0, data.batch, node_energy.view(-1))
+        energy = energy / self.avg_num_nodes
+        outputs['energy'] = energy
+
+        if self.regress_stress and self.regress_forces:
+            # Stress and forces prediction
+            grads = torch.autograd.grad(
+                [energy.sum()],
+                [data.pos, displacement],
+                create_graph=self.training,
+            )
+            forces = torch.neg(grads[0])
+            virial = grads[1].view(-1, 3, 3)
+            volume = torch.det(data.cell).abs().unsqueeze(-1)
+            stress = virial / volume.view(-1, 1, 1)
+            virial = torch.neg(virial)
+            outputs['forces'] = forces
+            outputs['stress'] = stress.view(-1, 9) * (~dens_batch_mask_tensor) # not predict stress during DeNS
+            data.cell = orig_cell
+        elif not self.regress_stress and self.regress_forces:
+            forces = (
+                -1
+                * torch.autograd.grad(
+                    energy.sum(), data.pos, create_graph=self.training
+                )[0]
+            )
+            outputs['forces'] = forces
+
+        # for DeNS
+        if self.regress_forces:
+            denoising_pos_vec = self.dens_block(
+                x,
+                source_atomic_numbers,
+                target_atomic_numbers,
+                edge_distance,
+                edge_index,
+                edge_envelope_weight
+            )
+            denoising_pos_vec = denoising_pos_vec.narrow(1, 1, 3)
+            denoising_pos_vec = denoising_pos_vec.view(-1, 3)
+            outputs['forces'] = forces * (~noise_mask_tensor) + denoising_pos_vec * noise_mask_tensor # not predict forces during DeNS
+
+        return outputs
+
+
+    @torch._dynamo.disable
+    def _generate_dens_data(self, data):
+        num_atoms = len(data.atomic_numbers)
+        if hasattr(data, "denoising_pos_forward") and data.denoising_pos_forward:
+            force_data = data.forces
+            if hasattr(data, "noise_mask"):
+                noise_mask_tensor = data.noise_mask.view(-1, 1)
+            else:
+                noise_mask_tensor = torch.ones((num_atoms, 1), dtype=torch.bool, device=self.device)
+            if hasattr(data, "dens_batch_mask"):
+                dens_batch_mask_tensor = data.dens_batch_mask.view(-1, 1)
+            else:
+                dens_batch_mask_tensor = torch.ones((len(data.natoms), 1), dtype=torch.bool, device=self.device)
+            dens_mask_tensor = torch.ones((1, 1), dtype=torch.bool, device=self.device)
+        else:
+            force_data = torch.zeros((num_atoms, 3), dtype=self.dtype, device=self.device)
+            noise_mask_tensor = torch.zeros((num_atoms, 1), dtype=torch.bool, device=self.device)
+            dens_batch_mask_tensor = torch.zeros((len(data.natoms), 1), dtype=torch.bool, device=self.device)
+            dens_mask_tensor = torch.zeros((1, 1), dtype=torch.bool, device=self.device)
+        force_sh = e3nn.o3.spherical_harmonics(
+            l=self.irreps_sh,
+            x=force_data,
+            normalize=True,
+            normalization='component'
+        )
+
+        return force_data, force_sh, noise_mask_tensor, dens_batch_mask_tensor, dens_mask_tensor
+
+
+    def _forward_dens_force_encoding(self, data):
+        force_data, force_sh, noise_mask_tensor, dens_batch_mask_tensor, dens_mask_tensor = self._generate_dens_data(data)
+        force_norm = force_data.norm(dim=-1, keepdim=True)
+        force_norm = force_norm / math.sqrt(3.0)
+        force_embedding = force_sh * force_norm
+        force_embedding = force_embedding.view(force_embedding.shape[0], -1, 1)
+        force_embedding = self.force_embedding(force_embedding)
+        noise_mask_tensor = noise_mask_tensor.view(-1, 1, 1)
+        force_embedding = force_embedding * noise_mask_tensor
+        noise_mask_tensor = noise_mask_tensor.view(-1, 1)
+
+        return force_embedding, noise_mask_tensor, dens_batch_mask_tensor, dens_mask_tensor

phonons.py

@@ -0,0 +1,103 @@
+"""Run finite-displacement phonons with Equiformer V3 and ASE."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+from ase.build import bulk
+from ase.io import read, write
+
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3Calculator,
+    relax_structure,
+    run_phonon_workflow,
+    write_workflow_result,
+)
+
+
+def default_checkpoint() -> str:
+    return str(
+        Path(__file__).resolve().parent
+        / "weight"
+        / "omat24-mptrj-salex_gradient.pt"
+    )
+
+
+def load_structure(path: str | None):
+    if path:
+        return read(path)
+    return bulk("Cu")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint")
+    parser.add_argument(
+        "--input",
+        help="periodic structure readable by ASE; defaults to built-in bulk Cu",
+    )
+    parser.add_argument(
+        "--relax",
+        action="store_true",
+        help="relax the cell and atomic positions before finite displacements",
+    )
+    parser.add_argument("--relax-fmax", type=float, default=0.01)
+    parser.add_argument("--relax-steps", type=int, default=200)
+    parser.add_argument("--supercell", type=int, nargs=3, default=(3, 3, 3))
+    parser.add_argument("--delta", type=float, default=0.01)
+    parser.add_argument("--bandpath", help="ASE band path, for example GXWKGL")
+    parser.add_argument("--band-points", type=int, default=100)
+    parser.add_argument("--dos-kpts", type=int, nargs=3, default=(10, 10, 10))
+    parser.add_argument("--dos-points", type=int, default=400)
+    parser.add_argument("--dos-width", type=float, default=0.001, help="eV")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--output-dir", default="outputs/phonons")
+    args = parser.parse_args()
+
+    checkpoint = args.checkpoint or default_checkpoint()
+    calculator = EquiformerV3Calculator.from_checkpoint(
+        checkpoint, device=args.device
+    )
+    output_dir = Path(args.output_dir)
+    structure = load_structure(args.input)
+    relaxation = None
+    if args.relax:
+        structure, relaxation = relax_structure(
+            structure,
+            calculator,
+            relax_cell=True,
+            fmax=args.relax_fmax,
+            steps=args.relax_steps,
+        )
+        output_dir.mkdir(parents=True, exist_ok=True)
+        write(output_dir / "relaxed.cif", structure)
+    result = run_phonon_workflow(
+        structure,
+        calculator,
+        output_dir,
+        supercell=args.supercell,
+        delta=args.delta,
+        bandpath=args.bandpath,
+        band_points=args.band_points,
+        dos_kpts=args.dos_kpts,
+        dos_points=args.dos_points,
+        dos_width_ev=args.dos_width,
+    )
+    result["checkpoint"] = str(checkpoint)
+    result["relaxation"] = relaxation
+    output = write_workflow_result(result, output_dir / "phonons.json")
+
+    print("formula:", result["formula"])
+    print("minimum phonon energy (eV):", result["minimum_band_energy_ev"])
+    print("result:", output)
+
+
+if __name__ == "__main__":
+    main()

prepare_mptrj_dataset.py

@@ -0,0 +1,456 @@
+"""Convert an MPtrj JSON file into ASE-LMDB shards for Equiformer V3."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import multiprocessing as mp
+import os
+import time
+from collections import deque
+from concurrent.futures import Future, ProcessPoolExecutor
+from itertools import islice
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Iterable, Iterator
+
+import ase
+import numpy as np
+import torch
+from ase import Atoms
+from ase.calculators.singlepoint import SinglePointCalculator
+from ase.db import connect
+
+from onescience.modules.func_utils.uma_graph.radius_graph_pbc import radius_graph_pbc
+
+
+def _to_atoms(identifier: str, record: dict) -> Atoms:
+    """Convert one MPtrj record while preserving FairChem label conventions."""
+
+    structure = record["structure"]
+    sites = structure["sites"]
+    numbers = [
+        ase.data.atomic_numbers[site["species"][0]["element"]] for site in sites
+    ]
+    positions = [site["xyz"] for site in sites]
+    atoms = Atoms(
+        numbers,
+        positions,
+        cell=structure["lattice"]["matrix"],
+        pbc=True,
+    )
+
+    # MPtrj stress is reported in kbar; retain the upstream ASE sign and units.
+    stress = np.asarray(record["stress"], dtype=np.float32)
+    stress = stress * (-0.1 * ase.units.GPa)
+    energy = record["uncorrected_total_energy"]
+    atoms.calc = SinglePointCalculator(
+        atoms,
+        energy=energy,
+        free_energy=energy,
+        forces=record["force"],
+        stress=stress,
+    )
+    atoms.info["sid"] = identifier
+    return atoms
+
+
+def _has_no_isolated_atoms(
+    atoms: Atoms, cutoff: float, max_neighbors: int
+) -> bool:
+    # Match the upstream AtomsToGraphs + radius_graph_pbc(..., True) path.
+    data = SimpleNamespace(
+        pos=torch.as_tensor(atoms.positions, dtype=torch.float32),
+        cell=torch.as_tensor(atoms.cell.array, dtype=torch.float32).view(1, 3, 3),
+        pbc=torch.as_tensor(atoms.pbc, dtype=torch.bool).view(1, 3),
+        natoms=torch.tensor([len(atoms)], dtype=torch.long),
+    )
+    edge_index, _, _ = radius_graph_pbc(
+        data,
+        cutoff,
+        max_neighbors,
+        True,
+        pbc=data.pbc[0],
+    )
+    counts = torch.bincount(edge_index[1], minlength=len(atoms))
+    return bool(torch.all(counts > 0))
+
+
+def _init_worker() -> None:
+    """Keep each conversion worker single-threaded to avoid CPU oversubscription."""
+
+    os.environ["OMP_NUM_THREADS"] = "1"
+    os.environ["MKL_NUM_THREADS"] = "1"
+    torch.set_num_threads(1)
+
+
+def _worker_ready() -> None:
+    """Start the process pool before the parent opens an ASE database."""
+
+
+def _convert_batch(
+    batch: list[tuple[str, dict]], cutoff: float, max_neighbors: int
+) -> list[Atoms | None]:
+    converted: list[Atoms | None] = []
+    for identifier, record in batch:
+        atoms = _to_atoms(identifier, record)
+        converted.append(
+            atoms if _has_no_isolated_atoms(atoms, cutoff, max_neighbors) else None
+        )
+    return converted
+
+
+def _iter_batches(
+    records: Iterable[tuple[str, dict]], batch_size: int
+) -> Iterator[list[tuple[str, dict]]]:
+    iterator = iter(records)
+    while batch := list(islice(iterator, batch_size)):
+        yield batch
+
+
+def _convert_batches_ordered(
+    records: Iterable[tuple[str, dict]],
+    cutoff: float,
+    max_neighbors: int,
+    workers: int,
+    batch_size: int,
+    executor: ProcessPoolExecutor | None,
+) -> Iterator[list[Atoms | None]]:
+    batches = _iter_batches(records, batch_size)
+    if executor is None:
+        for batch in batches:
+            yield _convert_batch(batch, cutoff, max_neighbors)
+        return
+
+    pending: deque[Future[list[Atoms | None]]] = deque()
+    max_pending = workers * 2
+    for batch in batches:
+        pending.append(
+            executor.submit(_convert_batch, batch, cutoff, max_neighbors)
+        )
+        if len(pending) >= max_pending:
+            yield pending.popleft().result()
+    while pending:
+        yield pending.popleft().result()
+
+
+class _JSONStream:
+    """Incrementally decode JSON values without loading the full MPtrj file."""
+
+    def __init__(self, path: Path, chunk_size: int = 1 << 20):
+        self.handle = path.open("r", encoding="utf-8")
+        self.chunk_size = chunk_size
+        self.buffer = ""
+        self.decoder = json.JSONDecoder()
+        self.eof = False
+
+    def close(self) -> None:
+        self.handle.close()
+
+    def _fill(self) -> None:
+        if not self.eof:
+            chunk = self.handle.read(self.chunk_size)
+            if chunk:
+                self.buffer += chunk
+            else:
+                self.eof = True
+
+    def _skip_whitespace(self) -> None:
+        while True:
+            stripped = self.buffer.lstrip()
+            if stripped:
+                self.buffer = stripped
+                return
+            self._fill()
+            if not self.buffer and self.eof:
+                raise EOFError("unexpected end of JSON input")
+
+    def _take(self, token: str) -> None:
+        self._skip_whitespace()
+        if not self.buffer.startswith(token):
+            raise ValueError(f"expected {token!r} in MPtrj JSON")
+        self.buffer = self.buffer[len(token) :]
+
+    def value(self):
+        """Read one complete JSON value, filling until the decoder succeeds."""
+
+        self._skip_whitespace()
+        while True:
+            try:
+                value, end = self.decoder.raw_decode(self.buffer)
+            except json.JSONDecodeError:
+                if self.eof:
+                    raise
+                self._fill()
+                continue
+            self.buffer = self.buffer[end:]
+            return value
+
+    def iter_group_records(self) -> Iterator[tuple[str, dict]]:
+        """Yield ``(record_id, record)`` pairs from the current group object."""
+
+        self._take("{")
+        self._skip_whitespace()
+        if self.buffer.startswith("}"):
+            self.buffer = self.buffer[1:]
+            return
+        while True:
+            record_id = self.value()
+            if not isinstance(record_id, str):
+                raise ValueError("MPtrj record identifiers must be strings")
+            self._take(":")
+            record = self.value()
+            if not isinstance(record, dict):
+                raise ValueError("MPtrj records must be JSON objects")
+            yield record_id, record
+            self._skip_whitespace()
+            if self.buffer.startswith(","):
+                self.buffer = self.buffer[1:]
+                continue
+            self._take("}")
+            return
+
+    def iter_records(self) -> Iterator[tuple[str, dict]]:
+        """Yield all records from the two-level MPtrj root object."""
+
+        self._take("{")
+        self._skip_whitespace()
+        if self.buffer.startswith("}"):
+            self.buffer = self.buffer[1:]
+            return
+        while True:
+            group_id = self.value()
+            if not isinstance(group_id, str):
+                raise ValueError("MPtrj group identifiers must be strings")
+            self._take(":")
+            for record_id, record in self.iter_group_records():
+                # Upstream discards the outer group key and uses this ID as sid.
+                yield record_id, record
+            self._skip_whitespace()
+            if self.buffer.startswith(","):
+                self.buffer = self.buffer[1:]
+                continue
+            self._take("}")
+            while not self.eof:
+                self._fill()
+            if self.buffer.strip():
+                raise ValueError("trailing data after MPtrj JSON root object")
+            self.buffer = ""
+            return
+
+
+def _count_records(
+    input_path: Path,
+    max_records: int | None,
+    progress_every: int,
+) -> int:
+    """Count records without retaining the full upstream JSON object in memory."""
+
+    stream = _JSONStream(input_path)
+    started = time.monotonic()
+    try:
+        records = stream.iter_records()
+        if max_records is not None:
+            records = islice(records, max_records)
+        count = 0
+        for _ in records:
+            count += 1
+            if progress_every and count % progress_every == 0:
+                elapsed = max(time.monotonic() - started, 1e-9)
+                print(
+                    f"counted {count} input records; "
+                    f"rate={count / elapsed:.1f} records/s",
+                    flush=True,
+                )
+        print(
+            f"counted {count} input records; starting conversion",
+            flush=True,
+        )
+        return count
+    finally:
+        stream.close()
+
+
+def _shard_sizes(total_records: int, shards: int) -> list[int]:
+    """Return the same contiguous, remainder-first split used upstream."""
+
+    chunk_size, remainder = divmod(total_records, shards)
+    return [
+        chunk_size + (1 if shard_index < remainder else 0)
+        for shard_index in range(shards)
+    ]
+
+
+def convert(
+    input_path: Path,
+    output_dir: Path,
+    cutoff: float,
+    max_neighbors: int,
+    shards: int,
+    max_records: int | None = None,
+    progress_every: int = 10_000,
+    workers: int = 1,
+    batch_size: int = 16,
+) -> None:
+    if shards < 1:
+        raise ValueError("shards must be positive")
+    if max_records is not None and max_records < 1:
+        raise ValueError("max_records must be positive")
+    if progress_every < 0:
+        raise ValueError("progress_every cannot be negative")
+    if workers < 1:
+        raise ValueError("workers must be positive")
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive")
+    output_dir.mkdir(parents=True, exist_ok=True)
+    existing = sorted(output_dir.iterdir())
+    if existing:
+        raise FileExistsError(
+            f"refusing to write into non-empty output directory {output_dir}; "
+            "choose a new directory or clear it explicitly"
+        )
+    total_records = _count_records(input_path, max_records, progress_every)
+    shard_sizes = _shard_sizes(total_records, shards)
+    atom_counts: list[int] = []
+    written = 0
+    examined = 0
+    next_progress = progress_every
+    conversion_started = time.monotonic()
+    executor: ProcessPoolExecutor | None = None
+    if workers > 1:
+        executor = ProcessPoolExecutor(
+            max_workers=workers,
+            mp_context=mp.get_context("fork"),
+            initializer=_init_worker,
+        )
+        # ProcessPoolExecutor launches all workers together on its first submit.
+        executor.submit(_worker_ready).result()
+    stream = _JSONStream(input_path)
+    try:
+        records = stream.iter_records()
+        if max_records is not None:
+            records = islice(records, max_records)
+
+        for shard_index, shard_size in enumerate(shard_sizes):
+            # Zero padding preserves numeric shard order in AseDBDataset, whose
+            # directory discovery sorts paths lexicographically.
+            database = connect(
+                str(output_dir / f"data_{shard_index:05d}.aselmdb")
+            )
+            shard_written = 0
+            try:
+                shard_records = islice(records, shard_size)
+                for converted in _convert_batches_ordered(
+                    shard_records,
+                    cutoff,
+                    max_neighbors,
+                    workers,
+                    batch_size,
+                    executor,
+                ):
+                    for atoms in converted:
+                        examined += 1
+                        if atoms is not None:
+                            database.write(atoms, data=atoms.info)
+                            atom_counts.append(len(atoms))
+                            shard_written += 1
+                            written += 1
+                        if progress_every and examined >= next_progress:
+                            elapsed = max(
+                                time.monotonic() - conversion_started, 1e-9
+                            )
+                            rate = examined / elapsed
+                            remaining = total_records - examined
+                            eta_seconds = remaining / rate if rate else float("inf")
+                            print(
+                                f"processed {examined}/{total_records} records; "
+                                f"wrote {written}; filtered {examined - written}; "
+                                f"rate={rate:.1f} records/s; "
+                                f"eta={eta_seconds / 60:.1f} min",
+                                flush=True,
+                            )
+                            next_progress += progress_every
+            finally:
+                database.close()
+            print(
+                f"finished shard {shard_index}: examined {shard_size}; "
+                f"wrote {shard_written}",
+                flush=True,
+            )
+    finally:
+        stream.close()
+        if executor is not None:
+            executor.shutdown(wait=True, cancel_futures=True)
+
+    if examined != total_records:
+        raise RuntimeError(
+            f"MPtrj input changed while converting: counted {total_records} "
+            f"records but read {examined}"
+        )
+
+    np.savez(
+        output_dir / "metadata.npz",
+        natoms=np.asarray(atom_counts, dtype=np.int64),
+    )
+    print(f"wrote {written} structures to {output_dir} ({shards} shards)")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input", type=Path, required=True, help="MPtrj JSON file")
+    parser.add_argument(
+        "--output", type=Path, required=True, help="ASE-LMDB output directory"
+    )
+    parser.add_argument("--cutoff", type=float, default=6.0)
+    parser.add_argument("--max-neighbors", type=int, default=1000)
+    parser.add_argument("--shards", type=int, default=15)
+    parser.add_argument(
+        "--max-records",
+        type=int,
+        default=None,
+        help="optional maximum input records to examine for bounded validation",
+    )
+    parser.add_argument(
+        "--progress-every",
+        type=int,
+        default=10_000,
+        help="print progress after this many input records; zero disables it",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="parallel CPU conversion workers; results remain ordered",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=16,
+        help="records sent to each worker task",
+    )
+    args = parser.parse_args()
+    if args.shards < 1:
+        parser.error("--shards must be positive")
+    if args.max_records is not None and args.max_records < 1:
+        parser.error("--max-records must be positive")
+    if args.progress_every < 0:
+        parser.error("--progress-every cannot be negative")
+    if args.workers < 1:
+        parser.error("--workers must be positive")
+    if args.batch_size < 1:
+        parser.error("--batch-size must be positive")
+    convert(
+        args.input,
+        args.output,
+        args.cutoff,
+        args.max_neighbors,
+        args.shards,
+        args.max_records,
+        args.progress_every,
+        args.workers,
+        args.batch_size,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/create_finetune_dataset.py

@@ -0,0 +1,186 @@
+from __future__ import annotations
+
+import argparse
+import glob
+import logging
+import multiprocessing as mp
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+from ase.db import connect
+from ase.io import read
+from tqdm import tqdm
+
+from onescience.datapipes.materials.custom_stack.storage import AseDBDataset
+logging.basicConfig(level=logging.INFO)
+
+
+def compute_normalizer_and_linear_reference(train_path, num_workers):
+    """
+    Given a path to an ASE database file, compute the normalizer value and linear
+    reference coefficients. These are used to normalize energies and forces during
+    training. For large datasets, compute this for only a subset of the data.
+    """
+    global dataset
+    dataset = AseDBDataset({"src": str(train_path)})
+
+    sample_indices = random.sample(range(len(dataset)), min(100000, len(dataset)))
+    with mp.Pool(num_workers) as pool:
+        outputs = list(
+            tqdm(
+                pool.imap(extract_energy_and_forces, sample_indices),
+                total=len(sample_indices),
+                desc="Computing normalizer values.",
+            )
+        )
+        atomic_numbers = [x[0] for x in outputs]
+        energies = [x[1] for x in outputs]
+        forces = np.array([force for x in outputs for force in x[2]])
+        force_rms = np.sqrt(np.mean(np.square(forces)))
+        assert (
+            np.isfinite(force_rms).all()
+        ), "We found non-finite values in the forces, please check your input data!"
+        coeff = compute_lin_ref(atomic_numbers, energies)
+    return force_rms, coeff
+
+
+def extract_energy_and_forces(idx):
+    """
+    Extract energy and forces from an ASE atoms object at a given index in the dataset.
+    """
+    atoms = dataset.get_atoms(idx)
+    atomic_numbers = atoms.get_atomic_numbers()
+    energy = atoms.get_potential_energy()
+    forces = atoms.get_forces()
+    n_atoms = len(atoms)
+
+    fixed_idx = np.zeros(n_atoms)
+    if hasattr(atoms, "constraints"):
+        from ase.constraints import FixAtoms
+
+        for constraint in atoms.constraints:
+            if isinstance(constraint, FixAtoms):
+                fixed_idx[constraint.index] = 1
+
+    mask = fixed_idx == 0
+    forces = forces[mask]
+
+    return atomic_numbers, energy, forces
+
+
+def compute_lin_ref(atomic_numbers, energies):
+    """
+    Compute linear reference coefficients given atomic numbers and energies.
+    """
+    features = [np.bincount(x, minlength=100).astype(int) for x in atomic_numbers]
+
+    X = np.vstack(features)
+    y = energies
+
+    coeff = np.linalg.lstsq(X, y, rcond=None)[0]
+    assert np.isfinite(
+        coeff
+    ).all(), "Found some non-finite values while computing element references, please check/sanitize your inputs before proceeding"
+    return coeff.tolist()
+
+
+def write_ase_db(mp_arg):
+    """
+    Write ASE atoms objects to an ASE database file. This function is designed to be
+    run in parallel using multiprocessing.
+    """
+    db_file, file_list, worker_id = mp_arg
+
+    successful = []
+    failed = []
+    natoms = []
+    with connect(str(db_file)) as db:
+        for file in tqdm(file_list, position=worker_id):
+            atoms_list = read(file, ":")
+            for i, atoms in enumerate(atoms_list):
+                try:
+                    assert (
+                        atoms.calc is not None
+                    ), "No calculator attached to atoms object."
+                    assert "energy" in atoms.calc.results, "Missing energy result"
+                    assert "forces" in atoms.calc.results, "Missing forces result"
+                    db.write(atoms, data=atoms.info)
+                    natoms.append(len(atoms))
+                    successful.append(f"{file},{i}")
+                except AssertionError as e:
+                    failed.append(f"{file},{i}: {e!s}")
+
+    return db_file, natoms, successful, failed
+
+
+def launch_processing(data_dir, output_dir, num_workers):
+    """
+    Driver script to launch processing of ASE atoms files into an ASE database.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    input_files = [
+        f
+        for f in glob.glob(os.path.join(data_dir, "**/*"), recursive=True)
+        if os.path.isfile(f)
+    ]
+    chunked_files = np.array_split(input_files, num_workers)
+    db_files = [output_dir / f"data.{i:04d}.aselmdb" for i in range(num_workers)]
+    mp_args = [(db_files[i], chunked_files[i], i) for i in range(num_workers)]
+    with mp.Pool(num_workers) as pool:
+        outputs = pool.map(write_ase_db, mp_args)
+
+    # Log results
+    natoms = []
+    for output in outputs:
+        db_file, _natoms, successful, failed = output
+        natoms.extend(_natoms)
+        log_file = db_file.with_suffix(".log")
+        failed_file = db_file.with_suffix(".failed")
+
+        with open(log_file, "w") as log:
+            log.write("\n".join(successful))
+        with open(failed_file, "w") as failed_log:
+            failed_log.write("\n".join(failed))
+    np.savez_compressed(output_dir / "metadata.npz", natoms=natoms)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--train-dir",
+        type=str,
+        required=True,
+        help="Directory of ASE atoms objects to convert for training.",
+    )
+    parser.add_argument(
+        "--val-dir",
+        type=str,
+        required=True,
+        help="Directory of ASE atoms objects to convert for validation.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        required=True,
+        help="Output directory to save required finetuning artifacts.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=8,
+        help="Number of parallel workers for processing files.",
+    )
+    args = parser.parse_args()
+
+    # Launch processing for training data
+    train_path = args.output_dir / "train"
+    launch_processing(args.train_dir, train_path, args.num_workers)
+    force_rms, linref_coeff = compute_normalizer_and_linear_reference(
+        train_path, args.num_workers
+    )
+    val_path = args.output_dir / "val"
+    launch_processing(args.val_dir, val_path, args.num_workers)
+
+    logging.info(f"Generated dataset at {args.output_dir}")

scripts/create_uma_finetune_dataset.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from pathlib import Path
+
+import yaml
+
+from onescience.utils.uma.scripts.create_finetune_dataset import (
+    compute_normalizer_and_linear_reference,
+    launch_processing,
+)
+from onescience.utils.uma.units.mlip_unit.api.inference import UMATask
+
+
+logging.basicConfig(level=logging.INFO)
+
+TEMPLATE_DIR = Path("configs")
+DATA_YAML_DIR = Path("data")
+REGRESSION_LABEL_TO_TASK_YAML = {
+    "e": DATA_YAML_DIR / Path("uma_conserving_data_task_energy.yaml"),
+    "ef": DATA_YAML_DIR / Path("uma_conserving_data_task_energy_force.yaml"),
+    "efs": DATA_YAML_DIR / Path("uma_conserving_data_task_energy_force_stress.yaml"),
+}
+UMA_SM_FINETUNE_YAML = Path("uma_sm_finetune_template.yaml")
+
+
+def create_yaml(
+    train_path: str,
+    val_path: str,
+    force_rms: float,
+    linref_coeff: list,
+    output_dir: str,
+    dataset_name: str,
+    regression_tasks: str,
+    # base_model_name: str,
+):
+    data_task_yaml = TEMPLATE_DIR / REGRESSION_LABEL_TO_TASK_YAML[regression_tasks]
+    with open(data_task_yaml) as file:
+        template = yaml.safe_load(file)
+        template["dataset_name"] = dataset_name
+        template["normalizer_rmsd"] = force_rms
+        template["elem_refs"] = linref_coeff
+        template["train_dataset"]["splits"]["train"]["src"] = train_path
+        template["val_dataset"]["splits"]["val"]["src"] = val_path
+        # add extra large vaccum box for molecules
+        # if dataset_name == str(UMATask.OMOL):
+        #     template["train_dataset"]["a2g_args"]["molecule_cell_size"] = 1000.0
+        os.makedirs(output_dir / DATA_YAML_DIR, exist_ok=True)
+        with open(
+            output_dir / REGRESSION_LABEL_TO_TASK_YAML[regression_tasks], "w"
+        ) as yaml_file:
+            yaml.dump(template, yaml_file, default_flow_style=False, sort_keys=False)
+
+    uma_finetune_yaml = TEMPLATE_DIR / UMA_SM_FINETUNE_YAML
+    with open(uma_finetune_yaml) as file:
+        template_ft = yaml.safe_load(file)
+        #template_ft["base_model_name"] = base_model_name
+        template_ft["defaults"][0]["data"] = REGRESSION_LABEL_TO_TASK_YAML[
+            regression_tasks
+        ].stem
+    with open(output_dir / UMA_SM_FINETUNE_YAML, "w") as yaml_file:
+        yaml.dump(template_ft, yaml_file, default_flow_style=False, sort_keys=False)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--train-dir",
+        type=str,
+        required=True,
+        help="Directory of ASE atoms objects to convert for training.",
+    )
+    parser.add_argument(
+        "--val-dir",
+        type=str,
+        required=True,
+        help="Directory of ASE atoms objects to convert for validation.",
+    )
+    parser.add_argument(
+        "--uma-task",
+        type=str,
+        required=True,
+        choices=[t.value for t in UMATask],
+        help="choose a uma task to finetune",
+    )
+    parser.add_argument(
+        "--regression-tasks",
+        type=str,
+        choices=["e", "ef", "efs"],
+        required=True,
+        help="Choose to finetune based on regression task set (you must have the corresponding labels in your dataset), can be energy (e), energy+force (ef) or energy+force+stress(efs)",
+    )
+    #parser.add_argument(
+    #    "--base-model",
+    #    type=str,
+    #   default="uma-s-1",
+    #    help="Name of base uma model",
+    #)
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        required=True,
+        help="Output directory to save required finetuning artifacts.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=8,
+        help="Number of parallel workers for processing files.",
+    )
+    args = parser.parse_args()
+    assert not Path(
+        args.output_dir
+    ).exists(), f"{args.output_dir} can't already exist, please choose a different dir"
+
+    # Launch processing for training data
+    train_path = args.output_dir / "train"
+    launch_processing(args.train_dir, train_path, args.num_workers)
+    force_rms, linref_coeff = compute_normalizer_and_linear_reference(
+        train_path, args.num_workers
+    )
+    val_path = args.output_dir / "val"
+    launch_processing(args.val_dir, val_path, args.num_workers)
+
+    create_yaml(
+        train_path=str(train_path),
+        val_path=str(val_path),
+        force_rms=float(force_rms),
+        linref_coeff=linref_coeff,
+        output_dir=args.output_dir,
+        dataset_name=args.uma_task,
+        regression_tasks=args.regression_tasks,
+        # base_model_name=args.base_model,
+    )
+    logging.info(f"Generated dataset and data config yaml in {args.output_dir}")
+    logging.info(
+        f"To run finetuning, run fairchem -c {args.output_dir}/{UMA_SM_FINETUNE_YAML}"
+    )

single_point.py

@@ -0,0 +1,80 @@
+"""Run one Equiformer V3 energy, force, and stress prediction through ASE."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+from ase.build import bulk
+from ase.io import read
+
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3Calculator,
+    write_workflow_result,
+)
+
+
+def default_checkpoint() -> str:
+    return str(
+        Path(__file__).resolve().parent
+        / "weight"
+        / "omat24-mptrj-salex_gradient.pt"
+    )
+
+
+def load_structure(path: str | None):
+    if path:
+        return read(path)
+    # Keep the example runnable without requiring a separate structure file.
+    return bulk("Cu")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint")
+    parser.add_argument(
+        "--input",
+        help=(
+            "CIF, POSCAR, XYZ, trajectory, or another ASE-readable structure; "
+            "defaults to the built-in periodic Cu example"
+        ),
+    )
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--output", default="outputs/single_point.json")
+    args = parser.parse_args()
+
+    checkpoint = args.checkpoint or default_checkpoint()
+    atoms = load_structure(args.input)
+    atoms.calc = EquiformerV3Calculator.from_checkpoint(
+        checkpoint,
+        device=args.device,
+    )
+
+    result = {
+        "formula": atoms.get_chemical_formula(),
+        "natoms": len(atoms),
+        "input": str(Path(args.input).expanduser()) if args.input else None,
+        "input_source": args.input or "ASE bulk Cu default",
+        "checkpoint": str(Path(checkpoint).expanduser()),
+        "pbc": atoms.pbc.tolist(),
+        "cell_angstrom": atoms.cell.array.tolist(),
+        "energy_ev": float(atoms.get_potential_energy()),
+        "forces_ev_per_angstrom": atoms.get_forces().tolist(),
+        "stress_ev_per_angstrom_cubed_voigt": atoms.get_stress().tolist(),
+    }
+    output = write_workflow_result(result, args.output)
+
+    print("formula:", result["formula"])
+    print("atoms:", result["natoms"])
+    print("energy (eV):", result["energy_ev"])
+    print("result:", output)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,1585 @@
+"""Train, fine-tune, or resume an Equiformer V3 atomistic model.
+
+The input splits must be ASE DB or ASE-LMDB datasets containing calculator
+results. Checkpoints remain compatible with ``EquiformerV3Calculator`` while
+also carrying optimizer, scheduler, EMA, and progress state for resuming.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import os
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+os.environ.setdefault(
+    "ONESCIENCE_EQUIFORMER_V3_JD_PATH",
+    str(Path(__file__).resolve().parent / "weight" / "Jd.pt"),
+)
+
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from onescience.datapipes.materials.custom_stack import data_list_collater
+from onescience.datapipes.materials.custom_stack.base_dataset import (
+    Subset as MetadataSubset,
+)
+from onescience.datapipes.materials.custom_stack.storage.ase_datasets import (
+    AseDBDataset,
+)
+from onescience.modules.loss.uma_loss import DDPLoss
+from onescience.utils.equiformer_v3 import (
+    EquiformerV3CheckpointTransforms,
+    load_equiformer_v3_checkpoint,
+)
+from onescience.utils.uma.common.data_parallel import BalancedBatchSampler
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.normalization.element_references import (
+    LinearReferences,
+    create_element_references,
+    fit_linear_references,
+)
+from onescience.utils.uma.normalization.normalizer import (
+    create_normalizer,
+    fit_normalizers,
+)
+from onescience.utils.uma.scheduler import CosineLRLambda
+
+
+MODE_ALIASES = {
+    "scratch": "train_from_scratch",
+    "train": "train_from_scratch",
+    "train_from_scratch": "train_from_scratch",
+    "finetune": "init_from_checkpoint",
+    "fine_tune": "init_from_checkpoint",
+    "init_from_checkpoint": "init_from_checkpoint",
+    "resume": "resume_training",
+    "resume_training": "resume_training",
+}
+
+
+@dataclass(frozen=True)
+class DistributedContext:
+    """Runtime information for a normal process or a torchrun worker."""
+
+    rank: int = 0
+    world_size: int = 1
+    local_rank: int = 0
+
+    @property
+    def enabled(self) -> bool:
+        return self.world_size > 1
+
+    @property
+    def is_main(self) -> bool:
+        return self.rank == 0
+
+
+@dataclass(frozen=True)
+class LossSpec:
+    """One target loss from the YAML contract."""
+
+    name: str
+    function: str
+    coefficient: float
+    free_atoms_only: bool = False
+
+
+@dataclass(frozen=True)
+class DenoisingPosParams:
+    """Official Equiformer V3 DeNS position-corruption contract."""
+
+    enabled: bool = False
+    prob: float = 0.0
+    fixed_noise_std: bool = True
+    std: float = 0.025
+    corrupt_ratio: float | None = None
+    all_atoms: bool = False
+    min_num_atoms: int | None = None
+    strict_max_ratio: float | None = None
+    max_force_norm: float | None = None
+    max_stress_norm: float | None = None
+    max_mean_force_norm: float | None = None
+    coefficient: float = 1.0
+
+
+class ModelEMA:
+    """Exponential moving average of trainable model parameters."""
+
+    def __init__(self, model: torch.nn.Module, decay: float):
+        if not 0.0 < decay < 1.0:
+            raise ValueError("ema_decay must be between zero and one")
+        self.decay = float(decay)
+        self.shadow = {
+            name: parameter.detach().clone()
+            for name, parameter in model.named_parameters()
+            if parameter.requires_grad
+        }
+
+    @torch.no_grad()
+    def update(self, model: torch.nn.Module) -> None:
+        parameters = dict(model.named_parameters())
+        for name, value in self.shadow.items():
+            value.lerp_(parameters[name].detach(), 1.0 - self.decay)
+
+    @contextmanager
+    def apply(self, model: torch.nn.Module):
+        parameters = dict(model.named_parameters())
+        backup = {
+            name: parameters[name].detach().clone() for name in self.shadow
+        }
+        with torch.no_grad():
+            for name, value in self.shadow.items():
+                parameters[name].copy_(value)
+        try:
+            yield
+        finally:
+            with torch.no_grad():
+                for name, value in backup.items():
+                    parameters[name].copy_(value)
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "decay": self.decay,
+            "shadow": {
+                name: value.detach().cpu() for name, value in self.shadow.items()
+            },
+        }
+
+    def load_state_dict(self, state: dict[str, Any], device: torch.device) -> None:
+        self.decay = float(state["decay"])
+        if set(state["shadow"]) != set(self.shadow):
+            raise ValueError("EMA parameters do not match the resumed model")
+        self.shadow = {
+            name: value.to(device=device) for name, value in state["shadow"].items()
+        }
+
+
+def _init_distributed(device_name: str, backend: str) -> DistributedContext:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if world_size == 1:
+        if device_name.startswith("cuda") and torch.cuda.is_available():
+            torch.cuda.set_device(0)
+        return DistributedContext()
+    if not torch.distributed.is_available():
+        raise RuntimeError("torch.distributed is required for torchrun training")
+    rank = int(os.environ["RANK"])
+    local_rank = int(os.environ.get("LOCAL_RANK", rank))
+    if device_name.startswith("cuda"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("torchrun requested CUDA/DCU, but it is unavailable")
+        torch.cuda.set_device(local_rank)
+    torch.distributed.init_process_group(
+        backend=backend, rank=rank, world_size=world_size
+    )
+    return DistributedContext(rank, world_size, local_rank)
+
+
+def _close_distributed(context: DistributedContext) -> None:
+    if context.enabled and torch.distributed.is_initialized():
+        torch.distributed.barrier()
+        torch.distributed.destroy_process_group()
+
+
+def _loader(
+    path: str | list[str],
+    batch_size: int,
+    workers: int,
+    max_samples: int | None = None,
+    context: DistributedContext | None = None,
+    train: bool = False,
+    seed: int = 0,
+    max_atoms: int | None = None,
+    load_balancing: str | bool | None = "atoms",
+    load_balancing_on_error: str = "raise",
+    device: torch.device | None = None,
+) -> DataLoader:
+    """Build the shared OneScience FairChem-style ASE data loader."""
+
+    dataset = AseDBDataset(
+        {
+            "src": path,
+            "a2g_args": {
+                "r_edges": False,
+                "r_energy": True,
+                "r_forces": True,
+                "r_stress": True,
+            },
+        }
+    )
+    indices = list(range(len(dataset)))
+    if max_atoms is not None:
+        if not dataset.metadata_hasattr("natoms"):
+            raise ValueError("max_atoms requires dataset metadata.npz with natoms")
+        natoms = dataset.get_metadata("natoms", indices)
+        indices = [
+            index for index, count in zip(indices, natoms) if int(count) <= max_atoms
+        ]
+        if not indices:
+            raise ValueError(f"max_atoms={max_atoms} filtered every sample")
+    if max_samples is not None:
+        sample_count = min(max_samples, len(indices))
+        generator = torch.Generator().manual_seed(seed)
+        order = torch.randperm(len(indices), generator=generator)[:sample_count]
+        indices = [indices[index] for index in order.tolist()]
+    if len(indices) != len(dataset):
+        dataset = MetadataSubset(dataset, indices, metadata={})
+
+    context = context or DistributedContext()
+    if load_balancing:
+        batch_sampler = BalancedBatchSampler(
+            dataset,
+            batch_size=batch_size,
+            num_replicas=context.world_size,
+            rank=context.rank,
+            device=device,
+            seed=seed,
+            mode=load_balancing,
+            shuffle=train,
+            on_error=load_balancing_on_error,
+            drop_last=False,
+        )
+        return DataLoader(
+            dataset,
+            batch_sampler=batch_sampler,
+            num_workers=workers,
+            collate_fn=lambda items: data_list_collater(items, otf_graph=True),
+            generator=torch.Generator().manual_seed(seed),
+        )
+
+    sampler = None
+    if context.enabled:
+        sampler = DistributedSampler(
+            dataset,
+            num_replicas=context.world_size,
+            rank=context.rank,
+            shuffle=train,
+            drop_last=False,
+            seed=seed,
+        )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=sampler is None and train,
+        sampler=sampler,
+        num_workers=workers,
+        collate_fn=lambda items: data_list_collater(items, otf_graph=True),
+        generator=torch.Generator().manual_seed(seed),
+    )
+
+
+def _expand_path(value: str | list[str] | None) -> str | list[str] | None:
+    if value is None:
+        return None
+    if isinstance(value, list):
+        return [_expand_path(item) for item in value]
+    return os.path.expandvars(os.path.expanduser(str(value)))
+
+
+def _normalize_config(raw: dict[str, Any]) -> dict[str, Any]:
+    config = copy.deepcopy(raw)
+    mode = config.get("mode")
+    if mode is None:
+        mode = "resume_training" if config.get("resume") else None
+        mode = mode or (
+            "init_from_checkpoint"
+            if config.get("initialization_checkpoint") or config.get("checkpoint")
+            else "train_from_scratch"
+        )
+    try:
+        config["mode"] = MODE_ALIASES[str(mode)]
+    except KeyError as error:
+        raise ValueError(f"unsupported training mode: {mode!r}") from error
+
+    if config.get("checkpoint") and not config.get("initialization_checkpoint"):
+        config["initialization_checkpoint"] = config["checkpoint"]
+    for key in ("initialization_checkpoint", "resume", "train", "val", "output"):
+        config[key] = _expand_path(config.get(key))
+    config["transforms_checkpoint"] = _expand_path(
+        config.get("transforms_checkpoint")
+    )
+
+    optimizer = config.setdefault("optimizer", {})
+    optimizer.setdefault("name", "AdamW")
+    optimizer.setdefault("lr", config.get("lr", 5.0e-5))
+    optimizer.setdefault("weight_decay", config.get("weight_decay", 1.0e-3))
+    optimizer.setdefault("betas", config.get("betas", [0.9, 0.98]))
+    optimizer.setdefault("eps", config.get("eps", 1.0e-6))
+
+    scheduler = config.setdefault("scheduler", {})
+    scheduler.setdefault("name", "cosine")
+    scheduler.setdefault("warmup_factor", 0.0)
+    scheduler.setdefault("warmup_epochs", 0.1)
+    scheduler.setdefault("lr_min_factor", 0.01)
+
+    config.setdefault("device", "cuda")
+    config.setdefault("backend", "nccl")
+    config.setdefault("seed", 0)
+    config.setdefault("epochs", 1)
+    config.setdefault("batch_size", 1)
+    config.setdefault("eval_batch_size", config["batch_size"])
+    config.setdefault("workers", 0)
+    config.setdefault("grad_accumulation_steps", 1)
+    config.setdefault("log_every_n_steps", 0)
+    config.setdefault("log_every_n_validation_batches", 0)
+    config.setdefault("clip_grad_norm", 100.0)
+    config.setdefault("ema_decay", 0.999)
+    config.setdefault("amp", False)
+    config.setdefault("amp_dtype", "float16")
+    config.setdefault("amp_init_scale", 65536.0)
+    config.setdefault("load_balancing", "atoms")
+    config.setdefault("load_balancing_on_error", "raise")
+    config.setdefault("ddp_find_unused_parameters", False)
+    return config
+
+
+def _loss_specs(config: dict[str, Any]) -> list[LossSpec]:
+    raw = config.get("losses") or config.get("loss_functions")
+    if raw is None:
+        raw = {
+            "energy": {
+                "fn": "per_atom_mae",
+                "coefficient": config.get("energy_weight", 1.0),
+            },
+            "forces": {
+                "fn": "l2mae",
+                "coefficient": config.get("force_weight", 10.0),
+            },
+            "stress": {
+                "fn": "mae",
+                "coefficient": config.get("stress_weight", 0.0),
+            },
+        }
+    if isinstance(raw, list):
+        merged = {}
+        for item in raw:
+            merged.update(item)
+        raw = merged
+    specs = []
+    for name, values in raw.items():
+        coefficient = float(values.get("coefficient", values.get("weight", 1.0)))
+        if coefficient == 0.0:
+            continue
+        specs.append(
+            LossSpec(
+                name=name,
+                function=str(values.get("fn", values.get("function", "mae"))),
+                coefficient=coefficient,
+                free_atoms_only=bool(values.get("free_atoms_only", False)),
+            )
+        )
+    if not specs:
+        raise ValueError("at least one nonzero target loss is required")
+    unknown = {spec.name for spec in specs} - {"energy", "forces", "stress"}
+    if unknown:
+        raise ValueError(f"unsupported loss targets: {sorted(unknown)}")
+    return specs
+
+
+def _optional_float(value: Any) -> float | None:
+    return None if value is None else float(value)
+
+
+def _denoising_pos_params(config: dict[str, Any]) -> DenoisingPosParams:
+    values = config.get("denoising_pos_params") or {}
+    return DenoisingPosParams(
+        enabled=bool(config.get("use_denoising_pos", False)),
+        prob=float(values.get("prob", 0.0)),
+        fixed_noise_std=bool(values.get("fixed_noise_std", True)),
+        std=float(values.get("std", 0.025)),
+        corrupt_ratio=_optional_float(values.get("corrupt_ratio")),
+        all_atoms=bool(values.get("all_atoms", False)),
+        min_num_atoms=(
+            None
+            if values.get("min_num_atoms") is None
+            else int(values["min_num_atoms"])
+        ),
+        strict_max_ratio=_optional_float(values.get("strict_max_ratio")),
+        max_force_norm=_optional_float(values.get("max_force_norm")),
+        max_stress_norm=_optional_float(values.get("max_stress_norm")),
+        max_mean_force_norm=_optional_float(values.get("max_mean_force_norm")),
+        coefficient=float(config.get("denoising_pos_coefficient", 1.0)),
+    )
+
+
+def _transforms_from_config(config: dict[str, Any]) -> EquiformerV3CheckpointTransforms:
+    transform_config = config.get("transforms", {})
+    normalizers = {}
+    for name, values in transform_config.get("normalizers", {}).items():
+        values = copy.deepcopy(values)
+        if "file" in values:
+            values["file"] = _expand_path(values["file"])
+        normalizers[name] = create_normalizer(**values)
+    elementrefs = {}
+    for name, values in transform_config.get("element_references", {}).items():
+        values = copy.deepcopy(values)
+        if "values" in values:
+            elementrefs[name] = LinearReferences(
+                torch.as_tensor(values["values"], dtype=torch.float32)
+            )
+        else:
+            if "file" in values:
+                values["file"] = _expand_path(values["file"])
+            elementrefs[name] = create_element_references(**values)
+    return EquiformerV3CheckpointTransforms(normalizers, elementrefs)
+
+
+def _training_transforms(
+    config: dict[str, Any], checkpoint_path: str | Path | None
+) -> EquiformerV3CheckpointTransforms:
+    """Resolve target transforms independently from model initialization weights.
+
+    Initialization checkpoints carry the statistics used by their training
+    dataset.  Fine-tuning may instead point at a target-domain checkpoint or
+    override individual entries in ``transforms``.  Resume deliberately keeps
+    the source checkpoint transforms unchanged and is validated separately.
+    """
+
+    source = config.get("transforms_checkpoint") or checkpoint_path
+    if source:
+        if config.get("clear_checkpoint_transforms"):
+            transforms = EquiformerV3CheckpointTransforms()
+        else:
+            transforms = EquiformerV3CheckpointTransforms.from_checkpoint(source)
+    else:
+        transforms = EquiformerV3CheckpointTransforms()
+
+    overrides = _transforms_from_config(config)
+    for name, module in overrides.normalizers.items():
+        transforms.normalizers[name] = module
+    for name, module in overrides.elementrefs.items():
+        transforms.elementrefs[name] = module
+    return transforms
+
+
+def _construct_model(model_config: dict[str, Any]) -> torch.nn.Module:
+    import onescience.models.equiformer_v3  # noqa: F401
+
+    kwargs = copy.deepcopy(model_config)
+    name = kwargs.pop("name", None)
+    if name not in {"equiformer_v3", "equiformer_v3_dens"}:
+        raise ValueError(f"unsupported Equiformer V3 model name: {name!r}")
+    return registry.get_model_class(name)(**kwargs)
+
+
+def _checkpoint_document(path: str | Path) -> dict[str, Any]:
+    path = Path(path)
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    document = torch.load(path, map_location="cpu", weights_only=False)
+    if "config" not in document or "state_dict" not in document:
+        raise ValueError(f"invalid Equiformer V3 checkpoint: {path}")
+    return document
+
+
+def _reset_module(module: torch.nn.Module) -> None:
+    for child in module.modules():
+        if hasattr(child, "reset_parameters"):
+            child.reset_parameters()
+
+
+def _initialize_model(
+    config: dict[str, Any],
+) -> tuple[
+    torch.nn.Module,
+    EquiformerV3CheckpointTransforms,
+    dict[str, Any],
+    dict[str, Any] | None,
+]:
+    mode = config["mode"]
+    if mode == "train_from_scratch":
+        if not config.get("model"):
+            raise ValueError("train_from_scratch requires a model mapping")
+        model_config = copy.deepcopy(config["model"])
+        return (
+            _construct_model(model_config),
+            _training_transforms(config, None),
+            model_config,
+            None,
+        )
+
+    path = config.get("resume") if mode == "resume_training" else config.get(
+        "initialization_checkpoint"
+    )
+    if not path:
+        required = "resume" if mode == "resume_training" else "initialization_checkpoint"
+        raise ValueError(f"{mode} requires {required}")
+    document = _checkpoint_document(path)
+    source_model_config = copy.deepcopy(document["config"]["model"])
+    model_config = source_model_config | copy.deepcopy(config.get("model", {}))
+
+    if mode == "resume_training":
+        model = _construct_model(model_config)
+        state = document.get("training_state_dict", document["state_dict"])
+        model.load_state_dict(state, strict=True)
+    elif model_config == source_model_config:
+        model = load_equiformer_v3_checkpoint(path)
+    else:
+        model = _construct_model(model_config)
+        source_state = document["state_dict"]
+        excluded = tuple(config.get("exclude_initialization_prefixes", []))
+        compatible = {
+            key.removeprefix("_orig_mod."): value
+            for key, value in source_state.items()
+            if not key.removeprefix("_orig_mod.").startswith(excluded)
+            and key.removeprefix("_orig_mod.") in model.state_dict()
+            and model.state_dict()[key.removeprefix("_orig_mod.")].shape == value.shape
+        }
+        model.load_state_dict(compatible, strict=False)
+        print(
+            f"initialized {len(compatible)}/{len(model.state_dict())} model tensors "
+            f"from {path}",
+            flush=True,
+        )
+    if config.get("reset_energy_head"):
+        _reset_module(model.energy_block)
+
+    if mode == "resume_training":
+        transforms = EquiformerV3CheckpointTransforms.from_checkpoint(path)
+    else:
+        transforms = _training_transforms(config, path)
+    return model, transforms, model_config, document
+
+
+def _target_tensor(name: str, batch) -> torch.Tensor:
+    if not hasattr(batch, name):
+        raise ValueError(f"the batch does not contain required {name} labels")
+    return getattr(batch, name)
+
+
+def _masked_tensors(
+    prediction: torch.Tensor,
+    target: torch.Tensor,
+    spec: LossSpec,
+    batch,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if spec.name == "forces" and spec.free_atoms_only:
+        if not hasattr(batch, "fixed"):
+            raise ValueError("free_atoms_only requires a fixed atom mask")
+        mask = batch.fixed.reshape(-1) == 0
+        prediction = prediction[mask]
+        target = target[mask]
+    return prediction, target
+
+
+def _loss_functions(specs: list[LossSpec]) -> dict[str, DDPLoss]:
+    """Build the same DDP-aware reductions used by the official trainer."""
+
+    return {
+        spec.name: DDPLoss(spec.function, reduction="mean") for spec in specs
+    }
+
+
+def _loss(
+    prediction: dict[str, torch.Tensor],
+    batch,
+    specs: list[LossSpec],
+    transforms: EquiformerV3CheckpointTransforms,
+    loss_functions: dict[str, DDPLoss] | None = None,
+    dens_params: DenoisingPosParams | None = None,
+) -> tuple[torch.Tensor, dict[str, float]]:
+    loss_functions = loss_functions or _loss_functions(specs)
+    total = next(iter(prediction.values())).new_zeros(())
+    components = {}
+    for spec in specs:
+        if spec.name not in prediction:
+            raise ValueError(f"the model did not return required {spec.name} output")
+        target = _target_tensor(spec.name, batch)
+        normalized = transforms.normalize_target(
+            spec.name, target, prediction[spec.name], batch
+        )
+        if spec.name == "forces" and _is_dens_batch(batch):
+            if dens_params is None:
+                raise RuntimeError("DeNS batch requires denoising parameters")
+            pred = prediction[spec.name]
+            noise_mask = batch.noise_mask.reshape(-1, 1).bool()
+            denoising_target = batch.noise_vec.to(pred) / dens_params.std
+            hybrid_target = torch.where(noise_mask, denoising_target, normalized)
+            selection = torch.ones(
+                pred.shape[0], dtype=torch.bool, device=pred.device
+            )
+            if spec.free_atoms_only:
+                if not hasattr(batch, "fixed"):
+                    raise ValueError("free_atoms_only requires a fixed atom mask")
+                selection = batch.fixed.reshape(-1) == 0
+                if dens_params.all_atoms:
+                    selection = selection | noise_mask.reshape(-1)
+            if not bool(selection.any()):
+                raise RuntimeError("DeNS batch has no atoms selected for force loss")
+
+            atomwise = torch.linalg.vector_norm(pred - hybrid_target, dim=-1)
+            coefficients = torch.where(
+                noise_mask.reshape(-1),
+                atomwise.new_full(atomwise.shape, dens_params.coefficient),
+                atomwise.new_full(atomwise.shape, spec.coefficient),
+            )
+            value = (atomwise[selection] * coefficients[selection]).mean()
+            total = total + value
+            force_mask = selection & ~noise_mask.reshape(-1)
+            dens_mask = selection & noise_mask.reshape(-1)
+            if bool(force_mask.any()):
+                components[f"{spec.name}_{spec.function}"] = float(
+                    atomwise[force_mask].mean().detach()
+                )
+            if bool(dens_mask.any()):
+                components["denoising_pos_l2mae"] = float(
+                    atomwise[dens_mask].mean().detach()
+                )
+            components["forces_dens_hybrid_l2mae"] = float(value.detach())
+            continue
+        pred, normalized = _masked_tensors(
+            prediction[spec.name], normalized, spec, batch
+        )
+        value = loss_functions[spec.name](pred, normalized, natoms=batch.natoms)
+        total = total + spec.coefficient * value
+        components[f"{spec.name}_{spec.function}"] = float(value.detach())
+    return total, components
+
+
+@torch.no_grad()
+def _physical_metrics(
+    prediction: dict[str, torch.Tensor],
+    batch,
+    specs: list[LossSpec],
+    transforms: EquiformerV3CheckpointTransforms,
+    dens_params: DenoisingPosParams | None = None,
+) -> dict[str, tuple[float, int]]:
+    metrics = {}
+    for spec in specs:
+        if spec.name == "forces" and _is_dens_batch(batch):
+            if dens_params is None:
+                raise RuntimeError("DeNS batch requires denoising parameters")
+            prediction_tensor = prediction[spec.name].detach()
+            noise_mask = batch.noise_mask.reshape(-1).bool()
+            selection = torch.ones_like(noise_mask)
+            if spec.free_atoms_only:
+                selection = batch.fixed.reshape(-1) == 0
+                if dens_params.all_atoms:
+                    selection = selection | noise_mask
+            force_mask = selection & ~noise_mask
+            dens_mask = selection & noise_mask
+
+            physical_force = transforms.denormalize_prediction(
+                spec.name, prediction_tensor, batch
+            )
+            force_error = physical_force[force_mask] - batch.forces[force_mask]
+            dens_prediction = prediction_tensor[dens_mask] * dens_params.std
+            dens_error = dens_prediction - batch.noise_vec[dens_mask]
+            metrics["denoising_force_mae"] = (
+                float(force_error.abs().sum()),
+                force_error.numel(),
+            )
+            metrics["denoising_force_l2mae"] = (
+                float(torch.linalg.vector_norm(force_error, dim=-1).sum()),
+                force_error.shape[0],
+            )
+            metrics["denoising_pos_mae"] = (
+                float(dens_error.abs().sum()),
+                dens_error.numel(),
+            )
+            metrics["denoising_pos_l2mae"] = (
+                float(torch.linalg.vector_norm(dens_error, dim=-1).sum()),
+                dens_error.shape[0],
+            )
+            metrics["dens_corrupted_atom_fraction"] = (
+                float(dens_mask.sum()),
+                int(selection.sum()),
+            )
+            continue
+        physical = transforms.denormalize_prediction(
+            spec.name, prediction[spec.name].detach(), batch
+        )
+        target = _target_tensor(spec.name, batch).reshape_as(physical)
+        physical, target = _masked_tensors(physical, target, spec, batch)
+        error = physical - target
+        absolute_error = error.abs()
+        metrics[f"{spec.name}_mae"] = (
+            float(absolute_error.sum()),
+            absolute_error.numel(),
+        )
+        if spec.name == "energy":
+            shape = (-1,) + (1,) * (error.ndim - 1)
+            per_atom = error / batch.natoms.to(error).reshape(shape)
+            metrics["energy_per_atom_mae"] = (
+                float(per_atom.abs().sum()),
+                per_atom.numel(),
+            )
+        elif error.ndim >= 2:
+            vector_error = torch.linalg.vector_norm(error, dim=-1)
+            metrics[f"{spec.name}_l2mae"] = (
+                float(vector_error.sum()),
+                vector_error.numel(),
+            )
+    return metrics
+
+
+def _reduce_metrics(
+    sums: dict[str, tuple[float, int]],
+    device: torch.device,
+    context: DistributedContext,
+) -> dict[str, float]:
+    if not sums:
+        raise RuntimeError("the dataset contains no samples")
+    names = sorted(sums)
+    if context.enabled:
+        rank_names: list[list[str] | None] = [None] * context.world_size
+        torch.distributed.all_gather_object(rank_names, names)
+        names = sorted(
+            {
+                name
+                for gathered_names in rank_names
+                if gathered_names is not None
+                for name in gathered_names
+            }
+        )
+    values = torch.tensor(
+        [
+            *(sums.get(name, (0.0, 0))[0] for name in names),
+            *(sums.get(name, (0.0, 0))[1] for name in names),
+        ],
+        device=device,
+        dtype=torch.float64,
+    )
+    if context.enabled:
+        torch.distributed.all_reduce(values, op=torch.distributed.ReduceOp.SUM)
+    split = len(names)
+    reduced = {}
+    for index, name in enumerate(names):
+        count = values[split + index].item()
+        if count <= 0:
+            continue
+        reduced[name] = values[index].item() / count
+    return reduced
+
+
+def _collect_batch_metrics(
+    sums: dict[str, tuple[float, int]],
+    loss: torch.Tensor,
+    components: dict[str, float],
+    physical: dict[str, tuple[float, int]],
+) -> None:
+    values = {
+        "loss": (float(loss.detach()), 1),
+        **{
+            f"normalized_{name}": (value, 1) for name, value in components.items()
+        },
+        **physical,
+    }
+    for name, (total, count) in values.items():
+        previous_total, previous_count = sums.get(name, (0.0, 0))
+        sums[name] = previous_total + total, previous_count + count
+
+
+def _unwrap(model: torch.nn.Module) -> torch.nn.Module:
+    return model.module if isinstance(model, DistributedDataParallel) else model
+
+
+def _is_dens_batch(batch) -> bool:
+    value = getattr(batch, "denoising_pos_forward", False)
+    if torch.is_tensor(value):
+        return bool(value.reshape(-1)[0].item())
+    return bool(value)
+
+
+def _graph_max(
+    values: torch.Tensor, batch_index: torch.Tensor, graph_count: int
+) -> torch.Tensor:
+    result = values.new_full((graph_count,), float("-inf"))
+    return result.scatter_reduce_(
+        0, batch_index, values, reduce="amax", include_self=True
+    )
+
+
+def _apply_graph_filter(
+    graph_mask: torch.Tensor,
+    dens_batch_mask: torch.Tensor,
+    noise_mask: torch.Tensor,
+    noise_vec: torch.Tensor,
+    batch_index: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    dens_batch_mask = dens_batch_mask & graph_mask
+    atom_mask = graph_mask[batch_index]
+    noise_mask = noise_mask & atom_mask
+    noise_vec = noise_vec * atom_mask.reshape(-1, 1)
+    return dens_batch_mask, noise_mask, noise_vec
+
+
+def _add_gaussian_noise_to_position(batch, params: DenoisingPosParams):
+    """Apply the official Equiformer V3 DeNS corruption to one collated batch."""
+
+    graph_count = int(batch.natoms.numel())
+    batch_index = batch.batch.long()
+    noise_vec = torch.empty_like(batch.pos).normal_(mean=0.0, std=params.std)
+    if params.corrupt_ratio is None:
+        noise_mask = torch.ones(
+            batch.pos.shape[0], dtype=torch.bool, device=batch.pos.device
+        )
+    else:
+        noise_mask = (
+            torch.rand(
+                batch.pos.shape[0], dtype=batch.pos.dtype, device=batch.pos.device
+            )
+            < params.corrupt_ratio
+        )
+    noise_vec = noise_vec * noise_mask.reshape(-1, 1)
+    dens_batch_mask = torch.ones(
+        graph_count, dtype=torch.bool, device=batch.pos.device
+    )
+
+    if hasattr(batch, "skip_dens"):
+        graph_mask = ~batch.skip_dens.reshape(-1).bool()
+        dens_batch_mask, noise_mask, noise_vec = _apply_graph_filter(
+            graph_mask, dens_batch_mask, noise_mask, noise_vec, batch_index
+        )
+    if params.min_num_atoms is not None:
+        graph_mask = batch.natoms >= params.min_num_atoms
+        dens_batch_mask, noise_mask, noise_vec = _apply_graph_filter(
+            graph_mask, dens_batch_mask, noise_mask, noise_vec, batch_index
+        )
+    if params.strict_max_ratio is not None:
+        corrupted = batch.pos.new_zeros(graph_count)
+        corrupted.index_add_(0, batch_index, noise_mask.to(batch.pos.dtype))
+        graph_mask = corrupted <= batch.natoms.to(corrupted) * params.strict_max_ratio
+        dens_batch_mask, noise_mask, noise_vec = _apply_graph_filter(
+            graph_mask, dens_batch_mask, noise_mask, noise_vec, batch_index
+        )
+    if params.max_force_norm is not None:
+        graph_mask = _graph_max(
+            torch.linalg.vector_norm(batch.forces, dim=-1),
+            batch_index,
+            graph_count,
+        ) <= params.max_force_norm
+        dens_batch_mask, noise_mask, noise_vec = _apply_graph_filter(
+            graph_mask, dens_batch_mask, noise_mask, noise_vec, batch_index
+        )
+    if params.max_stress_norm is not None:
+        graph_mask = (
+            torch.linalg.vector_norm(batch.stress.reshape(graph_count, -1), dim=-1)
+            <= params.max_stress_norm
+        )
+        dens_batch_mask, noise_mask, noise_vec = _apply_graph_filter(
+            graph_mask, dens_batch_mask, noise_mask, noise_vec, batch_index
+        )
+    if params.max_mean_force_norm is not None:
+        force_sum = batch.forces.new_zeros((graph_count, batch.forces.shape[-1]))
+        force_sum.index_add_(0, batch_index, batch.forces)
+        graph_mask = (
+            torch.linalg.vector_norm(force_sum, dim=-1)
+            <= params.max_mean_force_norm
+        )
+        dens_batch_mask, noise_mask, noise_vec = _apply_graph_filter(
+            graph_mask, dens_batch_mask, noise_mask, noise_vec, batch_index
+        )
+
+    if params.all_atoms:
+        position_mask = torch.ones_like(noise_mask)
+    else:
+        if not hasattr(batch, "fixed"):
+            raise ValueError("DeNS with all_atoms=false requires a fixed atom mask")
+        position_mask = batch.fixed.reshape(-1) == 0
+    batch.pos = batch.pos + noise_vec * position_mask.reshape(-1, 1)
+    batch.noise_vec = noise_vec
+    batch.noise_mask = noise_mask
+    batch.denoising_pos_forward = True
+    batch.dens_batch_mask = dens_batch_mask
+    return batch
+
+
+def _should_apply_dens(
+    params: DenoisingPosParams,
+    device: torch.device,
+    context: DistributedContext,
+) -> bool:
+    if not params.enabled or params.prob <= 0.0:
+        return False
+    decision = torch.rand((), device=device) < params.prob
+    if context.enabled:
+        torch.distributed.broadcast(decision, src=0)
+    return bool(decision.item())
+
+
+def _run_train_epoch(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    device: torch.device,
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler.LambdaLR,
+    ema: ModelEMA | None,
+    specs: list[LossSpec],
+    loss_functions: dict[str, DDPLoss],
+    transforms: EquiformerV3CheckpointTransforms,
+    context: DistributedContext,
+    scaler: torch.GradScaler | None,
+    amp_dtype: torch.dtype,
+    grad_accumulation_steps: int,
+    clip_grad_norm: float | None,
+    global_step: int,
+    max_steps: int | None,
+    epoch: int,
+    log_every_n_steps: int,
+    dens_params: DenoisingPosParams,
+) -> tuple[dict[str, float], int]:
+    model.train()
+    optimizer.zero_grad(set_to_none=True)
+    sums: dict[str, tuple[float, int]] = {}
+    progress_sums: dict[str, tuple[float, int]] = {}
+    progress_start_step = global_step
+    pending = 0
+    completed_updates = 0
+    skipped_optimizer_steps = 0
+    updates_per_epoch = _updates_per_epoch(len(loader), grad_accumulation_steps)
+    batches_to_process = updates_per_epoch * grad_accumulation_steps
+    for index, batch in enumerate(loader):
+        if index >= batches_to_process:
+            break
+        batch = batch.to(device)
+        if _should_apply_dens(dens_params, device, context):
+            batch = _add_gaussian_noise_to_position(batch, dens_params)
+        synchronize_gradients = pending + 1 == grad_accumulation_steps
+        with _gradient_sync_context(model, synchronize_gradients):
+            with torch.autocast(
+                device_type=device.type,
+                enabled=scaler is not None,
+                dtype=amp_dtype,
+            ):
+                prediction = model(batch)
+                loss, components = _loss(
+                    prediction,
+                    batch,
+                    specs,
+                    transforms,
+                    loss_functions,
+                    dens_params,
+                )
+            physical = _physical_metrics(
+                prediction, batch, specs, transforms, dens_params
+            )
+            backward_loss = loss / grad_accumulation_steps
+            if scaler is None:
+                backward_loss.backward()
+            else:
+                scaler.scale(backward_loss).backward()
+        _collect_batch_metrics(sums, loss, components, physical)
+        _collect_batch_metrics(progress_sums, loss, components, physical)
+        pending += 1
+        if pending != grad_accumulation_steps:
+            continue
+        if clip_grad_norm:
+            if scaler is not None:
+                scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad_norm)
+        if scaler is None:
+            optimizer.step()
+            optimizer_step_succeeded = True
+        else:
+            previous_scale = float(scaler.get_scale())
+            scaler.step(optimizer)
+            scaler.update()
+            optimizer_step_succeeded = _amp_step_succeeded(
+                previous_scale, float(scaler.get_scale())
+            )
+        optimizer.zero_grad(set_to_none=True)
+        pending = 0
+        if not optimizer_step_succeeded:
+            skipped_optimizer_steps += 1
+            continue
+        scheduler.step()
+        if ema is not None:
+            ema.update(_unwrap(model))
+        global_step += 1
+        completed_updates += 1
+        reached_max_steps = max_steps is not None and global_step >= max_steps
+        should_log = _should_log_progress(
+            global_step,
+            completed_updates,
+            updates_per_epoch,
+            log_every_n_steps,
+            reached_max_steps,
+        )
+        if should_log:
+            window_metrics = _reduce_metrics(progress_sums, device, context)
+            if context.is_main:
+                print(
+                    json.dumps(
+                        {
+                            "event": "train_progress",
+                            "epoch": epoch,
+                            "epoch_step": completed_updates,
+                            "epoch_steps": updates_per_epoch,
+                            "global_step": global_step,
+                            "lr": float(optimizer.param_groups[0]["lr"]),
+                            "window_steps": global_step - progress_start_step,
+                            "window_metrics": window_metrics,
+                        },
+                        sort_keys=True,
+                    ),
+                    flush=True,
+                )
+            progress_sums = {}
+            progress_start_step = global_step
+        if reached_max_steps:
+            break
+    metrics = _reduce_metrics(sums, device, context)
+    metrics["lr"] = float(optimizer.param_groups[0]["lr"])
+    metrics["skipped_optimizer_steps"] = float(skipped_optimizer_steps)
+    return metrics, global_step
+
+
+def _amp_step_succeeded(previous_scale: float, current_scale: float) -> bool:
+    """A decreasing GradScaler scale means optimizer.step was skipped."""
+
+    return current_scale >= previous_scale
+
+
+@contextmanager
+def _gradient_sync_context(
+    model: torch.nn.Module, synchronize_gradients: bool
+):
+    """Delay DDP reduction until the final microbatch in an update."""
+
+    if synchronize_gradients or not hasattr(model, "no_sync"):
+        yield
+        return
+    with model.no_sync():
+        yield
+
+
+def _updates_per_epoch(loader_batches: int, grad_accumulation_steps: int) -> int:
+    """Return the upstream trainer's number of complete optimizer updates."""
+
+    updates = loader_batches // grad_accumulation_steps
+    if updates < 1:
+        raise ValueError(
+            "grad_accumulation_steps exceeds the number of training batches; "
+            "reduce it or provide more training samples"
+        )
+    return updates
+
+
+def _should_log_progress(
+    global_step: int,
+    completed: int,
+    total: int,
+    interval: int,
+    reached_limit: bool = False,
+) -> bool:
+    """Log periodic progress plus the final update or batch in a phase."""
+
+    if interval <= 0:
+        return False
+    return global_step % interval == 0 or completed == total or reached_limit
+
+
+def _run_validation(
+    model: torch.nn.Module,
+    loader: DataLoader,
+    device: torch.device,
+    specs: list[LossSpec],
+    loss_functions: dict[str, DDPLoss],
+    transforms: EquiformerV3CheckpointTransforms,
+    context: DistributedContext,
+    amp: bool,
+    amp_dtype: torch.dtype,
+    epoch: int,
+    log_every_n_batches: int,
+) -> dict[str, float]:
+    # Gradient models derive forces/stress from energy, so validation must keep
+    # autograd enabled even though no parameter update is performed.
+    model.eval()
+    sums: dict[str, tuple[float, int]] = {}
+    total_batches = len(loader)
+    for batch_index, batch in enumerate(loader, start=1):
+        batch = batch.to(device)
+        with torch.autocast(
+            device_type=device.type,
+            enabled=amp,
+            dtype=amp_dtype,
+        ):
+            prediction = model(batch)
+            loss, components = _loss(
+                prediction, batch, specs, transforms, loss_functions
+            )
+        physical = _physical_metrics(prediction, batch, specs, transforms)
+        _collect_batch_metrics(sums, loss, components, physical)
+        if context.is_main and _should_log_progress(
+            batch_index,
+            batch_index,
+            total_batches,
+            log_every_n_batches,
+        ):
+            print(
+                json.dumps(
+                    {
+                        "event": "validation_progress",
+                        "epoch": epoch,
+                        "batch": batch_index,
+                        "batches": total_batches,
+                    },
+                    sort_keys=True,
+                ),
+                flush=True,
+            )
+    return _reduce_metrics(sums, device, context)
+
+
+def _cosine_scheduler(
+    optimizer: torch.optim.Optimizer,
+    scheduler_config: dict[str, Any],
+    steps_per_epoch: int,
+    epochs: int,
+    max_steps: int | None,
+) -> torch.optim.lr_scheduler.LambdaLR:
+    if scheduler_config.get("name", "cosine").lower() not in {
+        "cosine",
+        "lambdalr",
+    }:
+        raise ValueError("only the official cosine LambdaLR scheduler is supported")
+    del max_steps
+    total_steps = max(1, steps_per_epoch * epochs)
+    warmup_steps = int(
+        float(scheduler_config.get("warmup_epochs", 0.0)) * steps_per_epoch
+    )
+    # Official full runs always have at least one warmup update. Keep bounded
+    # smoke configurations away from the upstream zero-step division edge case.
+    warmup_steps = max(1, min(warmup_steps, total_steps))
+    warmup_factor = float(scheduler_config.get("warmup_factor", 0.0))
+    minimum = float(scheduler_config.get("lr_min_factor", 0.01))
+    lr_lambda = CosineLRLambda(
+        warmup_epochs=warmup_steps,
+        warmup_factor=warmup_factor,
+        epochs=total_steps,
+        lr_min_factor=minimum,
+    )
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+def _state_dict_cpu(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    return {name: value.detach().cpu() for name, value in state.items()}
+
+
+def _inference_state_dict(model: torch.nn.Module, ema: ModelEMA | None) -> dict:
+    state = _state_dict_cpu(model.state_dict())
+    if ema is not None:
+        for name, value in ema.shadow.items():
+            state[name] = value.detach().cpu()
+    return state
+
+
+def _save_checkpoint(
+    output: Path,
+    model: torch.nn.Module,
+    transforms: EquiformerV3CheckpointTransforms,
+    model_config: dict[str, Any],
+    training_config: dict[str, Any],
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler.LambdaLR,
+    ema: ModelEMA | None,
+    epoch: int,
+    global_step: int,
+    history: list[dict[str, Any]],
+    source_document: dict[str, Any] | None,
+    scaler: torch.GradScaler | None = None,
+) -> None:
+    output.parent.mkdir(parents=True, exist_ok=True)
+    source_metadata = copy.deepcopy((source_document or {}).get("metadata", {}))
+    source_metadata.update(
+        {
+            "onescience_equiformer_v3_history": history,
+            "training_mode": training_config["mode"],
+            "source_checkpoint": training_config.get("initialization_checkpoint"),
+            "resume_checkpoint": training_config.get("resume"),
+            "global_step": global_step,
+            "ema_decay": ema.decay if ema is not None else None,
+            "amp": scaler is not None,
+        }
+    )
+    document = {
+        "config": {"model": copy.deepcopy(model_config), "training": training_config},
+        "normalizers": {
+            name: _state_dict_cpu(module.state_dict())
+            for name, module in transforms.normalizers.items()
+        },
+        "elementrefs": {
+            name: _state_dict_cpu(module.state_dict())
+            for name, module in transforms.elementrefs.items()
+        },
+        "state_dict": _inference_state_dict(model, ema),
+        "training_state_dict": _state_dict_cpu(model.state_dict()),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "scheduler_state_dict": scheduler.state_dict(),
+        "ema_state_dict": ema.state_dict() if ema is not None else None,
+        "amp_state_dict": scaler.state_dict() if scaler is not None else None,
+        "training_state": {"epoch": epoch, "global_step": global_step},
+        "metadata": source_metadata,
+    }
+    torch.save(document, output)
+    history_path = output.with_name(output.name + ".history.json")
+    history_path.write_text(json.dumps(history, indent=2, sort_keys=True) + "\n")
+
+
+def _fit_transforms(
+    config: dict[str, Any],
+    transforms: EquiformerV3CheckpointTransforms,
+    dataset,
+) -> None:
+    if config.get("fit_element_references"):
+        fitted = fit_linear_references(
+            targets=["energy"],
+            dataset=dataset,
+            batch_size=config["batch_size"],
+            num_batches=config.get("fit_statistics_batches"),
+            num_workers=config["workers"],
+            log_metrics=False,
+            shuffle=False,
+        )
+        transforms.elementrefs["energy"] = fitted["energy"]
+    requested = config.get("fit_normalizers")
+    if requested:
+        targets = (
+            [spec.name for spec in _loss_specs(config)]
+            if requested is True
+            else list(requested)
+        )
+        fitted = fit_normalizers(
+            targets=targets,
+            dataset=dataset,
+            batch_size=config["batch_size"],
+            num_batches=config.get("fit_statistics_batches"),
+            num_workers=config["workers"],
+            shuffle=False,
+            element_references=dict(transforms.elementrefs),
+        )
+        for name, normalizer in fitted.items():
+            transforms.normalizers[name] = normalizer
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", required=True, help="training YAML path")
+    parser.add_argument("--mode", choices=sorted(MODE_ALIASES))
+    parser.add_argument("--checkpoint", dest="initialization_checkpoint")
+    parser.add_argument("--transforms-checkpoint")
+    parser.add_argument("--resume")
+    parser.add_argument("--train")
+    parser.add_argument("--val")
+    parser.add_argument("--output")
+    parser.add_argument("--device")
+    parser.add_argument("--epochs", type=int)
+    parser.add_argument("--max-steps", type=int)
+    parser.add_argument("--batch-size", type=int)
+    parser.add_argument("--eval-batch-size", type=int)
+    parser.add_argument("--workers", type=int)
+    parser.add_argument("--max-train-samples", type=int)
+    parser.add_argument("--max-val-samples", type=int)
+    parser.add_argument("--max-atoms", type=int)
+    parser.add_argument("--log-every-n-steps", type=int)
+    parser.add_argument("--log-every-n-validation-batches", type=int)
+    parser.add_argument("--seed", type=int)
+    parser.add_argument(
+        "--amp", action=argparse.BooleanOptionalAction, default=None
+    )
+    parser.add_argument("--amp-dtype", choices=("float16", "bfloat16"))
+    parser.add_argument(
+        "--clear-checkpoint-transforms",
+        action=argparse.BooleanOptionalAction,
+        default=None,
+    )
+    return parser.parse_args()
+
+
+def _load_config(args: argparse.Namespace) -> dict[str, Any]:
+    with Path(args.config).expanduser().open(encoding="utf-8") as handle:
+        raw = yaml.safe_load(handle) or {}
+    for key, value in vars(args).items():
+        if key != "config" and value is not None:
+            raw[key] = value
+    return _normalize_config(raw)
+
+
+def _validate_config(config: dict[str, Any]) -> None:
+    missing = [key for key in ("train", "val", "output") if not config.get(key)]
+    if missing:
+        raise ValueError("missing required config fields: " + ", ".join(missing))
+    for name in ("epochs", "batch_size", "eval_batch_size", "grad_accumulation_steps"):
+        if int(config[name]) < 1:
+            raise ValueError(f"{name} must be positive")
+    if config.get("max_steps") is not None and int(config["max_steps"]) < 1:
+        raise ValueError("max_steps must be positive")
+    if config.get("max_atoms") is not None and int(config["max_atoms"]) < 1:
+        raise ValueError("max_atoms must be positive")
+    for name in ("log_every_n_steps", "log_every_n_validation_batches"):
+        if int(config.get(name, 0)) < 0:
+            raise ValueError(f"{name} must be non-negative")
+    if config["amp_dtype"] not in {"float16", "bfloat16"}:
+        raise ValueError("amp_dtype must be float16 or bfloat16")
+    if config["mode"] == "resume_training" and (
+        config.get("transforms_checkpoint")
+        or config.get("transforms")
+        or config.get("clear_checkpoint_transforms")
+    ):
+        raise ValueError(
+            "resume_training restores transforms from the resume checkpoint; "
+            "remove transforms_checkpoint/transforms overrides"
+        )
+    specs = _loss_specs(config)
+    dens_params = _denoising_pos_params(config)
+    if dens_params.enabled:
+        if not dens_params.fixed_noise_std:
+            raise ValueError("the official DeNS trainer requires fixed_noise_std=true")
+        if not 0.0 <= dens_params.prob <= 1.0:
+            raise ValueError("denoising_pos_params.prob must be between zero and one")
+        if dens_params.std <= 0.0:
+            raise ValueError("denoising_pos_params.std must be positive")
+        for name, value in (
+            ("corrupt_ratio", dens_params.corrupt_ratio),
+            ("strict_max_ratio", dens_params.strict_max_ratio),
+        ):
+            if value is not None and not 0.0 <= value <= 1.0:
+                raise ValueError(
+                    f"denoising_pos_params.{name} must be between zero and one"
+                )
+        if dens_params.min_num_atoms is not None and dens_params.min_num_atoms < 1:
+            raise ValueError("denoising_pos_params.min_num_atoms must be positive")
+        if dens_params.coefficient <= 0.0:
+            raise ValueError("denoising_pos_coefficient must be positive")
+        force_specs = [spec for spec in specs if spec.name == "forces"]
+        if len(force_specs) != 1 or force_specs[0].function != "l2mae":
+            raise ValueError("DeNS requires one forces loss using l2mae")
+        if config["mode"] == "train_from_scratch":
+            model_config = config.get("model") or {}
+            if model_config.get("name") != "equiformer_v3_dens":
+                raise ValueError("DeNS requires model.name=equiformer_v3_dens")
+            if not model_config.get("direct_prediction", False):
+                raise ValueError("DeNS pre-training requires direct_prediction=true")
+
+
+def main() -> None:
+    args = _parse_args()
+    try:
+        config = _load_config(args)
+        _validate_config(config)
+    except ValueError as error:
+        raise SystemExit(str(error)) from error
+    if config["device"].startswith("cuda") and not torch.cuda.is_available():
+        raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
+
+    context = _init_distributed(config["device"], config["backend"])
+    try:
+        device = (
+            torch.device(f"cuda:{context.local_rank}")
+            if config["device"].startswith("cuda")
+            else torch.device(config["device"])
+        )
+        torch.manual_seed(int(config["seed"]) + context.rank)
+        dens_params = _denoising_pos_params(config)
+        model, transforms, model_config, source_document = _initialize_model(config)
+        model = model.to(device)
+
+        train_loader = _loader(
+            config["train"],
+            config["batch_size"],
+            config["workers"],
+            config.get("max_train_samples"),
+            context,
+            train=True,
+            seed=config["seed"],
+            max_atoms=config.get("max_atoms"),
+            load_balancing=config.get("load_balancing"),
+            load_balancing_on_error=config["load_balancing_on_error"],
+            device=device,
+        )
+        val_loader = _loader(
+            config["val"],
+            config["eval_batch_size"],
+            config["workers"],
+            config.get("max_val_samples"),
+            context,
+            train=False,
+            seed=config["seed"] + 1,
+            max_atoms=config.get("eval_max_atoms"),
+            load_balancing=config.get("load_balancing"),
+            load_balancing_on_error=config["load_balancing_on_error"],
+            device=device,
+        )
+        if config["mode"] != "resume_training":
+            _fit_transforms(config, transforms, train_loader.dataset)
+        transforms = transforms.to(device)
+
+        optimizer_config = config["optimizer"]
+        if optimizer_config["name"].lower() != "adamw":
+            raise ValueError("only the official AdamW optimizer is supported")
+        optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=float(optimizer_config["lr"]),
+            weight_decay=float(optimizer_config["weight_decay"]),
+            betas=tuple(float(value) for value in optimizer_config["betas"]),
+            eps=float(optimizer_config["eps"]),
+        )
+        updates_per_epoch = _updates_per_epoch(
+            len(train_loader), int(config["grad_accumulation_steps"])
+        )
+        scheduler = _cosine_scheduler(
+            optimizer,
+            config["scheduler"],
+            updates_per_epoch,
+            int(config["epochs"]),
+            config.get("max_steps"),
+        )
+        if config["amp"] and device.type != "cuda":
+            raise ValueError("amp requires a CUDA/DCU device")
+        amp_dtype = getattr(torch, config["amp_dtype"])
+        scaler = (
+            torch.GradScaler(
+                "cuda", init_scale=float(config["amp_init_scale"])
+            )
+            if config["amp"]
+            else None
+        )
+        ema = (
+            ModelEMA(model, float(config["ema_decay"]))
+            if config.get("ema_decay")
+            else None
+        )
+
+        start_epoch = 0
+        global_step = 0
+        history: list[dict[str, Any]] = []
+        if config["mode"] == "resume_training":
+            state = source_document.get("training_state", {})
+            start_epoch = int(state.get("epoch", -1)) + 1
+            global_step = int(state.get("global_step", 0))
+            optimizer.load_state_dict(source_document["optimizer_state_dict"])
+            scheduler.load_state_dict(source_document["scheduler_state_dict"])
+            if scaler is not None and source_document.get("amp_state_dict") is not None:
+                scaler.load_state_dict(source_document["amp_state_dict"])
+            if ema is not None and source_document.get("ema_state_dict") is not None:
+                ema.load_state_dict(source_document["ema_state_dict"], device)
+            history = list(
+                source_document.get("metadata", {}).get(
+                    "onescience_equiformer_v3_history", []
+                )
+            )
+
+        if context.enabled:
+            find_unused_parameters = bool(
+                config["ddp_find_unused_parameters"]
+            )
+            if context.is_main:
+                print(
+                    json.dumps(
+                        {
+                            "event": "ddp_setup",
+                            "find_unused_parameters": find_unused_parameters,
+                        },
+                        sort_keys=True,
+                    ),
+                    flush=True,
+                )
+            model = DistributedDataParallel(
+                model,
+                device_ids=[context.local_rank] if device.type == "cuda" else None,
+                output_device=context.local_rank if device.type == "cuda" else None,
+                find_unused_parameters=find_unused_parameters,
+            )
+        specs = _loss_specs(config)
+        loss_functions = _loss_functions(specs)
+        output = Path(config["output"])
+
+        for epoch in range(start_epoch, int(config["epochs"])):
+            if hasattr(train_loader.batch_sampler, "set_epoch"):
+                train_loader.batch_sampler.set_epoch(epoch)
+            elif isinstance(train_loader.sampler, DistributedSampler):
+                train_loader.sampler.set_epoch(epoch)
+            train_metrics, global_step = _run_train_epoch(
+                model,
+                train_loader,
+                device,
+                optimizer,
+                scheduler,
+                ema,
+                specs,
+                loss_functions,
+                transforms,
+                context,
+                scaler,
+                amp_dtype,
+                int(config["grad_accumulation_steps"]),
+                config.get("clip_grad_norm"),
+                global_step,
+                config.get("max_steps"),
+                epoch,
+                int(config["log_every_n_steps"]),
+                dens_params,
+            )
+            base_model = _unwrap(model)
+            if ema is None:
+                val_metrics = _run_validation(
+                    model,
+                    val_loader,
+                    device,
+                    specs,
+                    loss_functions,
+                    transforms,
+                    context,
+                    scaler is not None,
+                    amp_dtype,
+                    epoch,
+                    int(config["log_every_n_validation_batches"]),
+                )
+            else:
+                with ema.apply(base_model):
+                    val_metrics = _run_validation(
+                        model,
+                        val_loader,
+                        device,
+                        specs,
+                        loss_functions,
+                        transforms,
+                        context,
+                        scaler is not None,
+                        amp_dtype,
+                        epoch,
+                        int(config["log_every_n_validation_batches"]),
+                    )
+            record = {
+                "epoch": epoch,
+                "global_step": global_step,
+                "train": train_metrics,
+                "val": val_metrics,
+            }
+            if context.is_main:
+                history.append(record)
+                print(json.dumps(record, sort_keys=True), flush=True)
+                _save_checkpoint(
+                    output,
+                    base_model,
+                    transforms,
+                    model_config,
+                    config,
+                    optimizer,
+                    scheduler,
+                    ema,
+                    epoch,
+                    global_step,
+                    history,
+                    source_document,
+                    scaler,
+                )
+                print(f"saved checkpoint: {output}", flush=True)
+            if config.get("max_steps") is not None and global_step >= int(
+                config["max_steps"]
+            ):
+                break
+    finally:
+        _close_distributed(context)
+
+
+if __name__ == "__main__":
+    main()

weight/Jd.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4059c45be246dcb6c49c545670b65c56550eb0c2e7a9c92b4b50a92d370dbe2
+size 21697

weight/mptrj_gradient.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59c6c23573a3b05b347662f473209d1bb1ccb5b85f6624a8f87072e7c397addf
+size 485354773

weight/omat24-mptrj-salex_gradient.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:429ccded98163122e7ba588d78e2441653f37f3e091e106c432807fe373c8f98
+size 486243861

weight/omat24_direct.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd874ab2698a391a03a5083b8729c286c676d76ca584437ff5e550696609a96d
+size 559867247

weight/omat24_gradient.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c4b9139c1a14094ed724c3addd584ce07afaf0032e1ec77490d2c60ed8b8a2b
+size 486243797