Upload folder using huggingface_hub

README.md

@@ -0,0 +1,169 @@
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">UMA</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+UMA (Universal Materials Interaction Model) is a universal machine-learning interatomic potential model for materials and catalysis systems, built on equivariant graph neural networks. It predicts energy and forces for atomic structures.
+
+
+# Model Description
+
+UMA is based on an equivariant graph neural network architecture and is trained on diverse materials and catalysis datasets such as OC20, OC22, OC25, OMat, OMOL, ODAC, and OMC. It performs energy and force prediction and structure optimization for catalytic adsorption, inorganic materials, molecular systems, MOFs, and related scenarios.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :---: |
+| OC20 energy and force fine-tuning | Train the UMA model using the standard configuration to read OC20 fine-tuning data |
+| OC22 oxide catalysis fine-tuning | Train the UMA model using the standard configuration to read OC22 fine-tuning data (1P2 only) |
+| OC25 (electro-)catalysis fine-tuning | Train the UMA model using the standard configuration to read OC25 fine-tuning data (1P2 only) |
+| OMat inorganic materials fine-tuning | Train the UMA model using the standard configuration to read OMat fine-tuning data |
+| OMOL molecules + polymers fine-tuning | Train the UMA model using the standard configuration to read OMOL fine-tuning data |
+| ODAC MOFs fine-tuning | Train the UMA model using the standard configuration to read ODAC fine-tuning data |
+| OMC molecular crystals fine-tuning | Train the UMA model using the standard configuration to read OMC fine-tuning data |
+| Training workflow pre-check | Check configuration, data paths, run scripts, and checkpoint placement |
+| Catalytic adsorption system modeling | Refer to OC20/OC22/OC25 tasks for adsorption/catalytic surface system training and inference |
+| Inference script reference | Use upstream inference examples for crystal relaxation, adsorption system relaxation, or molecular MD adaptation |
+| Custom data migration | Convert ASE-readable structures into UMA fine-tuning data and replace the training and validation paths |
+
+
+
+# Usage Instructions
+
+## 1. Using OneCode
+
+You can try out 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**
+
+- GPU or DCU is recommended for fine-tuning training.
+- CPU can be used for configuration and data-path pre-checks; formal UMA training is not recommended on CPU.
+- DCU users need to install DTK in advance. DTK 25.04.2 or above, or the OneScience-recommended version matching the current cluster, is suggested.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/UMA --local_dir ./UMA
+cd UMA
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is also supported
+pip install onescience[matchem-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is also supported
+pip install onescience[matchem-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+
+### Training Dataset Description
+
+This repository does not include built-in training data. The following uses **OC20 fine-tuning** as an example; the procedures for OC22, OC25, OMat, OMOL, ODAC, OMC, and other tasks are the same—just replace `--uma-task` and the data path accordingly.
+
+
+```bash
+modelscope download --dataset OneScience/oc20 --local_dir ./data
+```
+
+
+**Data Format Conversion**
+
+The downloaded raw data is usually in `.extxyz` format and must first be converted to ASE-lmdb format using `scripts/create_uma_finetune_dataset.py`, which also computes `elem_refs` and `normalizer_rmsd`. The script supports the following tasks:
+
+| Task | Description |
+| --- | --- |
+| `oc20` | Catalysis (example) |
+| `oc22` | Oxide catalysis (1P2 only) |
+| `oc25` | (Electro-)catalysis (1P2 only) |
+| `omat` | Inorganic materials |
+| `omol` | Molecules + polymers |
+| `odac` | MOFs |
+| `omc` | Molecular crystals |
+
+
+```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
+```
+
+After conversion, the following are generated:
+
+```text
+data/oc20_finetune/
+├── train/                 # ASE-lmdb training data
+├── val/                   # ASE-lmdb validation data
+└── data/                  # Generated data configuration yaml
+    └── uma_conserving_data_task_energy_force.yaml
+```
+
+Then use `scripts/update_demo_config.py` to update the generated `elem_refs`, `normalizer_rmsd`, and data paths into the demo configuration file:
+
+```bash
+python scripts/update_demo_config.py --demo-config demo/configs/oc20_ef_4dcu.yaml
+```
+
+`demo/run.sh` automatically uses the repository root as `ONESCIENCE_DATASETS_DIR`, so relative paths in the configuration file will match automatically.
+
+### Training Weights
+
+This repository already includes the rotation basis file `weight/Jd.pt`. UMA pre-trained checkpoints (e.g., `uma-s-1p1_converted.pt`) need to be downloaded from the fairchem official repository, converted to UMA format, and placed at (upload coming soon):
+
+```text
+weight/uma-s-1p1_converted.pt
+```
+
+- fairchem official repository: https://github.com/facebookresearch/fairchem
+
+`demo/run.sh` and the example scripts under `inference/` will automatically detect `weight/Jd.pt` and set `ONESCIENCE_UMA_JD_PATH`.
+
+
+### Fine-tuning
+
+```bash
+bash demo/run.sh --config demo/configs/oc20_ef_4dcu.yaml
+```
+
+### Inference
+```bash
+python inference/run_molecular_md.py
+```
+You need to specify the pre-trained weights path yourself.
+
+## OneScience Official Information
+
+| 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 UMA example code comes from the matchem example implementation in the OneScience project and references the upstream fairchem project (https://github.com/facebookresearch/fairchem). The upstream fairchem repository software is released under the [MIT License](https://fair-chem.github.io/core/install.html#license); fairchem model checkpoints and datasets may carry their own separate licenses, so please follow the corresponding instructions when using them.
+- If you use UMA fine-tuning or inference results in scientific research, we recommend citing the UMA/related universal materials interaction model methods, fairchem/OneScience project information, and the sources of the datasets actually used.

configs/data/uma_conserving_data_task_energy.yaml

@@ -0,0 +1,61 @@
+dataset_name: ??
+elem_refs: ??
+normalizer_rmsd: ??
+
+train_dataset:
+  splits:
+    train:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+val_dataset:
+  splits:
+    val:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+regress_stress: False
+pass_through_head_outputs: False
+
+heads:
+  energy:
+    module: onescience.modules.head.uma_head.MLP_Energy_Head
+
+tasks_list:
+- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+  name: energy
+  level: system
+  property: energy
+  loss_fn:
+    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+    coefficient: 20
+  out_spec:
+    dim: [1]
+    dtype: float32
+  normalizer:
+    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  element_references:
+    _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+    element_references:
+      _target_: torch.DoubleTensor
+      _args_:
+        - ${data.elem_refs}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - per_atom_mae

configs/data/uma_conserving_data_task_energy_force.yaml

@@ -0,0 +1,86 @@
+dataset_name: ??
+elem_refs: ??
+normalizer_rmsd: ??
+
+train_dataset:
+  splits:
+    train:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+val_dataset:
+  splits:
+    val:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+regress_stress: False
+pass_through_head_outputs: True
+
+heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+
+tasks_list:
+- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+  name: energy
+  level: system
+  property: energy
+  loss_fn:
+    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+    coefficient: 20
+  out_spec:
+    dim: [1]
+    dtype: float32
+  normalizer:
+    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  element_references:
+    _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+    element_references:
+      _target_: torch.DoubleTensor
+      _args_:
+        - ${data.elem_refs}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - per_atom_mae
+- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+  name: forces
+  level: atom
+  property: forces
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  loss_fn:
+    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.L2NormLoss
+    reduction: mean
+    coefficient: 2
+  out_spec:
+    dim: [3]
+    dtype: float32
+  normalizer:
+    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error

configs/data/uma_conserving_data_task_energy_force_stress.yaml

@@ -0,0 +1,107 @@
+dataset_name: ??
+elem_refs: ??
+normalizer_rmsd: ??
+
+train_dataset:
+  splits:
+    train:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+val_dataset:
+  splits:
+    val:
+      src: ??
+  format: ase_db
+  transforms:
+    common_transform:
+      dataset_name: ${data.dataset_name}
+    stress_reshape_transform:
+      dataset_name: ${data.dataset_name}
+
+regress_stress: True
+pass_through_head_outputs: True
+
+heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+
+tasks_list:
+- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+  name: energy
+  level: system
+  property: energy
+  loss_fn:
+    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+    coefficient: 20
+  out_spec:
+    dim: [1]
+    dtype: float32
+  normalizer:
+    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  element_references:
+    _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+    element_references:
+      _target_: torch.DoubleTensor
+      _args_:
+        - ${data.elem_refs}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - per_atom_mae
+- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+  name: forces
+  level: atom
+  property: forces
+  train_on_free_atoms: True
+  eval_on_free_atoms: True
+  loss_fn:
+    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.L2NormLoss
+    reduction: mean
+    coefficient: 2
+  out_spec:
+    dim: [3]
+    dtype: float32
+  normalizer:
+    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+- _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+  name: stress
+  level: system
+  property: stress
+  loss_fn:
+    _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.MAELoss
+    reduction: mean
+    coefficient: 1
+  out_spec:
+    dim: [1, 9]
+    dtype: float32
+  normalizer:
+    _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+    mean: 0.0
+    rmsd: ${data.normalizer_rmsd}
+  datasets:
+    - ${data.dataset_name}
+  metrics:
+    - mae

configs/uma_sm_finetune_template.yaml

@@ -0,0 +1,115 @@
+defaults:
+  - data: ??
+  - _self_
+
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL
+    ranks_per_node: 1
+    num_nodes: 1
+  debug: True
+  run_dir: ./uma_finetune_runs/
+  run_name: uma_finetune
+  logger:
+    _target_: onescience.utils.uma.common.logger.WandBSingletonLogger.init_wandb
+    _partial_: true
+    entity: example
+    project: uma_finetune
+
+#base_model_name: uma-s-1p1
+# 300 is effectively infinite neighbors, but if you are memory strained, using 100 is usually just as good
+max_neighbors: 300
+# choose to either run for integer number of epochs or steps, only 1 can be step, the other must be null
+epochs: 1
+steps: null
+# customize batch size depending on your system
+batch_size: 2
+lr: 4e-4
+weight_decay: 1e-3
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: 1000
+
+
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config: { sampling: {type: temperature, temperature: 1.0} }
+
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: True
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config: { sampling: {type: temperature, temperature: 1.0} }
+
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: False
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: "${oc.env:ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt"
+      overrides:
+        backbone:
+          otf_graph: True
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: False
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+    - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+      checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+      max_saved_checkpoints: 5
+    - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/_parse_config.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""UMA demo YAML 配置解析器 - 把 YAML 转换为训练启动命令与相关变量。
+
+用法:
+    python _parse_config.py <config.yaml> <action> [hydra_config_path]
+
+action:
+    name           打印实验名称
+    command        打印训练启动命令 (python/torchrun/srun 风格)
+                   需要额外传入 hydra_config_path，作为 train.py -c 的目标
+    env            打印环境变量 export 语句
+    env-args       打印 env_setup.sh 的参数 (conda_env module1 module2 ...)
+    data-files     打印 preflight_check.sh 要检查的路径
+                   (checkpoint_location + data.train_dataset.splits.*.src + data.val_dataset.splits.*.src)
+    slurm          打印 SLURM 配置变量
+    hydra-config   把 YAML 剥离 demo meta 后的部分(纯 hydra 内容)打印到 stdout
+                   run.sh 会把它重定向到 outputs/xxx_ts/hydra_config.yaml
+
+顶层 meta 键（只给 demo 工具用，不进 hydra）：
+    name, description, launch, env, slurm, nccl
+
+其余顶层键（data, job, runner, reducer, train_dataset, ...）作为 hydra 内容。
+"""
+
+import os
+import sys
+
+import yaml
+
+
+# 这些顶层键只给 demo 工具用，在写 hydra_config 时会被剥离
+META_KEYS = {"name", "description", "launch", "env", "slurm", "nccl"}
+
+
+def load_config(path):
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def get_train_py_path(config_path):
+    """根据 demo/configs/xxx.yaml 的位置推算 UMA/train.py 路径。"""
+    # config_path: .../UMA/demo/configs/xxx.yaml
+    demo_dir = os.path.dirname(os.path.dirname(os.path.abspath(config_path)))
+    uma_dir = os.path.dirname(demo_dir)
+    return os.path.join(uma_dir, "train.py")
+
+
+# -----------------------------------------------------------------------------
+# action 实现
+# -----------------------------------------------------------------------------
+
+def print_name(cfg):
+    print(cfg.get("name", "uma_train"))
+
+
+def print_command(cfg, config_path, hydra_config_path):
+    """拼 python/torchrun/srun train.py -c <hydra_config_path>。"""
+    if not hydra_config_path:
+        print(
+            "[ERROR] command 动作需要提供 hydra_config_path 作为第三个参数",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    launch = cfg.get("launch", {}) or {}
+    num_nodes = launch.get("num_nodes", 1)
+    num_gpus = launch.get("num_gpus", 1)
+    launcher = launch.get("launcher", "python")
+
+    train_py = get_train_py_path(config_path)
+
+    if num_nodes > 1:
+        # 多节点: sbatch 内会用 srun 包裹 torchrun, 命令本体是 torchrun + train.py
+        # 注意: master_addr/master_port/node_rank 在 run.sh 生成的 submit.sh 里
+        # 基于 $SLURM_* 变量导出, 这里占位用 ${MASTER_ADDR} 等 shell 变量
+        print(
+            "torchrun \\\n"
+            f"  --nnodes={num_nodes} \\\n"
+            "  --node_rank=${SLURM_NODEID} \\\n"
+            f"  --nproc_per_node={num_gpus} \\\n"
+            "  --rdzv_id=${SLURM_JOB_ID} \\\n"
+            "  --rdzv_backend=c10d \\\n"
+            "  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} \\\n"
+            f"  {train_py} \\\n"
+            f"  -c {hydra_config_path}"
+        )
+    elif launcher == "torchrun" or num_gpus > 1:
+        # 单节点多卡: torchrun
+        print(
+            "torchrun \\\n"
+            "  --nnodes=1 \\\n"
+            f"  --nproc_per_node={num_gpus} \\\n"
+            f"  {train_py} \\\n"
+            f"  -c {hydra_config_path}"
+        )
+    else:
+        # 单卡: 直接 python
+        print(f"python {train_py} -c {hydra_config_path}")
+
+
+def print_env(cfg):
+    """打印环境变量 export 语句。"""
+    launch = cfg.get("launch", {}) or {}
+    num_nodes = launch.get("num_nodes", 1)
+    num_gpus = launch.get("num_gpus", 1)
+
+    omp = launch.get("omp_num_threads", 1)
+    print(f"export OMP_NUM_THREADS={omp}")
+
+    if num_gpus > 1:
+        devices = ",".join(str(i) for i in range(num_gpus))
+        print(f"export HIP_VISIBLE_DEVICES={devices}")
+    elif num_gpus == 1:
+        print("export HIP_VISIBLE_DEVICES=0")
+
+    # 多节点 NCCL 配置
+    if num_nodes > 1:
+        nccl = cfg.get("nccl", {}) or {}
+        print("export HSA_FORCE_FINE_GRAIN_PCIE=1")
+        if nccl.get("socket_ifname"):
+            print(f"export NCCL_SOCKET_IFNAME={nccl['socket_ifname']}")
+        if nccl.get("ib_hca"):
+            print(f"export NCCL_IB_HCA={nccl['ib_hca']}")
+        if nccl.get("proto"):
+            print(f"export NCCL_PROTO={nccl['proto']}")
+
+
+def print_env_args(cfg):
+    """打印 env_setup.sh 的参数: conda_env module1 module2 ..."""
+    env = cfg.get("env", {}) or {}
+    conda_env = env.get("conda_env", "chem")
+    modules = env.get("modules", []) or []
+    parts = [conda_env] + list(modules)
+    print(" ".join(parts))
+
+
+def _dig_src_paths(node, out):
+    """递归找出 splits 下所有 src 字段。"""
+    if isinstance(node, dict):
+        if "splits" in node and isinstance(node["splits"], dict):
+            for split_name, split_val in node["splits"].items():
+                if isinstance(split_val, dict) and "src" in split_val:
+                    out.append(split_val["src"])
+        for v in node.values():
+            _dig_src_paths(v, out)
+    elif isinstance(node, list):
+        for item in node:
+            _dig_src_paths(item, out)
+
+
+def _dig_checkpoint_locations(node, out):
+    """递归找出 checkpoint_location 字段。"""
+    if isinstance(node, dict):
+        if "checkpoint_location" in node and isinstance(
+            node["checkpoint_location"], str
+        ):
+            out.append(node["checkpoint_location"])
+        for v in node.values():
+            _dig_checkpoint_locations(v, out)
+    elif isinstance(node, list):
+        for item in node:
+            _dig_checkpoint_locations(item, out)
+
+
+def print_data_files(cfg):
+    """打印 preflight 要检查的路径: checkpoint + 数据 src。"""
+    paths = []
+    data = cfg.get("data", {})
+    _dig_src_paths(data, paths)
+    _dig_checkpoint_locations(cfg, paths)
+    # 去重保持顺序
+    seen = set()
+    for p in paths:
+        if p and p not in seen and not str(p).startswith("??"):
+            seen.add(p)
+            print(p)
+
+
+def print_slurm(cfg):
+    """打印 SLURM 配置变量。"""
+    launch = cfg.get("launch", {}) or {}
+    slurm = cfg.get("slurm", {}) or {}
+    num_nodes = launch.get("num_nodes", 1)
+    num_gpus = launch.get("num_gpus", 1)
+
+    name = cfg.get("name", "uma_train")
+    partition = slurm.get("partition", "k100ai")
+    time_limit = slurm.get("time", "12:00:00")
+    cpus = slurm.get("cpus_per_task", 128)
+
+    # UMA 的多节点策略: 每个节点 1 个 srun task, 由 torchrun 在该节点内再 spawn
+    # nproc_per_node=num_gpus 个进程. 因此无论单/多节点 NTASKS_PER_NODE 都是 1.
+    ntasks = 1
+    cpus_per_task = cpus
+    _ = num_gpus  # 保留 num_gpus 以便 GPUS_PER_NODE 正确填写
+
+    print(f"JOB_NAME={name}")
+    print(f"PARTITION={partition}")
+    print(f"NODES={num_nodes}")
+    print(f"NTASKS_PER_NODE={ntasks}")
+    print(f"CPUS_PER_TASK={cpus_per_task}")
+    print(f"GPUS_PER_NODE={num_gpus}")
+    print(f"TIME={time_limit}")
+
+
+def print_hydra_config(cfg):
+    """把 cfg 剥离 META_KEYS 后的部分用 YAML 输出。"""
+    hydra_only = {k: v for k, v in cfg.items() if k not in META_KEYS}
+    # sort_keys=False 保留原顺序; default_flow_style=False 强制 block 风格
+    yaml.safe_dump(
+        hydra_only,
+        sys.stdout,
+        sort_keys=False,
+        default_flow_style=False,
+        allow_unicode=True,
+    )
+
+
+# -----------------------------------------------------------------------------
+# main
+# -----------------------------------------------------------------------------
+
+def main():
+    if len(sys.argv) < 3:
+        print(__doc__, file=sys.stderr)
+        sys.exit(1)
+
+    config_path = sys.argv[1]
+    action = sys.argv[2]
+    hydra_config_path = sys.argv[3] if len(sys.argv) >= 4 else None
+
+    cfg = load_config(config_path)
+
+    actions = {
+        "name": lambda: print_name(cfg),
+        "command": lambda: print_command(cfg, config_path, hydra_config_path),
+        "env": lambda: print_env(cfg),
+        "env-args": lambda: print_env_args(cfg),
+        "data-files": lambda: print_data_files(cfg),
+        "slurm": lambda: print_slurm(cfg),
+        "hydra-config": lambda: print_hydra_config(cfg),
+    }
+
+    fn = actions.get(action)
+    if fn is None:
+        print(f"未知 action: {action}", file=sys.stderr)
+        print(f"可用: {', '.join(actions.keys())}", file=sys.stderr)
+        sys.exit(1)
+    fn()
+
+
+if __name__ == "__main__":
+    main()

demo/configs/_oc20_ef_4dcu_filled.yaml

@@ -0,0 +1,303 @@
+name: oc20_ef_4dcu
+description: UMA OC20 finetune - 1 node x 4 DCU (energy+forces)
+launch:
+  launcher: torchrun
+  num_nodes: 1
+  num_gpus: 4
+  omp_num_threads: 1
+env:
+  conda_env: chem
+  modules:
+  - sghpc-mpi-gcc/26.3
+slurm:
+  partition: k100ai
+  time: '12:00:00'
+  cpus_per_task: 128
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL
+    ranks_per_node: 4
+    num_nodes: 1
+  debug: true
+  run_dir: ./uma_finetune_runs/
+  run_name: oc20_ef_4dcu_smoke
+  logger: null
+max_neighbors: 300
+epochs: 1
+steps: 20
+batch_size: 2
+lr: 0.0004
+weight_decay: 0.001
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: 1000
+data:
+  dataset_name: oc20
+  elem_refs:
+  - -2.338112300549499e-17
+  - -3.397406033745951
+  - -3.552713678800501e-15
+  - -1.4883927423881005e-15
+  - -2.4424906541753444e-15
+  - -6.05474814945516
+  - -8.359419908238154
+  - -8.461445707766968
+  - -4.772501185489789
+  - -4.9960036108132044e-15
+  - -3.552713678800501e-15
+  - -1.505306152902931
+  - 6.217248937900877e-15
+  - -3.5935859953202867
+  - -5.298043899940335
+  - -5.289195691841476
+  - -4.586309292435159
+  - -3.145306223924418
+  - 8.43769498715119e-15
+  - -1.5534031951839165
+  - -2.452149990972268
+  - -6.698950860236158
+  - -7.784447374762749
+  - -8.35172207094275
+  - -8.475443777072595
+  - -7.835952917487167
+  - -7.177746464190568
+  - -6.276117296504732
+  - -4.906667819649391
+  - -2.9770431302950078
+  - -0.6870767989674853
+  - -2.570966565048506
+  - -4.112249499833631
+  - -4.501937385014317
+  - -3.7911090450737395
+  - 0.0
+  - 0.0
+  - -1.404014767200028
+  - -2.3340636233592837
+  - -6.977544301327938
+  - -8.429019781059974
+  - -9.558754587545973
+  - -9.914673095966668
+  - -9.262176302297178
+  - -8.344618489779986
+  - -6.81614547353467
+  - -4.823268946054286
+  - -1.9248732292365998
+  - -0.22855559486612576
+  - -2.0711980156857903
+  - -3.3848661082103417
+  - -3.693940774789746
+  - -3.0317412637444727
+  - 0.0
+  - 4.440892098500626e-16
+  - -1.69210743883596
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - -9.847846174118736
+  - -11.050291739766923
+  - -11.612898250816091
+  - -11.043957156699099
+  - -10.20788381032279
+  - -8.329494152492249
+  - -5.758362379765965
+  - -2.705571683868682
+  - 0.3244215988004441
+  - -1.801036139937328
+  - -2.98927510335205
+  - -3.3414587373794893
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  normalizer_rmsd: 0.3387841421237734
+  train_dataset:
+    splits:
+      train:
+        src: "${ONESCIENCE_DATASETS_DIR}/oc20/s2ef_200k_uncompressed"
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  val_dataset:
+    splits:
+      val:
+        src: "${ONESCIENCE_DATASETS_DIR}/oc20/s2ef_val_id_uncompressed"
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  regress_stress: false
+  pass_through_head_outputs: true
+  heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+  tasks_list:
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: energy
+    level: system
+    property: energy
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+      coefficient: 20
+    out_spec:
+      dim:
+      - 1
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    element_references:
+      _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+      element_references:
+        _target_: torch.DoubleTensor
+        _args_:
+        - ${data.elem_refs}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - per_atom_mae
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: forces
+    level: atom
+    property: forces
+    train_on_free_atoms: true
+    eval_on_free_atoms: true
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.L2NormLoss
+      reduction: mean
+      coefficient: 2
+    out_spec:
+      dim:
+      - 3
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: true
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: false
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: "${ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt"
+      overrides:
+        backbone:
+          otf_graph: true
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: false
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+  - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+    checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+    max_saved_checkpoints: 5
+  - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/configs/oc20_ef_16dcu.yaml

@@ -0,0 +1,270 @@
+# ============================================================
+# UMA OC20 微调示例 - 2 节点 x 8 DCU, energy + forces
+# ============================================================
+#
+# 使用前置:
+#   1. 已运行 scripts/create_uma_finetune_dataset.py 转换数据集
+#      (或手动将 ASE-lmdb 格式的 train/ val/ 放到合适路径)
+#   2. 已运行 scripts/convert_model.py 得到可微调的 checkpoint .pt
+#   3. 将 data.dataset_name / elem_refs / normalizer_rmsd 填为数据生成脚本输出的值
+#
+# 运行:
+#   bash run.sh --config configs/oc20_ef_16dcu.yaml --submit     # 多节点只能 sbatch
+#   bash run.sh --config configs/oc20_ef_16dcu.yaml --dry-run    # 预览
+# ============================================================
+
+name: oc20_ef_16dcu
+description: "UMA OC20 finetune - 2 nodes x 8 DCU (energy+forces)"
+
+# ============================================================
+# 启动配置 (控制 python / torchrun / srun 选择)
+# ============================================================
+launch:
+  launcher: torchrun
+  num_nodes: 2
+  num_gpus: 8
+  omp_num_threads: 1
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: chem
+  modules:
+    - sghpc-mpi-gcc/26.3
+
+# ============================================================
+# SLURM 配置 (仅 --submit 模式使用)
+# ============================================================
+slurm:
+  partition: k100ai
+  time: "48:00:00"
+  cpus_per_task: 128
+
+# ============================================================
+# 多节点 NCCL 通信配置
+# ============================================================
+nccl:
+  socket_ifname: ib0
+  ib_hca: mlx5_0
+  proto: Simple
+
+# ============================================================
+# ---------- 下面是 hydra 配置内容 (直接传给 train.py) ----------
+# 注意: 本文件的 name/launch/env/slurm/nccl 会被 run.sh 在写出
+# outputs/xxx_ts/hydra_config.yaml 时剥离, 不会进入 hydra 配置。
+# ============================================================
+
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL                    # 由 run.sh 驱动 srun+torchrun, 不让 train.py submitit
+    ranks_per_node: 8
+    num_nodes: 2
+  debug: True
+  run_dir: ./uma_finetune_runs/
+  run_name: oc20_ef_16dcu
+  logger:
+    _target_: onescience.utils.uma.common.logger.WandBSingletonLogger.init_wandb
+    _partial_: true
+    entity: example
+    project: uma_finetune
+
+max_neighbors: 300
+epochs: 1
+steps: null
+batch_size: 2
+lr: 4.0e-4
+weight_decay: 1.0e-3
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: 1000
+
+# ============================================================
+# 数据配置 (来自 configs/data/uma_conserving_data_task_energy_force.yaml)
+# 下面 4 项必须由用户填写:
+#   dataset_name    : 例如 "oc20" / "omat" / "omol" (与 base model 预训练任务对应)
+#   elem_refs       : 每元素参考能量列表 (长度 118), 由 create_finetune_dataset.py 生成
+#   normalizer_rmsd : 力的 RMSD 归一化常数, 由 create_finetune_dataset.py 生成
+#   train/val src   : ASE-lmdb 数据目录的绝对/相对路径
+# ============================================================
+data:
+  dataset_name: ??                 # 请填写, e.g. "oc20"
+  elem_refs: ??                    # 请填写, 列表
+  normalizer_rmsd: ??              # 请填写, float
+
+  train_dataset:
+    splits:
+      train:
+        src: "${ONESCIENCE_DATASETS_DIR}/oc20/s2ef_200k_uncompressed"
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+
+  val_dataset:
+    splits:
+      val:
+        src: "${ONESCIENCE_DATASETS_DIR}/oc20/s2ef_val_id_uncompressed"
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+
+  regress_stress: False
+  pass_through_head_outputs: True
+
+  heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+
+  tasks_list:
+    - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+      name: energy
+      level: system
+      property: energy
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+        loss_fn:
+          _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+        coefficient: 20
+      out_spec:
+        dim: [1]
+        dtype: float32
+      normalizer:
+        _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+        mean: 0.0
+        rmsd: ${data.normalizer_rmsd}
+      element_references:
+        _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+        element_references:
+          _target_: torch.DoubleTensor
+          _args_:
+            - ${data.elem_refs}
+      datasets:
+        - ${data.dataset_name}
+      metrics:
+        - mae
+        - per_atom_mae
+    - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+      name: forces
+      level: atom
+      property: forces
+      train_on_free_atoms: True
+      eval_on_free_atoms: True
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+        loss_fn:
+          _target_: onescience.modules.loss.uma_loss.L2NormLoss
+        reduction: mean
+        coefficient: 2
+      out_spec:
+        dim: [3]
+        dtype: float32
+      normalizer:
+        _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+        mean: 0.0
+        rmsd: ${data.normalizer_rmsd}
+      datasets:
+        - ${data.dataset_name}
+      metrics:
+        - mae
+        - cosine_similarity
+        - magnitude_error
+
+# ============================================================
+# 数据流与训练 Runner (来自 configs/uma_sm_finetune_template.yaml)
+# ============================================================
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: True
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: False
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: "${ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt"
+      overrides:
+        backbone:
+          otf_graph: True
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: False
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+    - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+      checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+      max_saved_checkpoints: 5
+    - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/configs/oc20_ef_4dcu.yaml

@@ -0,0 +1,303 @@
+name: oc20_ef_4dcu
+description: UMA OC20 finetune - 1 node x 4 DCU (energy+forces)
+launch:
+  launcher: torchrun
+  num_nodes: 1
+  num_gpus: 4
+  omp_num_threads: 1
+env:
+  conda_env: uv311_mat
+  modules:
+  - sghpc-mpi-gcc/26.3
+slurm:
+  partition: hpctest01
+  time: '2:00:00'
+  cpus_per_task: 64
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL
+    ranks_per_node: 4
+    num_nodes: 1
+  debug: true
+  run_dir: ./uma_finetune_runs/
+  run_name: oc20_ef_4dcu
+  logger: null
+max_neighbors: 300
+epochs: 1
+steps: null
+batch_size: 2
+lr: 0.0004
+weight_decay: 0.001
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: null
+data:
+  dataset_name: oc20
+  elem_refs:
+  - -9.01836389963216e-17
+  - -3.384793156087369
+  - -3.552713678800501e-15
+  - -1.4210854715202004e-14
+  - -3.552713678800501e-15
+  - -6.096600969433027
+  - -8.288307365078008
+  - -8.377401871630358
+  - -4.73277382959691
+  - -4.218847493575595e-15
+  - 8.881784197001252e-15
+  - -1.5758719946770752
+  - 3.552713678800501e-15
+  - -3.605510955946432
+  - -5.341968801697664
+  - -5.26720503097145
+  - -4.566219949129778
+  - -3.0961819892844145
+  - 0.0
+  - -1.5790045471096414
+  - -2.4721953289673877
+  - -6.6807503507860915
+  - -7.74796841292303
+  - -8.403945117874883
+  - -8.507943997230633
+  - -7.859727128191719
+  - -7.202700747006435
+  - -6.283218856635957
+  - -4.8972399918354865
+  - -2.9438503901650144
+  - -0.6766002592520115
+  - -2.577391165910655
+  - -4.06962014738207
+  - -4.497309369632559
+  - -3.780665648345848
+  - -5.329070518200751e-15
+  - -1.3322676295501878e-15
+  - -1.4256648977370752
+  - -2.412526481733912
+  - -7.014442857905434
+  - -8.512272383088028
+  - -9.592355889199172
+  - -9.910578824883942
+  - -9.267018354580234
+  - -8.354773494931816
+  - -6.854409258193916
+  - -4.840002859370445
+  - -1.918019107887611
+  - -0.23821365616560497
+  - -2.0561553805596358
+  - -3.3842132848566258
+  - -3.667988196267832
+  - -3.032988525381068
+  - 2.220446049250313e-16
+  - 0.0
+  - -1.6559264334420458
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - -9.875111813171198
+  - -11.052180114376393
+  - -11.621851352916064
+  - -11.117519848695258
+  - -10.219707389902464
+  - -8.357987841676932
+  - -5.796390731438155
+  - -2.701232714681403
+  - 0.32034288848399406
+  - -1.8259293180072989
+  - -2.991135014870159
+  - -3.3375851657056166
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  normalizer_rmsd: 0.3184592617813517
+  train_dataset:
+    splits:
+      train:
+        src: ${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/train
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  val_dataset:
+    splits:
+      val:
+        src: ${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/val
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  regress_stress: false
+  pass_through_head_outputs: true
+  heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+  tasks_list:
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: energy
+    level: system
+    property: energy
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+      coefficient: 20
+    out_spec:
+      dim:
+      - 1
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    element_references:
+      _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+      element_references:
+        _target_: torch.DoubleTensor
+        _args_:
+        - ${data.elem_refs}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - per_atom_mae
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: forces
+    level: atom
+    property: forces
+    train_on_free_atoms: true
+    eval_on_free_atoms: true
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.L2NormLoss
+      reduction: mean
+      coefficient: 2
+    out_spec:
+      dim:
+      - 3
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: true
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: false
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: ${ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt
+      overrides:
+        backbone:
+          otf_graph: true
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: false
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+  - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+    checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+    max_saved_checkpoints: 5
+  - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/configs/oc20_ef_8dcu.yaml

@@ -0,0 +1,263 @@
+# ============================================================
+# UMA OC20 微调示例 - 单节点 8 DCU, energy + forces
+# ============================================================
+#
+# 使用前置:
+#   1. 已运行 scripts/create_uma_finetune_dataset.py 转换数据集
+#      (或手动将 ASE-lmdb 格式的 train/ val/ 放到合适路径)
+#   2. 已运行 scripts/convert_model.py 得到可微调的 checkpoint .pt
+#   3. 将 data.dataset_name / elem_refs / normalizer_rmsd 填为数据生成脚本输出的值
+#
+# 运行:
+#   bash run.sh --config configs/oc20_ef_8dcu.yaml              # 直接跑
+#   bash run.sh --config configs/oc20_ef_8dcu.yaml --submit     # sbatch
+#   bash run.sh --config configs/oc20_ef_8dcu.yaml --dry-run    # 预览
+# ============================================================
+
+name: oc20_ef_8dcu
+description: "UMA OC20 finetune - 1 node x 8 DCU (energy+forces)"
+
+# ============================================================
+# 启动配置 (控制 python / torchrun / srun 选择)
+# ============================================================
+launch:
+  launcher: torchrun
+  num_nodes: 1
+  num_gpus: 8
+  omp_num_threads: 1
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: chem
+  modules:
+    - sghpc-mpi-gcc/26.3
+
+# ============================================================
+# SLURM 配置 (仅 --submit 模式使用)
+# ============================================================
+slurm:
+  partition: k100ai
+  time: "12:00:00"
+  cpus_per_task: 128
+
+# ============================================================
+# ---------- 下面是 hydra 配置内容 (直接传给 train.py) ----------
+# 注意: 本文件的 name/launch/env/slurm/nccl 会被 run.sh 在写出
+# outputs/xxx_ts/hydra_config.yaml 时剥离, 不会进入 hydra 配置。
+# ============================================================
+
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL                    # 由 run.sh 驱动 torchrun, 不让 train.py submitit
+    ranks_per_node: 8
+    num_nodes: 1
+  debug: True
+  run_dir: ./uma_finetune_runs/
+  run_name: oc20_ef_8dcu
+  logger:
+    _target_: onescience.utils.uma.common.logger.WandBSingletonLogger.init_wandb
+    _partial_: true
+    entity: example
+    project: uma_finetune
+
+max_neighbors: 300
+epochs: 1
+steps: null
+batch_size: 2
+lr: 4.0e-4
+weight_decay: 1.0e-3
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: 1000
+
+# ============================================================
+# 数据配置 (来自 configs/data/uma_conserving_data_task_energy_force.yaml)
+# 下面 4 项必须由用户填写:
+#   dataset_name    : 例如 "oc20" / "omat" / "omol" (与 base model 预训练任务对应)
+#   elem_refs       : 每元素参考能量列表 (长度 118), 由 create_finetune_dataset.py 生成
+#   normalizer_rmsd : 力的 RMSD 归一化常数, 由 create_finetune_dataset.py 生成
+#   train/val src   : ASE-lmdb 数据目录的绝对/相对路径
+# ============================================================
+data:
+  dataset_name: ??                 # 请填写, e.g. "oc20"
+  elem_refs: ??                    # 请填写, 列表
+  normalizer_rmsd: ??              # 请填写, float
+
+  train_dataset:
+    splits:
+      train:
+        src: "${ONESCIENCE_DATASETS_DIR}/oc20/s2ef_200k_uncompressed"
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+
+  val_dataset:
+    splits:
+      val:
+        src: "${ONESCIENCE_DATASETS_DIR}/oc20/s2ef_val_id_uncompressed"
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+
+  regress_stress: False
+  pass_through_head_outputs: True
+
+  heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+
+  tasks_list:
+    - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+      name: energy
+      level: system
+      property: energy
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+        loss_fn:
+          _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+        coefficient: 20
+      out_spec:
+        dim: [1]
+        dtype: float32
+      normalizer:
+        _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+        mean: 0.0
+        rmsd: ${data.normalizer_rmsd}
+      element_references:
+        _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+        element_references:
+          _target_: torch.DoubleTensor
+          _args_:
+            - ${data.elem_refs}
+      datasets:
+        - ${data.dataset_name}
+      metrics:
+        - mae
+        - per_atom_mae
+    - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+      name: forces
+      level: atom
+      property: forces
+      train_on_free_atoms: True
+      eval_on_free_atoms: True
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+        loss_fn:
+          _target_: onescience.modules.loss.uma_loss.L2NormLoss
+        reduction: mean
+        coefficient: 2
+      out_spec:
+        dim: [3]
+        dtype: float32
+      normalizer:
+        _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+        mean: 0.0
+        rmsd: ${data.normalizer_rmsd}
+      datasets:
+        - ${data.dataset_name}
+      metrics:
+        - mae
+        - cosine_similarity
+        - magnitude_error
+
+# ============================================================
+# 数据流与训练 Runner (来自 configs/uma_sm_finetune_template.yaml)
+# ============================================================
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: True
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: True
+    batch_size: ${batch_size}
+    shuffle: False
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: "${ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt"
+      overrides:
+        backbone:
+          otf_graph: True
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: False
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+    - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+      checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+      max_saved_checkpoints: 5
+    - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/outputs/oc20_ef_4dcu_20260721_114458/config.yaml

@@ -0,0 +1,303 @@
+name: oc20_ef_4dcu
+description: UMA OC20 finetune - 1 node x 4 DCU (energy+forces)
+launch:
+  launcher: torchrun
+  num_nodes: 1
+  num_gpus: 4
+  omp_num_threads: 1
+env:
+  conda_env: uv311_mat
+  modules:
+  - sghpc-mpi-gcc/26.3
+slurm:
+  partition: hpctest01
+  time: '2:00:00'
+  cpus_per_task: 64
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL
+    ranks_per_node: 4
+    num_nodes: 1
+  debug: true
+  run_dir: ./uma_finetune_runs/
+  run_name: oc20_ef_4dcu
+  logger: null
+max_neighbors: 300
+epochs: 1
+steps: null
+batch_size: 2
+lr: 0.0004
+weight_decay: 0.001
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: null
+data:
+  dataset_name: oc20
+  elem_refs:
+  - -9.01836389963216e-17
+  - -3.384793156087369
+  - -3.552713678800501e-15
+  - -1.4210854715202004e-14
+  - -3.552713678800501e-15
+  - -6.096600969433027
+  - -8.288307365078008
+  - -8.377401871630358
+  - -4.73277382959691
+  - -4.218847493575595e-15
+  - 8.881784197001252e-15
+  - -1.5758719946770752
+  - 3.552713678800501e-15
+  - -3.605510955946432
+  - -5.341968801697664
+  - -5.26720503097145
+  - -4.566219949129778
+  - -3.0961819892844145
+  - 0.0
+  - -1.5790045471096414
+  - -2.4721953289673877
+  - -6.6807503507860915
+  - -7.74796841292303
+  - -8.403945117874883
+  - -8.507943997230633
+  - -7.859727128191719
+  - -7.202700747006435
+  - -6.283218856635957
+  - -4.8972399918354865
+  - -2.9438503901650144
+  - -0.6766002592520115
+  - -2.577391165910655
+  - -4.06962014738207
+  - -4.497309369632559
+  - -3.780665648345848
+  - -5.329070518200751e-15
+  - -1.3322676295501878e-15
+  - -1.4256648977370752
+  - -2.412526481733912
+  - -7.014442857905434
+  - -8.512272383088028
+  - -9.592355889199172
+  - -9.910578824883942
+  - -9.267018354580234
+  - -8.354773494931816
+  - -6.854409258193916
+  - -4.840002859370445
+  - -1.918019107887611
+  - -0.23821365616560497
+  - -2.0561553805596358
+  - -3.3842132848566258
+  - -3.667988196267832
+  - -3.032988525381068
+  - 2.220446049250313e-16
+  - 0.0
+  - -1.6559264334420458
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - -9.875111813171198
+  - -11.052180114376393
+  - -11.621851352916064
+  - -11.117519848695258
+  - -10.219707389902464
+  - -8.357987841676932
+  - -5.796390731438155
+  - -2.701232714681403
+  - 0.32034288848399406
+  - -1.8259293180072989
+  - -2.991135014870159
+  - -3.3375851657056166
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  normalizer_rmsd: 0.3184592617813517
+  train_dataset:
+    splits:
+      train:
+        src: ${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/train
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  val_dataset:
+    splits:
+      val:
+        src: ${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/val
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  regress_stress: false
+  pass_through_head_outputs: true
+  heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+  tasks_list:
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: energy
+    level: system
+    property: energy
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+      coefficient: 20
+    out_spec:
+      dim:
+      - 1
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    element_references:
+      _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+      element_references:
+        _target_: torch.DoubleTensor
+        _args_:
+        - ${data.elem_refs}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - per_atom_mae
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: forces
+    level: atom
+    property: forces
+    train_on_free_atoms: true
+    eval_on_free_atoms: true
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.L2NormLoss
+      reduction: mean
+      coefficient: 2
+    out_spec:
+      dim:
+      - 3
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: true
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: false
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: ${ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt
+      overrides:
+        backbone:
+          otf_graph: true
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: false
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+  - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+    checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+    max_saved_checkpoints: 5
+  - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/outputs/oc20_ef_4dcu_20260721_114458/hydra_config.yaml

@@ -0,0 +1,288 @@
+job:
+  device_type: CUDA
+  scheduler:
+    mode: LOCAL
+    ranks_per_node: 4
+    num_nodes: 1
+  debug: true
+  run_dir: ./uma_finetune_runs/
+  run_name: oc20_ef_4dcu
+  logger: null
+max_neighbors: 300
+epochs: 1
+steps: null
+batch_size: 2
+lr: 0.0004
+weight_decay: 0.001
+evaluate_every_n_steps: 100
+checkpoint_every_n_steps: null
+data:
+  dataset_name: oc20
+  elem_refs:
+  - -9.01836389963216e-17
+  - -3.384793156087369
+  - -3.552713678800501e-15
+  - -1.4210854715202004e-14
+  - -3.552713678800501e-15
+  - -6.096600969433027
+  - -8.288307365078008
+  - -8.377401871630358
+  - -4.73277382959691
+  - -4.218847493575595e-15
+  - 8.881784197001252e-15
+  - -1.5758719946770752
+  - 3.552713678800501e-15
+  - -3.605510955946432
+  - -5.341968801697664
+  - -5.26720503097145
+  - -4.566219949129778
+  - -3.0961819892844145
+  - 0.0
+  - -1.5790045471096414
+  - -2.4721953289673877
+  - -6.6807503507860915
+  - -7.74796841292303
+  - -8.403945117874883
+  - -8.507943997230633
+  - -7.859727128191719
+  - -7.202700747006435
+  - -6.283218856635957
+  - -4.8972399918354865
+  - -2.9438503901650144
+  - -0.6766002592520115
+  - -2.577391165910655
+  - -4.06962014738207
+  - -4.497309369632559
+  - -3.780665648345848
+  - -5.329070518200751e-15
+  - -1.3322676295501878e-15
+  - -1.4256648977370752
+  - -2.412526481733912
+  - -7.014442857905434
+  - -8.512272383088028
+  - -9.592355889199172
+  - -9.910578824883942
+  - -9.267018354580234
+  - -8.354773494931816
+  - -6.854409258193916
+  - -4.840002859370445
+  - -1.918019107887611
+  - -0.23821365616560497
+  - -2.0561553805596358
+  - -3.3842132848566258
+  - -3.667988196267832
+  - -3.032988525381068
+  - 2.220446049250313e-16
+  - 0.0
+  - -1.6559264334420458
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - -9.875111813171198
+  - -11.052180114376393
+  - -11.621851352916064
+  - -11.117519848695258
+  - -10.219707389902464
+  - -8.357987841676932
+  - -5.796390731438155
+  - -2.701232714681403
+  - 0.32034288848399406
+  - -1.8259293180072989
+  - -2.991135014870159
+  - -3.3375851657056166
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  - 0.0
+  normalizer_rmsd: 0.3184592617813517
+  train_dataset:
+    splits:
+      train:
+        src: ${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/train
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  val_dataset:
+    splits:
+      val:
+        src: ${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/val
+    format: ase_db
+    transforms:
+      common_transform:
+        dataset_name: ${data.dataset_name}
+      stress_reshape_transform:
+        dataset_name: ${data.dataset_name}
+  regress_stress: false
+  pass_through_head_outputs: true
+  heads:
+    efs:
+      module: onescience.modules.head.uma_head.MLP_EFS_Head
+  tasks_list:
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: energy
+    level: system
+    property: energy
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.PerAtomMAELoss
+      coefficient: 20
+    out_spec:
+      dim:
+      - 1
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    element_references:
+      _target_: onescience.utils.uma.normalization.element_references.ElementReferences
+      element_references:
+        _target_: torch.DoubleTensor
+        _args_:
+        - ${data.elem_refs}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - per_atom_mae
+  - _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.Task
+    name: forces
+    level: atom
+    property: forces
+    train_on_free_atoms: true
+    eval_on_free_atoms: true
+    loss_fn:
+      _target_: onescience.modules.loss.uma_loss.DDPMTLoss
+      loss_fn:
+        _target_: onescience.modules.loss.uma_loss.L2NormLoss
+      reduction: mean
+      coefficient: 2
+    out_spec:
+      dim:
+      - 3
+      dtype: float32
+    normalizer:
+      _target_: onescience.utils.uma.normalization.normalizer.Normalizer
+      mean: 0.0
+      rmsd: ${data.normalizer_rmsd}
+    datasets:
+    - ${data.dataset_name}
+    metrics:
+    - mae
+    - cosine_similarity
+    - magnitude_error
+train_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.train_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+train_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${train_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: true
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+val_dataset:
+  _target_: onescience.datapipes.materials.custom_stack.storage.mt_concat_dataset.create_concat_dataset
+  dataset_configs:
+    finetune_dataset: ${data.val_dataset}
+  combined_dataset_config:
+    sampling:
+      type: temperature
+      temperature: 1.0
+eval_dataloader:
+  _target_: onescience.datapipes.materials.uma.dataloader_builder.get_dataloader
+  dataset: ${val_dataset}
+  batch_sampler_fn:
+    _target_: onescience.utils.uma.common.data_parallel.BalancedBatchSampler
+    _partial_: true
+    batch_size: ${batch_size}
+    shuffle: false
+    seed: 0
+  num_workers: 0
+  collate_fn:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.mt_collater_adapter
+    tasks: ${data.tasks_list}
+    exclude_keys: []
+runner:
+  _target_: onescience.utils.uma.components.train.train_runner.TrainEvalRunner
+  train_dataloader: ${train_dataloader}
+  eval_dataloader: ${eval_dataloader}
+  train_eval_unit:
+    _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.MLIPTrainEvalUnit
+    job_config: ${job}
+    tasks: ${data.tasks_list}
+    model:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit.initialize_finetuning_model
+      checkpoint_location: ${ONESCIENCE_MODELS_DIR}/weight/uma-s-1p1_converted.pt
+      overrides:
+        backbone:
+          otf_graph: true
+          max_neighbors: ${max_neighbors}
+          regress_stress: ${data.regress_stress}
+          always_use_pbc: false
+        pass_through_head_outputs: ${data.pass_through_head_outputs}
+      heads: ${data.heads}
+    optimizer_fn:
+      _target_: torch.optim.AdamW
+      _partial_: true
+      lr: ${lr}
+      weight_decay: ${weight_decay}
+    cosine_lr_scheduler_fn:
+      _target_: onescience.utils.uma.units.mlip_unit.mlip_unit._get_consine_lr_scheduler
+      _partial_: true
+      warmup_factor: 0.2
+      warmup_epochs: 0.01
+      lr_min_factor: 0.01
+      epochs: ${epochs}
+      steps: ${steps}
+    print_every: 10
+    clip_grad_norm: 100
+  max_epochs: ${epochs}
+  max_steps: ${steps}
+  evaluate_every_n_steps: ${evaluate_every_n_steps}
+  callbacks:
+  - _target_: onescience.utils.uma.components.train.train_runner.TrainCheckpointCallback
+    checkpoint_every_n_steps: ${checkpoint_every_n_steps}
+    max_saved_checkpoints: 5
+  - _target_: torchtnt.framework.callbacks.TQDMProgressBar

demo/outputs/oc20_ef_4dcu_20260721_114458/models/pretrained_models.json

@@ -0,0 +1,56 @@
+{
+    "uma-s-1": {
+        "filename": "uma-s-1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "uma-s-1p1": {
+        "filename": "uma-s-1p1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "uma-m-1p1": {
+        "filename": "uma-m-1p1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-md-direct-all-omol": {
+        "filename": "esen_md_direct_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-sm-conserving-all-omol": {
+        "filename": "esen_sm_conserving_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-sm-direct-all-omol": {
+        "filename": "esen_sm_direct_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    }
+}

demo/run.sh

@@ -0,0 +1,300 @@
+#!/bin/bash
+# run.sh - UMA 训练统一入口脚本
+#
+# 用法:
+#   bash run.sh --config configs/oc20_ef_4dcu.yaml              # 直接运行训练
+#   bash run.sh --config configs/oc20_ef_4dcu.yaml --submit     # 提交 SLURM 作业
+#   bash run.sh --config configs/oc20_ef_4dcu.yaml --dry-run    # 仅打印命令，不执行
+#
+set -euo pipefail
+
+# ============================================================
+# 参数解析
+# ============================================================
+CONFIG=""
+SUBMIT=false
+DRY_RUN=false
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --config)
+            CONFIG="$2"; shift 2 ;;
+        --config=*)
+            CONFIG="${1#*=}"; shift ;;
+        --submit)
+            SUBMIT=true; shift ;;
+        --dry-run)
+            DRY_RUN=true; shift ;;
+        -h|--help)
+            echo "用法: bash run.sh --config <config.yaml> [--submit] [--dry-run]"
+            echo ""
+            echo "选项:"
+            echo "  --config <file>   YAML 配置文件路径 (必需)"
+            echo "  --submit          生成 SLURM 脚本并提交作业"
+            echo "  --dry-run         仅打印训练命令，不执行"
+            exit 0 ;;
+        *)
+            echo "[ERROR] 未知参数: $1"
+            exit 1 ;;
+    esac
+done
+
+if [ -z "$CONFIG" ]; then
+    echo "[ERROR] 请指定配置文件: bash run.sh --config configs/xxx.yaml"
+    exit 1
+fi
+
+if [ ! -f "$CONFIG" ]; then
+    echo "[ERROR] 配置文件不存在: $CONFIG"
+    exit 1
+fi
+
+# ============================================================
+# 路径设置
+# ============================================================
+DEMO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$DEMO_DIR/.." && pwd)"
+CONFIG_ABS="$(cd "$(dirname "$CONFIG")" && pwd)/$(basename "$CONFIG")"
+PARSE_PY="$DEMO_DIR/_parse_config.py"
+
+# 如果用户没有设置 ONESCIENCE_DATASETS_DIR，默认指向仓库根目录
+export ONESCIENCE_DATASETS_DIR="${ONESCIENCE_DATASETS_DIR:-$REPO_ROOT}"
+# 如果用户没有设置 ONESCIENCE_MODELS_DIR，默认也指向仓库根目录（用于定位 weight/ 下权重）
+export ONESCIENCE_MODELS_DIR="${ONESCIENCE_MODELS_DIR:-$REPO_ROOT}"
+
+# ============================================================
+# 解析 meta 部分 (不依赖 hydra_config 路径)
+# ============================================================
+EXP_NAME=$(python3 "$PARSE_PY" "$CONFIG_ABS" name)
+ENV_EXPORTS=$(python3 "$PARSE_PY" "$CONFIG_ABS" env)
+ENV_ARGS=$(python3 "$PARSE_PY" "$CONFIG_ABS" env-args)
+DATA_FILES=$(python3 "$PARSE_PY" "$CONFIG_ABS" data-files)
+
+# ============================================================
+# 预生成输出目录与 hydra_config.yaml（dry-run 也生成，便于预览）
+# ============================================================
+TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
+HYDRA_CONFIG="$OUTPUT_DIR/hydra_config.yaml"
+
+# dry-run 不实际创建目录
+if $DRY_RUN; then
+    DRY_HYDRA_TMP=$(mktemp -t "uma_hydra_XXXXXX.yaml")
+    HYDRA_CONFIG="$DRY_HYDRA_TMP"
+    python3 "$PARSE_PY" "$CONFIG_ABS" hydra-config > "$HYDRA_CONFIG"
+else
+    mkdir -p "$OUTPUT_DIR"
+    cp "$CONFIG_ABS" "$OUTPUT_DIR/config.yaml"
+    python3 "$PARSE_PY" "$CONFIG_ABS" hydra-config > "$HYDRA_CONFIG"
+    # UMA/calculate/pretrained_mlip.py 使用 os.getcwd()+"/models/pretrained_models.json"
+    # 定位预训练模型清单, 而 run.sh 会 cd 到 OUTPUT_DIR 再启动训练, 所以在 OUTPUT_DIR
+    # 下放一个指回 UMA/models-json 的软链 (软链名仍为 models), 保证 cwd 相对路径能找到 json。
+    UMA_ROOT_DIR="$(cd "$DEMO_DIR/.." && pwd)"
+    if [ -d "$UMA_ROOT_DIR/models-json" ] && [ ! -e "$OUTPUT_DIR/models" ]; then
+        ln -s "$UMA_ROOT_DIR/models-json" "$OUTPUT_DIR/models"
+    fi
+fi
+
+# 现在 hydra_config 路径就绪, 拼 command
+TRAIN_CMD=$(python3 "$PARSE_PY" "$CONFIG_ABS" command "$HYDRA_CONFIG")
+
+# ============================================================
+# Dry-run 模式
+# ============================================================
+if $DRY_RUN; then
+    echo "========================================="
+    echo "Dry-run: $EXP_NAME"
+    echo "Config:  $CONFIG_ABS"
+    echo "Output:  $OUTPUT_DIR (未创建)"
+    echo "Hydra:   $HYDRA_CONFIG (临时)"
+    echo "========================================="
+    echo ""
+    echo "# 环境变量:"
+    echo "$ENV_EXPORTS"
+    echo ""
+    echo "# 训练命令:"
+    echo "$TRAIN_CMD"
+    echo ""
+    echo "# env_setup.sh 参数: $ENV_ARGS"
+    echo "# 预检路径:"
+    echo "$DATA_FILES" | sed 's/^/#   /'
+    echo ""
+    echo "# hydra_config.yaml 预览 (前 40 行):"
+    head -n 40 "$HYDRA_CONFIG" | sed 's/^/#   /'
+    rm -f "$HYDRA_CONFIG"
+    exit 0
+fi
+
+# ============================================================
+# SLURM 提交模式
+# ============================================================
+if $SUBMIT; then
+    SLURM_VARS=$(python3 "$PARSE_PY" "$CONFIG_ABS" slurm)
+    eval "$SLURM_VARS"
+
+    SLURM_SCRIPT="$OUTPUT_DIR/submit.sh"
+
+    # 生成 SLURM header
+    sed -e "s|{{JOB_NAME}}|${JOB_NAME}|g" \
+        -e "s|{{PARTITION}}|${PARTITION}|g" \
+        -e "s|{{NODES}}|${NODES}|g" \
+        -e "s|{{NTASKS_PER_NODE}}|${NTASKS_PER_NODE}|g" \
+        -e "s|{{CPUS_PER_TASK}}|${CPUS_PER_TASK}|g" \
+        -e "s|{{GPUS_PER_NODE}}|${GPUS_PER_NODE}|g" \
+        -e "s|{{TIME}}|${TIME}|g" \
+        "$DEMO_DIR/templates/slurm_header.template" > "$SLURM_SCRIPT"
+
+    # 环境初始化
+    cat >> "$SLURM_SCRIPT" << 'SETUP_BLOCK'
+
+# 环境初始化
+SETUP_BLOCK
+    cat >> "$SLURM_SCRIPT" << 'ENV_BLOCK'
+set +u
+if [ -n "${UMA_ENV_SCRIPT:-}" ] && [ -f "${UMA_ENV_SCRIPT}" ]; then
+    source "${UMA_ENV_SCRIPT}"
+else
+    echo "[WARN] UMA_ENV_SCRIPT 未设置或文件不存在，跳过环境初始化。请自行确保 conda/matchem 环境已激活。"
+fi
+set -u
+ENV_BLOCK
+
+    # 预检
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# 预检" >> "$SLURM_SCRIPT"
+    DATA_ARGS=""
+    while IFS= read -r line; do
+        [ -z "$line" ] && continue
+        DATA_ARGS="$DATA_ARGS \"$line\""
+    done <<< "$DATA_FILES"
+    echo "bash $DEMO_DIR/templates/preflight_check.sh $DATA_ARGS" >> "$SLURM_SCRIPT"
+
+    # 工作目录与环境变量
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# 工作目录" >> "$SLURM_SCRIPT"
+    echo "cd $OUTPUT_DIR" >> "$SLURM_SCRIPT"
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# 环境变量" >> "$SLURM_SCRIPT"
+    echo "$ENV_EXPORTS" >> "$SLURM_SCRIPT"
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# 将仓库根目录加入 PYTHONPATH，确保能 import 本地 model 包" >> "$SLURM_SCRIPT"
+    echo "export PYTHONPATH=\"$REPO_ROOT:\${PYTHONPATH:-}\"" >> "$SLURM_SCRIPT"
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# 屏蔽 PyTorch NCCL C++ INFO 日志" >> "$SLURM_SCRIPT"
+    echo "export TORCH_CPP_LOG_LEVEL=WARNING" >> "$SLURM_SCRIPT"
+    echo "export NCCL_DEBUG=ERROR" >> "$SLURM_SCRIPT"
+    echo "export GLOG_minloglevel=1" >> "$SLURM_SCRIPT"
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# AMD DCU: 避免 RCCL \"Missing HSA_FORCE_FINE_GRAIN_PCIE\" 警告" >> "$SLURM_SCRIPT"
+    echo "export HSA_FORCE_FINE_GRAIN_PCIE=1" >> "$SLURM_SCRIPT"
+    echo "" >> "$SLURM_SCRIPT"
+    echo "# UMA 依赖: 清理 rocblas tensile 路径, 避免拉到错版本" >> "$SLURM_SCRIPT"
+    echo "unset ROCBLAS_TENSILE_LIBPATH" >> "$SLURM_SCRIPT"
+
+    # 多节点特殊处理: 设置 MASTER_ADDR/PORT, 用 srun 包裹
+    if [ "$NODES" -gt 1 ]; then
+        cat >> "$SLURM_SCRIPT" << 'MULTI_NODE'
+
+# 多节点分布式设置
+nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
+nodes_array=($nodes)
+export MASTER_ADDR=${nodes_array[0]}
+export MASTER_PORT=29504
+echo "SLURM_NNODES=$SLURM_NNODES"
+echo "MASTER_ADDR=$MASTER_ADDR"
+echo "MASTER_PORT=$MASTER_PORT"
+
+# srun 启动分布式训练 (每节点 1 个 task, torchrun 在节点内 spawn 多个 rank)
+MULTI_NODE
+        echo "srun --nodes=\$SLURM_NNODES --ntasks=\$SLURM_NNODES \\" >> "$SLURM_SCRIPT"
+        # 把 TRAIN_CMD 原样追加 (其中的 \${SLURM_*} / \${MASTER_*} 会在每节点上展开)
+        echo "  $TRAIN_CMD" >> "$SLURM_SCRIPT"
+    else
+        echo "" >> "$SLURM_SCRIPT"
+        echo "# 训练命令" >> "$SLURM_SCRIPT"
+        echo "$TRAIN_CMD" >> "$SLURM_SCRIPT"
+    fi
+
+    echo "========================================="
+    echo "SLURM 脚本已生成: $SLURM_SCRIPT"
+    echo "hydra 配置已生成: $HYDRA_CONFIG"
+    echo "配置快照已保存: $OUTPUT_DIR/config.yaml"
+    echo "========================================="
+    echo ""
+    echo "提交作业..."
+    sbatch "$SLURM_SCRIPT"
+    exit 0
+fi
+
+# ============================================================
+# 直接运行模式
+# ============================================================
+echo "========================================="
+echo "实验: $EXP_NAME"
+echo "配置: $CONFIG_ABS"
+echo "输出: $OUTPUT_DIR"
+echo "Hydra: $HYDRA_CONFIG"
+echo "========================================="
+
+# 环境初始化
+set +u
+if [ -n "${UMA_ENV_SCRIPT:-}" ] && [ -f "${UMA_ENV_SCRIPT}" ]; then
+    source "${UMA_ENV_SCRIPT}"
+else
+    echo "[WARN] UMA_ENV_SCRIPT 未设置或文件不存在，跳过环境初始化。请自行确保 conda/matchem 环境已激活。"
+fi
+set -u
+
+# 预检
+DATA_ARGS=""
+while IFS= read -r line; do
+    [ -z "$line" ] && continue
+    DATA_ARGS="$DATA_ARGS \"$line\""
+done <<< "$DATA_FILES"
+eval "bash $DEMO_DIR/templates/preflight_check.sh $DATA_ARGS"
+
+# 设置环境变量
+eval "$ENV_EXPORTS"
+
+# 将仓库根目录加入 PYTHONPATH，确保能 import 本地 model 包
+export PYTHONPATH="$REPO_ROOT:${PYTHONPATH:-}"
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+UMA_JD_PATH="$REPO_ROOT/weight/Jd.pt"
+if [ -f "$UMA_JD_PATH" ]; then
+    export ONESCIENCE_UMA_JD_PATH="$UMA_JD_PATH"
+    echo "[OK] 找到 Jd.pt: $UMA_JD_PATH"
+else
+    echo "[WARN] 未找到 Jd.pt: $UMA_JD_PATH，如训练/推理报错请检查"
+fi
+
+# 屏蔽 PyTorch NCCL C++ INFO 日志
+export TORCH_CPP_LOG_LEVEL=WARNING
+export NCCL_DEBUG=ERROR
+export GLOG_minloglevel=1
+
+# AMD DCU
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+unset ROCBLAS_TENSILE_LIBPATH
+
+# 切换到输出目录
+cd "$OUTPUT_DIR"
+
+# 执行训练
+# 执行训练（实时合并日志）
+echo "========================================="
+echo "开始训练..."
+echo "========================================="
+
+export PYTHONUNBUFFERED=1
+MERGED_LOG="${OUTPUT_DIR}/train_merged.out"
+: > "${MERGED_LOG}"
+echo "[LOG] merged realtime log: ${MERGED_LOG}"
+
+set +e
+eval "$TRAIN_CMD" 2>&1 | tee -a "${MERGED_LOG}"
+TRAIN_EXIT=${PIPESTATUS[0]}
+set -e
+
+exit "${TRAIN_EXIT}"
+

demo/templates/env_setup.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# env_setup.sh - shared environment setup for UMA demo
+# Usage: source templates/env_setup.sh <conda_env> <module1> [module2 ...]
+
+CONDA_ENV="${1:?Please provide conda environment name}"
+shift
+
+# System scripts (bashrc/conda/module) may reference unset vars.
+# Temporarily relax bash strict mode during bootstrap.
+set +eu
+
+module purge
+source ~/.bashrc
+conda activate "$CONDA_ENV"
+
+for mod in "$@"; do
+    module load "$mod"
+done
+
+set -eu
+
+# Load shared dataset/model roots from repo env.sh if it exists
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+DEMO_DIR="$(dirname "$SCRIPT_DIR")"
+REPO_ROOT="$(cd "$DEMO_DIR/../.." && pwd)"
+if [ -f "$REPO_ROOT/env.sh" ]; then
+    source "$REPO_ROOT/env.sh"
+else
+    # Fallback: point dataset/model roots to repo root
+    export ONESCIENCE_DATASETS_DIR="${ONESCIENCE_DATASETS_DIR:-$REPO_ROOT}"
+    export ONESCIENCE_MODELS_DIR="${ONESCIENCE_MODELS_DIR:-$REPO_ROOT}"
+fi
+
+# UMA-specific runtime asset: Jd.pt
+# 如仓库根目录 weight/ 下存在 Jd.pt，会自动设置；否则仅打印警告，不影响其他推理。
+export ONESCIENCE_UMA_JD_PATH="${ONESCIENCE_MODELS_DIR}/weight/Jd.pt"
+
+if [ -f "${ONESCIENCE_UMA_JD_PATH}" ]; then
+    echo "[OK] 文件: ${ONESCIENCE_UMA_JD_PATH}"
+else
+    echo "[WARN] 文件不存在: ${ONESCIENCE_UMA_JD_PATH}"
+fi
+

demo/templates/preflight_check.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+# preflight_check.sh - UMA 训练前预检脚本
+# 用法: bash templates/preflight_check.sh [path1] [path2] ...
+#
+# 每个 path 可以是:
+#   - 普通文件 (如 checkpoint .pt)
+#   - 目录 (如 ASE-lmdb 的 train/ val/)
+# 检查项:
+#   1. ONESCIENCE_DATASETS_DIR 环境变量 (如果 env.sh 已 source)
+#   2. 每个路径是否存在
+#   3. DCU/GPU 状态
+
+ERRORS=0
+
+echo "========================================="
+echo "Node: $(hostname)"
+echo "Time: $(date)"
+echo "========================================="
+
+# 检查路径是否存在 (文件或目录均可)
+for p in "$@"; do
+    [ -z "$p" ] && continue
+    expanded=$(eval echo "$p")
+    if [ ! -e "$expanded" ]; then
+        echo "[ERROR] 路径不存在: $expanded"
+        ERRORS=$((ERRORS + 1))
+    elif [ -d "$expanded" ]; then
+        echo "[OK] 目录: $expanded"
+    else
+        echo "[OK] 文件: $expanded"
+    fi
+done
+
+# 打印 DCU/GPU 状态
+echo "-----------------------------------------"
+echo ">>> DCU/GPU 状态:"
+if command -v hy-smi &>/dev/null; then
+    hy-smi
+elif command -v rocm-smi &>/dev/null; then
+    rocm-smi
+elif command -v nvidia-smi &>/dev/null; then
+    nvidia-smi
+else
+    echo "[WARN] 未找到 GPU 监控工具 (hy-smi/rocm-smi/nvidia-smi)"
+fi
+echo "-----------------------------------------"
+
+if [ $ERRORS -gt 0 ]; then
+    echo "[FATAL] 预检发现 $ERRORS 个错误，终止训练"
+    exit 1
+fi
+
+echo "[OK] 预检通过"

demo/templates/slurm_header.template

@@ -0,0 +1,9 @@
+#!/bin/bash
+#SBATCH -J {{JOB_NAME}}
+#SBATCH -p {{PARTITION}}
+#SBATCH --nodes={{NODES}}
+#SBATCH --ntasks-per-node={{NTASKS_PER_NODE}}
+#SBATCH --cpus-per-task={{CPUS_PER_TASK}}
+#SBATCH --gres=dcu:{{GPUS_PER_NODE}}
+#SBATCH --output=%j.log
+#SBATCH --time={{TIME}}

inference/batch_inference_with_dataloader.py

@@ -0,0 +1,50 @@
+from torch.utils.data import DataLoader
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+import os
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from onescience.datapipes.materials.custom_stack.core.atomic_data import (
+    atomicdata_list_to_batch,
+)
+from onescience.datapipes.materials.custom_stack.storage.ase_datasets import AseDBDataset
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+
+
+def main() -> None:
+    # Update these two paths before running.
+    db_path = "../dataset/omat24/val/rattled-300-subsampled/data.aselmdb"
+    checkpoint_path = "../weight/uma-s-1p1_converted.pt"
+
+    dataset = AseDBDataset(
+        config={
+            "src": db_path,
+            "a2g_args": {"task_name": "omat"},
+        }
+    )
+
+    loader = DataLoader(
+        dataset,
+        batch_size=16,
+        collate_fn=atomicdata_list_to_batch,
+    )
+
+    predictor = load_predict_unit(checkpoint_path, device="cuda")
+
+    for i, batch in enumerate(loader):
+        preds = predictor.predict(batch)
+
+        for j in range(len(preds["energy"])):
+            energy = preds["energy"][j].item()
+            forces = preds["forces"][batch.batch == j].cpu().numpy()
+
+            print(f"\\n[Batch {i} | Structure {j}]")
+            print("Predicted energy:", energy)
+            print("Predicted forces:\\n", forces)
+
+
+if __name__ == "__main__":
+    main()

inference/calculate_spin_gap.py

@@ -0,0 +1,65 @@
+from ase.build import molecule
+from ase.io import write
+from ase.optimize import LBFGS
+import numpy as np
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+import os
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+from onescience.utils.uma.calculate.ase_calculator import FAIRChemCalculator
+
+# ===== 路径：改成你的本地检查点 =====
+CKPT = "../weight/uma-s-1p1_converted.pt"
+
+# ===== 选项：是否在各自自旋态下弛豫几何（绝热差） =====
+RELAX = False
+FMAX = 0.05
+STEPS = 200
+
+# ===== 载入预测器（GPU）与构造计算器 =====
+predictor = load_predict_unit(CKPT, device="cuda")
+
+def make_state(mol_name: str, mult: int, charge: int = 0):
+    atoms = molecule(mol_name)
+    atoms.pbc = False
+    # 关键：同时写入两个键，避免读取不到
+    atoms.info["spin"] = mult
+    atoms.info["spin_multiplicity"] = mult
+    atoms.info["charge"] = charge
+    atoms.calc = FAIRChemCalculator(predictor, task_name="omol")
+    return atoms
+
+# 单重态 / 三重态 CH2
+singlet = make_state("CH2_s1A1d", mult=1, charge=0)
+triplet = make_state("CH2_s3B1d", mult=3, charge=0)
+
+# （可选）各自自旋面上弛豫
+if RELAX:
+    for tag, atoms in [("singlet", singlet), ("triplet", triplet)]:
+        opt = LBFGS(atoms, logfile=f"{tag}_opt.log")
+        opt.run(fmax=FMAX, steps=STEPS)
+        write(f"{tag}_final.xyz", atoms)
+
+# 计算单点能（或弛豫后能量）
+E_s = singlet.get_potential_energy()
+E_t = triplet.get_potential_energy()
+
+gap_ev = E_t - E_s
+gap_kcal = gap_ev * 23.060543
+
+print(f"Singlet energy (CH2 1A1):  {E_s:.6f} eV")
+print(f"Triplet energy (CH2 3B1):  {E_t:.6f} eV")
+print(f"Gap (T - S):               {gap_ev:.6f} eV  ({gap_kcal:.2f} kcal/mol)")
+
+# 保存结构，便于复查
+write("CH2_singlet.xyz", singlet)
+write("CH2_triplet.xyz", triplet)
+
+# 打印自旋/电荷确认，确保不会再触发“未设置自旋”的警告
+print("Singlet info:", {k: singlet.info[k] for k in ("spin", "spin_multiplicity", "charge")})
+print("Triplet info:", {k: triplet.info[k] for k in ("spin", "spin_multiplicity", "charge")})

inference/generate_batches_at_runtime.py

@@ -0,0 +1,49 @@
+import os
+
+# 自动定位 UMA 旋转基文件 Jd.pt（如果在仓库根目录 weight/ 下存在）
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from ase.build import bulk
+
+from onescience.datapipes.materials.custom_stack.core.atomic_data import (
+    AtomicData,
+    atomicdata_list_to_batch,
+)
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+
+# 构建多个结构，可替换为 molecule() 或 slab(...)
+atoms_list = [
+    bulk("Pt"),
+    bulk("Cu"),
+    bulk("NaCl", crystalstructure="rocksalt", a=2.0),
+]
+
+# 转换为 AtomicData 并赋予任务名
+atomic_data_list = [
+    AtomicData.from_ase(atoms, task_name="omat") for atoms in atoms_list
+]
+
+# 合并成一个 batch
+batch = atomicdata_list_to_batch(atomic_data_list)
+
+# 加载模型（默认从仓库根目录的 weight/ 下读取）
+checkpoint_path = os.environ.get(
+    "UMA_CHECKPOINT_PATH",
+    os.path.join(_REPO_ROOT, "weight", "uma-s-1p1_converted.pt"),
+)
+predictor = load_predict_unit(checkpoint_path, device="cuda")
+
+# 执行推理
+preds = predictor.predict(batch)
+
+# 输出每个结构的能量和原子力
+for i, atoms in enumerate(atoms_list):
+    energy = preds["energy"][i].item()
+    forces = preds["forces"][batch.batch == i].cpu().numpy()
+
+    print(f"\nStructure #{i + 1}: {atoms.get_chemical_formula()}")
+    print("Predicted energy:", energy)
+    print("Predicted forces:\n", forces)

inference/inference_heterogeneous_batches.py

@@ -0,0 +1,50 @@
+from ase.build import bulk, molecule, fcc100, add_adsorbate
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+import os
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+from onescience.datapipes.materials.custom_stack.core.atomic_data import AtomicData, atomicdata_list_to_batch
+
+# 1. 创建异构结构
+h2o = molecule("H2O")
+h2o.info.update({"charge": 0, "spin": 1})
+
+pt = bulk("Pt")
+
+slab = fcc100("Cu", (3, 3, 3), vacuum=8, periodic=True)
+adsorbate = molecule("CO")
+add_adsorbate(slab, adsorbate, 2.0, "bridge")
+
+# 2. 结构转为 AtomicData，并指定不同 task_name
+atomic_data_list = [
+    AtomicData.from_ase(
+        h2o, task_name="omol", r_data_keys=["spin", "charge"], molecule_cell_size=12
+    ),
+    AtomicData.from_ase(pt, task_name="omat"),
+    AtomicData.from_ase(slab, task_name="oc20"),
+]
+
+# 3. 合成 batch
+batch = atomicdata_list_to_batch(atomic_data_list)
+
+# 4. 加载 UMA 模型
+predictor = load_predict_unit(
+    "../weight/uma-s-1p1_converted.pt", device="cuda"#替换为你的检查点路径
+)
+
+# 5. 执行联合推理
+preds = predictor.predict(batch)
+
+# 6. 输出每个结构结果
+for i in range(len(preds["energy"])):
+    energy = preds["energy"][i].item()
+    forces = preds["forces"][batch.batch == i].cpu().numpy()
+
+    print(f"\n[Structure {i}]")
+    print("Predicted energy:", energy)
+    print("Predicted forces:\n", forces)

inference/models/pretrained_models.json

@@ -0,0 +1,56 @@
+{
+    "uma-s-1": {
+        "filename": "uma-s-1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "uma-s-1p1": {
+        "filename": "uma-s-1p1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "uma-m-1p1": {
+        "filename": "uma-m-1p1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-md-direct-all-omol": {
+        "filename": "esen_md_direct_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-sm-conserving-all-omol": {
+        "filename": "esen_sm_conserving_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-sm-direct-all-omol": {
+        "filename": "esen_sm_direct_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    }
+}

inference/relax_adsorbate_on_slab.py

@@ -0,0 +1,67 @@
+import time, torch, numpy as np
+from ase.build import fcc100, add_adsorbate, molecule
+from ase.optimize import LBFGS
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+import os
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+from onescience.utils.uma.calculate.ase_calculator import FAIRChemCalculator
+
+# === 工具函数 ===
+def _sync():
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+def clear_cache(atoms):
+    # 避免命中 ASE/Calculator 的上次结果缓存，确保测到真实推理时间
+    if hasattr(atoms, "calc") and hasattr(atoms.calc, "results"):
+        atoms.calc.results.clear()
+
+# === 加载模型 ===
+predictor = load_predict_unit(
+    "../weight/uma-s-1p1_converted.pt",
+    device="cuda"
+)
+calc = FAIRChemCalculator(predictor, task_name="oc20")
+
+# === 构建体系 ===
+slab = fcc100("Cu", (3, 3, 3), vacuum=8, periodic=True)
+adsorbate = molecule("CO")
+add_adsorbate(slab, adsorbate, 2.0, "bridge")
+slab.calc = calc
+
+# === 预热（把模型/张量搬到显存，避免首轮开销影响统计）===
+_ = slab.get_potential_energy(); _sync()
+_ = slab.get_forces();            _sync()
+
+# === 计时：LBFGS 优化（总耗时 & 平均每步）===
+t0 = time.perf_counter()
+opt = LBFGS(slab, logfile=None)  # 如需日志，改成 "lbfgs.log"
+opt.run(fmax=0.05, steps=100)
+_sync()
+lbfgs_total = time.perf_counter() - t0
+steps_done = getattr(opt, "nsteps", None) or 100
+
+# === 计时：单次能量 & 力 推理延迟 ===
+clear_cache(slab); t0 = time.perf_counter()
+energy = slab.get_potential_energy(); _sync()
+t_energy = time.perf_counter() - t0
+
+clear_cache(slab); t0 = time.perf_counter()
+forces = slab.get_forces(); _sync()
+t_forces = time.perf_counter() - t0
+
+# === 打印结果 ===
+print(f"Predicted energy: {energy}")
+print(f"Predicted forces shape: {forces.shape}")
+# 如需打印具体力矩阵，取消下一行注释（可能很长）
+# print(forces)
+
+print(f"[E] latency: {t_energy*1e3:.2f} ms/call")
+print(f"[F] latency: {t_forces*1e3:.2f} ms/call")
+print(f"[LBFGS] total: {lbfgs_total:.3f} s, per-step: {lbfgs_total/steps_done*1e3:.2f} ms/step")

inference/relax_inorganic_crystal.py

@@ -0,0 +1,58 @@
+from ase.build import bulk
+from ase.optimize import LBFGS
+from ase.filters import FrechetCellFilter
+from ase.io import write
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+import os
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+from onescience.utils.uma.calculate.ase_calculator import FAIRChemCalculator
+import numpy as np
+
+# === 本地检查点（改成你的实际路径） ===
+ckpt = "../weight/uma-s-1p1_converted.pt"
+
+# === 加载预测器与计算器 ===
+predictor = load_predict_unit(ckpt, device="cuda")   # GPU 加速
+calc = FAIRChemCalculator(predictor, task_name="omat")
+
+# === 建 Fe 体相并绑定计算器 ===
+atoms = bulk("Fe")
+atoms.calc = calc
+
+# === 原子+晶胞联合弛豫 ===
+filt = FrechetCellFilter(atoms)        # 同时优化原子位置与晶胞
+opt = LBFGS(filt, logfile="relax.log") # 记录优化日志到文件
+opt.run(fmax=0.05, steps=100)
+
+# === 计算与打印结果 ===
+E = atoms.get_potential_energy()
+F = atoms.get_forces()
+fmax = np.sqrt((F**2).sum(axis=1)).max()
+cell = atoms.get_cell()
+a, b, c, alpha, beta, gamma = atoms.get_cell_lengths_and_angles()
+vol = atoms.get_volume()
+
+print(f"Final energy: {E:.6f} eV  (per atom: {E/len(atoms):.6f} eV/atom)")
+print(f"Final fmax: {fmax:.6f} eV/Å")
+print(f"Volume: {vol:.6f} Å^3")
+try:
+    stress = atoms.get_stress()  # Voigt, eV/Å^3
+    print("Final stress (Voigt, eV/Å^3):", stress)
+except Exception as e:
+    print("Stress not available:", e)
+
+print("Cell (Å):\n", cell)
+print(f"Lengths/angles: a={a:.4f} b={b:.4f} c={c:.4f}  "
+      f"alpha={alpha:.2f} beta={beta:.2f} gamma={gamma:.2f}")
+print("Positions (Å):\n", atoms.get_positions())
+
+# === 保存结果到文件 ===
+write("Fe_optimized.cif", atoms)       # 方便后续可视化/复现实验
+write("Fe_optimized.traj", atoms)      # ASE 轨迹
+np.savetxt("Fe_forces.txt", F)         # 保存力矩阵

inference/run_molecular_md.py

@@ -0,0 +1,76 @@
+from ase import units
+from ase.io import write, Trajectory, read
+from ase.md.langevin import Langevin
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation
+from ase.build import molecule
+from ase.md import MDLogger
+import numpy as np
+
+# 自动定位 UMA 旋转基文件 Jd.pt
+import os
+_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_JD_PATH = os.path.join(_REPO_ROOT, "weight", "Jd.pt")
+if os.path.isfile(_JD_PATH):
+    os.environ.setdefault("ONESCIENCE_UMA_JD_PATH", _JD_PATH)
+
+from onescience.utils.uma.units.mlip_unit import load_predict_unit
+from onescience.utils.uma.calculate.ase_calculator import FAIRChemCalculator
+
+# === 本地检查点路径（改成你的实际路径）===
+ckpt = "../weight/uma-s-1p1_converted.pt"
+
+# === 加载预测器与计算器 ===
+predictor = load_predict_unit(ckpt, device="cuda")     # GPU加速
+calc = FAIRChemCalculator(predictor, task_name="omol") # 分子任务
+
+# === 构建体系 ===
+atoms = molecule("H2O")  # 可替换为其它分子
+atoms.calc = calc
+
+# === 速度初始化（400 K），去线动量/角动量 ===
+T0 = 400  # K
+MaxwellBoltzmannDistribution(atoms, temperature_K=T0)
+Stationary(atoms)       # 去整体平动
+ZeroRotation(atoms)     # 去整体转动
+
+# === 动力学设置：Langevin ===
+dt = 0.1 * units.fs
+gamma = 0.001 / units.fs
+dyn = Langevin(atoms, timestep=dt, temperature_K=T0, friction=gamma)
+
+# === 轨迹与日志 ===
+traj = Trajectory("my_md.traj", "w", atoms)            # ASE 原生轨迹
+dyn.attach(traj.write, interval=1)                     # 每步写一帧
+logger = MDLogger(dyn, atoms, "md.log", header=True, stress=False, peratom=False)
+dyn.attach(logger, interval=10)                        # 每10步记录一次
+
+# 也打印到屏幕（可选）
+def printer():
+    epot = atoms.get_potential_energy()
+    ekin = atoms.get_kinetic_energy()
+    Tinst = 2.0 * ekin / (3 * len(atoms) * units.kB)
+    print(f"Step {dyn.nsteps:5d}  Epot={epot: .6f} eV  Ekin={ekin: .6f} eV  T={Tinst: .1f} K")
+dyn.attach(printer, interval=50)
+
+# === 运行MD ===
+steps = 1000
+dyn.run(steps=steps)
+
+# === 结束后输出关键信息与保存 ===
+Epot = atoms.get_potential_energy()
+Ekin = atoms.get_kinetic_energy()
+Tfinal = 2.0 * Ekin / (3 * len(atoms) * units.kB)
+print("\n== MD Finished ==")
+print(f"Final Epot = {Epot:.6f} eV,  Ekin = {Ekin:.6f} eV,  T = {Tfinal:.2f} K")
+print("COM =", atoms.get_center_of_mass())
+
+# 保存最终结构（xyz 与 cif）
+write("final.xyz", atoms)
+try:
+    write("final.cif", atoms)  # 分子一般无PBC，cif可有可无
+except Exception as e:
+    print("Save CIF skipped:", e)
+
+# 如需导出整段轨迹为 .xyz（多帧），取消注释以下两行
+# frames = read("my_md.traj", ":")
+# write("my_md_all.xyz", frames)

model/__init__.py

@@ -0,0 +1,30 @@
+
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+from model.base import HydraModel, HydraModelV2
+
+try:
+    __version__ = version("onescience.utils.uma")
+except PackageNotFoundError:
+    # package is not installed
+    __version__ = ""
+
+__all__ = [
+    "HydraModel",
+    "HydraModelV2",
+    "eSCNMDBackbone",
+    "eSCNMDMoeBackbone",
+]
+
+
+def __getattr__(name: str):
+    if name == "eSCNMDBackbone":
+        from model.uma_escn_md import eSCNMDBackbone
+        return eSCNMDBackbone
+    if name == "eSCNMDMoeBackbone":
+        from model.uma_escn_moe import eSCNMDMoeBackbone
+        return eSCNMDMoeBackbone
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

model/base.py

@@ -0,0 +1,222 @@
+
+
+from __future__ import annotations
+
+import copy
+import logging
+from abc import ABCMeta, abstractmethod
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.common.utils import (
+    load_model_and_weights_from_checkpoint,
+)
+
+if TYPE_CHECKING:
+    from onescience.datapipes.materials.custom_stack.core.atomic_data import AtomicData
+
+
+class HeadInterface(metaclass=ABCMeta):
+    @property
+    def use_amp(self):
+        return False
+
+    @abstractmethod
+    def forward(
+        self, data: AtomicData, emb: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        """Head forward.
+
+        Arguments
+        ---------
+        data: AtomicData
+            Atomic systems as input
+        emb: dict[str->torch.Tensor]
+            Embeddings of the input as generated by the backbone
+
+        Returns
+        -------
+        outputs: dict[str->torch.Tensor]
+            Return one or more targets generated by this head
+        """
+        return
+
+
+class BackboneInterface(metaclass=ABCMeta):
+    @abstractmethod
+    def forward(self, data: AtomicData) -> dict[str, torch.Tensor]:
+        """Backbone forward.
+
+        Arguments
+        ---------
+        data: AtomicData
+            Atomic systems as input
+
+        Returns
+        -------
+        embedding: dict[str->torch.Tensor]
+            Return backbone embeddings for the given input
+        """
+        return
+
+
+@registry.register_model("hydra")
+class HydraModel(nn.Module):
+    def __init__(
+        self,
+        backbone: dict | None = None,
+        heads: dict | None = None,
+        finetune_config: dict | None = None,
+        otf_graph: bool = True,
+        pass_through_head_outputs: bool = False,
+        freeze_backbone: bool = False,
+    ):
+        super().__init__()
+        self.device = None
+        self.otf_graph = otf_graph
+        # This is required for hydras with models that have multiple outputs per head, since we will deprecate
+        # the old config system at some point, this will prevent the need to make major modifications to the trainer
+        # because they all expect the name of the outputs directly instead of the head_name.property_name
+        self.pass_through_head_outputs = pass_through_head_outputs
+
+        # if finetune_config is provided, then attempt to load the model from the given finetune checkpoint
+        starting_model = None
+        if finetune_config is not None:
+            # Make it hard to sneak more fields into finetuneconfig
+            assert (
+                len(set(finetune_config.keys()) - {"starting_checkpoint", "override"})
+                == 0
+            )
+            starting_model: HydraModel = load_model_and_weights_from_checkpoint(
+                finetune_config["starting_checkpoint"]
+            )
+            logging.info(
+                f"Found and loaded fine-tuning checkpoint: {finetune_config['starting_checkpoint']} (Note we are NOT loading the training state from this checkpoint, only parts of the model and weights)"
+            )
+            assert isinstance(
+                starting_model, HydraModel
+            ), "Can only finetune starting from other hydra models!"
+            # TODO this is a bit hacky to overrride attrs in the backbone
+            if "override" in finetune_config:
+                for key, value in finetune_config["override"].items():
+                    setattr(starting_model.backbone, key, value)
+
+        if backbone is not None:
+            backbone = copy.deepcopy(backbone)
+            backbone_model_name = backbone.pop("model")
+            self.backbone: BackboneInterface = registry.get_model_class(
+                backbone_model_name
+            )(
+                **backbone,
+            )
+        elif starting_model is not None:
+            self.backbone = starting_model.backbone
+            logging.info(
+                f"User did not specify a backbone, using the backbone from the starting checkpoint {self.backbone}"
+            )
+        else:
+            raise RuntimeError(
+                "Backbone not specified and not found in the starting checkpoint"
+            )
+
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+        if heads is not None:
+            heads = copy.deepcopy(heads)
+            # Iterate through outputs_cfg and create heads
+            self.output_heads: dict[str, HeadInterface] = {}
+
+            head_names_sorted = sorted(heads.keys())
+            assert len(set(head_names_sorted)) == len(
+                head_names_sorted
+            ), "Head names must be unique!"
+            for head_name in head_names_sorted:
+                head_config = heads[head_name]
+                if "module" not in head_config:
+                    raise ValueError(
+                        f"{head_name} head does not specify module to use for the head"
+                    )
+
+                module_name = head_config.pop("module")
+                self.output_heads[head_name] = registry.get_model_class(module_name)(
+                    self.backbone,
+                    **head_config,
+                )
+
+            self.output_heads = torch.nn.ModuleDict(self.output_heads)
+        elif starting_model is not None:
+            self.output_heads = starting_model.output_heads
+            logging.info(
+                f"User did not specify heads, using the output heads from the starting checkpoint {self.output_heads}"
+            )
+        else:
+            raise RuntimeError(
+                "Heads not specified and not found in the starting checkpoint"
+            )
+
+    def forward(self, data: AtomicData):
+        # lazily get device from input to use with amp, at least one input must be a tensor to figure out it's device
+        if not self.device:
+            device_from_tensors = {
+                x.device.type for x in data.values() if isinstance(x, torch.Tensor)
+            }
+            assert (
+                len(device_from_tensors) == 1
+            ), f"all inputs must be on the same device, found the following devices {device_from_tensors}"
+            self.device = device_from_tensors.pop()
+
+        emb = self.backbone(data)
+        # Predict all output properties for all structures in the batch for now.
+        out = {}
+        for k in self.output_heads:
+            with torch.autocast(
+                device_type=self.device, enabled=self.output_heads[k].use_amp
+            ):
+                if self.pass_through_head_outputs:
+                    out.update(self.output_heads[k](data, emb))
+                else:
+                    out[k] = self.output_heads[k](data, emb)
+
+        return out
+
+
+class HydraModelV2(nn.Module):
+    def __init__(
+        self,
+        backbone: BackboneInterface,
+        heads: dict[str, HeadInterface],
+        freeze_backbone: bool = False,
+    ):
+        super().__init__()
+        self.backbone = backbone
+        self.output_heads = torch.nn.ModuleDict(heads)
+        self.device = None
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+    def forward(self, data):
+        # lazily get device from input to use with amp, at least one input must be a tensor to figure out it's device
+        if not self.device:
+            device_from_tensors = {
+                x.device.type for x in data.values() if isinstance(x, torch.Tensor)
+            }
+            assert (
+                len(device_from_tensors) == 1
+            ), f"all inputs must be on the same device, found the following devices {device_from_tensors}"
+            self.device = device_from_tensors.pop()
+
+        emb = self.backbone(data)
+        # Predict all output properties for all structures in the batch for now.
+        out = {}
+        for k in self.output_heads:
+            with torch.autocast(
+                device_type=self.device, enabled=self.output_heads[k].use_amp
+            ):
+                out[k] = self.output_heads[k](data, emb)
+        return out

model/models/__init__.py

@@ -0,0 +1 @@
+"""compat aliases"""\n

model/models/base.py

@@ -0,0 +1 @@
+from model.base import *

model/models/uma/__init__.py

@@ -0,0 +1 @@
+"""compat aliases"""

model/models/uma/escn_md.py

@@ -0,0 +1 @@
+from model.uma_escn_md import *

model/models/uma/escn_moe.py

@@ -0,0 +1 @@
+from model.uma_escn_moe import *

model/uma_escn_md.py

@@ -0,0 +1,621 @@
+
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn as nn
+from torch.profiler import record_function
+
+from onescience.utils.uma.common import gp_utils
+from onescience.utils.uma.common.distutils import get_device_for_local_rank
+from onescience.utils.uma.common.registry import registry
+from onescience.utils.uma.common.utils import conditional_grad
+from onescience.modules.func_utils.uma_graph.compute import generate_graph
+from onescience.modules.func_utils.uma_path_utils import resolve_jd_path
+from onescience.modules.head.uma_head import (
+    Linear_Energy_Head,
+    Linear_Force_Head,
+    MLP_EFS_Head,
+    MLP_Energy_Head,
+    MLP_Stress_Head,
+    compose_tensor,
+)
+from onescience.modules.func_utils.uma_rotation import (
+    init_edge_rot_mat,
+    rotation_to_wigner,
+)
+from onescience.modules.func_utils.uma_rotation_cuda_graph import RotMatWignerCudaGraph
+from onescience.modules.func_utils.uma_so3 import CoefficientMapping, SO3_Grid
+from onescience.modules.embedding.uma_embedding import (
+    ChgSpinEmbedding,
+    DatasetEmbedding,
+    EdgeDegreeEmbedding,
+)
+from onescience.modules.layer.uma_layer_norm import (
+    EquivariantLayerNormArray,
+    EquivariantLayerNormArraySphericalHarmonics,
+    EquivariantRMSNormArraySphericalHarmonics,
+    EquivariantRMSNormArraySphericalHarmonicsV2,
+    get_normalization_layer,
+)
+from onescience.modules.func_utils.uma_mole_utils import MOLEInterface
+from onescience.modules.layer.uma_radial import GaussianSmearing
+from onescience.modules.layer.uma_so3_layers import SO3_Linear
+
+from onescience.modules.block.uma_escn_md_block import eSCNMD_Block
+
+ESCNMD_DEFAULT_EDGE_CHUNK_SIZE = 1024 * 128
+
+
+@registry.register_model("escnmd_backbone")
+class eSCNMDBackbone(nn.Module, MOLEInterface):
+    def __init__(
+        self,
+        max_num_elements: int = 100,
+        sphere_channels: int = 128,
+        lmax: int = 2,
+        mmax: int = 2,
+        grid_resolution: int | None = None,
+        num_sphere_samples: int = 128,  # NOTE not used
+        # NOTE: graph construction related, to remove
+        otf_graph: bool = False,
+        max_neighbors: int = 300,
+        use_pbc: bool = True,  # deprecated
+        use_pbc_single: bool = True,  # deprecated
+        cutoff: float = 5.0,
+        edge_channels: int = 128,
+        distance_function: str = "gaussian",
+        num_distance_basis: int = 512,
+        direct_forces: bool = True,
+        regress_forces: bool = True,
+        regress_stress: bool = False,
+        # escnmd specific
+        num_layers: int = 2,
+        hidden_channels: int = 128,
+        norm_type: str = "rms_norm_sh",
+        act_type: str = "gate",
+        ff_type: str = "grid",
+        activation_checkpointing: bool = False,
+        chg_spin_emb_type: str = "pos_emb",
+        cs_emb_grad: bool = False,
+        dataset_emb_grad: bool = False,
+        dataset_list: list[str] | None = None,
+        use_dataset_embedding: bool = True,
+        use_cuda_graph_wigner: bool = False,
+        radius_pbc_version: int = 1,
+        always_use_pbc: bool = True,
+        jd_path: str | None = None,
+    ):
+        super().__init__()
+        self.max_num_elements = max_num_elements
+        self.lmax = lmax
+        self.mmax = mmax
+        self.sphere_channels = sphere_channels
+        self.grid_resolution = grid_resolution
+        self.num_sphere_samples = num_sphere_samples
+        # set this True if we want to ALWAYS use pbc for internal graph gen
+        # despite what's in the input data this only affects when otf_graph is True
+        # in this mode, the user must be responsible for providing a large vaccum box
+        # for aperiodic systems
+        self.always_use_pbc = always_use_pbc
+
+        # energy conservation related
+        self.regress_forces = regress_forces
+        self.direct_forces = direct_forces
+        self.regress_stress = regress_stress
+
+        # NOTE: graph construction related, to remove, except for cutoff
+        self.otf_graph = otf_graph
+        self.max_neighbors = max_neighbors
+        self.radius_pbc_version = radius_pbc_version
+        self.enforce_max_neighbors_strictly = False
+
+        activation_checkpoint_chunk_size = None
+        if activation_checkpointing:
+            # The size of edge blocks to use in activation checkpointing
+            activation_checkpoint_chunk_size = ESCNMD_DEFAULT_EDGE_CHUNK_SIZE
+
+        # related to charge spin dataset system embedding
+        self.chg_spin_emb_type = chg_spin_emb_type
+        self.cs_emb_grad = cs_emb_grad
+        self.dataset_emb_grad = dataset_emb_grad
+        self.dataset_list = dataset_list
+        self.use_dataset_embedding = use_dataset_embedding
+        self.use_cuda_graph_wigner = use_cuda_graph_wigner
+        assert (
+            self.dataset_list
+        ), "the dataset list is empty, please add it to the model backbone config"
+
+        # rotation utils
+        Jd_list = torch.load(resolve_jd_path(jd_path))
+        for l in range(self.lmax + 1):
+            self.register_buffer(f"Jd_{l}", Jd_list[l])
+        self.sph_feature_size = int((self.lmax + 1) ** 2)
+        self.mappingReduced = CoefficientMapping(self.lmax, self.mmax)
+
+        # lmax_lmax for node, lmax_mmax for edge
+        self.SO3_grid = nn.ModuleDict()
+        self.SO3_grid["lmax_lmax"] = SO3_Grid(
+            self.lmax, self.lmax, resolution=grid_resolution, rescale=True
+        )
+        self.SO3_grid["lmax_mmax"] = SO3_Grid(
+            self.lmax, self.mmax, resolution=grid_resolution, rescale=True
+        )
+
+        # atom embedding
+        self.sphere_embedding = nn.Embedding(
+            self.max_num_elements, self.sphere_channels
+        )
+
+        # charge / spin embedding
+        self.charge_embedding = ChgSpinEmbedding(
+            self.chg_spin_emb_type,
+            "charge",
+            self.sphere_channels,
+            grad=self.cs_emb_grad,
+        )
+        self.spin_embedding = ChgSpinEmbedding(
+            self.chg_spin_emb_type,
+            "spin",
+            self.sphere_channels,
+            grad=self.cs_emb_grad,
+        )
+
+        # dataset embedding
+        if self.use_dataset_embedding:
+            self.dataset_embedding = DatasetEmbedding(
+                self.sphere_channels,
+                grad=self.dataset_emb_grad,
+                dataset_list=self.dataset_list,
+            )
+            # mix charge, spin, dataset embeddings
+            self.mix_csd = nn.Linear(3 * self.sphere_channels, self.sphere_channels)
+        else:
+            # mix charge, spin
+            self.mix_csd = nn.Linear(2 * self.sphere_channels, self.sphere_channels)
+
+        # edge distance embedding
+        self.cutoff = cutoff
+        self.edge_channels = edge_channels
+        self.distance_function = distance_function
+        self.num_distance_basis = num_distance_basis
+
+        if self.distance_function == "gaussian":
+            self.distance_expansion = GaussianSmearing(
+                0.0,
+                self.cutoff,
+                self.num_distance_basis,
+                2.0,
+            )
+        else:
+            raise ValueError("Unknown distance function")
+
+        # equivariant initial embedding
+        self.source_embedding = nn.Embedding(self.max_num_elements, self.edge_channels)
+        self.target_embedding = nn.Embedding(self.max_num_elements, self.edge_channels)
+        nn.init.uniform_(self.source_embedding.weight.data, -0.001, 0.001)
+        nn.init.uniform_(self.target_embedding.weight.data, -0.001, 0.001)
+
+        self.edge_channels_list = [
+            self.num_distance_basis + 2 * self.edge_channels,
+            self.edge_channels,
+            self.edge_channels,
+        ]
+
+        self.edge_degree_embedding = EdgeDegreeEmbedding(
+            sphere_channels=self.sphere_channels,
+            lmax=self.lmax,
+            mmax=self.mmax,
+            max_num_elements=self.max_num_elements,
+            edge_channels_list=self.edge_channels_list,
+            rescale_factor=5.0,  # NOTE: sqrt avg degree
+            cutoff=self.cutoff,
+            mappingReduced=self.mappingReduced,
+            activation_checkpoint_chunk_size=activation_checkpoint_chunk_size,
+        )
+
+        self.num_layers = num_layers
+        self.hidden_channels = hidden_channels
+        self.norm_type = norm_type
+        self.act_type = act_type
+        self.ff_type = ff_type
+
+        # Initialize the blocks for each layer
+        self.blocks = nn.ModuleList()
+        for _ in range(self.num_layers):
+            block = eSCNMD_Block(
+                self.sphere_channels,
+                self.hidden_channels,
+                self.lmax,
+                self.mmax,
+                self.mappingReduced,
+                self.SO3_grid,
+                self.edge_channels_list,
+                self.cutoff,
+                self.norm_type,
+                self.act_type,
+                self.ff_type,
+                activation_checkpoint_chunk_size=activation_checkpoint_chunk_size,
+            )
+            self.blocks.append(block)
+
+        self.norm = get_normalization_layer(
+            self.norm_type,
+            lmax=self.lmax,
+            num_channels=self.sphere_channels,
+        )
+
+        self.rot_mat_wigner_cuda = None  # lazily initialize this
+        coefficient_index = self.SO3_grid["lmax_lmax"].mapping.coefficient_idx(
+            self.lmax, self.mmax
+        )
+        self.register_buffer("coefficient_index", coefficient_index, persistent=False)
+
+    def _get_rotmat_and_wigner(
+        self, edge_distance_vecs: torch.Tensor, use_cuda_graph: bool
+    ):
+        Jd_buffers = [
+            getattr(self, f"Jd_{l}").type(edge_distance_vecs.dtype)
+            for l in range(self.lmax + 1)
+        ]
+
+        if use_cuda_graph:
+            if self.rot_mat_wigner_cuda is None:
+                self.rot_mat_wigner_cuda = RotMatWignerCudaGraph()
+            with record_function("obtain rotmat wigner cudagraph"):
+                edge_rot_mat, wigner, wigner_inv = (
+                    self.rot_mat_wigner_cuda.get_rotmat_and_wigner(
+                        edge_distance_vecs, Jd_buffers
+                    )
+                )
+        else:
+            with record_function("obtain rotmat wigner original"):
+                edge_rot_mat = init_edge_rot_mat(
+                    edge_distance_vecs, rot_clip=(not self.direct_forces)
+                )
+                wigner = rotation_to_wigner(
+                    edge_rot_mat,
+                    0,
+                    self.lmax,
+                    Jd_buffers,
+                    rot_clip=(not self.direct_forces),
+                )
+                wigner_inv = torch.transpose(wigner, 1, 2).contiguous()
+
+        # select subset of coefficients we are using
+        if self.mmax != self.lmax:
+            wigner = wigner.index_select(1, self.coefficient_index)
+            wigner_inv = wigner_inv.index_select(2, self.coefficient_index)
+
+        wigner_and_M_mapping = torch.einsum(
+            "mk,nkj->nmj", self.mappingReduced.to_m, wigner
+        )
+        wigner_and_M_mapping_inv = torch.einsum(
+            "njk,mk->njm", wigner_inv, self.mappingReduced.to_m
+        )
+        return edge_rot_mat, wigner_and_M_mapping, wigner_and_M_mapping_inv
+
+    def _get_displacement_and_cell(self, data_dict):
+        ###############################################################
+        # gradient-based forces/stress
+        ###############################################################
+        displacement = None
+        orig_cell = None
+        if self.regress_stress and not self.direct_forces:
+            displacement = torch.zeros(
+                (3, 3),
+                dtype=data_dict["pos"].dtype,
+                device=data_dict["pos"].device,
+            )
+            num_batch = len(data_dict["natoms"])
+            displacement = displacement.view(-1, 3, 3).expand(num_batch, 3, 3)
+            displacement.requires_grad = True
+            symmetric_displacement = 0.5 * (
+                displacement + displacement.transpose(-1, -2)
+            )
+            if data_dict["pos"].requires_grad is False:
+                data_dict["pos"].requires_grad = True
+            data_dict["pos_original"] = data_dict["pos"]
+            data_dict["pos"] = data_dict["pos"] + torch.bmm(
+                data_dict["pos"].unsqueeze(-2),
+                torch.index_select(symmetric_displacement, 0, data_dict["batch"]),
+            ).squeeze(-2)
+
+            orig_cell = data_dict["cell"]
+            data_dict["cell"] = data_dict["cell"] + torch.bmm(
+                data_dict["cell"], symmetric_displacement
+            )
+
+        if (
+            not self.regress_stress
+            and self.regress_forces
+            and not self.direct_forces
+            and data_dict["pos"].requires_grad is False
+        ):
+            data_dict["pos"].requires_grad = True
+        return displacement, orig_cell
+
+    def csd_embedding(self, charge, spin, dataset):
+        with record_function("charge spin dataset embeddings"):
+            # Add charge, spin, and dataset embeddings
+            chg_emb = self.charge_embedding(charge)
+            spin_emb = self.spin_embedding(spin)
+            if self.use_dataset_embedding:
+                assert dataset is not None
+                dataset_emb = self.dataset_embedding(dataset)
+                return torch.nn.SiLU()(
+                    self.mix_csd(torch.cat((chg_emb, spin_emb, dataset_emb), dim=1))
+                )
+            return torch.nn.SiLU()(self.mix_csd(torch.cat((chg_emb, spin_emb), dim=1)))
+
+    def _generate_graph(self, data_dict):
+        if self.otf_graph:
+            pbc = None
+            if self.always_use_pbc:
+                pbc = torch.ones(len(data_dict), 3, dtype=torch.bool)
+            else:
+                assert (
+                    "pbc" in data_dict
+                ), "Since always_use_pbc is False, pbc conditions must be supplied by the input data"
+                pbc = data_dict["pbc"]
+            assert (
+                pbc.all() or (~pbc).all()
+            ), "We can only accept pbc that is all true or all false"
+            logging.debug(f"Using radius graph gen version {self.radius_pbc_version}")
+            graph_dict = generate_graph(
+                data_dict,
+                cutoff=self.cutoff,
+                max_neighbors=self.max_neighbors,
+                enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
+                radius_pbc_version=self.radius_pbc_version,
+                pbc=pbc,
+            )
+        else:
+            # this assume edge_index is provided
+            assert (
+                "edge_index" in data_dict
+            ), "otf_graph is false, need to provide edge_index as input!"
+            cell_per_edge = data_dict["cell"].repeat_interleave(
+                data_dict["nedges"], dim=0
+            )
+            shifts = torch.einsum(
+                "ij,ijk->ik",
+                data_dict["cell_offsets"].to(cell_per_edge.dtype),
+                cell_per_edge,
+            )
+            edge_distance_vec = (
+                data_dict["pos"][data_dict["edge_index"][0]]
+                - data_dict["pos"][data_dict["edge_index"][1]]
+                + shifts
+            )  # [n_edges, 3]
+            # pylint: disable=E1102
+            edge_distance = torch.linalg.norm(
+                edge_distance_vec, dim=-1, keepdim=False
+            )  # [n_edges, 1]
+
+            graph_dict = {
+                "edge_index": data_dict["edge_index"],
+                "edge_distance": edge_distance,
+                "edge_distance_vec": edge_distance_vec,
+                "node_offset": 0,
+            }
+
+        if gp_utils.initialized():
+            graph_dict = self._init_gp_partitions(
+                graph_dict, data_dict["atomic_numbers_full"]
+            )
+            # create partial atomic numbers and batch tensors for GP
+            node_partition = graph_dict["node_partition"]
+            data_dict["atomic_numbers"] = data_dict["atomic_numbers_full"][
+                node_partition
+            ]
+            data_dict["batch"] = data_dict["batch_full"][node_partition]
+        else:
+            graph_dict["node_offset"] = 0
+            graph_dict["edge_distance_vec_full"] = graph_dict["edge_distance_vec"]
+            graph_dict["edge_distance_full"] = graph_dict["edge_distance"]
+            graph_dict["edge_index_full"] = graph_dict["edge_index"]
+
+        return graph_dict
+
+    @conditional_grad(torch.enable_grad())
+    def forward(self, data_dict) -> dict[str, torch.Tensor]:
+        data_dict["atomic_numbers"] = data_dict["atomic_numbers"].long()
+        data_dict["atomic_numbers_full"] = data_dict["atomic_numbers"]
+        data_dict["batch_full"] = data_dict["batch"]
+
+        csd_mixed_emb = self.csd_embedding(
+            charge=data_dict["charge"],
+            spin=data_dict["spin"],
+            dataset=data_dict.get("dataset", None),
+        )
+
+        self.set_MOLE_coefficients(
+            atomic_numbers_full=data_dict["atomic_numbers_full"],
+            batch_full=data_dict["batch_full"],
+            csd_mixed_emb=csd_mixed_emb,
+        )
+
+        with record_function("get_displacement_and_cell"):
+            displacement, orig_cell = self._get_displacement_and_cell(data_dict)
+
+        with record_function("generate_graph"):
+            graph_dict = self._generate_graph(data_dict)
+
+        if graph_dict["edge_index"].numel() == 0:
+            raise ValueError(
+                f"No edges found in input system, this means either you have a single atom in the system or the atoms are farther apart than the radius cutoff of the model of {self.cutoff} Angstroms. We don't know how to handle this case. Check the positions of system: {data_dict['pos']}"
+            )
+
+        with record_function("obtain wigner"):
+            (edge_rot_mat, wigner_and_M_mapping_full, wigner_and_M_mapping_inv_full) = (
+                self._get_rotmat_and_wigner(
+                    graph_dict["edge_distance_vec_full"],
+                    use_cuda_graph=self.use_cuda_graph_wigner
+                    and "cuda" in get_device_for_local_rank()
+                    and not self.training,
+                )
+            )
+            # As a sanity check this should all be 0, dist, 0 (dist = scalar distance)
+            # rotated_ones = torch.bmm(edge_rot_mat, graph_dict["edge_distance_vec"].unsqueeze(-1)).squeeze(-1)
+            if gp_utils.initialized():
+                wigner_and_M_mapping = wigner_and_M_mapping_full[
+                    graph_dict["edge_partition"]
+                ]
+                wigner_and_M_mapping_inv = wigner_and_M_mapping_inv_full[
+                    graph_dict["edge_partition"]
+                ]
+            else:
+                wigner_and_M_mapping = wigner_and_M_mapping_full
+                wigner_and_M_mapping_inv = wigner_and_M_mapping_inv_full
+
+        ###############################################################
+        # Initialize node embeddings
+        ###############################################################
+
+        # Init per node representations using an atomic number based embedding
+        with record_function("atom embedding"):
+            x_message = torch.zeros(
+                data_dict["atomic_numbers"].shape[0],
+                self.sph_feature_size,
+                self.sphere_channels,
+                device=data_dict["pos"].device,
+                dtype=data_dict["pos"].dtype,
+            )
+            x_message[:, 0, :] = self.sphere_embedding(data_dict["atomic_numbers"])
+
+        sys_node_embedding = csd_mixed_emb[data_dict["batch"]]
+        x_message[:, 0, :] = x_message[:, 0, :] + sys_node_embedding
+
+        ###
+        # Hook to allow MOLE
+        ###
+        self.set_MOLE_sizes(
+            nsystems=csd_mixed_emb.shape[0],
+            batch_full=data_dict["batch_full"],
+            edge_index=graph_dict["edge_index"],
+        )
+        self.log_MOLE_stats()
+
+        # edge degree embedding
+        with record_function("edge embedding"):
+            edge_distance_embedding = self.distance_expansion(
+                graph_dict["edge_distance"]
+            )
+            source_embedding = self.source_embedding(
+                data_dict["atomic_numbers_full"][graph_dict["edge_index"][0]]
+            )
+            target_embedding = self.target_embedding(
+                data_dict["atomic_numbers_full"][graph_dict["edge_index"][1]]
+            )
+            x_edge = torch.cat(
+                (edge_distance_embedding, source_embedding, target_embedding), dim=1
+            )
+            x_message = self.edge_degree_embedding(
+                x_message,
+                x_edge,
+                graph_dict["edge_distance"],
+                graph_dict["edge_index"],
+                wigner_and_M_mapping_inv,
+                graph_dict["node_offset"],
+            )
+
+        ###############################################################
+        # Update spherical node embeddings
+        ###############################################################
+        for i in range(self.num_layers):
+            with record_function(f"message passing {i}"):
+                x_message = self.blocks[i](
+                    x_message,
+                    x_edge,
+                    graph_dict["edge_distance"],
+                    graph_dict["edge_index"],
+                    wigner_and_M_mapping,
+                    wigner_and_M_mapping_inv,
+                    sys_node_embedding=sys_node_embedding,
+                    node_offset=graph_dict["node_offset"],
+                )
+
+        # Final layer norm
+        x_message = self.norm(x_message)
+        out = {
+            "node_embedding": x_message,
+            "displacement": displacement,
+            "orig_cell": orig_cell,
+            "batch": data_dict["batch"],
+        }
+        return out
+
+    def _init_gp_partitions(self, graph_dict, atomic_numbers_full):
+        """Graph Parallel
+        This creates the required partial tensors for each rank given the full tensors.
+        The tensors are split on the dimension along the node index using node_partition.
+        """
+        edge_index = graph_dict["edge_index"]
+        edge_distance = graph_dict["edge_distance"]
+        edge_distance_vec_full = graph_dict["edge_distance_vec"]
+
+        node_partition = torch.tensor_split(
+            torch.arange(len(atomic_numbers_full)).to(atomic_numbers_full.device),
+            gp_utils.get_gp_world_size(),
+        )[gp_utils.get_gp_rank()]
+
+        assert (
+            node_partition.numel() > 0
+        ), "Looks like there is no atoms in this graph paralell partition. Cannot proceed"
+        edge_partition = torch.where(
+            torch.logical_and(
+                edge_index[1] >= node_partition.min(),
+                edge_index[1] <= node_partition.max(),  # TODO: 0 or 1?
+            )
+        )[0]
+
+        # full versions of data
+        graph_dict["edge_distance_vec_full"] = edge_distance_vec_full
+        graph_dict["edge_distance_full"] = edge_distance
+        graph_dict["edge_index_full"] = edge_index
+        graph_dict["edge_partition"] = edge_partition
+        graph_dict["node_partition"] = node_partition
+
+        # gp versions of data
+        graph_dict["edge_index"] = edge_index[:, edge_partition]
+        graph_dict["edge_distance"] = edge_distance[edge_partition]
+        graph_dict["edge_distance_vec"] = edge_distance_vec_full[edge_partition]
+        graph_dict["node_offset"] = node_partition.min().item()
+
+        return graph_dict
+
+    @property
+    def num_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    @torch.jit.ignore
+    def no_weight_decay(self) -> set:
+        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.Linear,
+                    SO3_Linear,
+                    torch.nn.LayerNorm,
+                    EquivariantLayerNormArray,
+                    EquivariantLayerNormArraySphericalHarmonics,
+                    EquivariantRMSNormArraySphericalHarmonics,
+                    EquivariantRMSNormArraySphericalHarmonicsV2,
+                ),
+            ):
+                for parameter_name, _ in module.named_parameters():
+                    if (
+                        isinstance(module, (torch.nn.Linear, SO3_Linear))
+                        and "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)

model/uma_escn_moe.py

@@ -0,0 +1,184 @@
+
+
+from __future__ import annotations
+
+import logging
+import warnings
+
+import numpy as np
+import torch
+import torch.nn as nn
+from matplotlib import pyplot as plt
+
+from onescience.modules.head.uma_head import (
+    DatasetSpecificMoEWrapper,
+    DatasetSpecificSingleHeadWrapper,
+)
+from onescience.utils.uma.common.registry import registry
+from model.uma_escn_md import eSCNMDBackbone
+from onescience.modules.func_utils.uma_mole_utils import (
+    MOLEInterface,
+    convert_model_to_MOLE_model,
+    model_search_and_replace,
+    recursive_replace_so2_MOLE,
+    replace_MOLE_with_linear,
+)
+
+# This will catch the warning despite its C++ origin
+# torch.Tensor.index_reduce is in beta
+warnings.filterwarnings(
+    "ignore",
+    message="index_reduce\\(\\) is in beta",
+    category=UserWarning,
+)
+
+
+@registry.register_model("escnmd_moe_backbone")
+class eSCNMDMoeBackbone(eSCNMDBackbone, MOLEInterface):
+    def __init__(
+        self,
+        num_experts: int = 8,
+        moe_dropout: float = 0.0,
+        use_global_embedding: bool = False,  # obsolete
+        use_composition_embedding: bool = False,
+        moe_expert_coefficient_norm: str = "softmax",
+        act=torch.nn.SiLU,
+        layers_moe=None,
+        moe_layer_type: str = "pytorch",
+        moe_single: bool = False,
+        moe_type: str = "so2",
+        model_version: float = 1.0,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.parent_kwargs = kwargs
+        self.num_experts = num_experts
+        self.model_version = model_version
+        if num_experts > 0:
+            convert_model_to_MOLE_model(
+                model=self,
+                num_experts=num_experts,
+                mole_dropout=moe_dropout,
+                mole_expert_coefficient_norm=moe_expert_coefficient_norm,
+                act=act,
+                layers_mole=layers_moe,
+                use_composition_embedding=use_composition_embedding,
+                mole_layer_type=moe_layer_type,
+                mole_single=moe_single,
+                mole_type=moe_type,
+            )
+
+    def merge_MOLE_model(self, data):
+        if self.num_experts == 0:
+            return self
+        data["atomic_numbers"] = data["atomic_numbers"].long()
+        csd_mixed_emb = self.csd_embedding(
+            charge=data["charge"],
+            spin=data["spin"],
+            dataset=data["dataset"],
+        )
+        self.set_MOLE_coefficients(
+            atomic_numbers_full=data["atomic_numbers"],
+            batch_full=data["batch"],
+            csd_mixed_emb=csd_mixed_emb,
+        )
+        if self.mole_type != "so2":
+            raise ValueError("Only mole_type=so2 supported for merging")
+
+        model_search_and_replace(
+            self, recursive_replace_so2_MOLE, replace_MOLE_with_linear
+        )
+
+        # drop moe parameters from merged model
+        self.routing_mlp = None
+        self.composition_embedding = None
+        self.num_experts = 0
+
+        # create a new non moe model and load weights into there
+        new_model = eSCNMDBackbone(**self.parent_kwargs)
+        new_model.load_state_dict(self.state_dict())
+        new_model.eval()
+        return new_model
+
+    def set_MOLE_coefficients(self, atomic_numbers_full, batch_full, csd_mixed_emb):
+        if self.num_experts == 0:
+            return
+        with torch.autocast(device_type=atomic_numbers_full.device.type, enabled=False):
+            embeddings = []
+            if self.use_composition_embedding:
+                composition_by_atom = self.composition_embedding(atomic_numbers_full)
+                composition = composition_by_atom.new_zeros(
+                    csd_mixed_emb.shape[0],
+                    self.sphere_channels,
+                ).index_reduce_(
+                    0,
+                    batch_full,
+                    composition_by_atom,
+                    reduce="mean",
+                    include_self=np.isclose(self.model_version, 1.0).item(),
+                )
+                embeddings.append(composition.unsqueeze(0))
+            embeddings.append(csd_mixed_emb[None])
+
+            expert_mixing_coefficients_before_norm = self.routing_mlp(
+                torch.vstack(embeddings)
+                .transpose(0, 1)
+                .reshape(csd_mixed_emb.shape[0], -1)
+            )
+            self.global_mole_tensors.expert_mixing_coefficients = (
+                self.mole_expert_coefficient_norm(
+                    self.mole_dropout(expert_mixing_coefficients_before_norm)
+                )
+            )
+
+    def set_MOLE_sizes(self, nsystems, batch_full, edge_index):
+        if self.num_experts == 0:
+            return
+        with torch.autocast(device_type=batch_full.device.type, enabled=False):
+            # Generate edge mix_size routing each edge in this instance (GP or not)
+            # using its local edge and batch routing
+
+            # Local edge_index is 2xN where [1,:] is the target node, the target node does not
+            # have the gp offset applied, which means we need to lookup in the full batch_full
+            # _, mix_size = torch.unique(data.batch_full[edge_index[1]], return_counts=True)
+            mole_sizes = torch.zeros(
+                nsystems,  # data.natoms.shape[0],
+                dtype=torch.int,
+                device=batch_full[edge_index[1]].device,
+            ).scatter_(0, batch_full[edge_index[1]], 1, reduce="add")
+
+            self.global_mole_tensors.mole_sizes = mole_sizes.cpu()
+
+    def log_MOLE_stats(self):
+        if not self.training or self.num_experts == 0:
+            return
+        if not hasattr(self, "fig"):
+            self.fig, self.axs = plt.subplots(2, 1)
+        with torch.no_grad():
+            if self.counter % 500 == 0:
+                logging.info(
+                    f"{self.counter }: Expert variance: "
+                    + ",".join(
+                        [
+                            f"{x:.2e}"
+                            for x in self.global_mole_tensors.expert_mixing_coefficients.var(
+                                axis=0
+                            ).tolist()
+                        ]
+                    )
+                )
+                logging.info(
+                    f"{self.counter }: Expert mean: "
+                    + ",".join(
+                        [
+                            f"{x:.2e}"
+                            for x in self.global_mole_tensors.expert_mixing_coefficients.mean(
+                                axis=0
+                            ).tolist()
+                        ]
+                    )
+                )
+                self.fig.tight_layout()
+                self.plot_ready = True
+
+        self.counter += 1

models-json/pretrained_models.json

@@ -0,0 +1,56 @@
+{
+    "uma-s-1": {
+        "filename": "uma-s-1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "uma-s-1p1": {
+        "filename": "uma-s-1p1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "uma-m-1p1": {
+        "filename": "uma-m-1p1.pt",
+        "repo_id": "facebook/UMA",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-md-direct-all-omol": {
+        "filename": "esen_md_direct_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-sm-conserving-all-omol": {
+        "filename": "esen_sm_conserving_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    },
+    "esen-sm-direct-all-omol": {
+        "filename": "esen_sm_direct_all.pt",
+        "repo_id": "facebook/OMol25",
+        "subfolder": "checkpoints",
+        "atom_refs": {
+            "subfolder": "references",
+            "filename": "iso_atom_elem_refs.yaml"
+        }
+    }
+}

scripts/convert_model.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+import argparse
+import pickle
+from pathlib import Path
+import torch
+from omegaconf import OmegaConf, DictConfig, ListConfig
+
+ROOT = Path(__file__).resolve().parents[2]
+BACKBONE_MOE = "model.uma_escn_moe"
+BACKBONE_MD = "model.uma_escn_md"
+BACKBONE_BASE = "model.base"
+
+RULES = [
+    ("onescience.models.UMA.models.base", BACKBONE_BASE),
+    ("onescience.models.UMA.base", BACKBONE_BASE),
+    ("onescience.models.UMA.models.uma.escn_md", BACKBONE_MD),
+    ("onescience.models.UMA.models.uma.escn_moe", BACKBONE_MOE),
+    ("onescience.models.UMA.uma.escn_md", BACKBONE_MD),
+    ("onescience.models.UMA.uma.escn_moe", BACKBONE_MOE),
+    ("onescience.models.UMA.uma_escn_md", BACKBONE_MD),
+    ("onescience.models.UMA.uma_escn_moe", BACKBONE_MOE),
+    ("onescience.utils.uma.models.uma.escn_md", BACKBONE_MD),
+    ("onescience.utils.uma.models.uma.escn_moe", BACKBONE_MOE),
+    ("fairchem.core.models.base", BACKBONE_BASE),
+    ("fairchem.core.models.uma.escn_md", BACKBONE_MD),
+    ("fairchem.core.models.uma.escn_moe", BACKBONE_MOE),
+    ("onescience.models.UMA.units", "onescience.utils.uma.units"),
+    ("onescience.models.UMA.modules.head", "onescience.modules.head.uma_head"),
+    ("onescience.models.UMA.modules.loss", "onescience.modules.loss.uma_loss"),
+    ("onescience.models.UMA.common", "onescience.utils.uma.common"),
+    ("onescience.utils.uma.modules.head", "onescience.modules.head.uma_head"),
+    ("onescience.utils.uma.modules.loss", "onescience.modules.loss.uma_loss"),
+    ("fairchem.core.units", "onescience.utils.uma.units"),
+    ("fairchem.core.modules.head", "onescience.modules.head.uma_head"),
+    ("fairchem.core.modules.loss", "onescience.modules.loss.uma_loss"),
+    ("fairchem.core.components", "onescience.utils.uma.components"),
+    ("fairchem.core.common", "onescience.utils.uma.common"),
+    ("fairchem.core", "onescience.utils.uma"),
+    # 兜底：修复之前被错误转换的中间状态
+    ("onescience.modules.normalization", "onescience.utils.uma.normalization"),
+]
+
+LEGACY_PREFIX = (
+    "onescience.models.UMA.models.",
+    "onescience.models.UMA.uma.",
+    "onescience.utils.uma.models.",
+    "onescience.utils.uma.modules.head.",
+    "onescience.utils.uma.modules.loss.",
+    "fairchem.core.models.",
+    "fairchem.core.modules.head.",
+    "fairchem.core.modules.loss.",
+    "onescience.modules.normalization.",
+)
+
+def remap(s: str) -> str:
+    prev = None
+    cur = s
+    while cur != prev:
+        prev = cur
+        for old, new in RULES:
+            if cur == old or cur.startswith(old + "."):
+                cur = new + cur[len(old):]
+                break
+        # 兜底：修复之前错误转换出的中间状态
+        if cur.startswith("onescience.modules.loss.") and not cur.startswith("onescience.modules.loss.uma_loss."):
+            cur = "onescience.modules.loss.uma_loss." + cur[len("onescience.modules.loss."):]
+        if cur.startswith("onescience.modules.head.") and not cur.startswith("onescience.modules.head.uma_head."):
+            cur = "onescience.modules.head.uma_head." + cur[len("onescience.modules.head."):]
+    return cur
+
+def to_py(x):
+    if isinstance(x, (DictConfig, ListConfig)):
+        return OmegaConf.to_container(x, resolve=False)
+    return x
+
+def walk(x):
+    if isinstance(x, str):
+        return remap(x)
+    if isinstance(x, dict):
+        return {k: walk(v) for k, v in x.items()}
+    if isinstance(x, list):
+        return [walk(v) for v in x]
+    return x
+
+def find_legacy(x, path="root"):
+    out = []
+    if isinstance(x, str):
+        if x.startswith(LEGACY_PREFIX):
+            out.append((path, x))
+    elif isinstance(x, dict):
+        for k, v in x.items():
+            out.extend(find_legacy(v, f"{path}.{k}"))
+    elif isinstance(x, list):
+        for i, v in enumerate(x):
+            out.extend(find_legacy(v, f"{path}[{i}]"))
+    return out
+
+class CustomUnpickler(pickle.Unpickler):
+    def find_class(self, module_name, class_name):
+        mapped = remap(module_name)
+        try:
+            return super().find_class(mapped, class_name)
+        except ModuleNotFoundError:
+            return super().find_class(module_name, class_name)
+
+class CustomPickleModule:
+    Unpickler = CustomUnpickler
+
+def gf(obj, k):
+    return obj[k] if isinstance(obj, dict) else getattr(obj, k)
+
+def sf(obj, k, v):
+    if isinstance(obj, dict):
+        obj[k] = v
+    else:
+        setattr(obj, k, v)
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("src")
+    ap.add_argument("--out", required=True)
+    args = ap.parse_args()
+
+    print(f"[*] loading: {args.src}")
+    try:
+        obj = torch.load(args.src, map_location="cpu", pickle_module=CustomPickleModule, weights_only=False)
+    except TypeError:
+        obj = torch.load(args.src, map_location="cpu", pickle_module=CustomPickleModule)
+
+    mc = OmegaConf.create(walk(to_py(gf(obj, "model_config"))))
+    tc = OmegaConf.create(walk(to_py(gf(obj, "tasks_config"))))
+    sf(obj, "model_config", mc)
+    sf(obj, "tasks_config", tc)
+
+    cfg_py = OmegaConf.to_container(mc, resolve=False)
+    bad = find_legacy(cfg_py)
+    if bad:
+        print("[FATAL] legacy paths still exist in model_config:")
+        for p, v in bad[:20]:
+            print(" ", p, "=", v)
+        raise SystemExit(2)
+
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(obj, str(out))
+    print("[OK] saved:", out)
+    print("[OK] model _target_:", mc.get("_target_", "<none>"))
+    print("[OK] backbone model:", mc.get("backbone", {}).get("model"))
+
+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}"
+    )

scripts/update_demo_config.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+"""用 create_uma_finetune_dataset.py 生成的 data yaml 更新 demo 配置文件。"""
+
+import argparse
+from pathlib import Path
+
+import yaml
+
+
+def load_yaml(path: Path):
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def save_yaml(path: Path, data):
+    with open(path, "w") as f:
+        yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
+
+
+def update_demo_config(demo_config_path: Path, generated_data_path: Path):
+    demo_cfg = load_yaml(demo_config_path)
+    gen_data = load_yaml(generated_data_path)
+
+    # 生成的 yaml 直接就是 data 部分的内容（没有顶层 data 键）
+    demo_cfg["data"] = gen_data
+
+    # 保持 train/val 路径使用 ONESCIENCE_DATASETS_DIR 变量（数据放在 UMA/data/ 下）
+    demo_cfg["data"]["train_dataset"]["splits"]["train"][
+        "src"
+    ] = "${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/train"
+    demo_cfg["data"]["val_dataset"]["splits"]["val"][
+        "src"
+    ] = "${ONESCIENCE_DATASETS_DIR}/data/oc20_finetune/val"
+
+    save_yaml(demo_config_path, demo_cfg)
+    print(f"[OK] 已更新 {demo_config_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--demo-config", type=Path, required=True)
+    parser.add_argument(
+        "--generated-data",
+        type=Path,
+        default=Path("data/oc20_finetune/data/uma_conserving_data_task_energy_force.yaml"),
+    )
+    args = parser.parse_args()
+    update_demo_config(args.demo_config, args.generated_data)

train.py

@@ -0,0 +1,600 @@
+from __future__ import annotations
+import argparse
+import logging
+import os
+import random
+import tempfile
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Optional
+
+# import clusterscope
+try:
+    import clusterscope
+except Exception:
+    clusterscope = None
+import hydra 
+import numpy as np
+import torch
+from omegaconf import OmegaConf
+from omegaconf.errors import InterpolationKeyError
+
+from onescience.utils.uma.common import gp_utils
+
+if TYPE_CHECKING:
+    from omegaconf import DictConfig
+
+    from onescience.utils.uma.components.reducer import Reducer
+    from onescience.utils.uma.components.runner import Runner
+
+from submitit import AutoExecutor
+from submitit.core.utils import JobPaths, cloudpickle_dump
+from submitit.helpers import Checkpointable, DelayedSubmission
+from submitit.slurm.slurm import SlurmJobEnvironment
+
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.uma.common import distutils
+from onescience.monitoring.uma.logger import WandBSingletonLogger
+from onescience.monitoring.uma.runtime_logging import setup_logging
+from onescience.utils.uma.common.utils import (
+    get_commit_hash,
+    get_timestamp_uid,
+    setup_env_vars,
+)
+
+# this effects the cli only since the actual job will be run in subprocesses or remoe
+logging.basicConfig(level=logging.INFO)
+
+
+ALLOWED_TOP_LEVEL_KEYS = {"job", "runner", "reducer"}
+
+LOG_DIR_NAME = "logs"
+CHECKPOINT_DIR_NAME = "checkpoints"
+RESULTS_DIR = "results"
+CONFIG_FILE_NAME = "canonical_config.yaml"
+PREEMPTION_STATE_DIR_NAME = "preemption_state"
+
+
+class SchedulerType(str, Enum):
+    LOCAL = "local"
+    SLURM = "slurm"
+
+
+class DeviceType(str, Enum):
+    CPU = "cpu"
+    CUDA = "cuda"
+
+
+class RunType(str, Enum):
+    RUN = "run"
+    REDUCE = "reduce"
+
+
+class DistributedInitMethod(str, Enum):
+    TCP = "tcp"
+    FILE = "file"
+
+
+@dataclass
+class SlurmConfig:
+    mem_gb: int = 80
+    timeout_hr: int = 168
+    cpus_per_task: int = 8
+    partition: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    qos: Optional[str] = None  # omegaconf in python 3.9 does not backport annotations
+    account: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    additional_parameters: Optional[dict] = None  # 字典格式，用于存储环境变量和配置
+
+
+@dataclass
+class SchedulerConfig:
+    mode: SchedulerType = SchedulerType.LOCAL
+    distributed_init_method: DistributedInitMethod = DistributedInitMethod.TCP
+    ranks_per_node: int = 1
+    num_nodes: int = 1
+    num_array_jobs: int = 1
+
+    # 新增：仅 LOCAL+elastic 多节点时使用（也可不填，见下文）
+    rdzv_backend: str = "c10d"
+    rdzv_endpoint: Optional[str] = None
+    run_id: str = field(default_factory=lambda: f"run_{uuid.uuid4().hex[:8]}")
+
+    slurm: SlurmConfig = field(default_factory=lambda: SlurmConfig)
+
+
+@dataclass
+class SlurmEnv:
+    # reflects the job_id given by submitit (slurm id with array job id and array task id if they exist)
+    job_id: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    # reflects SLURM_JOB_ID only
+    raw_job_id: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    # SLURM_ARRAY_JOB_ID
+    array_job_id: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    # SLURM_ARRAY_TASK_ID
+    array_task_id: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    # reflects SLURM_RESTART_COUNT env variable
+    restart_count: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+
+
+@dataclass
+class Metadata:
+    # read-only metadata about the job, not user inputs
+    commit: str
+    log_dir: str
+    checkpoint_dir: str
+    results_dir: str
+    config_path: str
+    preemption_checkpoint_dir: str
+    cluster_name: str
+    array_job_num: int = 0
+    slurm_env: SlurmEnv = field(default_factory=lambda: SlurmEnv())
+
+
+@dataclass
+class JobConfig:
+    run_name: str = field(
+        default_factory=lambda: get_timestamp_uid() + uuid.uuid4().hex.upper()[0:4]
+    )
+    timestamp_id: str = field(default_factory=lambda: get_timestamp_uid())
+    run_dir: str = field(default_factory=lambda: tempfile.TemporaryDirectory().name)
+    device_type: DeviceType = DeviceType.CUDA
+    debug: bool = False
+    scheduler: SchedulerConfig = field(default_factory=lambda: SchedulerConfig)
+    logger: Optional[dict] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    seed: int = 0
+    deterministic: bool = False
+    runner_state_path: Optional[str] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    # read-only metadata about the job, not user inputs
+    metadata: Optional[Metadata] = (
+        None  # omegaconf in python 3.9 does not backport annotations
+    )
+    graph_parallel_group_size: Optional[int] = None
+
+    def __post_init__(self) -> None:
+        self.run_dir = os.path.abspath(self.run_dir)
+        try:
+            try:
+                cluster = clusterscope.cluster()
+            except Exception as e:
+                # DCU 或受限环境下，clusterscope 可能调用 nvidia-smi 失败
+                cluster = os.environ.get("ONESCIENCE_CLUSTER", "sghpc")
+                if os.environ.get("LOCAL_RANK", "0") == "0":
+                    print(
+                            f"[WARN] clusterscope.cluster() failed "
+                            f"({type(e).__name__}: {e}); fallback to '{cluster}'"
+                            )
+        except RuntimeError:
+            cluster = ""
+        self.metadata = Metadata(
+            commit=get_commit_hash(),
+            log_dir=os.path.join(self.run_dir, self.timestamp_id, LOG_DIR_NAME),
+            checkpoint_dir=os.path.join(
+                self.run_dir, self.timestamp_id, CHECKPOINT_DIR_NAME
+            ),
+            results_dir=os.path.join(self.run_dir, self.timestamp_id, RESULTS_DIR),
+            config_path=os.path.join(self.run_dir, self.timestamp_id, CONFIG_FILE_NAME),
+            preemption_checkpoint_dir=os.path.join(
+                self.run_dir,
+                self.timestamp_id,
+                CHECKPOINT_DIR_NAME,
+                PREEMPTION_STATE_DIR_NAME,
+            ),
+            cluster_name=cluster,
+        )
+
+
+def _set_seeds(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+def _set_deterministic_mode() -> None:
+    # this is required for full cuda deterministic mode
+    logging.info("Setting deterministic mode!")
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    torch.use_deterministic_algorithms(True)
+
+
+def _get_slurm_env() -> SlurmEnv:
+    slurm_job_env = SlurmJobEnvironment()
+    try:
+        slurm_env = SlurmEnv(
+            job_id=slurm_job_env.job_id,
+            raw_job_id=slurm_job_env.raw_job_id,
+            array_job_id=slurm_job_env.array_job_id,
+            array_task_id=slurm_job_env.array_task_id,
+            restart_count=os.environ.get("SLURM_RESTART_COUNT"),
+        )
+    except KeyError:
+        # slurm environment variables are undefined, running locally
+        slurm_env = SlurmEnv()
+
+    return slurm_env
+def remove_runner_state_from_submission(log_folder: str, job_id: str) -> None:
+    # (HACK) Decouple the job from the runner state by manually modifying it
+    # this ensures the saved runner state is not re-submitted in the event of a node failure
+    # ie: if the job was started at state t=T, a requeue during node failure would resubmit the job
+    # starting at state t=T again without calling the checkpoint callback, losing all progress in between.
+    job_path = JobPaths(folder=log_folder, job_id=job_id)
+    if os.path.isfile(job_path.submitted_pickle):
+        submission_obj = DelayedSubmission.load(job_path.submitted_pickle)
+        submission_obj.args[0].job.runner_state_path = None
+        cloudpickle_dump(submission_obj, job_path.submitted_pickle)
+class Submitit(Checkpointable):
+    def __init__(self) -> None:
+        self.config = None
+        self.runner = None
+        self.reducer = None
+
+    def __call__(
+        self, dict_config: DictConfig, run_type: RunType = RunType.RUN
+    ) -> None:
+        self.config = dict_config
+        self.run_type = run_type
+        # modify the config metadata to add slurm info if they exist
+        self.config.job.metadata.slurm_env = _get_slurm_env()
+
+        setup_env_vars()
+        setup_logging()
+        dist_config = map_job_config_to_dist_config(self.config.job)
+        logging.info("Setting up distributed backend...")
+        distutils.setup(dist_config)
+        distutils.synchronize()
+        if (
+            distutils.is_master()
+            and self.config.job.scheduler.mode == SchedulerType.SLURM
+        ):
+            # this pickle file is shared across all processes so can only modify this on the main rank
+            remove_runner_state_from_submission(
+                dict_config.job.metadata.log_dir,
+                self.config.job.metadata.slurm_env.job_id,
+            )
+
+        if self.config.job.graph_parallel_group_size is not None:
+            logging.info("Setting up graph parallel...")
+            gp_utils.setup_graph_parallel_groups(
+                self.config.job.graph_parallel_group_size,
+                dist_config["distributed_backend"],
+            )
+
+        self._init_logger()
+        print(f"rank:{distutils.get_rank()}")
+
+        _set_seeds(self.config.job.seed)
+        if self.config.job.deterministic:
+            _set_deterministic_mode()
+
+        if run_type == RunType.RUN:
+            logging.info("Calling runner.run() ...")
+            self.runner: Runner = hydra.utils.instantiate(self.config.runner)
+            self.runner.job_config = self.config.job
+            # must call resume state AFTER the runner has been initialized
+            self.runner.load_state(self.config.job.runner_state_path)
+            self.runner.run()
+        elif run_type == RunType.REDUCE:
+            logging.info("Calling reducer.reduce() ...")
+            self.reducer: Reducer = hydra.utils.instantiate(self.config.reducer)
+            self.reducer.job_config = self.config.job
+            self.reducer.runner_config = self.config.runner
+            # must call resume state AFTER the runner has been initialized
+            self.reducer.load_state(self.config.job.runner_state_path)
+            self.reducer.reduce()
+        else:
+            raise ValueError(f"run type {run_type} is not recognized!")
+
+        distutils.cleanup()
+
+    def _init_logger(self) -> None:
+        if (
+            self.config.job.logger
+            and distutils.is_master()
+            and not self.config.job.debug
+            and self.config.job.metadata.array_job_num == 0
+        ):
+            # get a partial function from the config and instantiate wandb with it
+            # currently code assumes that we only use the WandBSingletonLogger
+            logger_initializer = hydra.utils.instantiate(self.config.job.logger)
+            simple_config = OmegaConf.to_container(
+                self.config, resolve=True, throw_on_missing=True
+            )
+            logger_initializer(
+                config=simple_config,
+                run_id=self.config.job.timestamp_id,
+                run_name=self.config.job.run_name,
+                log_dir=self.config.job.metadata.log_dir,
+            )
+
+    def checkpoint(self, *args, **kwargs) -> DelayedSubmission:
+        logging.error("Submitit checkpointing callback is triggered")
+        save_path = self.config.job.metadata.preemption_checkpoint_dir
+        cfg_copy = self.config.copy()
+        # only assign if the save was successful
+        cfg_copy.job.runner_state_path = None
+
+        if (
+            self.run_type == RunType.RUN
+            and self.runner.save_state(save_path, is_preemption=True)
+        ) or (
+            self.run_type == RunType.REDUCE
+            and self.reducer.save_state(save_path, is_preemption=True)
+        ):
+            cfg_copy.job.runner_state_path = save_path
+
+        if WandBSingletonLogger.initialized():
+            WandBSingletonLogger.get_instance().mark_preempting()
+        logging.info(
+            f"Submitit checkpointing callback is completed, resuming with use the following state: {save_path}"
+        )
+        return DelayedSubmission(Submitit(), cfg_copy)
+
+
+def map_job_config_to_dist_config(job_cfg: JobConfig) -> dict:
+    scheduler_config = job_cfg.scheduler
+    return {
+        "world_size": scheduler_config.num_nodes * scheduler_config.ranks_per_node,
+        "distributed_backend": (
+            "gloo" if job_cfg.device_type == DeviceType.CPU else "nccl"
+        ),
+        "submit": scheduler_config.mode == SchedulerType.SLURM,
+        "cpu": job_cfg.device_type == DeviceType.CPU,
+        "init_method": scheduler_config.distributed_init_method,
+        # for distributed shared file initialization
+        "shared_file_dir": os.path.join(job_cfg.run_dir, job_cfg.timestamp_id),
+        "array_job_num": job_cfg.metadata.array_job_num,
+    }
+
+
+def get_canonical_config(config: DictConfig) -> DictConfig:
+    # manually initialize metadata, because OmegaConf currently doesn't call __post_init__ on dataclasses
+    job = OmegaConf.to_object(config.job)
+    job.__post_init__()
+    config.job = job
+    # check that each key other than the allowed top level keys are used in config
+    # find all top level keys are not in the allowed set
+    all_keys = set(config.keys()).difference(ALLOWED_TOP_LEVEL_KEYS)
+    used_keys = set()
+    for key in all_keys:
+        # make a copy of all keys except the key in question
+        copy_cfg = OmegaConf.create({k: v for k, v in config.items() if k != key})
+        try:
+            OmegaConf.resolve(copy_cfg)
+        except InterpolationKeyError:
+            # if this error is thrown, this means the key was actually required
+            used_keys.add(key)
+
+    unused_keys = all_keys.difference(used_keys)
+    if unused_keys != set():
+        raise ValueError(
+            f"Found unused keys in the config: {unused_keys}, please remove them!, only keys other than {ALLOWED_TOP_LEVEL_KEYS} or ones that are used as variables are allowed."
+        )
+
+    # resolve the config to fully replace the variables and delete all top level keys except for the ALLOWED_TOP_LEVEL_KEYS
+    for _k in ("ONESCIENCE_DATASETS_DIR", "ONESCIENCE_MODELS_DIR"):
+        if OmegaConf.select(config, _k, default=None) is None:
+            _v = os.environ.get(_k)
+            if _v:
+                OmegaConf.update(config, _k, _v, force_add=True)
+    OmegaConf.resolve(config)
+    return OmegaConf.create(
+        {k: v for k, v in config.items() if k in ALLOWED_TOP_LEVEL_KEYS}
+    )
+
+
+def get_hydra_config_from_yaml(
+    config_yml: str, overrides_args: list[str]
+) -> DictConfig:
+    # Load the configuration from the file
+    os.environ["HYDRA_FULL_ERROR"] = "1"
+    config_directory = os.path.dirname(os.path.abspath(config_yml))
+    config_name = os.path.basename(config_yml)
+    hydra.initialize_config_dir(config_directory, version_base="1.1")
+    cfg = hydra.compose(config_name=config_name, overrides=overrides_args)
+    # merge default structured config with initialized job object
+    cfg = OmegaConf.merge({"job": OmegaConf.structured(JobConfig)}, cfg)
+    # canonicalize config (remove top level keys that just used replacing variables)
+    return get_canonical_config(cfg)
+
+
+def _runner_wrapper(config: DictConfig, run_type: RunType = RunType.RUN):
+    # This is needed when using elastic_launch for local runs since it looks for
+    # the __name__ attribute of the function, Submitit.__call__ does not have one
+    Submitit()(config, run_type)
+
+
+def _run_under_torchrun(cfg: DictConfig) -> None:
+    """在 torchrun / srun+torchrun 启动的 worker 进程里运行训练。
+
+    调用方已经是 rank N 之一（LOCAL_RANK 已由 torchrun 注入），因此这里
+    绝对不能再走 submitit 或 elastic_launch，而是直接：
+      1) setup_env_vars / setup_logging
+      2) DistributedManager.initialize() (读取 RANK/LOCAL_RANK/WORLD_SIZE 等 env 变量)
+      3) gp_utils (可选)
+      4) _set_seeds + deterministic
+      5) WandB 日志 (rank0)
+      6) hydra.utils.instantiate(cfg.runner) + runner.run()
+      7) DistributedManager.cleanup()
+    """
+    setup_env_vars()
+    setup_logging()
+
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    logging.info(
+        f"Init distributed from ENV: RANK={os.environ.get('RANK')}, "
+        f"LOCAL_RANK={os.environ.get('LOCAL_RANK')}, WORLD_SIZE={world_size}, "
+        f"MASTER={os.environ.get('MASTER_ADDR')}:{os.environ.get('MASTER_PORT')}"
+    )
+
+    DistributedManager.initialize()
+
+    if cfg.job.graph_parallel_group_size is not None:
+        backend = "gloo" if cfg.job.device_type == DeviceType.CPU else "nccl"
+        gp_utils.setup_graph_parallel_groups(
+            cfg.job.graph_parallel_group_size, backend
+        )
+
+    _set_seeds(cfg.job.seed)
+    if cfg.job.deterministic:
+        _set_deterministic_mode()
+
+    dm = DistributedManager()
+    if cfg.job.logger and dm.rank == 0 and not cfg.job.debug:
+        logger_initializer = hydra.utils.instantiate(cfg.job.logger)
+        simple_config = OmegaConf.to_container(
+            cfg, resolve=True, throw_on_missing=True
+        )
+        logger_initializer(
+            config=simple_config,
+            run_id=cfg.job.timestamp_id,
+            run_name=cfg.job.run_name,
+            log_dir=cfg.job.metadata.log_dir,
+        )
+
+    runner = hydra.utils.instantiate(cfg.runner)
+    runner.job_config = cfg.job
+    runner.load_state(cfg.job.runner_state_path)
+    runner.run()
+
+    DistributedManager.cleanup()
+
+
+def main(
+    args: argparse.Namespace | None = None, override_args: list[str] | None = None
+):
+    if args is None:
+        parser = argparse.ArgumentParser()
+        parser.add_argument("-c", "--config", type=str, required=True)
+        args, override_args = parser.parse_known_args()
+
+    cfg = get_hydra_config_from_yaml(args.config, override_args)
+
+    log_dir = cfg.job.metadata.log_dir
+    os.makedirs(cfg.job.run_dir, exist_ok=True)
+    os.makedirs(log_dir, exist_ok=True)
+
+    OmegaConf.save(cfg, cfg.job.metadata.config_path)
+    logging.info(f"saved canonical config to {cfg.job.metadata.config_path}")
+
+    scheduler_cfg = cfg.job.scheduler
+
+    # ------------------------------------------------------------------
+    # torchrun / srun+torchrun worker 分支
+    # LOCAL_RANK 由 torchrun 注入, 说明我们已经在 rank N 的子进程中,
+    # 不应再走 submitit / elastic_launch (否则会递归 spawn N*N 个进程)。
+    # demo/run.sh 会把 LOCAL 调度与 torchrun 启动器组合使用, 进入这里。
+    # ------------------------------------------------------------------
+    if (
+        os.environ.get("LOCAL_RANK") is not None
+        and scheduler_cfg.mode == SchedulerType.LOCAL
+    ):
+        _run_under_torchrun(cfg)
+        return
+
+    logging.info(f"Running fairchemv2 cli with {cfg}")
+    if scheduler_cfg.mode == SchedulerType.SLURM:  # Run on cluster
+        assert (
+            os.getenv("SLURM_SUBMIT_HOST") is None
+        ), "SLURM DID NOT SUBMIT JOB!! Please do not submit jobs from an active slurm job (srun or otherwise)"
+        executor = AutoExecutor(folder=log_dir, slurm_max_num_timeout=3)
+        executor.update_parameters(
+            name=cfg.job.run_name,
+            mem_gb=scheduler_cfg.slurm.mem_gb,
+            timeout_min=scheduler_cfg.slurm.timeout_hr * 60,
+            slurm_partition=scheduler_cfg.slurm.partition,
+            # gpus_per_node=scheduler_cfg.ranks_per_node,
+            gpus_per_node=None,  # 不设置 gpus-per-node
+            cpus_per_task=scheduler_cfg.slurm.cpus_per_task,
+            tasks_per_node=scheduler_cfg.ranks_per_node,
+            nodes=scheduler_cfg.num_nodes,
+            slurm_qos=scheduler_cfg.slurm.qos,
+            slurm_account=scheduler_cfg.slurm.account,
+            # 新增：把 YAML 里的 additional_parameters 传给 sbatch
+            slurm_additional_parameters=getattr(
+                scheduler_cfg.slurm, "additional_parameters", None
+            ),
+        )
+        if scheduler_cfg.num_array_jobs == 1:
+            job = executor.submit(Submitit(), cfg)
+            logging.info(
+                f"Submitted job id: {cfg.job.timestamp_id}, slurm id: {job.job_id}, logs: {cfg.job.metadata.log_dir}"
+            )
+            jobs = [job]
+        elif scheduler_cfg.num_array_jobs > 1:
+            executor.update_parameters(
+                slurm_array_parallelism=scheduler_cfg.num_array_jobs,
+            )
+
+            jobs = []
+            with executor.batch():
+                for job_number in range(scheduler_cfg.num_array_jobs):
+                    _cfg = cfg.copy()
+                    _cfg.job.metadata.array_job_num = job_number
+                    job = executor.submit(Submitit(), _cfg)
+                    jobs.append(job)
+            logging.info(f"Submitted {len(jobs)} jobs: {jobs[0].job_id.split('_')[0]}")
+
+        if "reducer" in cfg:
+            job_id = jobs[0].job_id.split("_")[0]
+            executor.update_parameters(
+                name=f"{cfg.job.run_name}_reduce",
+                # set a single node, or do we want the same config as the Runner or a separate JobConfig
+                nodes=1,
+                slurm_dependency=f"afterok:{job_id}",
+                slurm_additional_parameters={
+                    "kill-on-invalid-dep": "yes"
+                },  # kill the reducer if run fails
+            )
+            executor.submit(Submitit(), cfg, RunType.REDUCE)
+    else:
+        from torch.distributed.launcher.api import LaunchConfig, elastic_launch
+
+        if scheduler_cfg.num_nodes > 1:
+            cfg.job.scheduler.num_nodes = 1
+            logging.warning(
+                f"You cannot use more than one node (scheduler_cfg.num_nodes={scheduler_cfg.num_nodes}) in LOCAL mode, over-riding to 1 node"
+            )
+        if scheduler_cfg.ranks_per_node > 1:
+            logging.info(
+                f"Running in local mode with {scheduler_cfg.ranks_per_node} ranks using device_type:{cfg.job.device_type}"
+            )
+            launch_config = LaunchConfig(
+                min_nodes=1,
+                max_nodes=1,
+                nproc_per_node=scheduler_cfg.ranks_per_node,
+                rdzv_backend="c10d",
+                max_restarts=0,
+            )
+            elastic_launch(launch_config, _runner_wrapper)(cfg)
+            if "reducer" in cfg:
+                elastic_launch(launch_config, _runner_wrapper)(cfg, RunType.REDUCE)
+        else:
+            logging.info("Running in local mode without elastic launch")
+            distutils.setup_env_local()
+            Submitit()(cfg)
+            if "reducer" in cfg:
+                Submitit()(cfg, RunType.REDUCE)
+
+if __name__ == "__main__":
+    main()

weight/Jd.pt

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