README.md

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