【MIIT program】Add UMA interatomic potential support by yqf66 · Pull Request #290 · PaddlePaddle/PaddleMaterials

yqf66 · 2026-06-02T08:48:02Z

Summary

This PR adds UMA support to interatomic_potentials and documents the migration validation results following the checklist in #258.

Main additions:

Add ppmat.models.uma with the Paddle implementation of the UMA/eSCN backbone, direct energy/force heads, task/dataset embedding path, rotation utilities, and ASELMDB dataset adapters.
Add UMA configs for OMat24 rattled-500 S2EF, OC20 S2EF, budget-scale validation, and multi-task short-run validation.
Add OC20 S2EF ASELMDB conversion and budget data preparation scripts.
Add UMA README with data preparation, training/evaluation commands, validation scope, checkpoints, and medium-scale experiment metrics.
Update the interatomic_potentials model matrix to include UMA.

Validation

Code/config checks:

python -m py_compile $(find ppmat/models/uma interatomic_potentials/configs/uma -name '*.py' -print)
Parsed all 7 UMA YAML configs with yaml.safe_load.
git diff --check
UMA import check in a clean clone after building the existing MatterSim Cython extension:
from ppmat.models.uma.single_dataset import UMAAseDBDataset
and from ppmat.models.uma.escn_md import UMASingleTaskModel.

Training/evaluation runs completed on a single RTX 4090 with PaddlePaddle 3.3.0 and AMP O1:

Experiment	Config	Train / Val / Test	Epochs	Best checkpoint
OMat24 rattled-500	`uma_omat24_r500_budget_s2ef.yaml`	50k / 5k / 5k	3	`output/uma_omat24_r500_budget_s2ef/checkpoints/best.pdparams`
OC20 S2EF	`uma_oc20_50k_budget_s2ef.yaml`	50k / 5k / 5k	2	`output/uma_oc20_50k_budget_s2ef/checkpoints/best.pdparams`
Multi-task short-run	`uma_multitask_budget_smoke.yaml`	small 5-task proxy dataset	1	`output/uma_multitask_budget_smoke/checkpoints/best.pdparams`

Training loss convergence:

Experiment	Train energy(loss)	Train forces(loss)	Train loss(loss)
OMat24 epoch 1 -> 2 -> 3	6.627979 -> 4.304747 -> 3.176400	0.483082 -> 0.310235 -> 0.257879	80.772244 -> 52.354506 -> 39.500356
OC20 epoch 1 -> 2	51.754490 -> 19.917693	0.181425 -> 0.178013	522.987643 -> 204.517324

Final evaluation:

Experiment	Val energy(metric)	Val forces(metric)	Val loss(loss)	Test energy(metric)	Test forces(metric)	Test loss(loss)
OMat24 rattled-500	11.780332	0.412024	127.072415	11.963885	0.331205	128.924666
OC20 S2EF	16.703999	0.182028	172.312066	16.996420	0.161219	175.239833
Multi-task short-run	237.356260	0.493668	2385.158658	183.573112	1.271681	1850.638359

Notes

The OMat24/OC20 runs are intended as medium-scale migration validation, not full original benchmark reproduction.
The multi-task short-run config uses real OC20 samples and OMat24 subsets as proxy datasets for multiple UMA task names. It validates task-name propagation, task embedding selection, forward/loss/backward, eval/test, and checkpoint saving, but does not claim scientific benchmark quality for the real OMol/ODAC/OMC domains.
The included .pt files under ppmat/models/uma are small precomputed rotation/Wigner coefficient tensors required by the UMA implementation, not trained model checkpoints.

leeleolay

辛苦补充数据集和paddle模型权重可下载链接，我提供url可补充进来

leeleolay · 2026-06-09T12:09:44Z

+除 PaddleMaterials 常规依赖外，UMA 数据读取还需要：
+
+```bash
+pip install ase lmdb e3nn omegaconf


套件已有的requirements 是否已经有了相关的配置，另外e3nn目前没有paddle版本，这里是否合理？套件内有部分实现

leeleolay · 2026-06-09T12:13:15Z

辛苦参考已有model的readme 格式，本次模型复现的指标可在pr中回复

leeleolay · 2026-06-09T12:15:10Z

这里的注释？

leeleolay · 2026-06-09T12:15:45Z

辛苦检查下所有的注释，并且在每个文件补充paddle 权利声明

yqf66 · 2026-06-17T08:00:42Z

已根据意见做如下修改：

README 已参考现有 interatomic potential model 的格式重写，调整为 Abstract / Datasets / Models / Results / Training / Evaluation / Testing / Citations 结构；详细训练收敛过程和迁移验证指标不再作为冗长实验记录放在 README 中。
依赖说明已修正：ase、lmdb、omegaconf 已在 PaddleMaterials requirements.txt 中维护，README 不再提供额外 pip install 命令；同时移除了外部 e3nn 依赖，UMA 中涉及 FromS2Grid、ToS2Grid、o3 的代码已切换到套件内 ppmat.models.common.e3nn 实现。
已检查 UMA 新增源码文件，并为 ppmat/models/uma/**/*.py、interatomic_potentials/configs/uma/prepare_oc20_s2ef_aselmdb.py 和 prepare_budget_data.sh 补充 PaddlePaddle Apache 版权声明；原 fairchem/Meta 来源声明保留在文件内部。
README 中已补充 OMat24/OC20 的官方原始数据来源链接先做占位。Paddle 官方数据集/模型权重 URL 这边我会先将本次实验使用的数据子集、ckpt 和 log 稍后上传到百度网盘，提供给您后再补充到 README 的 Results 表中。

leeleolay

基本未按照套件内已有各模块组织逻辑进行复现，包括dataset，model的组织方式

leeleolay · 2026-06-18T01:44:51Z

该文件可删除或移动到/test/uma下

leeleolay · 2026-06-18T01:45:57Z

同上，这类数据预处理过程性文件与套件核心功能（模型训推无关）

leeleolay · 2026-06-18T01:53:03Z

按照abstarct（包括模型图），dataset，model（模型原理），results组织该页面
在interatomic_potentials/README.md中添加该模型链接

leeleolay · 2026-06-18T01:55:55Z

+| Weight decay | `1e-3` |
+| Loss weights | energy 10.0, forces 30.0 |
+
+`ppmat/models/uma/Jd.pt` stores precomputed Wigner-d coefficients used by UMA


这里是预先设定的参数吗？为什么是硬编码逻辑，pt是torch格式吧？
另外readme这里的模型介绍不保留具体的模型参数说明，只介绍模型的通用原理

leeleolay · 2026-06-18T01:56:37Z

+Single-GPU training:
+
+```bash
+python interatomic_potentials/train.py -c interatomic_potentials/configs/uma/uma_omat24_r500_s2ef.yaml


严格按照已有格式

leeleolay · 2026-06-18T02:08:03Z

+              include_forces: True
+              include_stress: False
+              select_args:
+                limit: 2


leeleolay · 2026-06-18T02:08:48Z

+              src: ./data/oc20/uma_aselmdb/val/val.aselmdb
+              dataset_name: oc20-s2ef-200k-val
+              task_name: oc20
+              include_energy: True


通过grapg_convert可以构建吗

leeleolay · 2026-06-18T02:09:17Z

其他config同上

leeleolay · 2026-06-18T02:09:56Z

使用graph convert，build等方式构建数据

leeleolay · 2026-06-18T02:18:06Z

参考套件内的dataset格式修改

leeleolay

辛苦参考该指南 #258

yqf66 · 2026-06-22T12:39:54Z

已将 UMA dataset 从 ppmat/models/uma 迁移到 ppmat/datasets/uma_dataset.py，collator 迁移到 ppmat/datasets/collate_fn.py。这样 dataset 读取、batch 组装和 model 实现边界更清晰，也对齐套件内已有 dataset/model 的组织方式。
新增 UMAGraphConverter，并在正式 config 中通过 Global.graph_converter 和 Dataset.*.build_graph_cfg 接入图构建流程。现在 UMA 的数据流是：dataset 读取结构与标签，graph converter 构建邻接边和周期边偏移，model 接收构建好的图字段进行 forward，而不是把图构建逻辑隐藏在模型内部。
UMA config 已改为 PaddleMaterials 现有风格，不再使用动态 __import__。正式 configs/uma/ 目录只保留 OMat24 和 OC20 两个训练配置；此前带有 limit: 2、多任务 proxy、budget 验证的小规模配置已移动到 test/uma/，避免把 smoke-test 参数放在正式模型配置里。
过程性数据准备脚本已从正式 configs/uma/ 移到 test/uma/。这些脚本仅用于复现实验数据子集和 smoke validation，不属于模型训推的核心入口。
关于 Jd 文件：它是 UMA/eSCN 旋转模块使用的预计算 Wigner-d 系数表，用于构造球谐/旋转相关的变换矩阵；它不是训练得到的模型权重，也不是 checkpoint。此前文件名为 Jd.pt，来源格式是 torch 序列化文件，并且代码中保留了 torch fallback 加载逻辑。为避免引入 torch 格式和额外依赖，已转换为 Paddle 可直接加载的 Jd.pdparams，并删除 torch fallback，当前仅使用 paddle.load 读取该预计算常量。
README 已按 Abstract / Datasets / Models / Results 重新整理，并补充模型流程图。模型原理部分只保留 UMA/eSCN 的通用机制说明，不再展开具体 K4L2、channel、cutoff 等实现参数；具体训练配置保留在 YAML 中。

leeleolay · 2026-06-22T13:01:40Z

+`Dataset.*.build_graph_cfg`, and training/evaluation runs through
+`interatomic_potentials/train.py`.
+
+```mermaid


修改为论文模型插图

leeleolay · 2026-06-22T13:09:27Z

        )


+class UMASingleCollator:


默认的不能用吗？

leeleolay · 2026-06-22T13:14:40Z

其他的我再详细review下，数据的部分得提供链接吧？需支持一键训练，还有权重，需支持一键推理

yqf66 · 2026-06-29T02:14:44Z

README 中原先的 mermaid 自绘流程图已删除。由于论文/官方实现没有提供可直接引用的完整模型结构图，这里不再放自绘图，避免造成误导。
关于 collator：UMA 这里不能直接使用默认 collator。默认 collator 会对 tensor 做 stack，要求同一个 batch 内各字段 shape 一致；但 UMA 的 atomistic sample 每个结构原子数不同，pos、atomic_numbers、forces 等 node-level 字段长度不同，edge_index 也需要根据 batch 内累计原子数做 offset 平移。因此需要 UMA 专用的 batch 组装逻辑，将 node-level 字段按节点维 concat，并生成 batch 索引，同时修正 edge_index。
已将 UMASingleCollator 精简为一层薄封装，实际逻辑复用 ppmat.datasets.uma_dataset.uma_data_list_to_batch，并删除了 uma_dataset.py 中重复的 collator 类，避免重复实现。
数据集、ckpt链接：https://pan.baidu.com/s/1CvgDvm4duZtQqwW5v8lvlQ?pwd=qkfi 提取码: qkfi

paddle-bot · 2026-07-13T09:14:59Z

Thanks for your contribution!

leeleolay · 2026-07-16T13:35:44Z

数据集链接在readme里只介绍，不写具体链接，自动down的逻辑在datasetli实现，辛苦再检查

leeleolay · 2026-07-16T13:37:30Z

ai生成的代码容易有太多helper，建议从已有的mp20dataset 或者类似的mptrj等复制过来重新修改。需要再精简该数据集。只给模型传pgl封装等ndarray

tensor在model forward内改拜年形式

leeleolay · 2026-07-16T13:39:03Z

这个文件位置不对，

leeleolay · 2026-07-16T13:44:29Z

关于 collator：UMA 这里不能直接使用默认 collator。默认 collator 会对 tensor 做 stack，要求同一个 batch 内各字段 shape 一致；但 UMA 的 atomistic sample 每个结构原子数不同，pos、atomic_numbers、forces 等 node-level 字段长度不同，edge_index 也需要根据 batch 内累计原子数做 offset 平移。因此需要 UMA 专用的 batch 组装逻辑，将 node-level 字段按节点维 concat，并生成 batch 索引，同时修正 edge_index。
已将 UMASingleCollator 精简为一层薄封装，实际逻辑复用 ppmat.datasets.uma_dataset.uma_data_list_to_batch，并删除了 uma_dataset.py 中重复的 collator 类，避免重复实现。

molecular_graph 是否已经验证确实跑不了呢，辛苦比较下chgnet的实现逻辑

leeleolay

model这块的代码辛苦再优化下，ai修改的代码，help太多，逻辑冗余，抽象出太多多余的函数，仅保留必要的函数和重要的逻辑部分，不做过多的工程冗余，辛苦参考已有的模型进行优化和精简

leeleolay · 2026-07-16T13:45:23Z

+    def _to_python(self, value: Any):
+        if paddle.is_tensor(value):
+            return value.numpy()
+        if isinstance(value, dict):
+            return {k: self._to_python(v) for k, v in value.items()}
+        if isinstance(value, list):
+            return [self._to_python(v) for v in value]
+        return value


leeleolay · 2026-08-11T09:03:17Z

重构了部分基础组件，移动了推理器的位置，辛苦基于新的开发和尝试，解决下冲突 @yqf66

yqf66 · 2026-08-12T06:11:01Z

通过网盘分享的文件：UMA
链接: https://pan.baidu.com/s/1rLa2KEtlxSDCiLEzotfNDw?pwd=g83s 提取码: g83s
--来自百度网盘超级会员v1的分享

leeleolay · 2026-08-19T10:42:17Z

https://paddle-org.bj.bcebos.com/paddlematerials/checkpoints/interatomic_potentials/uma/uma_checkpoints.zip 建议权重名参考模型注册表里的命名方式

leeleolay · 2026-08-19T12:28:50Z

https://paddle-org.bj.bcebos.com/paddlematerials/datasets/UMA/uma_datasets.zip

leeleolay added the MIIT Program label Jun 9, 2026

leeleolay changed the title ~~Add UMA interatomic potential support~~ 【MIIT program】Add UMA interatomic potential support Jun 9, 2026

leeleolay reviewed Jun 9, 2026

View reviewed changes

leeleolay reviewed Jun 18, 2026

View reviewed changes

leeleolay reviewed Jun 22, 2026

View reviewed changes

paddle-bot Bot added the contributor External developers label Jul 13, 2026

leeleolay reviewed Jul 16, 2026

View reviewed changes

yqf66 added 7 commits August 21, 2026 15:46

Add UMA interatomic potential support

aa665e2

Address UMA review comments

c741d27

Refactor UMA integration to PaddleMaterials structure

79c784f

Address UMA README and collator review

787d74f

Simplify UMA for native single-task integration

acee689

Document UMA full-training results

252697e

Address latest UMA review feedback

75e49fd

yqf66 force-pushed the add-uma branch from 1ea9405 to 75e49fd Compare August 21, 2026 07:56

		)


		class UMASingleCollator:

Uh oh!

Conversation

yqf66 commented Jun 2, 2026

Summary

Validation

Notes

Uh oh!

leeleolay left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yqf66 commented Jun 17, 2026

Uh oh!

leeleolay left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leeleolay left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yqf66 commented Jun 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leeleolay commented Jun 22, 2026

Uh oh!

yqf66 commented Jun 29, 2026

Uh oh!

paddle-bot Bot commented Jul 13, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leeleolay commented Jul 16, 2026

Uh oh!

leeleolay left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

leeleolay commented Aug 11, 2026

Uh oh!

yqf66 commented Aug 12, 2026

Uh oh!

leeleolay commented Aug 19, 2026

leeleolay left a comment •

edited

Loading

leeleolay left a comment •

edited

Loading