OneScience-Group/Equiformer_v3

Equiformer V3

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

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

hf download --model OneScience-Group/Equiformer_v3 --local-dir ./Equiformer_v3
cd Equiformer_v3

Install the Runtime Environment

DCU environment

# 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

# 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 and the FAIR Chemistry OC20 documentation 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.

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:

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:

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:

${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:

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:

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:

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 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:

python single_point.py --device cuda --output outputs/single_point.json

Formation energy:

python formation_energy.py --device cuda --output outputs/formation_energy.json

Elastic tensor:

python elastic.py --relax --device cuda --output outputs/elastic.json

Phonons:

python phonons.py \
  --supercell 3 3 3 \
  --bandpath GXWKGL \
  --device cuda \
  --output-dir outputs/phonons

Custom structures or weights:

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.

@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}
}