Upload folder using huggingface_hub

README.md

@@ -0,0 +1,124 @@
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MACE</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+MACE is a machine-learning interatomic potential (MLIP) model for molecular and materials systems, built on E(3)-equivariant graph neural networks. It predicts energy and forces for atomic structures.
+
+Paper: *MACE: Higher order equivariant message passing neural networks for fast and accurate force fields*  
+Reference implementation: https://github.com/ACEsuit/mace
+
+# Model Description
+
+MACE is based on an E(3)-equivariant graph neural network architecture and is trained on HDF5/XYZ format data. It performs energy and force prediction and structure optimization for molecular and materials systems.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :--- |
+| Interatomic potential training | Train a MACE model using standard configurations that read HDF5/XYZ data |
+| Distributed training pre-check | Check multi-GPU/multi-node training configuration, data paths, and statistics consistency |
+| Validation-set evaluation | Output energy- and force-related error metrics on the validation set during training |
+| Custom data migration | Replace the existing HDF5/XYZ data and statistics files with your own data |
+| Environment connectivity check | Use the pre-check script to verify the OneScience matchem environment, pyyaml, h5py, and data readability |
+
+# 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 training.
+- CPU can be used for import and small-configuration connectivity checks; full training will be slow.
+- 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/MACE --local_dir ./mace
+cd mace
+```
+
+### 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 Data Description
+
+This repository does not include built-in training data. Taking the DMC starter dataset as an example, download it from ModelScope and place it under `data/` in the repository root:
+
+```bash
+modelscope download --dataset OneScience/DMC --local_dir ./data
+```
+
+After downloading, the data path will be `data/data/DMC/`. `scripts/demo/run.sh` automatically uses the repository root as `ONESCIENCE_DATASETS_DIR`, so there is no need to manually set this variable to match the paths in the configuration file.
+
+For other configurations (e.g., `ani1x_8dcu.yaml`, `water_*.yaml`, etc.), download the corresponding datasets and adjust the data paths in the YAML file.
+
+
+### Training
+
+Single-GPU:
+
+```bash
+bash scripts/demo/run.sh --config scripts/demo/configs/DMC.yaml
+```
+
+Multi-GPU:
+
+```bash
+# Taking 8 GPUs as an example; launch.launcher in the config should be torchrun
+bash scripts/demo/run.sh --config scripts/demo/configs/ani1x_8dcu.yaml
+```
+
+SLURM submission:
+
+```bash
+bash scripts/demo/run.sh --config scripts/demo/configs/ani1x_8dcu.yaml --submit
+```
+
+### Training Weights
+
+This repository currently does not include built-in trained weights. Weights can be obtained through the training steps above.
+
+# 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 MACE-related code comes from the matchem example implementation in the OneScience project and references the upstream MACE project (https://github.com/ACEsuit/mace). The upstream MACE code is released under the [MIT License](https://github.com/ACEsuit/mace/blob/main/LICENSE).
+- If you use MACE training results in scientific research, we recommend citing the original MACE paper, the relevant OneScience project information, and the sources of the datasets actually used.

model/__init__.py

@@ -0,0 +1,14 @@
+"""MACE model definitions."""
+
+from __future__ import annotations
+
+from .__version__ import __version__
+from .mace import AtomicDipolesMACE, EnergyDipolesMACE, MACE, ScaleShiftMACE
+
+__all__ = [
+    "__version__",
+    "AtomicDipolesMACE",
+    "EnergyDipolesMACE",
+    "MACE",
+    "ScaleShiftMACE",
+]

model/__version__.py

@@ -0,0 +1,3 @@
+__version__ = "0.3.13"
+
+__all__ = ["__version__"]

model/mace.py

@@ -0,0 +1,945 @@
+
+# Implementation of MACE models and other models based E(3)-Equivariant MPNNs
+
+from typing import Any, Callable, Dict, List, Optional, Type, Union
+
+import numpy as np
+import torch
+from e3nn import o3
+from e3nn.util.jit import compile_mode
+
+from onescience.modules.layer.mace_radial import ZBLBasis
+#from onescience.utils.mace.tools.scatter import scatter_sum
+from onescience.datapipes.materials.tools.scatter import scatter_sum
+
+from onescience.modules.block.mace_block import (
+    AtomicEnergiesBlock,
+    EquivariantProductBasisBlock,
+    InteractionBlock,
+    LinearDipoleReadoutBlock,
+    LinearNodeEmbeddingBlock,
+    LinearReadoutBlock,
+    NonLinearDipoleReadoutBlock,
+    NonLinearReadoutBlock,
+    RadialEmbeddingBlock,
+    ScaleShiftBlock,
+)
+from onescience.modules.func_utils.mace_func_utils import (
+    compute_fixed_charge_dipole,
+    get_atomic_virials_stresses,
+    get_edge_vectors_and_lengths,
+    get_outputs,
+    get_symmetric_displacement,
+    prepare_graph,
+)
+
+# pylint: disable=C0302
+
+
+@compile_mode("script")
+class MACE(torch.nn.Module):
+    def __init__(
+        self,
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        max_ell: int,
+        interaction_cls: Type[InteractionBlock],
+        interaction_cls_first: Type[InteractionBlock],
+        num_interactions: int,
+        num_elements: int,
+        hidden_irreps: o3.Irreps,
+        MLP_irreps: o3.Irreps,
+        atomic_energies: np.ndarray,
+        avg_num_neighbors: float,
+        atomic_numbers: List[int],
+        correlation: Union[int, List[int]],
+        gate: Optional[Callable],
+        pair_repulsion: bool = False,
+        distance_transform: str = "None",
+        radial_MLP: Optional[List[int]] = None,
+        radial_type: Optional[str] = "bessel",
+        heads: Optional[List[str]] = None,
+        cueq_config: Optional[Dict[str, Any]] = None,
+        lammps_mliap: Optional[bool] = False,
+    ):
+        super().__init__()
+        self.register_buffer(
+            "atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int64)
+        )
+        self.register_buffer(
+            "r_max", torch.tensor(r_max, dtype=torch.get_default_dtype())
+        )
+        self.register_buffer(
+            "num_interactions", torch.tensor(num_interactions, dtype=torch.int64)
+        )
+        if heads is None:
+            heads = ["Default"]
+        self.heads = heads
+        if isinstance(correlation, int):
+            correlation = [correlation] * num_interactions
+        self.lammps_mliap = lammps_mliap
+        # Embedding
+        node_attr_irreps = o3.Irreps([(num_elements, (0, 1))])
+        node_feats_irreps = o3.Irreps([(hidden_irreps.count(o3.Irrep(0, 1)), (0, 1))])
+        self.node_embedding = LinearNodeEmbeddingBlock(
+            irreps_in=node_attr_irreps,
+            irreps_out=node_feats_irreps,
+            cueq_config=cueq_config,
+        )
+        self.radial_embedding = RadialEmbeddingBlock(
+            r_max=r_max,
+            num_bessel=num_bessel,
+            num_polynomial_cutoff=num_polynomial_cutoff,
+            radial_type=radial_type,
+            distance_transform=distance_transform,
+        )
+        edge_feats_irreps = o3.Irreps(f"{self.radial_embedding.out_dim}x0e")
+        if pair_repulsion:
+            self.pair_repulsion_fn = ZBLBasis(p=num_polynomial_cutoff)
+            self.pair_repulsion = True
+
+        sh_irreps = o3.Irreps.spherical_harmonics(max_ell)
+        num_features = hidden_irreps.count(o3.Irrep(0, 1))
+        interaction_irreps = (sh_irreps * num_features).sort()[0].simplify()
+        self.spherical_harmonics = o3.SphericalHarmonics(
+            sh_irreps, normalize=True, normalization="component"
+        )
+        if radial_MLP is None:
+            radial_MLP = [64, 64, 64]
+        # Interactions and readout
+        self.atomic_energies_fn = AtomicEnergiesBlock(atomic_energies)
+
+        inter = interaction_cls_first(
+            node_attrs_irreps=node_attr_irreps,
+            node_feats_irreps=node_feats_irreps,
+            edge_attrs_irreps=sh_irreps,
+            edge_feats_irreps=edge_feats_irreps,
+            target_irreps=interaction_irreps,
+            hidden_irreps=hidden_irreps,
+            avg_num_neighbors=avg_num_neighbors,
+            radial_MLP=radial_MLP,
+            cueq_config=cueq_config,
+        )
+        self.interactions = torch.nn.ModuleList([inter])
+
+        # Use the appropriate self connection at the first layer for proper E0
+        use_sc_first = False
+        if "Residual" in str(interaction_cls_first):
+            use_sc_first = True
+
+        node_feats_irreps_out = inter.target_irreps
+        prod = EquivariantProductBasisBlock(
+            node_feats_irreps=node_feats_irreps_out,
+            target_irreps=hidden_irreps,
+            correlation=correlation[0],
+            num_elements=num_elements,
+            use_sc=use_sc_first,
+            cueq_config=cueq_config,
+        )
+        self.products = torch.nn.ModuleList([prod])
+
+        self.readouts = torch.nn.ModuleList()
+        self.readouts.append(
+            LinearReadoutBlock(
+                hidden_irreps, o3.Irreps(f"{len(heads)}x0e"), cueq_config
+            )
+        )
+
+        for i in range(num_interactions - 1):
+            if i == num_interactions - 2:
+                hidden_irreps_out = str(
+                    hidden_irreps[0]
+                )  # Select only scalars for last layer
+            else:
+                hidden_irreps_out = hidden_irreps
+            inter = interaction_cls(
+                node_attrs_irreps=node_attr_irreps,
+                node_feats_irreps=hidden_irreps,
+                edge_attrs_irreps=sh_irreps,
+                edge_feats_irreps=edge_feats_irreps,
+                target_irreps=interaction_irreps,
+                hidden_irreps=hidden_irreps_out,
+                avg_num_neighbors=avg_num_neighbors,
+                radial_MLP=radial_MLP,
+                cueq_config=cueq_config,
+            )
+            self.interactions.append(inter)
+            prod = EquivariantProductBasisBlock(
+                node_feats_irreps=interaction_irreps,
+                target_irreps=hidden_irreps_out,
+                correlation=correlation[i + 1],
+                num_elements=num_elements,
+                use_sc=True,
+                cueq_config=cueq_config,
+            )
+            self.products.append(prod)
+            if i == num_interactions - 2:
+                self.readouts.append(
+                    NonLinearReadoutBlock(
+                        hidden_irreps_out,
+                        (len(heads) * MLP_irreps).simplify(),
+                        gate,
+                        o3.Irreps(f"{len(heads)}x0e"),
+                        len(heads),
+                        cueq_config,
+                    )
+                )
+            else:
+                self.readouts.append(
+                    LinearReadoutBlock(
+                        hidden_irreps, o3.Irreps(f"{len(heads)}x0e"), cueq_config
+                    )
+                )
+
+    def forward(
+        self,
+        data: Dict[str, torch.Tensor],
+        training: bool = False,
+        compute_force: bool = True,
+        compute_virials: bool = False,
+        compute_stress: bool = False,
+        compute_displacement: bool = False,
+        compute_hessian: bool = False,
+        compute_edge_forces: bool = False,
+        compute_atomic_stresses: bool = False,
+        lammps_mliap: bool = False,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        # Setup
+        ctx = prepare_graph(
+            data,
+            compute_virials=compute_virials,
+            compute_stress=compute_stress,
+            compute_displacement=compute_displacement,
+            lammps_mliap=lammps_mliap,
+        )
+        is_lammps = ctx.is_lammps
+        num_atoms_arange = ctx.num_atoms_arange
+        num_graphs = ctx.num_graphs
+        displacement = ctx.displacement
+        positions = ctx.positions
+        vectors = ctx.vectors
+        lengths = ctx.lengths
+        cell = ctx.cell
+        node_heads = ctx.node_heads
+        interaction_kwargs = ctx.interaction_kwargs
+        lammps_natoms = interaction_kwargs.lammps_natoms
+        lammps_class = interaction_kwargs.lammps_class
+
+        # Atomic energies
+        node_e0 = self.atomic_energies_fn(data["node_attrs"])[
+            num_atoms_arange, node_heads
+        ]
+        e0 = scatter_sum(
+            src=node_e0, index=data["batch"], dim=0, dim_size=num_graphs
+        )  # [n_graphs, n_heads]
+        # Embeddings
+        node_feats = self.node_embedding(data["node_attrs"])
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(
+            lengths, data["node_attrs"], data["edge_index"], self.atomic_numbers
+        )
+        if hasattr(self, "pair_repulsion"):
+            pair_node_energy = self.pair_repulsion_fn(
+                lengths, data["node_attrs"], data["edge_index"], self.atomic_numbers
+            )
+            if is_lammps:
+                pair_node_energy = pair_node_energy[: lammps_natoms[0]]
+            pair_energy = scatter_sum(
+                src=pair_node_energy, index=data["batch"], dim=-1, dim_size=num_graphs
+            )  # [n_graphs,]
+        else:
+            pair_node_energy = torch.zeros_like(node_e0)
+            pair_energy = torch.zeros_like(e0)
+
+        # Interactions
+        energies = [e0, pair_energy]
+        node_energies_list = [node_e0, pair_node_energy]
+        node_feats_concat: List[torch.Tensor] = []
+
+        for i, (interaction, product, readout) in enumerate(
+            zip(self.interactions, self.products, self.readouts)
+        ):
+            node_attrs_slice = data["node_attrs"]
+            if is_lammps and i > 0:
+                node_attrs_slice = node_attrs_slice[: lammps_natoms[0]]
+            node_feats, sc = interaction(
+                node_attrs=node_attrs_slice,
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=data["edge_index"],
+                first_layer=(i == 0),
+                lammps_class=lammps_class,
+                lammps_natoms=lammps_natoms,
+            )
+            if is_lammps and i == 0:
+                node_attrs_slice = node_attrs_slice[: lammps_natoms[0]]
+            node_feats = product(
+                node_feats=node_feats, sc=sc, node_attrs=node_attrs_slice
+            )
+            node_feats_concat.append(node_feats)
+            node_es = readout(node_feats, node_heads)[num_atoms_arange, node_heads]
+            energy = scatter_sum(node_es, data["batch"], dim=0, dim_size=num_graphs)
+            energies.append(energy)
+            node_energies_list.append(node_es)
+
+        contributions = torch.stack(energies, dim=-1)
+        total_energy = torch.sum(contributions, dim=-1)
+        node_energy = torch.sum(torch.stack(node_energies_list, dim=-1), dim=-1)
+        node_feats_out = torch.cat(node_feats_concat, dim=-1)
+        node_energy = node_e0.double() + pair_node_energy.double()
+
+        forces, virials, stress, hessian, edge_forces = get_outputs(
+            energy=total_energy,
+            positions=positions,
+            displacement=displacement,
+            vectors=vectors,
+            cell=cell,
+            training=training,
+            compute_force=compute_force,
+            compute_virials=compute_virials,
+            compute_stress=compute_stress,
+            compute_hessian=compute_hessian,
+            compute_edge_forces=compute_edge_forces,
+        )
+
+        atomic_virials: Optional[torch.Tensor] = None
+        atomic_stresses: Optional[torch.Tensor] = None
+        if compute_atomic_stresses and edge_forces is not None:
+            atomic_virials, atomic_stresses = get_atomic_virials_stresses(
+                edge_forces=edge_forces,
+                edge_index=data["edge_index"],
+                vectors=vectors,
+                num_atoms=positions.shape[0],
+                batch=data["batch"],
+                cell=cell,
+            )
+        return {
+            "energy": total_energy,
+            "node_energy": node_energy,
+            "contributions": contributions,
+            "forces": forces,
+            "edge_forces": edge_forces,
+            "virials": virials,
+            "stress": stress,
+            "atomic_virials": atomic_virials,
+            "atomic_stresses": atomic_stresses,
+            "displacement": displacement,
+            "hessian": hessian,
+            "node_feats": node_feats_out,
+        }
+
+
+@compile_mode("script")
+class ScaleShiftMACE(MACE):
+    def __init__(
+        self,
+        atomic_inter_scale: float,
+        atomic_inter_shift: float,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.scale_shift = ScaleShiftBlock(
+            scale=atomic_inter_scale, shift=atomic_inter_shift
+        )
+
+    def forward(
+        self,
+        data: Dict[str, torch.Tensor],
+        training: bool = False,
+        compute_force: bool = True,
+        compute_virials: bool = False,
+        compute_stress: bool = False,
+        compute_displacement: bool = False,
+        compute_hessian: bool = False,
+        compute_edge_forces: bool = False,
+        compute_atomic_stresses: bool = False,
+        lammps_mliap: bool = False,
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        # Setup
+        ctx = prepare_graph(
+            data,
+            compute_virials=compute_virials,
+            compute_stress=compute_stress,
+            compute_displacement=compute_displacement,
+            lammps_mliap=lammps_mliap,
+        )
+
+        is_lammps = ctx.is_lammps
+        num_atoms_arange = ctx.num_atoms_arange
+        num_graphs = ctx.num_graphs
+        displacement = ctx.displacement
+        positions = ctx.positions
+        vectors = ctx.vectors
+        lengths = ctx.lengths
+        cell = ctx.cell
+        node_heads = ctx.node_heads
+        interaction_kwargs = ctx.interaction_kwargs
+        lammps_natoms = interaction_kwargs.lammps_natoms
+        lammps_class = interaction_kwargs.lammps_class
+
+        # Atomic energies
+        node_e0 = self.atomic_energies_fn(data["node_attrs"])[
+            num_atoms_arange, node_heads
+        ]
+        e0 = scatter_sum(
+            src=node_e0, index=data["batch"], dim=0, dim_size=num_graphs
+        )  # [n_graphs, num_heads]
+
+        # Embeddings
+        node_feats = self.node_embedding(data["node_attrs"])
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(
+            lengths, data["node_attrs"], data["edge_index"], self.atomic_numbers
+        )
+
+        if hasattr(self, "pair_repulsion"):
+            pair_node_energy = self.pair_repulsion_fn(
+                lengths, data["node_attrs"], data["edge_index"], self.atomic_numbers
+            )
+            if is_lammps:
+                pair_node_energy = pair_node_energy[: lammps_natoms[0]]
+        else:
+            pair_node_energy = torch.zeros_like(node_e0)
+
+        # Interactions
+        node_es_list = [pair_node_energy]
+        node_feats_list: List[torch.Tensor] = []
+
+        for i, (interaction, product, readout) in enumerate(
+            zip(self.interactions, self.products, self.readouts)
+        ):
+            node_attrs_slice = data["node_attrs"]
+            if is_lammps and i > 0:
+                node_attrs_slice = node_attrs_slice[: lammps_natoms[0]]
+            node_feats, sc = interaction(
+                node_attrs=node_attrs_slice,
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=data["edge_index"],
+                first_layer=(i == 0),
+                lammps_class=lammps_class,
+                lammps_natoms=lammps_natoms,
+            )
+            if is_lammps and i == 0:
+                node_attrs_slice = node_attrs_slice[: lammps_natoms[0]]
+            node_feats = product(
+                node_feats=node_feats, sc=sc, node_attrs=node_attrs_slice
+            )
+            node_feats_list.append(node_feats)
+            node_es_list.append(
+                readout(node_feats, node_heads)[num_atoms_arange, node_heads]
+            )
+
+        node_feats_out = torch.cat(node_feats_list, dim=-1)
+        node_inter_es = torch.sum(torch.stack(node_es_list, dim=0), dim=0)
+        node_inter_es = self.scale_shift(node_inter_es, node_heads)
+        inter_e = scatter_sum(node_inter_es, data["batch"], dim=-1, dim_size=num_graphs)
+
+        total_energy = e0 + inter_e
+        node_energy = node_e0.clone().double() + node_inter_es.clone().double()
+
+        forces, virials, stress, hessian, edge_forces = get_outputs(
+            energy=inter_e,
+            positions=positions,
+            displacement=displacement,
+            vectors=vectors,
+            cell=cell,
+            training=training,
+            compute_force=compute_force,
+            compute_virials=compute_virials,
+            compute_stress=compute_stress,
+            compute_hessian=compute_hessian,
+            compute_edge_forces=compute_edge_forces or compute_atomic_stresses,
+        )
+
+        atomic_virials: Optional[torch.Tensor] = None
+        atomic_stresses: Optional[torch.Tensor] = None
+        if compute_atomic_stresses and edge_forces is not None:
+            atomic_virials, atomic_stresses = get_atomic_virials_stresses(
+                edge_forces=edge_forces,
+                edge_index=data["edge_index"],
+                vectors=vectors,
+                num_atoms=positions.shape[0],
+                batch=data["batch"],
+                cell=cell,
+            )
+        return {
+            "energy": total_energy,
+            "node_energy": node_energy,
+            "interaction_energy": inter_e,
+            "forces": forces,
+            "edge_forces": edge_forces,
+            "virials": virials,
+            "stress": stress,
+            "atomic_virials": atomic_virials,
+            "atomic_stresses": atomic_stresses,
+            "hessian": hessian,
+            "displacement": displacement,
+            "node_feats": node_feats_out,
+        }
+
+
+@compile_mode("script")
+class AtomicDipolesMACE(torch.nn.Module):
+    def __init__(
+        self,
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        max_ell: int,
+        interaction_cls: Type[InteractionBlock],
+        interaction_cls_first: Type[InteractionBlock],
+        num_interactions: int,
+        num_elements: int,
+        hidden_irreps: o3.Irreps,
+        MLP_irreps: o3.Irreps,
+        avg_num_neighbors: float,
+        atomic_numbers: List[int],
+        correlation: int,
+        gate: Optional[Callable],
+        atomic_energies: Optional[
+            None
+        ],  # Just here to make it compatible with energy models, MUST be None
+        radial_type: Optional[str] = "bessel",
+        radial_MLP: Optional[List[int]] = None,
+        cueq_config: Optional[Dict[str, Any]] = None,  # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.register_buffer(
+            "atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int64)
+        )
+        self.register_buffer("r_max", torch.tensor(r_max, dtype=torch.float64))
+        self.register_buffer(
+            "num_interactions", torch.tensor(num_interactions, dtype=torch.int64)
+        )
+        assert atomic_energies is None
+
+        # Embedding
+        node_attr_irreps = o3.Irreps([(num_elements, (0, 1))])
+        node_feats_irreps = o3.Irreps([(hidden_irreps.count(o3.Irrep(0, 1)), (0, 1))])
+        self.node_embedding = LinearNodeEmbeddingBlock(
+            irreps_in=node_attr_irreps, irreps_out=node_feats_irreps
+        )
+        self.radial_embedding = RadialEmbeddingBlock(
+            r_max=r_max,
+            num_bessel=num_bessel,
+            num_polynomial_cutoff=num_polynomial_cutoff,
+            radial_type=radial_type,
+        )
+        edge_feats_irreps = o3.Irreps(f"{self.radial_embedding.out_dim}x0e")
+
+        sh_irreps = o3.Irreps.spherical_harmonics(max_ell)
+        num_features = hidden_irreps.count(o3.Irrep(0, 1))
+        interaction_irreps = (sh_irreps * num_features).sort()[0].simplify()
+        self.spherical_harmonics = o3.SphericalHarmonics(
+            sh_irreps, normalize=True, normalization="component"
+        )
+        if radial_MLP is None:
+            radial_MLP = [64, 64, 64]
+
+        # Interactions and readouts
+        inter = interaction_cls_first(
+            node_attrs_irreps=node_attr_irreps,
+            node_feats_irreps=node_feats_irreps,
+            edge_attrs_irreps=sh_irreps,
+            edge_feats_irreps=edge_feats_irreps,
+            target_irreps=interaction_irreps,
+            hidden_irreps=hidden_irreps,
+            avg_num_neighbors=avg_num_neighbors,
+            radial_MLP=radial_MLP,
+        )
+        self.interactions = torch.nn.ModuleList([inter])
+
+        # Use the appropriate self connection at the first layer
+        use_sc_first = False
+        if "Residual" in str(interaction_cls_first):
+            use_sc_first = True
+
+        node_feats_irreps_out = inter.target_irreps
+        prod = EquivariantProductBasisBlock(
+            node_feats_irreps=node_feats_irreps_out,
+            target_irreps=hidden_irreps,
+            correlation=correlation,
+            num_elements=num_elements,
+            use_sc=use_sc_first,
+        )
+        self.products = torch.nn.ModuleList([prod])
+
+        self.readouts = torch.nn.ModuleList()
+        self.readouts.append(LinearDipoleReadoutBlock(hidden_irreps, dipole_only=True))
+
+        for i in range(num_interactions - 1):
+            if i == num_interactions - 2:
+                assert (
+                    len(hidden_irreps) > 1
+                ), "To predict dipoles use at least l=1 hidden_irreps"
+                hidden_irreps_out = str(
+                    hidden_irreps[1]
+                )  # Select only l=1 vectors for last layer
+            else:
+                hidden_irreps_out = hidden_irreps
+            inter = interaction_cls(
+                node_attrs_irreps=node_attr_irreps,
+                node_feats_irreps=hidden_irreps,
+                edge_attrs_irreps=sh_irreps,
+                edge_feats_irreps=edge_feats_irreps,
+                target_irreps=interaction_irreps,
+                hidden_irreps=hidden_irreps_out,
+                avg_num_neighbors=avg_num_neighbors,
+                radial_MLP=radial_MLP,
+            )
+            self.interactions.append(inter)
+            prod = EquivariantProductBasisBlock(
+                node_feats_irreps=interaction_irreps,
+                target_irreps=hidden_irreps_out,
+                correlation=correlation,
+                num_elements=num_elements,
+                use_sc=True,
+            )
+            self.products.append(prod)
+            if i == num_interactions - 2:
+                self.readouts.append(
+                    NonLinearDipoleReadoutBlock(
+                        hidden_irreps_out, MLP_irreps, gate, dipole_only=True
+                    )
+                )
+            else:
+                self.readouts.append(
+                    LinearDipoleReadoutBlock(hidden_irreps, dipole_only=True)
+                )
+
+    def forward(
+        self,
+        data: Dict[str, torch.Tensor],
+        training: bool = False,  # pylint: disable=W0613
+        compute_force: bool = False,
+        compute_virials: bool = False,
+        compute_stress: bool = False,
+        compute_displacement: bool = False,
+        compute_edge_forces: bool = False, # pylint: disable=W0613
+        compute_atomic_stresses: bool = False, # pylint: disable=W0613
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        assert compute_force is False
+        assert compute_virials is False
+        assert compute_stress is False
+        assert compute_displacement is False
+        # Setup
+        data["node_attrs"].requires_grad_(True)
+        data["positions"].requires_grad_(True)
+        num_graphs = data["ptr"].numel() - 1
+
+        # Embeddings
+        node_feats = self.node_embedding(data["node_attrs"])
+        vectors, lengths = get_edge_vectors_and_lengths(
+            positions=data["positions"],
+            edge_index=data["edge_index"],
+            shifts=data["shifts"],
+        )
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(
+            lengths, data["node_attrs"], data["edge_index"], self.atomic_numbers
+        )
+
+        # Interactions
+        dipoles = []
+        for interaction, product, readout in zip(
+            self.interactions, self.products, self.readouts
+        ):
+            node_feats, sc = interaction(
+                node_attrs=data["node_attrs"],
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=data["edge_index"],
+            )
+            node_feats = product(
+                node_feats=node_feats,
+                sc=sc,
+                node_attrs=data["node_attrs"],
+            )
+            node_dipoles = readout(node_feats).squeeze(-1)  # [n_nodes,3]
+            dipoles.append(node_dipoles)
+
+        # Compute the dipoles
+        contributions_dipoles = torch.stack(
+            dipoles, dim=-1
+        )  # [n_nodes,3,n_contributions]
+        atomic_dipoles = torch.sum(contributions_dipoles, dim=-1)  # [n_nodes,3]
+        total_dipole = scatter_sum(
+            src=atomic_dipoles,
+            index=data["batch"],
+            dim=0,
+            dim_size=num_graphs,
+        )  # [n_graphs,3]
+        baseline = compute_fixed_charge_dipole(
+            charges=data["charges"],
+            positions=data["positions"],
+            batch=data["batch"],
+            num_graphs=num_graphs,
+        )  # [n_graphs,3]
+        total_dipole = total_dipole + baseline
+
+        output = {
+            "dipole": total_dipole,
+            "atomic_dipoles": atomic_dipoles,
+        }
+        return output
+
+
+@compile_mode("script")
+class EnergyDipolesMACE(torch.nn.Module):
+    def __init__(
+        self,
+        r_max: float,
+        num_bessel: int,
+        num_polynomial_cutoff: int,
+        max_ell: int,
+        interaction_cls: Type[InteractionBlock],
+        interaction_cls_first: Type[InteractionBlock],
+        num_interactions: int,
+        num_elements: int,
+        hidden_irreps: o3.Irreps,
+        MLP_irreps: o3.Irreps,
+        avg_num_neighbors: float,
+        atomic_numbers: List[int],
+        correlation: int,
+        gate: Optional[Callable],
+        atomic_energies: Optional[np.ndarray],
+        radial_MLP: Optional[List[int]] = None,
+        cueq_config: Optional[Dict[str, Any]] = None,  # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.register_buffer(
+            "atomic_numbers", torch.tensor(atomic_numbers, dtype=torch.int64)
+        )
+        self.register_buffer("r_max", torch.tensor(r_max, dtype=torch.float64))
+        self.register_buffer(
+            "num_interactions", torch.tensor(num_interactions, dtype=torch.int64)
+        )
+        # Embedding
+        node_attr_irreps = o3.Irreps([(num_elements, (0, 1))])
+        node_feats_irreps = o3.Irreps([(hidden_irreps.count(o3.Irrep(0, 1)), (0, 1))])
+        self.node_embedding = LinearNodeEmbeddingBlock(
+            irreps_in=node_attr_irreps, irreps_out=node_feats_irreps
+        )
+        self.radial_embedding = RadialEmbeddingBlock(
+            r_max=r_max,
+            num_bessel=num_bessel,
+            num_polynomial_cutoff=num_polynomial_cutoff,
+        )
+        edge_feats_irreps = o3.Irreps(f"{self.radial_embedding.out_dim}x0e")
+
+        sh_irreps = o3.Irreps.spherical_harmonics(max_ell)
+        num_features = hidden_irreps.count(o3.Irrep(0, 1))
+        interaction_irreps = (sh_irreps * num_features).sort()[0].simplify()
+        self.spherical_harmonics = o3.SphericalHarmonics(
+            sh_irreps, normalize=True, normalization="component"
+        )
+        if radial_MLP is None:
+            radial_MLP = [64, 64, 64]
+        # Interactions and readouts
+        self.atomic_energies_fn = AtomicEnergiesBlock(atomic_energies)
+
+        inter = interaction_cls_first(
+            node_attrs_irreps=node_attr_irreps,
+            node_feats_irreps=node_feats_irreps,
+            edge_attrs_irreps=sh_irreps,
+            edge_feats_irreps=edge_feats_irreps,
+            target_irreps=interaction_irreps,
+            hidden_irreps=hidden_irreps,
+            avg_num_neighbors=avg_num_neighbors,
+            radial_MLP=radial_MLP,
+        )
+        self.interactions = torch.nn.ModuleList([inter])
+
+        # Use the appropriate self connection at the first layer
+        use_sc_first = False
+        if "Residual" in str(interaction_cls_first):
+            use_sc_first = True
+
+        node_feats_irreps_out = inter.target_irreps
+        prod = EquivariantProductBasisBlock(
+            node_feats_irreps=node_feats_irreps_out,
+            target_irreps=hidden_irreps,
+            correlation=correlation,
+            num_elements=num_elements,
+            use_sc=use_sc_first,
+        )
+        self.products = torch.nn.ModuleList([prod])
+
+        self.readouts = torch.nn.ModuleList()
+        self.readouts.append(LinearDipoleReadoutBlock(hidden_irreps, dipole_only=False))
+
+        for i in range(num_interactions - 1):
+            if i == num_interactions - 2:
+                assert (
+                    len(hidden_irreps) > 1
+                ), "To predict dipoles use at least l=1 hidden_irreps"
+                hidden_irreps_out = str(
+                    hidden_irreps[:2]
+                )  # Select scalars and l=1 vectors for last layer
+            else:
+                hidden_irreps_out = hidden_irreps
+            inter = interaction_cls(
+                node_attrs_irreps=node_attr_irreps,
+                node_feats_irreps=hidden_irreps,
+                edge_attrs_irreps=sh_irreps,
+                edge_feats_irreps=edge_feats_irreps,
+                target_irreps=interaction_irreps,
+                hidden_irreps=hidden_irreps_out,
+                avg_num_neighbors=avg_num_neighbors,
+                radial_MLP=radial_MLP,
+            )
+            self.interactions.append(inter)
+            prod = EquivariantProductBasisBlock(
+                node_feats_irreps=interaction_irreps,
+                target_irreps=hidden_irreps_out,
+                correlation=correlation,
+                num_elements=num_elements,
+                use_sc=True,
+            )
+            self.products.append(prod)
+            if i == num_interactions - 2:
+                self.readouts.append(
+                    NonLinearDipoleReadoutBlock(
+                        hidden_irreps_out, MLP_irreps, gate, dipole_only=False
+                    )
+                )
+            else:
+                self.readouts.append(
+                    LinearDipoleReadoutBlock(hidden_irreps, dipole_only=False)
+                )
+
+    def forward(
+        self,
+        data: Dict[str, torch.Tensor],
+        training: bool = False,
+        compute_force: bool = True,
+        compute_virials: bool = False,
+        compute_stress: bool = False,
+        compute_displacement: bool = False,
+        compute_edge_forces: bool = False, # pylint: disable=W0613
+        compute_atomic_stresses: bool = False, # pylint: disable=W0613
+    ) -> Dict[str, Optional[torch.Tensor]]:
+        # Setup
+        data["node_attrs"].requires_grad_(True)
+        data["positions"].requires_grad_(True)
+        num_graphs = data["ptr"].numel() - 1
+        num_atoms_arange = torch.arange(data["positions"].shape[0])
+        displacement = torch.zeros(
+            (num_graphs, 3, 3),
+            dtype=data["positions"].dtype,
+            device=data["positions"].device,
+        )
+        if compute_virials or compute_stress or compute_displacement:
+            (
+                data["positions"],
+                data["shifts"],
+                displacement,
+            ) = get_symmetric_displacement(
+                positions=data["positions"],
+                unit_shifts=data["unit_shifts"],
+                cell=data["cell"],
+                edge_index=data["edge_index"],
+                num_graphs=num_graphs,
+                batch=data["batch"],
+            )
+
+        # Atomic energies
+        node_e0 = self.atomic_energies_fn(data["node_attrs"])[
+            num_atoms_arange, data["head"][data["batch"]]
+        ]
+        e0 = scatter_sum(
+            src=node_e0, index=data["batch"], dim=-1, dim_size=num_graphs
+        )  # [n_graphs,]
+
+        # Embeddings
+        node_feats = self.node_embedding(data["node_attrs"])
+        vectors, lengths = get_edge_vectors_and_lengths(
+            positions=data["positions"],
+            edge_index=data["edge_index"],
+            shifts=data["shifts"],
+        )
+        edge_attrs = self.spherical_harmonics(vectors)
+        edge_feats = self.radial_embedding(
+            lengths, data["node_attrs"], data["edge_index"], self.atomic_numbers
+        )
+
+        # Interactions
+        energies = [e0]
+        node_energies_list = [node_e0]
+        dipoles = []
+        for interaction, product, readout in zip(
+            self.interactions, self.products, self.readouts
+        ):
+            node_feats, sc = interaction(
+                node_attrs=data["node_attrs"],
+                node_feats=node_feats,
+                edge_attrs=edge_attrs,
+                edge_feats=edge_feats,
+                edge_index=data["edge_index"],
+            )
+            node_feats = product(
+                node_feats=node_feats,
+                sc=sc,
+                node_attrs=data["node_attrs"],
+            )
+            node_out = readout(node_feats).squeeze(-1)  # [n_nodes, ]
+            # node_energies = readout(node_feats).squeeze(-1)  # [n_nodes, ]
+            node_energies = node_out[:, 0]
+            energy = scatter_sum(
+                src=node_energies, index=data["batch"], dim=-1, dim_size=num_graphs
+            )  # [n_graphs,]
+            energies.append(energy)
+            node_dipoles = node_out[:, 1:]
+            dipoles.append(node_dipoles)
+
+        # Compute the energies and dipoles
+        contributions = torch.stack(energies, dim=-1)
+        total_energy = torch.sum(contributions, dim=-1)  # [n_graphs, ]
+        node_energy_contributions = torch.stack(node_energies_list, dim=-1)
+        node_energy = torch.sum(node_energy_contributions, dim=-1)  # [n_nodes, ]
+        contributions_dipoles = torch.stack(
+            dipoles, dim=-1
+        )  # [n_nodes,3,n_contributions]
+        atomic_dipoles = torch.sum(contributions_dipoles, dim=-1)  # [n_nodes,3]
+        total_dipole = scatter_sum(
+            src=atomic_dipoles,
+            index=data["batch"].unsqueeze(-1),
+            dim=0,
+            dim_size=num_graphs,
+        )  # [n_graphs,3]
+        baseline = compute_fixed_charge_dipole(
+            charges=data["charges"],
+            positions=data["positions"],
+            batch=data["batch"],
+            num_graphs=num_graphs,
+        )  # [n_graphs,3]
+        total_dipole = total_dipole + baseline
+
+        forces, virials, stress, _, _ = get_outputs(
+            energy=total_energy,
+            positions=data["positions"],
+            displacement=displacement,
+            cell=data["cell"],
+            training=training,
+            compute_force=compute_force,
+            compute_virials=compute_virials,
+            compute_stress=compute_stress,
+        )
+
+        output = {
+            "energy": total_energy,
+            "node_energy": node_energy,
+            "contributions": contributions,
+            "forces": forces,
+            "virials": virials,
+            "stress": stress,
+            "displacement": displacement,
+            "dipole": total_dipole,
+            "atomic_dipoles": atomic_dipoles,
+        }
+        return output

model/py.typed

@@ -0,0 +1 @@
+

scripts/demo/README.md

@@ -0,0 +1,118 @@
+# MACE 训练 Demo
+
+## 快速上手
+
+### 1. 选择配置
+
+`configs/` 目录下提供了多个预定义的实验配置：
+
+| 配置文件                        | 数据集       | GPU 数 | 说明             |
+| ------------------------------- | ------------ | ------ | ---------------- |
+| `DMC.yaml`                      | DMC 溶剂 XTB | 1      | 入门示例，最简单 |
+| `water_1dcu.yaml`               | Water        | 1      | 单卡训练         |
+| `water_4dcu.yaml`               | Water        | 4      | 4 卡分布式训练   |
+| `water_8dcu.yaml`               | Water        | 8      | 8 卡分布式训练   |
+| `ani1x_8dcu.yaml`               | ANI-1x       | 8      | 分布式训练       |
+| `nanotube_l0_8dcu.yaml`         | 碳纳米管     | 8      | max_L=0          |
+| `nanotube_l2_8dcu.yaml`         | 碳纳米管     | 8      | max_L=2          |
+| `nanotube_l2_16dcu.yaml`        | 碳纳米管     | 2x8    | 多节点分布式     |
+
+### 2. 运行训练
+
+```bash
+# 方式一：直接运行（交互式或在已分配的 SLURM 节点上）
+bash run.sh --config configs/DMC.yaml
+
+# 方式二：提交 SLURM 作业
+bash run.sh --config configs/DMC.yaml --submit
+
+# 方式三：预览命令（不执行）
+bash run.sh --config configs/DMC.yaml --dry-run
+```
+
+### 3. 查看输出
+
+训练输出自动保存到 `outputs/{实验名}_{时间戳}/` 目录，包含：
+- 模型 checkpoint
+- 训练日志
+- 当次使用的配置快照 (`config.yaml`)
+
+## 创建自定义实验
+
+1. 复制一个最接近的配置文件：
+   ```bash
+   cp configs/DMC.yaml configs/my_experiment.yaml
+   ```
+
+2. 编辑 YAML 中的参数（只需改参数值，不用碰任何 shell 脚本）
+
+3. 运行：
+   ```bash
+   bash run.sh --config configs/my_experiment.yaml
+   ```
+
+## YAML 配置字段说明
+
+### `train_args` - 训练参数
+
+所有字段直接映射为 `train.py` 的命令行参数。布尔值 `true` 转为标志参数（如 `swa: true` -> `--swa`），`false` 则跳过。
+
+常用参数：
+
+| 参数           | 说明             | 示例                          |
+| -------------- | ---------------- | ----------------------------- |
+| `model`        | 模型类型         | `MACE`                        |
+| `r_max`        | 截断半径 (A)     | `4.0` - `6.0`                 |
+| `num_channels` | 通道数           | `64`, `256`                   |
+| `max_L`        | 最大角动量量子数 | `0`, `2`                      |
+| `batch_size`   | 训练批大小       | `2` - `128`                   |
+| `E0s`          | 原子参考能量     | `average`, `isolated`, 显式字典 |
+| `swa`          | 启用随机权重平均 | `true`                        |
+| `ema`          | 启用指数移动平均 | `true`                        |
+| `distributed`  | 启用分布式训练   | `true` (多卡时自动添加)       |
+
+### `launch` - 启动配置
+
+| 参数                            | 说明       | 启动方式                       |
+| ------------------------------- | ---------- | ------------------------------ |
+| `num_nodes: 1, num_gpus: 1`    | 单卡       | `python train.py`              |
+| `num_nodes: 1, num_gpus: N`    | 单节点多卡 | `torchrun --nproc_per_node=N`  |
+| `num_nodes: M, num_gpus: N`    | 多节点     | `srun` (需 --submit)           |
+
+### `env` - 环境配置
+
+| 参数        | 说明                    |
+| ----------- | ----------------------- |
+| `conda_env` | conda 环境名            |
+| `modules`   | 需要加载的 module 列表  |
+
+### `slurm` - SLURM 作业配置
+
+| 参数            | 说明           |
+| --------------- | -------------- |
+| `partition`     | SLURM 分区     |
+| `time`          | 作业时间限制   |
+| `cpus_per_task` | CPU 核心数     |
+
+### `nccl` - 多节点通信配置（可选）
+
+| 参数            | 说明                |
+| --------------- | ------------------- |
+| `socket_ifname` | InfiniBand 网卡名   |
+| `ib_hca`        | IB HCA 设备名       |
+| `proto`         | NCCL 协议           |
+
+## 目录结构
+
+```
+demo/
+  run.sh                  # 统一入口脚本
+  _parse_config.py        # 配置解析器（内部使用）
+  README.md               # 本文件
+  configs/                # 实验配置
+  templates/              # 脚本模板
+    env_setup.sh          # 环境初始化
+    preflight_check.sh    # 训练前预检
+    slurm_header.template # SLURM header 模板
+  outputs/                # 训练输出（自动创建）
+```

scripts/demo/_parse_config.py

@@ -0,0 +1,304 @@
+#!/usr/bin/env python3
+"""Parse MACE demo YAML and print derived values for run.sh.
+
+Usage:
+    python _parse_config.py <config.yaml> <action>
+
+Actions:
+    command      Print launch command
+    env          Print environment exports
+    env-args     Print arguments for env_setup.sh
+    data-files   Print input files for preflight check
+    slurm        Print SLURM variables
+    name         Print experiment name
+"""
+
+from __future__ import annotations
+
+import os
+import shlex
+import sys
+from typing import Any
+
+import yaml
+
+
+def load_config(path: str) -> dict[str, Any]:
+    with open(path, encoding="utf-8") as f:
+        return yaml.safe_load(f) or {}
+
+
+def _needs_quoting(value: str) -> bool:
+    special = set(" {}()[]|&;'\"\\!`<>*?~#")
+    return any(c in special for c in value)
+
+
+def _quote_value(value: Any) -> str:
+    s = str(value)
+    if "$" in s:
+        escaped = s.replace("\\", "\\\\").replace('"', '\\"')
+        return f'"{escaped}"'
+    return shlex.quote(s)
+
+
+def _build_cli_args(arg_map: dict[str, Any], auto_name: str | None = None) -> list[str]:
+    args: list[str] = []
+    if auto_name and "name" not in arg_map:
+        args.append(f"--name={_quote_value(auto_name)}")
+
+    for key, value in arg_map.items():
+        if isinstance(value, bool):
+            if value:
+                args.append(f"--{key}")
+            continue
+        if isinstance(value, (dict, list)):
+            args.append(f"--{key}={_quote_value(value)}")
+            continue
+
+        s = str(value)
+        if _needs_quoting(s):
+            args.append(f"--{key}={_quote_value(s)}")
+        else:
+            args.append(f"--{key}={s}")
+    return args
+
+
+def _format_command_with_args(prefix: str, args: list[str]) -> str:
+    if not args:
+        return prefix
+    return f"{prefix} \
+  " + " \
+  ".join(args)
+
+
+def _demo_dir(config_path: str) -> str:
+    return os.path.dirname(os.path.dirname(os.path.abspath(config_path)))
+
+
+def _mace_dir(config_path: str) -> str:
+    return os.path.dirname(_demo_dir(config_path))
+
+
+def _get_train_py_path(config_path: str) -> str:
+    # train.py 放在 scripts/ 目录下，即 demo 的父目录
+    return os.path.join(_mace_dir(config_path), "train.py")
+
+
+def _resolve_script_path(config_path: str, script_value: str) -> str:
+    if os.path.isabs(script_value):
+        return script_value
+
+    config_dir = os.path.dirname(os.path.abspath(config_path))
+    demo_dir = _demo_dir(config_path)
+    mace_dir = _mace_dir(config_path)
+
+    candidates = [
+        os.path.join(config_dir, script_value),
+        os.path.join(demo_dir, script_value),
+        os.path.join(mace_dir, script_value),
+        os.path.join(mace_dir, "scripts", script_value),
+    ]
+    for c in candidates:
+        if os.path.exists(c):
+            return c
+
+    # Return the most likely path even if it does not exist; preflight will catch it.
+    return os.path.join(mace_dir, "scripts", script_value)
+
+
+def _is_entrypoint_mode(cfg: dict[str, Any]) -> bool:
+    entrypoint = cfg.get("entrypoint")
+    return isinstance(entrypoint, dict) and bool(entrypoint.get("script"))
+
+
+def _build_train_command(cfg: dict[str, Any], config_path: str) -> str:
+    launch = cfg.get("launch", {}) or {}
+    num_nodes = int(launch.get("num_nodes", 1))
+    num_gpus = int(launch.get("num_gpus", 1))
+    launcher = launch.get("launcher", "python")
+
+    train_py = _get_train_py_path(config_path)
+    train_args = _build_cli_args(
+        cfg.get("train_args", {}) or {},
+        auto_name=cfg.get("name"),
+    )
+
+    if num_nodes > 1 and "--distributed" not in train_args:
+        train_args.append("--distributed")
+    if launcher == "torchrun" and "--distributed" not in train_args:
+        train_args.append("--distributed")
+
+    if num_nodes > 1:
+        return _format_command_with_args(f"python {train_py}", train_args)
+    if launcher == "torchrun":
+        prefix = (
+            "torchrun \
+  --nnodes=1 \
+  "
+            f"--nproc_per_node={num_gpus} \
+  {train_py}"
+        )
+        return _format_command_with_args(prefix, train_args)
+    return _format_command_with_args(f"python {train_py}", train_args)
+
+
+def _build_entrypoint_command(cfg: dict[str, Any], config_path: str) -> str:
+    entrypoint = cfg.get("entrypoint", {}) or {}
+    script_path = _resolve_script_path(config_path, str(entrypoint["script"]))
+    args = _build_cli_args(entrypoint.get("args", {}) or {})
+    return _format_command_with_args(f"python {script_path}", args)
+
+
+def print_command(cfg: dict[str, Any], config_path: str) -> None:
+    if _is_entrypoint_mode(cfg):
+        print(_build_entrypoint_command(cfg, config_path))
+    else:
+        print(_build_train_command(cfg, config_path))
+
+
+def print_env(cfg: dict[str, Any]) -> None:
+    launch = cfg.get("launch", {}) or {}
+    num_nodes = int(launch.get("num_nodes", 1))
+    num_gpus = int(launch.get("num_gpus", 1))
+
+    omp = int(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")
+
+    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']}")
+
+    # 用户自定义额外环境变量（可覆盖上面自动生成的变量）
+    extra_env = cfg.get("extra_env", {})
+    for key, value in extra_env.items():
+        if value is None or value == "":
+            # 空值时只导出 key，保留当前 shell 中的值
+            print(f"export {key}")
+        else:
+            print(f"export {key}={_quote_value(str(value))}")
+
+
+def print_env_args(cfg: dict[str, Any]) -> None:
+    env = cfg.get("env", {}) or {}
+    conda_env = env.get("conda_env", "chem")
+    modules = env.get("modules", []) or []
+    parts = [str(conda_env)] + [str(m) for m in modules]
+    print(" ".join(parts))
+
+
+def _print_unique(paths: list[str]) -> None:
+    seen: set[str] = set()
+    for path in paths:
+        if not path:
+            continue
+        if path in seen:
+            continue
+        seen.add(path)
+        print(path)
+
+
+def print_data_files(cfg: dict[str, Any], config_path: str) -> None:
+    explicit = cfg.get("preflight_files")
+    if isinstance(explicit, list):
+        _print_unique([str(x) for x in explicit])
+        return
+
+    if _is_entrypoint_mode(cfg):
+        entrypoint = cfg.get("entrypoint", {}) or {}
+        files: list[str] = []
+        data_files = entrypoint.get("data_files", [])
+        if isinstance(data_files, list):
+            files.extend(str(x) for x in data_files)
+        script_value = entrypoint.get("script")
+        if script_value:
+            files.append(_resolve_script_path(config_path, str(script_value)))
+        _print_unique(files)
+        return
+
+    train_args = cfg.get("train_args", {}) or {}
+    files: list[str] = []
+    for key in (
+        "train_file",
+        "valid_file",
+        "test_file",
+        "statistics_file",
+        "foundation_model",
+        "pt_train_file",
+    ):
+        val = train_args.get(key)
+        if val:
+            files.append(str(val))
+    _print_unique(files)
+
+
+def print_slurm(cfg: dict[str, Any]) -> None:
+    launch = cfg.get("launch", {}) or {}
+    slurm = cfg.get("slurm", {}) or {}
+    num_nodes = int(launch.get("num_nodes", 1))
+    num_gpus = int(launch.get("num_gpus", 1))
+
+    name = cfg.get("name", "mace_train")
+    partition = slurm.get("partition", "newlarge")
+    time_limit = slurm.get("time", "8:00:00")
+    cpus = int(slurm.get("cpus_per_task", 64))
+
+    if num_nodes > 1:
+        ntasks = num_gpus
+        cpus_per_task = max(1, cpus // max(1, num_gpus))
+    else:
+        ntasks = 1
+        cpus_per_task = cpus
+
+    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_name(cfg: dict[str, Any]) -> None:
+    print(cfg.get("name", "mace_train"))
+
+
+def main() -> None:
+    if len(sys.argv) < 3:
+        print(__doc__, file=sys.stderr)
+        sys.exit(1)
+
+    config_path = sys.argv[1]
+    action = sys.argv[2]
+    cfg = load_config(config_path)
+
+    actions = {
+        "command": lambda: print_command(cfg, config_path),
+        "env": lambda: print_env(cfg),
+        "env-args": lambda: print_env_args(cfg),
+        "data-files": lambda: print_data_files(cfg, config_path),
+        "slurm": lambda: print_slurm(cfg),
+        "name": lambda: print_name(cfg),
+    }
+
+    fn = actions.get(action)
+    if fn is None:
+        print(f"Unknown action: {action}", file=sys.stderr)
+        print(f"Available: {', '.join(actions.keys())}", file=sys.stderr)
+        sys.exit(1)
+    fn()
+
+
+if __name__ == "__main__":
+    main()

scripts/demo/configs/DMC.yaml

@@ -0,0 +1,65 @@
+# ============================================================
+# DMC 溶剂 XTB 训练配置
+# 入门级示例：单 GPU，小数据集，参数最少
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: DMC
+
+# 实验描述
+description: "DMC solvent XTB - 单 GPU 入门示例 (200 训练帧)"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，
+  # 下载 ModelScope DMC 数据集到 ./data 后即可使用 可能需根据实际路径进行调整
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/data/DMC/solvent_xtb_train_200.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/data/DMC/solvent_xtb_test.xyz"
+  valid_fraction: 0.10           # 从训练集中划分 10% 作为验证集
+  energy_key: energy_xtb         # XYZ 文件中能量字段名
+  forces_key: forces_xtb         # XYZ 文件中力字段名
+  E0s: average                   # 原子参考能量方式: average / isolated / 显式字典
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 4.0                     # 截断半径 (Angstrom)
+  batch_size: 10                 # 训练批大小
+  max_num_epochs: 100            # 最大训练轮数
+  eval_interval: 10              # 每 10 个 epoch 评估并打印一次
+
+  # --- SWA (Stochastic Weight Averaging) ---
+  swa: true                      # 启用 SWA
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: python               # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 8                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest01            # SLURM 分区
+  time: "2:00:00"                # 作业时间限制
+  cpus_per_task: 128              # CPU 核心数

scripts/demo/configs/ani1x_8dcu.yaml

@@ -0,0 +1,84 @@
+# ============================================================
+# ANI-1x 8 DCU 训练配置
+# 单节点 8 DCU 分布式训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: ani1x_8dcu
+
+# 实验描述
+description: "ANI-1x - 8 DCU 分布式训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 64               # 通道数
+  max_L: 0                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/ani1x/ANI1x_cc_DFT_rc5_train"
+  valid_file: "${ONESCIENCE_DATASETS_DIR}/data/ani1x/ANI1x_cc_DFT_rc5_val"
+  statistics_file: "${ONESCIENCE_DATASETS_DIR}/data/ani1x/ANI1x_cc_DFT_rc5_statistics.json"
+  E0s: "{1: -13.62222753701504, 6: -1029.4130839658328, 7: -1484.8710358098756, 8: -2041.8396277138045}"
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 4.5                     # 截断半径 (Angstrom)
+  batch_size: 128                # 训练批大小
+  valid_batch_size: 128          # 验证批大小
+  max_num_epochs: 20            # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 40              # 能量的损失权重
+  weight_decay: 1e-7             # 权重衰减
+  clip_grad: 1.0                 # 梯度裁剪阈值
+  scheduler_patience: 20         # 学习率调度器耐心值
+  patience: 50                   # 早停耐心值
+  eval_interval: 1              # 每 10 个 epoch 评估并打印一次
+  num_workers: 32                # 数据加载线程数
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+  start_swa: 15                 # SWA 开始 epoch
+  swa_lr: 0.00025                # SWA 学习率
+  swa_forces_weight: 10          # SWA 力的损失权重
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                 # 保存模型到 CPU
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: torchrun             # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 8                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest02           # SLURM 分区
+  time: "8:00:00"                # 作业时间限制
+  cpus_per_task: 128             # CPU 核心数

scripts/demo/configs/nanotube_eval_l0_1dcu.yaml

@@ -0,0 +1,35 @@
+# ============================================================
+# Nanotube L0 evaluation (single DCU)
+# ============================================================
+name: nanotube_eval_l0_1dcu
+
+description: "Evaluate nanotube L0 model with eval_configs.py"
+
+entrypoint:
+  script: eval_configs.py
+  args:
+    configs: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large.xyz"
+    model: "${ONESCIENCE_MODELS_DIR}/mace/nanotube_l0_8dcu.model"
+    output: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large_eavl.xyz"
+    device: cuda
+    default_dtype: float64
+    batch_size: 2
+    info_prefix: "MACE_"
+  data_files:
+    - "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large.xyz"
+    - "${ONESCIENCE_MODELS_DIR}/mace/nanotube_l0_8dcu.model"
+
+launch:
+  launcher: python
+  num_nodes: 1
+  num_gpus: 1
+
+env:
+  conda_env: matchem
+  modules:
+    - sghpc-mpi-gcc/26.3
+
+slurm:
+  partition: hpctest01
+  time: "2:00:00"
+  cpus_per_task: 16

scripts/demo/configs/nanotube_finetune_l0_8dcu.yaml

@@ -0,0 +1,48 @@
+# ============================================================
+# Nanotube L0 finetuning (8 DCU, single node)
+# ============================================================
+name: nanotube_finetune_l0_8dcu
+
+description: "Finetune from nanotube L0 foundation model"
+
+train_args:
+  foundation_model: "${ONESCIENCE_MODELS_DIR}/mace/nanotube_l0_8dcu.model"
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_test.xyz"
+  valid_fraction: 0.60
+  multiheads_finetuning: false
+  E0s: average
+  energy_key: Energy
+  forces_key: forces
+
+  energy_weight: 1.0
+  forces_weight: 1.0
+  lr: 0.01
+  scaling: rms_forces_scaling
+  batch_size: 2
+  max_num_epochs: 10
+  num_samples_pt: 10
+
+  ema: true
+  swa: true
+  ema_decay: 0.99
+  amsgrad: true
+
+  default_dtype: float64
+  device: cuda
+  seed: 345
+
+launch:
+  launcher: torchrun
+  num_nodes: 1
+  num_gpus: 8
+
+env:
+  conda_env: matchem
+  modules:
+    - sghpc-mpi-gcc/26.3
+
+slurm:
+  partition: hpctest01
+  time: "2:00:00"
+  cpus_per_task: 128

scripts/demo/configs/nanotube_l0_8dcu.yaml

@@ -0,0 +1,83 @@
+# ============================================================
+# Nanotube L=0 训练配置
+# 单节点 8 DCU 分布式训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: nanotube_l0_8dcu
+
+# 实验描述
+description: "Carbon nanotube max_L=0 - 8 DCU 分布式训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 128              # 通道数
+  max_L: 0                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_test.xyz"
+  valid_fraction: 0.05           # 从训练集中划分 5% 作为验证集
+  E0s: average                   # 原子参考能量方式: average / isolated / 显式字典
+  energy_key: Energy             # XYZ 文件中能量字段名
+  forces_key: forces             # XYZ 文件中力字段名
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  enable_cueq: true              # 启用 cuequivariance + FastEq 加速
+  r_max: 6.0                     # 截断半径 (Angstrom)
+  batch_size: 4                  # 训练批大小
+  valid_batch_size: 8            # 验证批大小
+  max_num_epochs: 10             # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 10              # 能量的损失权重
+  scheduler_patience: 5          # 学习率调度器耐心值
+  patience: 15                   # 早停耐心值
+  eval_interval: 1               # 每 1 个 epoch 评估并打印一次
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+  start_swa: 7                   # SWA 开始 epoch
+  swa_forces_weight: 10          # SWA 力的损失权重
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                  # 保存检查点
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: torchrun             # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 8                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest01           # SLURM 分区
+  time: "2:00:00"                # 作业时间限制
+  cpus_per_task: 128             # CPU 核心数

scripts/demo/configs/nanotube_l2_16dcu.yaml

@@ -0,0 +1,114 @@
+# ============================================================
+# Nanotube L=2 多节点训练配置
+# 2 节点 x 8 DCU 多节点分布式训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: nanotube_l2_16dcu
+
+# 实验描述
+description: "Carbon nanotube max_L=2 - 2 节点 x 8 DCU 多节点分布式训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 256              # 通道数
+  max_L: 2                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_test.xyz"
+  valid_fraction: 0.05           # 从训练集中划分 5% 作为验证集
+  E0s: average                   # 原子参考能量方式: average / isolated / 显式字典
+  energy_key: Energy             # XYZ 文件中能量字段名
+  forces_key: forces             # XYZ 文件中力字段名
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 5.0                     # 截断半径 (Angstrom)
+  batch_size: 2                  # 训练批大小 (多节点时适当减小)
+  valid_batch_size: 4            # 验证批大小
+  max_num_epochs: 10             # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 10              # 能量的损失权重
+  scheduler_patience: 5          # 学习率调度器耐心值
+  patience: 15                   # 早停耐心值
+  eval_interval: 1               # 每 1 个 epoch 评估并打印一次
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+  start_swa: 7                   # SWA 开始 epoch
+  swa_forces_weight: 10          # SWA 力的损失权重
+
+  # --- 分布式 ---
+  distributed: true              # 启用分布式训练
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                  # 保存检查点
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: torchrun             # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 2                   # 节点数
+  num_gpus: 8                    # 每节点申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# 多节点 NCCL 通信配置
+# ============================================================
+nccl:
+  socket_ifname: ib0             # NCCL 网络接口
+  ib_hca: mlx5_0                 # InfiniBand HCA 设备
+  proto: Simple                  # NCCL 协议
+
+# ============================================================
+# 自定义环境变量（会覆盖自动生成的变量）
+# ============================================================
+extra_env:
+  NCCL_DEBUG: INFO1
+  NCCL_DEBUG_SUBSYS: ALL
+  NCCL_NET_GDR_LEVEL: 4
+  NCCL_NET_GDR_READ: 1
+  NCCL_PXN_DISABLE: 0
+  RCCL_PXN_GPU_BALANCE: 1
+  RCCL_NET_PLANE: "shca_0,shca_3|shca_1,shca_2"
+  NCCL_SOCKET_IFNAME: eno1
+  HIP_VISIBLE_DEVICES: "0,1,5,4,2,3,7,6"
+  NCCL_NET_PLUGIN: shca
+  NCCL_PLUGIN_P2P: ib
+  NCCL_WORK_FIFO_BYTES1: 4194304
+  LD_LIBRARY_PATH: ""            # 空值：保持当前 shell 中已有的值
+  SHCA_DEBUG_MASK: 0
+  NCCL_TOPO_FILE: ""             # 空值：保持当前 shell 中已有的值
+  UCX_IB_NUM_PATHS: 1
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest01           # SLURM 分区
+  time: "2:00:00"               # 作业时间限制
+  cpus_per_task: 128             # CPU 核心数

scripts/demo/configs/nanotube_l2_8dcu.yaml

@@ -0,0 +1,103 @@
+# ============================================================
+# Nanotube L=2 训练配置
+# 单节点 8 DCU 分布式训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: nanotube_l2_8dcu
+
+# 实验描述
+description: "Carbon nanotube max_L=2 - 8 DCU 分布式训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 128              # 通道数
+  max_L: 2                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_large.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_test.xyz"
+  valid_fraction: 0.05           # 从训练集中划分 5% 作为验证集
+  E0s: average                   # 原子参考能量方式: average / isolated / 显式字典
+  energy_key: Energy             # XYZ 文件中能量字段名
+  forces_key: forces             # XYZ 文件中力字段名
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 5.0                     # 截断半径 (Angstrom)
+  batch_size: 2                  # 训练批大小
+  valid_batch_size: 8            # 验证批大小
+  max_num_epochs: 10             # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 10              # 能量的损失权重
+  scheduler_patience: 5          # 学习率调度器耐心值
+  patience: 15                   # 早停耐心值
+  eval_interval: 1               # 每 1 个 epoch 评估并打印一次
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+  start_swa: 7                   # SWA 开始 epoch
+  swa_forces_weight: 10          # SWA 力的损失权重
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                  # 保存检查点
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: torchrun             # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 8                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# 自定义环境变量（会覆盖自动生成的变量）
+# ============================================================
+extra_env:
+  NCCL_DEBUG: INFO1
+  NCCL_DEBUG_SUBSYS: ALL
+  NCCL_NET_GDR_LEVEL: 4
+  NCCL_NET_GDR_READ: 1
+  NCCL_PXN_DISABLE: 0
+  RCCL_PXN_GPU_BALANCE: 1
+  RCCL_NET_PLANE: "shca_0,shca_3|shca_1,shca_2"
+  NCCL_SOCKET_IFNAME: eno1
+  HIP_VISIBLE_DEVICES: "0,1,5,4,2,3,7,6"
+  NCCL_NET_PLUGIN: shca
+  NCCL_PLUGIN_P2P: ib
+  NCCL_WORK_FIFO_BYTES1: 4194304
+  LD_LIBRARY_PATH: ""            # 空值：保持当前 shell 中已有的值
+  SHCA_DEBUG_MASK: 0
+  NCCL_TOPO_FILE: ""             # 空值：保持当前 shell 中已有的值
+  UCX_IB_NUM_PATHS: 1
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest02           # SLURM 分区
+  time: "8:00:00"                # 作业时间限制
+  cpus_per_task: 128             # CPU 核心数

scripts/demo/configs/nanotube_run_md_l0_1dcu.yaml

@@ -0,0 +1,40 @@
+# ============================================================
+# Nanotube L0 MD inference (single DCU)
+# ============================================================
+name: nanotube_run_md_l0_1dcu
+
+description: "Run MD using trained nanotube L0 model"
+
+entrypoint:
+  script: run_md.py
+  args:
+    init_conf: "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_init.xyz"
+    model: "${ONESCIENCE_MODELS_DIR}/mace/nanotube_l0_8dcu.model"
+    output_xyz: "md_result.xyz"
+    plot_file: "md_plot.png"
+    device: cuda
+    default_dtype: float32
+    temperature_k: 1200
+    time_step_fs: 1.0
+    friction: 0.1
+    save_interval: 10
+    steps: 1000
+    default_dtype: float64
+  data_files:
+    - "${ONESCIENCE_DATASETS_DIR}/data/nanotube/nanotube_init.xyz"
+    - "${ONESCIENCE_MODELS_DIR}/mace/nanotube_l0_8dcu.model"
+
+launch:
+  launcher: python
+  num_nodes: 1
+  num_gpus: 1
+
+env:
+  conda_env: matchem
+  modules:
+    - sghpc-mpi-gcc/26.3
+
+slurm:
+  partition: hpctest01
+  time: "2:00:00"
+  cpus_per_task: 16

scripts/demo/configs/water_1dcu.yaml

@@ -0,0 +1,80 @@
+# ============================================================
+# Water 单 DCU 训练配置
+# 单节点单卡训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: water_1dcu
+
+# 实验描述
+description: "Water 势能面 - 单 DCU 训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 64               # 通道数
+  max_L: 0                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_train.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_test.xyz"
+  valid_fraction: 0.05           # 从训练集中划分 5% 作为验证集
+  E0s: isolated                  # 原子参考能量方式: average / isolated / 显式字典
+  energy_key: TolEnergy          # XYZ 文件中能量字段名
+  forces_key: force              # XYZ 文件中力字段名
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 6.0                     # 截断半径 (Angstrom)
+  batch_size: 2                  # 训练批大小
+  valid_batch_size: 4            # 验证批大小
+  max_num_epochs: 10             # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 10              # 能量的损失权重
+  scheduler_patience: 15         # 学习率调度器耐心值
+  patience: 30                   # 早停耐心值
+  eval_interval: 1               # 每 1 个 epoch 评估并打印一次
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                  # 保存检查点
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: python               # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 1                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest02           # SLURM 分区
+  time: "8:00:00"                # 作业时间限制
+  cpus_per_task: 16              # CPU 核心数 (128 / 8 DCU = 16)

scripts/demo/configs/water_4dcu.yaml

@@ -0,0 +1,81 @@
+# ============================================================
+# Water 4 DCU 训练配置
+# 单节点 4 DCU 分布式训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: water_4dcu
+
+# 实验描述
+description: "Water 势能面 - 4 DCU 分布式训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 64               # 通道数
+  max_L: 0                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_train.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_test.xyz"
+  valid_fraction: 0.05           # 从训练集中划分 5% 作为验证集
+  E0s: isolated                  # 原子参考能量方式: average / isolated / 显式字典
+  energy_key: TolEnergy          # XYZ 文件中能量字段名
+  forces_key: force              # XYZ 文件中力字段名
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 6.0                     # 截断半径 (Angstrom)
+  batch_size: 2                  # 训练批大小
+  valid_batch_size: 4            # 验证批大小
+  max_num_epochs: 20            # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 10              # 能量的损失权重
+  scheduler_patience: 15         # 学习率调度器耐心值
+  patience: 30                   # 早停耐心值
+  eval_interval: 1               # 每 1 个 epoch 评估并打印一次
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+  start_swa: 15                  # SWA 开始 epoch
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                  # 保存检查点
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: torchrun             # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 4                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest02           # SLURM 分区
+  time: "8:00:00"                # 作业时间限制
+  cpus_per_task: 64             # CPU 核心数

scripts/demo/configs/water_8dcu.yaml

@@ -0,0 +1,81 @@
+# ============================================================
+# Water 8 DCU 训练配置
+# 单节点 8 DCU 分布式训练
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: water_8dcu
+
+# 实验描述
+description: "Water 势能面 - 8 DCU 分布式训练"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# 例如 model: MACE -> --model=MACE
+#      swa: true  -> --swa (布尔标志)
+# ============================================================
+train_args:
+  # --- 模型配置 ---
+  model: MACE                    # 模型类型
+  default_dtype: float64         # 数据精度
+  num_interactions: 2            # 交互层数
+  num_channels: 64               # 通道数
+  max_L: 0                       # 最大角动量
+  correlation: 3                 # 相关阶数
+
+  # --- 数据配置 ---
+  # 路径中的 ${ONESCIENCE_DATASETS_DIR} 默认指向仓库根目录，下载对应数据集到 ./data 后即可使用
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_train.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_test.xyz"
+  valid_fraction: 0.05           # 从训练集中划分 5% 作为验证集
+  E0s: isolated                  # 原子参考能量方式: average / isolated / 显式字典
+  energy_key: TolEnergy          # XYZ 文件中能量字段名
+  forces_key: force              # XYZ 文件中力字段名
+
+  # --- 训练超参 ---
+  seed: 123                      # 随机种子
+  device: cuda                   # 设备: cuda / cpu
+  r_max: 6.0                     # 截断半径 (Angstrom)
+  batch_size: 2                  # 训练批大小
+  valid_batch_size: 4            # 验证批大小
+  max_num_epochs: 20            # 最大训练轮数
+  forces_weight: 1000            # 力的损失权重
+  energy_weight: 10              # 能量的损失权重
+  scheduler_patience: 15         # 学习率调度器耐心值
+  patience: 30                   # 早停耐心值
+  eval_interval: 1               # 每 1 个 epoch 评估并打印一次
+  error_table: PerAtomMAE        # 误差表类型
+
+  # --- SWA / EMA ---
+  ema: true                      # 启用 EMA
+  swa: true                      # 启用 SWA
+  start_swa: 15                  # SWA 开始 epoch
+
+  # --- 其他 ---
+  restart_latest: true           # 从最新检查点重启
+  save_cpu: true                  # 保存检查点
+
+# ============================================================
+# 启动配置（控制 python / torchrun / srun 的选择）
+# ============================================================
+launch:
+  launcher: torchrun             # 启动方式: python / torchrun (多节点时自动用 srun)
+  num_nodes: 1                   # 节点数
+  num_gpus: 8                    # 申请的 DCU 卡数
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem                # conda 环境名
+  modules:
+    - sghpc-mpi-gcc/26.3         # 需要加载的 module
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest02           # SLURM 分区
+  time: "8:00:00"                # 作业时间限制
+  cpus_per_task: 128             # CPU 核心数

scripts/demo/configs/water_finetune_mp0_8dcu.yaml

@@ -0,0 +1,97 @@
+# ============================================================
+# Water 8 DCU 微调配置 —— 基于 MACE-MP-0 基础模型
+# 单节点 8 DCU 分布式微调
+# ============================================================
+
+# 实验名称（用于输出目录和模型文件命名）
+name: water_finetune_mp0_8dcu
+
+# 实验描述
+description: "Water 势能面 - 基于 MACE-MP-0 (medium) 微调 - 8 DCU"
+
+# ============================================================
+# train.py 参数
+# 以下所有字段将被转换为 train.py 的命令行参数
+# ============================================================
+train_args:
+  # --- 基础模型配置 ---
+  # 使用 mace-mp-0 medium 预训练模型（自动下载或加载本地缓存）
+  # 可选: small / medium / large / medium-mpa-0 / medium-omat-0 等
+  foundation_model: medium
+
+  # 保留基础模型的所有元素（不仅限于当前数据集出现的元素）
+  # 设置为 true 可防止元素嵌入被截断，有利于泛化
+  foundation_model_elements: true
+
+  # --- 模型配置 ---
+  # 当指定 foundation_model 时，以下参数通常会自动从基础模型继承
+  # 若需覆盖，可取消注释并修改
+  # model: MACE
+  # default_dtype: float64
+  # num_channels: 128
+  # max_L: 1
+  # num_interactions: 2
+  # correlation: 3
+
+  # --- 数据配置 ---
+  train_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_train.xyz"
+  test_file: "${ONESCIENCE_DATASETS_DIR}/data/water/water_test.xyz"
+  valid_fraction: 0.05
+  E0s: isolated
+  energy_key: TolEnergy
+  forces_key: force
+
+  # --- 训练超参（微调建议） ---
+  # 微调通常使用较小的学习率和较少的 epoch
+  seed: 123
+  device: cuda
+  r_max: 6.0
+  batch_size: 2
+  valid_batch_size: 4
+
+  # 微调 epoch 数通常较少（基础模型已具备先验知识）
+  max_num_epochs: 50
+
+  # 损失权重
+  forces_weight: 1000
+  energy_weight: 10
+
+  # 学习率相关（微调建议使用较小学习率）
+  lr: 0.001
+  scheduler_patience: 10
+  patience: 20
+  eval_interval: 1
+  error_table: PerAtomMAE
+
+  # --- SWA / EMA ---
+  ema: true
+  swa: true
+  start_swa: 35
+
+  # --- 其他 ---
+  restart_latest: true
+  save_cpu: true
+
+# ============================================================
+# 启动配置
+# ============================================================
+launch:
+  launcher: torchrun
+  num_nodes: 1
+  num_gpus: 8
+
+# ============================================================
+# 环境配置
+# ============================================================
+env:
+  conda_env: matchem
+  modules:
+    - sghpc-mpi-gcc/26.3
+
+# ============================================================
+# SLURM 配置（仅 --submit 模式使用）
+# ============================================================
+slurm:
+  partition: hpctest02
+  time: "4:00:00"
+  cpus_per_task: 128

scripts/demo/run.sh

@@ -0,0 +1,267 @@
+#!/bin/bash
+# run.sh - MACE 训练统一入口脚本
+#
+# 用法:
+#   bash run.sh --config configs/DMC.yaml              # 直接运行训练
+#   bash run.sh --config configs/DMC.yaml --submit      # 提交 SLURM 作业
+#   bash run.sh --config configs/DMC.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}"
+
+# ============================================================
+# 解析配置
+# ============================================================
+EXP_NAME=$(python3 "$PARSE_PY" "$CONFIG_ABS" name)
+TRAIN_CMD=$(python3 "$PARSE_PY" "$CONFIG_ABS" command)
+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)
+
+# ============================================================
+# Dry-run 模式
+# ============================================================
+if $DRY_RUN; then
+    TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+    OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
+    echo "========================================="
+    echo "Dry-run: $EXP_NAME"
+    echo "Config:  $CONFIG_ABS"
+    echo "Output:  $OUTPUT_DIR (未创建)"
+    echo "========================================="
+    echo ""
+    echo "# 环境变量:"
+    echo "$ENV_EXPORTS"
+    echo ""
+    echo "# 训练命令:"
+    echo "$TRAIN_CMD"
+    echo ""
+    echo "# env_setup.sh 参数: $ENV_ARGS"
+    echo "# 数据文件:"
+    echo "$DATA_FILES" | sed 's/^/#   /'
+    exit 0
+fi
+
+# ============================================================
+# 创建输出目录
+# ============================================================
+TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+OUTPUT_DIR="$DEMO_DIR/outputs/${EXP_NAME}_${TIMESTAMP}"
+mkdir -p "$OUTPUT_DIR"
+cp "$CONFIG_ABS" "$OUTPUT_DIR/config.yaml"
+
+# ============================================================
+# 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 "${MACE_ENV_SCRIPT:-}" ] && [ -f "${MACE_ENV_SCRIPT}" ]; then
+    source "${MACE_ENV_SCRIPT}"
+else
+    echo "[WARN] MACE_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 "# 屏蔽 e3nn FutureWarning 和 TorchScript UserWarning" >> "$SLURM_SCRIPT"
+    echo 'export PYTHONWARNINGS="ignore::FutureWarning:e3nn.o3._wigner,ignore::UserWarning:torch.jit._check"' >> "$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 '' >> "$SLURM_SCRIPT"
+    echo '# 屏蔽 glog INFO（ProcessGroupNCCL.cpp 初始化信息）' >> "$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"
+
+    # 多节点特殊处理
+    if [ "$NODES" -gt 1 ]; then
+        cat >> "$SLURM_SCRIPT" << 'MULTI_NODE'
+
+# 多节点分布式设置
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=29517
+export WORLD_SIZE=$SLURM_NTASKS
+
+echo "MASTER_ADDR: $MASTER_ADDR"
+echo "WORLD_SIZE: $WORLD_SIZE"
+
+# 使用 srun 启动分布式训练
+MULTI_NODE
+        # srun 包裹训练命令
+        echo "srun --export=ALL bash -c '" >> "$SLURM_SCRIPT"
+        echo "  export RANK=\$SLURM_PROCID" >> "$SLURM_SCRIPT"
+        echo "  export LOCAL_RANK=\$SLURM_LOCALID" >> "$SLURM_SCRIPT"
+        echo "  exec $TRAIN_CMD" >> "$SLURM_SCRIPT"
+        echo "'" >> "$SLURM_SCRIPT"
+    else
+        echo "" >> "$SLURM_SCRIPT"
+        echo "# 训练命令" >> "$SLURM_SCRIPT"
+        echo "$TRAIN_CMD" >> "$SLURM_SCRIPT"
+    fi
+
+    echo "========================================="
+    echo "SLURM 脚本已生成: $SLURM_SCRIPT"
+    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 "========================================="
+
+# 加载用户环境初始化脚本（路径通过 MACE_ENV_SCRIPT 指定）
+set +u
+if [ -n "${MACE_ENV_SCRIPT:-}" ] && [ -f "${MACE_ENV_SCRIPT}" ]; then
+    source "${MACE_ENV_SCRIPT}"
+else
+    echo "[WARN] MACE_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:-}"
+
+# 屏蔽 e3nn FutureWarning 和 TorchScript UserWarning
+export PYTHONWARNINGS="ignore::FutureWarning:e3nn.o3._wigner,ignore::UserWarning:torch.jit._check"
+
+# 屏蔽 PyTorch NCCL C++ INFO 日志
+export TORCH_CPP_LOG_LEVEL=WARNING
+export NCCL_DEBUG=ERROR
+
+# 屏蔽 glog INFO（ProcessGroupNCCL.cpp 初始化信息）
+export GLOG_minloglevel=1
+
+# AMD DCU: 避免 RCCL "Missing HSA_FORCE_FINE_GRAIN_PCIE" 警告
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+
+# 切换到输出目录
+cd "$OUTPUT_DIR"
+
+# 执行训练
+echo "========================================="
+echo "开始训练..."
+echo "========================================="
+if [ -n "${TRAIN_CMD:-}" ]; then
+    eval "$TRAIN_CMD"
+else
+    echo "[FATAL] TRAIN_CMD 为空，无法执行任务"
+    exit 1
+fi

scripts/demo/templates/env_setup.sh

@@ -0,0 +1,38 @@
+#!/bin/bash
+# env_setup.sh - 公共环境初始化脚本
+# 用法: source templates/env_setup.sh <conda_env> <module1> [module2 ...]
+#
+# 参数:
+#   $1        - conda 环境名 (如 chem)
+#   $2 ...    - 需要加载的 module 列表
+
+CONDA_ENV="${1:?请指定 conda 环境名}"
+shift
+
+# 系统脚本 (bashrc, conda, module) 可能引用未初始化变量或返回非零退出码，
+# 临时关闭 errexit 和 nounset 避免报错
+set +eu
+
+# 清理环境
+module purge
+source ~/.bashrc
+
+# 激活 conda 环境
+conda activate "$CONDA_ENV"
+
+
+# 加载指定 modules
+for mod in "$@"; do
+    module load "$mod"
+done
+
+# 恢复 errexit 和 nounset
+set -eu
+
+# 加载数据集路径
+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"
+fi

scripts/demo/templates/preflight_check.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# preflight_check.sh - 训练前预检脚本
+# 用法: bash templates/preflight_check.sh [data_file1] [data_file2] ...
+#
+# 检查项:
+#   1. ONESCIENCE_DATASETS_DIR 环境变量
+#   2. 数据文件是否存在
+#   3. DCU/GPU 状态
+#   4. 打印节点信息
+
+ERRORS=0
+
+echo "========================================="
+echo "Node: $(hostname)"
+echo "Time: $(date)"
+echo "========================================="
+
+# 检查数据集目录环境变量
+if [ -z "${ONESCIENCE_DATASETS_DIR:-}" ]; then
+    echo "[ERROR] ONESCIENCE_DATASETS_DIR 未设置，请先 source env.sh"
+    ERRORS=$((ERRORS + 1))
+fi
+
+# 检查数据文件是否存在
+for f in "$@"; do
+    # 跳过空参数
+    [ -z "$f" ] && continue
+    # 展开环境变量后检查
+    expanded=$(eval echo "$f")
+    if [ ! -e "$expanded" ]; then
+        echo "[ERROR] 数据文件不存在: $expanded"
+        ERRORS=$((ERRORS + 1))
+    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] 预检通过"

scripts/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}}

scripts/eval_configs.py

@@ -0,0 +1,187 @@
+# Script for evaluating configurations contained in an xyz file with a trained model
+
+import time
+import argparse
+
+import ase.io
+import numpy as np
+import torch
+from onescience.datapipes.materials.pyg_stack.core.utils import config_from_atoms
+from onescience.datapipes.materials.pyg_stack.core.atomic_data import AtomicData
+from onescience.datapipes.materials.tools import torch_geometric, torch_tools, utils
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument("--configs", help="path to XYZ configurations", required=True)
+    parser.add_argument("--model", help="path to model", required=True)
+    parser.add_argument("--output", help="output path", required=True)
+    parser.add_argument(
+        "--device",
+        help="select device",
+        type=str,
+        choices=["cpu", "cuda"],
+        default="cpu",
+    )
+    parser.add_argument(
+        "--default_dtype",
+        help="set default dtype",
+        type=str,
+        choices=["float32", "float64"],
+        default="float64",
+    )
+    parser.add_argument("--batch_size", help="batch size", type=int, default=64)
+    parser.add_argument(
+        "--compute_stress",
+        help="compute stress",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--return_contributions",
+        help="model outputs energy contributions for each body order, only supported for MACE, not ScaleShiftMACE",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--info_prefix",
+        help="prefix for energy, forces and stress keys",
+        type=str,
+        default="MACE_",
+    )
+    parser.add_argument(
+        "--head",
+        help="Model head used for evaluation",
+        type=str,
+        required=False,
+        default=None,
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    run(args)
+
+
+def run(args: argparse.Namespace) -> None:
+    torch_tools.set_default_dtype(args.default_dtype)
+    device = torch_tools.init_device(args.device)
+
+    # Load model
+    model = torch.load(f=args.model, map_location=args.device)
+    model = model.to(args.device)
+    model.eval()  # ✅ 保持 eval，不要关掉 requires_grad
+
+    # Load data and prepare input
+    atoms_list = ase.io.read(args.configs, index=":")
+    if args.head is not None:
+        for atoms in atoms_list:
+            atoms.info["head"] = args.head
+    configs = [config_from_atoms(atoms) for atoms in atoms_list]
+
+    z_table = utils.AtomicNumberTable([int(z) for z in model.atomic_numbers])
+
+    try:
+        heads = model.heads
+    except AttributeError:
+        heads = None
+
+    data_loader = torch_geometric.dataloader.DataLoader(
+        dataset=[
+            AtomicData.from_config(
+                config, z_table=z_table, cutoff=float(model.r_max), heads=heads
+            )
+            for config in configs
+        ],
+        batch_size=args.batch_size,
+        shuffle=False,
+        drop_last=False,
+    )
+
+    # Collect data
+    energies_list = []
+    contributions_list = []
+    stresses_list = []
+    forces_collection = []
+
+    total_time = 0.0
+    num_batches = 0
+    total_cfgs = len(configs)
+
+    start_total = time.time()
+    for batch in data_loader:
+        batch = batch.to(device)
+
+        if args.device == "cuda":
+            torch.cuda.synchronize()
+        start = time.time()
+
+        output = model(batch.to_dict(), compute_stress=args.compute_stress)
+
+        if args.device == "cuda":
+            torch.cuda.synchronize()
+        end = time.time()
+
+        total_time += (end - start)
+        num_batches += 1
+
+        energies_list.append(torch_tools.to_numpy(output["energy"]))
+        if args.compute_stress:
+            stresses_list.append(torch_tools.to_numpy(output["stress"]))
+        if args.return_contributions:
+            contributions_list.append(torch_tools.to_numpy(output["contributions"]))
+
+        forces = np.split(
+            torch_tools.to_numpy(output["forces"]),
+            indices_or_sections=batch.ptr[1:],
+            axis=0,
+        )
+        forces_collection.append(forces[:-1])  # drop last as its empty
+    end_total = time.time()
+
+    # ✅ 打印时间
+    if num_batches > 0 and total_cfgs > 0:
+        avg_time_batch = total_time / num_batches
+        avg_time_cfg = total_time / total_cfgs
+        print(f"Average inference time per batch: {avg_time_batch:.6f} s")
+        print(f"Average inference time per configuration: {avg_time_cfg:.6f} s")
+        print(f"Total inference time (only forward): {total_time:.6f} s")
+        print(f"Wall clock time (data+forward): {end_total - start_total:.6f} s")
+    else:
+        print("No batches were processed.")
+
+    # Process outputs
+    energies = np.concatenate(energies_list, axis=0)
+    forces_list = [
+        forces for forces_list in forces_collection for forces in forces_list
+    ]
+    assert len(atoms_list) == len(energies) == len(forces_list)
+    if args.compute_stress:
+        stresses = np.concatenate(stresses_list, axis=0)
+        assert len(atoms_list) == stresses.shape[0]
+
+    if args.return_contributions:
+        contributions = np.concatenate(contributions_list, axis=0)
+        assert len(atoms_list) == contributions.shape[0]
+
+    # Store data in atoms objects
+    for i, (atoms, energy, forces) in enumerate(zip(atoms_list, energies, forces_list)):
+        atoms.calc = None  # crucial
+        atoms.info[args.info_prefix + "energy"] = energy
+        atoms.arrays[args.info_prefix + "forces"] = forces
+
+        if args.compute_stress:
+            atoms.info[args.info_prefix + "stress"] = stresses[i]
+
+        if args.return_contributions:
+            atoms.info[args.info_prefix + "BO_contributions"] = contributions[i]
+
+    # Write atoms to output path
+    ase.io.write(args.output, images=atoms_list, format="extxyz")
+
+
+if __name__ == "__main__":
+    main()

scripts/fasteq/fasteq_install.sh

@@ -0,0 +1,88 @@
+#!/bin/bash
+# FastEq-hip 扩展安装脚本（OneScience MACE matchem 环境）
+# 前置条件：已安装好 OneScience 的 matchem 环境。
+# 本脚本安装 FastEq-hip 及 patch 后的 cuequivariance-torch，
+# 以实现 MACE 模型在 DCU 上的训练和推理加速。
+#
+# 使用方法：
+#   cd /public/home/easyscience2024/wangrui/onescience/examples/matchem/mace/fasteq 
+#   bash install.sh
+
+set -e
+
+# =============================================================================
+# 1. 环境初始化
+# =============================================================================
+# 必须先加载 module，再激活 conda；否则 sghpcdas 会覆盖 conda 环境变量。
+source ~/.bashrc
+
+module load sghpcdas/25.6
+module load sghpc-mpi-gcc/26.3
+
+conda activate matchem_opt
+
+# libgalaxyhip.so.5 位于 DTK 的 lib/ 和 hip/lib/ 下，但不在默认 LD_LIBRARY_PATH 中。
+# 不加此路径会导致导入 torch/fasteq 时报错：
+#   ImportError: libgalaxyhip.so.5: cannot open shared object file
+export LD_LIBRARY_PATH="/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/lib:/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/hip/lib:${LD_LIBRARY_PATH}"
+
+# DTK 的 cmake 配置查找 amd_comgr，但 CMAKE_PREFIX_PATH 默认为空，需要显式指定。
+export CMAKE_PREFIX_PATH="/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/dcc/comgr/lib64/cmake/amd_comgr:${CMAKE_PREFIX_PATH}"
+
+# FastEq-hip 编译需要显式声明后端为 HIP，否则默认回退到 CUDA 导致找不到 nvcc。
+export FASTEQ_BACKEND=hip
+export ROCM_PATH=/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4
+export HIP_PATH=/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4
+
+# CMake 3.26 的 HIP 编译器 ABI 检测阶段（try_compile）不会继承主项目的 cache 变量，
+# 但会读取 <Package>_DIR 环境变量。若不设置，会因找不到 amd_comgr 而报错。
+export amd_comgr_DIR=/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/dcc/comgr/lib64/cmake/amd_comgr
+
+# =============================================================================
+# 2. FastEq-hip 源码路径
+# =============================================================================
+# 默认路径：与 onescience 项目同级目录。
+# 可通过环境变量覆盖：FASTEQ_HIP_ROOT=/your/path bash install.sh
+FASTEQ_HIP_ROOT="${FASTEQ_HIP_ROOT}"
+
+if [[ ! -d "$FASTEQ_HIP_ROOT" ]]; then
+    echo "错误：FastEq-hip 源码未找到于 $FASTEQ_HIP_ROOT"
+    echo "请先克隆仓库，或通过 FASTEQ_HIP_ROOT 环境变量指定正确路径。"
+    exit 1
+fi
+
+echo ">>> FastEq-hip 源码路径: $FASTEQ_HIP_ROOT"
+
+# =============================================================================
+# 3. 安装 FastEq-hip
+# =============================================================================
+# PYTHONNOUSERSITE=1 避免使用 ~/.local 下的旧版 setuptools（59.8.0），
+# 该版本会导致 editable_wheel 构建失败（KeyError: 'unix_user'）。
+#
+# 注意：FastEq-hip 的 setup.py 中 HIP 后端的 extra_cmake_args 需补充两项：
+#   -DCMAKE_HIP_COMPILER=/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/llvm/bin/clang++
+#   -Damd_comgr_DIR=/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/dcc/comgr/lib64/cmake/amd_comgr
+# 原因：CMake 3.26 不支持将 hipcc wrapper 直接作为 CMAKE_HIP_COMPILER，
+# 必须使用 ROCm 的 clang++；且 try_compile 阶段不会继承 cache 变量中的 amd_comgr_DIR，
+# 需要通过环境变量（已在上方导出）或 cmake 参数显式传递。
+echo ">>> 正在安装 FastEq-hip（editable 模式）..."
+PYTHONNOUSERSITE=1 pip install -r "$FASTEQ_HIP_ROOT/requirements.txt"
+PYTHONNOUSERSITE=1 pip install -e "$FASTEQ_HIP_ROOT" --no-build-isolation
+
+# =============================================================================
+# 4. 安装 patch 后的 cuequivariance-torch
+# =============================================================================
+# cuequivariance==0.8.0 基础包已由 OneScience matchem 安装。
+# 此处只需覆盖安装本地子模块中的 HIP patch 版本。
+echo ">>> 正在安装 patched cuequivariance-torch（editable 模式）..."
+PYTHONNOUSERSITE=1 pip install -e "$FASTEQ_HIP_ROOT/3rdparty/cuequivariance_torch/cuequivariance_torch" --no-build-isolation
+
+# =============================================================================
+# 5. 验证安装
+# =============================================================================
+echo ">>> 验证导入..."
+python -c "import fasteq; print('  fasteq:', fasteq.__file__)"
+python -c "import cuequivariance_torch as cuet; print('  cuequivariance_torch:', cuet.__file__)"
+
+echo ">>> FastEq-hip 扩展安装成功！"
+echo ">>> 现在可以运行 MACE 训练，并传入 --enable_cueq 启用 FastEq 加速。"

scripts/fasteq/install.sh

@@ -0,0 +1,71 @@
+#!/bin/bash
+# FastEq-hip 扩展安装脚本（OneScience MACE matchem 环境）
+# 前置条件：已安装好 OneScience 的 matchem 环境。
+# 本脚本安装 FastEq-hip 及 patch 后的 cuequivariance-torch，
+# 以实现 MACE 模型在 DCU 上的训练和推理加速。
+#
+# 使用方法：
+#   cd /public/home/scnethpc2601/wr/onescience/examples/matchem/mace/demo/FastEq-hip
+#   bash install.sh
+
+set -e
+
+# =============================================================================
+# 1. 环境初始化
+# =============================================================================
+# 必须先加载 module，再激活 conda；否则 sghpcdas 会覆盖 conda 环境变量。
+module load sghpcdas/25.6
+module load sghpc-mpi-gcc/26.3
+
+source ~/.bashrc
+conda activate matchem
+
+# libgalaxyhip.so.5 位于 DTK 的 lib/ 和 hip/lib/ 下，但不在默认 LD_LIBRARY_PATH 中。
+# 不加此路径会导致导入 torch/fasteq 时报错：
+#   ImportError: libgalaxyhip.so.5: cannot open shared object file
+export LD_LIBRARY_PATH="/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/lib:/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/hip/lib:${LD_LIBRARY_PATH}"
+
+# DTK 的 cmake 配置查找 amd_comgr，但 CMAKE_PREFIX_PATH 默认为空，需要显式指定。
+export CMAKE_PREFIX_PATH="/public/software/sghpc_sdk.bak/Linux_x86_64/26.3/dtk/dtk-25.04.4/dcc/comgr/lib64/cmake/amd_comgr:${CMAKE_PREFIX_PATH}"
+
+# =============================================================================
+# 2. FastEq-hip 源码路径
+# =============================================================================
+# 默认路径：与 onescience 项目同级目录。
+# 可通过环境变量覆盖：FASTEQ_HIP_ROOT=/your/path bash install.sh
+FASTEQ_HIP_ROOT="${FASTEQ_HIP_ROOT:-/public/home/scnethpc2601/wr/project/FastEq-hip}"
+
+if [[ ! -d "$FASTEQ_HIP_ROOT" ]]; then
+    echo "错误：FastEq-hip 源码未找到于 $FASTEQ_HIP_ROOT"
+    echo "请先克隆仓库，或通过 FASTEQ_HIP_ROOT 环境变量指定正确路径。"
+    exit 1
+fi
+
+echo ">>> FastEq-hip 源码路径: $FASTEQ_HIP_ROOT"
+
+# =============================================================================
+# 3. 安装 FastEq-hip
+# =============================================================================
+# PYTHONNOUSERSITE=1 避免使用 ~/.local 下的旧版 setuptools（59.8.0），
+# 该版本会导致 editable_wheel 构建失败（KeyError: 'unix_user'）。
+echo ">>> 正在安装 FastEq-hip（editable 模式）..."
+PYTHONNOUSERSITE=1 pip install -r "$FASTEQ_HIP_ROOT/requirements.txt"
+PYTHONNOUSERSITE=1 pip install -e "$FASTEQ_HIP_ROOT" --no-build-isolation
+
+# =============================================================================
+# 4. 安装 patch 后的 cuequivariance-torch
+# =============================================================================
+# cuequivariance==0.8.0 基础包已由 OneScience matchem 安装。
+# 此处只需覆盖安装本地子模块中的 HIP patch 版本。
+echo ">>> 正在安装 patched cuequivariance-torch（editable 模式）..."
+PYTHONNOUSERSITE=1 pip install -e "$FASTEQ_HIP_ROOT/3rdparty/cuequivariance_torch/cuequivariance_torch" --no-build-isolation
+
+# =============================================================================
+# 5. 验证安装
+# =============================================================================
+echo ">>> 验证导入..."
+python -c "import fasteq; print('  fasteq:', fasteq.__file__)"
+python -c "import cuequivariance_torch as cuet; print('  cuequivariance_torch:', cuet.__file__)"
+
+echo ">>> FastEq-hip 扩展安装成功！"
+echo ">>> 现在可以运行 MACE 训练，并传入 --enable_cueq 启用 FastEq 加速。"

scripts/run_md.py

@@ -0,0 +1,112 @@
+from ase.io import read, write
+from ase import units
+from ase.md.langevin import Langevin
+from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution
+import random
+import os
+import time
+import numpy as np
+import matplotlib.pyplot as plt
+import argparse
+from onescience.utils.mace.calculators import MACECalculator
+
+
+def simpleMD(init_conf, temp, calc, fname, s, T, timestep_fs, friction, plot_file="md_plot.png"):
+    """
+    在服务器上运行的 MD 模拟，不依赖 Jupyter，结束后保存图片。
+    """
+    init_conf.set_calculator(calc)
+
+    # 初始化温度
+    random.seed(701)
+    MaxwellBoltzmannDistribution(init_conf, temperature_K=300)
+    Stationary(init_conf)
+    ZeroRotation(init_conf)
+
+    dyn = Langevin(
+        init_conf,
+        timestep_fs * units.fs,
+        temperature_K=temp,
+        friction=friction
+    )
+
+    time_fs = []
+    temperature = []
+    energies = []
+
+    # 删除之前存储的同名轨迹
+    if os.path.exists(fname):
+        os.system('rm -rfv ' + fname)
+
+    # 记录函数：仅存储数据，不实时绘图
+    def write_frame():
+        dyn.atoms.write(fname, append=True)
+        time_fs.append(dyn.get_time() / units.fs)
+        temperature.append(dyn.atoms.get_temperature())
+        energies.append(dyn.atoms.get_potential_energy() / len(dyn.atoms))
+
+        # 可选：每隔一段时间打印进度
+        if len(time_fs) % 10 == 0:   # 每10步打印一次
+            print(f"Step {len(time_fs)}, Time: {time_fs[-1]:.1f} fs, Temp: {temperature[-1]:.1f} K, E: {energies[-1]:.4f} eV/atom")
+
+    dyn.attach(write_frame, interval=s)
+
+    t0 = time.time()
+    dyn.run(T)
+    t1 = time.time()
+    print("MD 完成于 {0:.2f} 分钟!".format((t1 - t0) / 60))
+
+    # 模拟结束后绘制并保存图片
+    if time_fs and energies and temperature:
+        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 6), sharex='all')
+        ax1.plot(time_fs, energies, 'b-')
+        ax1.set_ylabel('E (eV/atom)')
+        ax2.plot(time_fs, temperature, 'r-')
+        ax2.set_ylabel('T (K)')
+        ax2.set_xlabel('Time (fs)')
+        plt.tight_layout()
+        plt.savefig(plot_file, dpi=150)
+        print(f"Plot saved to {plot_file}")
+        plt.close()   # 避免在无显示器的服务器上弹出窗口
+    else:
+        print("No data to plot.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--init_conf", type=str, required=True)
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument("--output_xyz", type=str, required=True)
+    parser.add_argument("--plot_file", type=str, required=True)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--default_dtype", type=str, default="float64")
+    parser.add_argument("--temperature_k", type=float, required=True)
+    parser.add_argument("--time_step_fs", type=float, default=1.0)
+    parser.add_argument("--friction", type=float, default=0.1)
+    parser.add_argument("--save_interval", type=int, default=10)
+    parser.add_argument("--steps", type=int, required=True)
+
+    args = parser.parse_args()
+
+    # 读取初始构型（使用传进来的绝对路径，不再硬编码相对路径）
+    init_conf = read(args.init_conf, ':')[0].copy()
+
+    # 加载MACE模型
+    mace_calc = MACECalculator(
+        model_paths=[args.model],
+        device=args.device,
+        default_dtype=args.default_dtype
+    )
+
+    # 启动MD
+    simpleMD(
+        init_conf=init_conf,
+        temp=args.temperature_k,
+        calc=mace_calc,
+        fname=args.output_xyz,
+        s=args.save_interval,
+        T=args.steps,
+        timestep_fs=args.time_step_fs,
+        friction=args.friction,
+        plot_file=args.plot_file
+    )

scripts/train.py

@@ -0,0 +1,1098 @@
+
+# Training script for MACE
+
+import ast
+import glob
+import importlib
+import json
+import logging
+import os
+import shutil
+import types
+from copy import deepcopy
+from pathlib import Path
+from typing import List, Optional
+
+import torch.distributed
+import torch.nn.functional
+from e3nn.util import jit
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import LBFGS
+from torch.utils.data import ConcatDataset
+from torch_ema import ExponentialMovingAverage
+
+import model as mace
+#from onescience.utils.mace import data, tools
+from onescience.utils.mace import tools
+from onescience.utils.mace.calculators.foundations_models import mace_mp, mace_off
+from onescience.utils.mace.cli.convert_cueq_e3nn import run as run_cueq_to_e3nn
+from onescience.utils.mace.cli.convert_e3nn_cueq import run as run_e3nn_to_cueq
+from onescience.utils.mace.cli.visualise_train import TrainingPlotter
+#from onescience.utils.mace.data import KeySpecification, update_keyspec_from_kwargs
+from onescience.datapipes.materials.pyg_stack.core.utils import KeySpecification, update_keyspec_from_kwargs
+#from onescience.utils.mace.tools import torch_geometric
+from onescience.datapipes.materials.tools import torch_geometric
+from onescience.utils.mace.tools.model_script_utils import configure_model
+import pdb
+import sys
+try:
+    import model
+    import onescience.utils.mace.tools
+except ImportError as e:
+    print(f"警告：无法导入 model 模块用于别名修复: {e}")
+
+
+# --- 在顶部添加这些导入 ---
+from onescience.datapipes.materials.pyg_stack.core.atomic_data import AtomicData
+from onescience.datapipes.materials.pyg_stack.storage.hdf5_dataset import HDF5Dataset, dataset_from_sharded_hdf5
+from onescience.datapipes.materials.pyg_stack.storage.text_dataset import TextDataset
+from onescience.utils.mace.tools.multihead_tools import (
+    HeadConfig,
+    assemble_mp_data,
+    dict_head_to_dataclass,
+    prepare_default_head,
+    prepare_pt_head,
+)
+from onescience.utils.mace.tools.run_train_utils import (
+    combine_datasets,
+    load_dataset_for_path,
+    normalize_file_paths,
+)
+from onescience.utils.mace.tools.scripts_utils import (
+    LRScheduler,
+    SubsetCollection,
+    check_path_ase_read,
+    convert_to_json_format,
+    dict_to_array,
+    extract_config_mace_model,
+    get_atomic_energies,
+    get_avg_num_neighbors,
+    get_config_type_weights,
+    get_dataset_from_xyz,
+    get_files_with_suffix,
+    get_loss_fn,
+    get_optimizer,
+    get_params_options,
+    get_swa,
+    print_git_commit,
+    remove_pt_head,
+    setup_wandb,
+)
+from onescience.utils.mace.tools.slurm_distributed import DistributedEnvironment
+from onescience.utils.mace.tools.tables_utils import create_error_table
+
+#from onescience.utils.mace.tools.utils import AtomicNumberTable
+from onescience.datapipes.materials.tools.utils import AtomicNumberTable
+
+def main() -> None:
+    """
+    This script runs the training/fine tuning for mace
+    """
+    args = tools.build_default_arg_parser().parse_args()
+    run(args)
+
+
+def run(args) -> None:
+    """
+    This script runs the training/fine tuning for mace
+    """
+
+    tag = tools.get_tag(name=args.name, seed=args.seed)
+    args, input_log_messages = tools.check_args(args)
+
+    # default keyspec to update using heads dictionary
+    args.key_specification = KeySpecification()
+    update_keyspec_from_kwargs(args.key_specification, vars(args))
+
+    if args.device == "xpu":
+        try:
+            import intel_extension_for_pytorch as ipex
+        except ImportError as e:
+            raise ImportError(
+                "Error: Intel extension for PyTorch not found, but XPU device was specified"
+            ) from e
+    if args.distributed:
+        try:
+            distr_env = DistributedEnvironment()
+        except Exception as e:  # pylint: disable=W0703
+            logging.error(f"Failed to initialize distributed environment: {e}")
+            return
+        world_size = distr_env.world_size
+        local_rank = distr_env.local_rank
+        rank = distr_env.rank
+        if rank == 0:
+            print(distr_env)
+        torch.distributed.init_process_group(backend="nccl")
+    else:
+        rank = int(0)
+
+    # Setup
+    tools.set_seeds(args.seed)
+    tools.setup_logger(level=args.log_level, tag=tag, directory=args.log_dir, rank=rank)
+    logging.info("===========VERIFYING SETTINGS===========")
+    for message, loglevel in input_log_messages:
+        logging.log(level=loglevel, msg=message)
+
+    if args.distributed:
+        torch.cuda.set_device(local_rank)
+        logging.info(f"!!!!!!!!!Set device {local_rank}")
+        logging.info(f"Process group initialized: {torch.distributed.is_initialized()}")
+        logging.info(f"Processes: {world_size}")
+
+    try:
+        logging.info(f"MACE version: {mace.__version__}")
+    except AttributeError:
+        logging.info("Cannot find MACE version, please install MACE via pip")
+    logging.debug(f"Configuration: {args}")
+
+    tools.set_default_dtype(args.default_dtype)
+    device = tools.init_device(args.device)
+    commit = print_git_commit()
+    model_foundation: Optional[torch.nn.Module] = None
+    foundation_model_avg_num_neighbors = 0
+    if args.foundation_model is not None:
+        if args.foundation_model in ["small", "medium", "large"]:
+            logging.info(
+                f"Using foundation model mace-mp-0 {args.foundation_model} as initial checkpoint."
+            )
+            calc = mace_mp(
+                model=args.foundation_model,
+                device=args.device,
+                default_dtype=args.default_dtype,
+            )
+            model_foundation = calc.models[0]
+        elif args.foundation_model in ["small_off", "medium_off", "large_off"]:
+            model_type = args.foundation_model.split("_")[0]
+            logging.info(
+                f"Using foundation model mace-off-2023 {model_type} as initial checkpoint. ASL license."
+            )
+            calc = mace_off(
+                model=model_type,
+                device=args.device,
+                default_dtype=args.default_dtype,
+            )
+            model_foundation = calc.models[0]
+        else:
+            # --- !! 插入这里的修复代码 !! ---
+            # (确保 'sys' 和 'model' 已经在文件顶部导入)
+            try:
+                logging.info(
+                    "Creating legacy 'mace' module aliases for foundation checkpoint loading..."
+                )
+                legacy_aliases = {
+                    "mace": "model",
+                    "mace.tools": "onescience.utils.mace.tools",
+                    "mace.data": "onescience.utils.mace.data",
+                    "mace.modules.models": "model.mace",
+                    "mace.modules.blocks": "onescience.modules.block.mace_block",
+                    "mace.modules.radial": "onescience.modules.layer.mace_radial",
+                    "mace.modules.loss": "onescience.modules.loss.mace_loss",
+                    "mace.modules.utils": "onescience.modules.func_utils.mace_func_utils",
+                    "mace.modules.irreps_tools": "onescience.modules.func_utils.mace_irreps_tools",
+                    "mace.modules.symmetric_contraction": "onescience.modules.equivariant.mace_symmetric_contraction",
+                    "mace.modules.wrapper_ops": "onescience.modules.equivariant.mace_wrapper_ops",
+                }
+                legacy_modules_pkg = sys.modules.setdefault(
+                    "mace.modules", types.ModuleType("mace.modules")
+                )
+                if not hasattr(legacy_modules_pkg, "__path__"):
+                    legacy_modules_pkg.__path__ = []  # type: ignore[attr-defined]
+
+                for legacy_name, new_name in legacy_aliases.items():
+                    module_obj = importlib.import_module(new_name)
+                    sys.modules[legacy_name] = module_obj
+                    if legacy_name.startswith("mace.modules."):
+                        setattr(
+                            legacy_modules_pkg,
+                            legacy_name.split(".")[-1],
+                            module_obj,
+                        )
+                logging.info("Legacy module aliases created successfully.")
+            except Exception as e:
+                logging.warning(f"Failed to create module alias: {e}")
+            # --- !! 修复代码结束 !! ---
+            
+            # 现在这个 torch.load 应该可以工作了
+            model_foundation = torch.load(
+                args.foundation_model, map_location=args.device
+            )
+            logging.info(
+                f"Using foundation model {args.foundation_model} as initial checkpoint."
+            )
+        args.r_max = model_foundation.r_max.item()
+        foundation_model_avg_num_neighbors = model_foundation.interactions[
+            0
+        ].avg_num_neighbors
+        if (
+            args.foundation_model not in ["small", "medium", "large"]
+            and args.pt_train_file is None
+        ):
+            logging.warning(
+                "Using multiheads finetuning with a foundation model that is not a Materials Project model, need to provied a path to a pretraining file with --pt_train_file."
+            )
+            args.multiheads_finetuning = False
+        if args.multiheads_finetuning:
+            assert (
+                args.E0s != "average"
+            ), "average atomic energies cannot be used for multiheads finetuning"
+            # check that the foundation model has a single head, if not, use the first head
+            if not args.force_mh_ft_lr:
+                logging.info(
+                    "Multihead finetuning mode, setting learning rate to 0.0001 and EMA to True. To use a different learning rate, set --force_mh_ft_lr=True."
+                )
+                args.lr = 0.0001
+                args.ema = True
+                args.ema_decay = 0.99999
+            logging.info(
+                "Using multiheads finetuning mode, setting learning rate to 0.0001 and EMA to True"
+            )
+            if hasattr(model_foundation, "heads"):
+                if len(model_foundation.heads) > 1:
+                    logging.warning(
+                        "Mutlihead finetuning with models with more than one head is not supported, using the first head as foundation head."
+                    )
+                    model_foundation = remove_pt_head(
+                        model_foundation, args.foundation_head
+                    )
+    else:
+        args.multiheads_finetuning = False
+
+    if args.heads is not None:
+        args.heads = ast.literal_eval(args.heads)
+        for _, head_dict in args.heads.items():
+            # priority is global args < head property_key values < head info_keys+arrays_keys
+            head_keyspec = deepcopy(args.key_specification)
+            update_keyspec_from_kwargs(head_keyspec, head_dict)
+            head_keyspec.update(
+                info_keys=head_dict.get("info_keys", {}),
+                arrays_keys=head_dict.get("arrays_keys", {}),
+            )
+            head_dict["key_specification"] = head_keyspec
+    else:
+        args.heads = prepare_default_head(args)
+    if args.multiheads_finetuning:
+        pt_keyspec = (
+            args.heads["pt_head"]["key_specification"]
+            if "pt_head" in args.heads
+            else deepcopy(args.key_specification)
+        )
+        args.heads["pt_head"] = prepare_pt_head(
+            args, pt_keyspec, foundation_model_avg_num_neighbors
+        )
+
+    logging.info("===========LOADING INPUT DATA===========")
+    heads = list(args.heads.keys())
+    logging.info(f"Using heads: {heads}")
+    logging.info("Using the key specifications to parse data:")
+    for name, head_dict in args.heads.items():
+        head_keyspec = head_dict["key_specification"]
+        logging.info(f"{name}: {head_keyspec}")
+
+    head_configs: List[HeadConfig] = []
+    for head, head_args in args.heads.items():
+        logging.info(f"=============    Processing head {head}     ===========")
+        head_config = dict_head_to_dataclass(head_args, head, args)
+
+        # Handle train_file and valid_file - normalize to lists
+        if hasattr(head_config, "train_file") and head_config.train_file is not None:
+            head_config.train_file = normalize_file_paths(head_config.train_file)
+        if hasattr(head_config, "valid_file") and head_config.valid_file is not None:
+            head_config.valid_file = normalize_file_paths(head_config.valid_file)
+        if hasattr(head_config, "test_file") and head_config.test_file is not None:
+            head_config.test_file = normalize_file_paths(head_config.test_file)
+
+        if (
+            head_config.statistics_file is not None
+            and head_config.head_name != "pt_head"
+        ):
+            with open(head_config.statistics_file, "r") as f:  # pylint: disable=W1514
+                statistics = json.load(f)
+            logging.info("Using statistics json file")
+            head_config.atomic_numbers = statistics["atomic_numbers"]
+            head_config.mean = statistics["mean"]
+            head_config.std = statistics["std"]
+            head_config.avg_num_neighbors = statistics["avg_num_neighbors"]
+            head_config.compute_avg_num_neighbors = False
+            if isinstance(statistics["atomic_energies"], str) and statistics[
+                "atomic_energies"
+            ].endswith(".json"):
+                with open(statistics["atomic_energies"], "r", encoding="utf-8") as f:
+                    atomic_energies = json.load(f)
+                head_config.E0s = atomic_energies
+                head_config.atomic_energies_dict = ast.literal_eval(atomic_energies)
+            else:
+                head_config.E0s = statistics["atomic_energies"]
+                head_config.atomic_energies_dict = ast.literal_eval(
+                    statistics["atomic_energies"]
+                )
+        if head_config.train_file == ["mp"]:
+            assert (
+                head_config.head_name == "pt_head"
+            ), "Only pt_head should use mp as train_file"
+            logging.info(
+                "Using the full Materials Project data for replay. You can construct a different subset using `fine_tuning_select.py` script."
+            )
+            collections = assemble_mp_data(args, head_config, tag)
+            head_config.collections = collections
+        elif any(check_path_ase_read(f) for f in head_config.train_file):
+            train_files_ase_list = [
+                f for f in head_config.train_file if check_path_ase_read(f)
+            ]
+            valid_files_ase_list = None
+            test_files_ase_list = None
+            if head_config.valid_file:
+                valid_files_ase_list = [
+                    f for f in head_config.valid_file if check_path_ase_read(f)
+                ]
+            if head_config.test_file:
+                test_files_ase_list = [
+                    f for f in head_config.test_file if check_path_ase_read(f)
+                ]
+            config_type_weights = get_config_type_weights(
+                head_config.config_type_weights
+            )
+            collections, atomic_energies_dict = get_dataset_from_xyz(
+                work_dir=os.path.join(args.output_dir,args.name),
+                train_path=train_files_ase_list,
+                valid_path=valid_files_ase_list,
+                valid_fraction=head_config.valid_fraction,
+                config_type_weights=config_type_weights,
+                test_path=test_files_ase_list,
+                seed=args.seed,
+                key_specification=head_config.key_specification,
+                head_name=head_config.head_name,
+                keep_isolated_atoms=head_config.keep_isolated_atoms,
+            )
+            head_config.collections = SubsetCollection(
+                train=collections.train,
+                valid=collections.valid,
+                tests=collections.tests,
+            )
+            head_config.atomic_energies_dict = atomic_energies_dict
+            logging.info(
+                f"Total number of configurations: train={len(collections.train)}, valid={len(collections.valid)}, "
+                f"tests=[{', '.join([name + ': ' + str(len(test_configs)) for name, test_configs in collections.tests])}],"
+            )
+        head_configs.append(head_config)
+
+    if all(
+        check_path_ase_read(head_config.train_file[0]) for head_config in head_configs
+    ):
+        size_collections_train = sum(
+            len(head_config.collections.train) for head_config in head_configs
+        )
+        size_collections_valid = sum(
+            len(head_config.collections.valid) for head_config in head_configs
+        )
+        if size_collections_train < args.batch_size:
+            logging.error(
+                f"Batch size ({args.batch_size}) is larger than the number of training data ({size_collections_train})"
+            )
+        if size_collections_valid < args.valid_batch_size:
+            logging.warning(
+                f"Validation batch size ({args.valid_batch_size}) is larger than the number of validation data ({size_collections_valid})"
+            )
+
+    if args.multiheads_finetuning:
+        logging.info(
+            "==================Using multiheads finetuning mode=================="
+        )
+        args.loss = "universal"
+
+        all_ase_readable = all(
+            all(check_path_ase_read(f) for f in head_config.train_file)
+            for head_config in head_configs
+        )
+        head_config_pt = filter(lambda x: x.head_name == "pt_head", head_configs)
+        head_config_pt = next(head_config_pt, None)
+        assert head_config_pt is not None, "Pretraining head not found"
+        if all_ase_readable:
+            ratio_pt_ft = size_collections_train / len(head_config_pt.collections.train)
+            if ratio_pt_ft < 0.1:
+                logging.warning(
+                    f"Ratio of the number of configurations in the training set and the in the pt_train_file is {ratio_pt_ft}, "
+                    f"increasing the number of configurations in the fine-tuning heads by {int(0.1 / ratio_pt_ft)}"
+                )
+                for head_config in head_configs:
+                    if head_config.head_name == "pt_head":
+                        continue
+                    head_config.collections.train += (
+                        head_config.collections.train * int(0.1 / ratio_pt_ft)
+                    )
+            logging.info(
+                f"Total number of configurations in pretraining: train={len(head_config_pt.collections.train)}, valid={len(head_config_pt.collections.valid)}"
+            )
+        else:
+            logging.debug(
+                "Using LMDB/HDF5 datasets for pretraining or fine-tuning - skipping ratio check"
+            )
+
+    # Atomic number table
+    # yapf: disable
+    for head_config in head_configs:
+        if head_config.atomic_numbers is None:
+            assert all(check_path_ase_read(f) for f in head_config.train_file), "Must specify atomic_numbers when using .h5 or .aselmdb train_file input"
+            z_table_head = tools.get_atomic_number_table_from_zs(
+                z
+                for configs in (head_config.collections.train, head_config.collections.valid)
+                for config in configs
+                for z in config.atomic_numbers
+            )
+            head_config.atomic_numbers = z_table_head.zs
+            head_config.z_table = z_table_head
+        else:
+            if head_config.statistics_file is None:
+                logging.info("Using atomic numbers from command line argument")
+            else:
+                logging.info("Using atomic numbers from statistics file")
+            zs_list = ast.literal_eval(head_config.atomic_numbers)
+            assert isinstance(zs_list, list)
+            z_table_head = tools.AtomicNumberTable(zs_list)
+            head_config.atomic_numbers = zs_list
+            head_config.z_table = z_table_head
+        # yapf: enable
+    all_atomic_numbers = set()
+    for head_config in head_configs:
+        all_atomic_numbers.update(head_config.atomic_numbers)
+    z_table = AtomicNumberTable(sorted(list(all_atomic_numbers)))
+    if args.foundation_model_elements and model_foundation:
+        z_table = AtomicNumberTable(sorted(model_foundation.atomic_numbers.tolist()))
+    logging.info(f"Atomic Numbers used: {z_table.zs}")
+
+    # Atomic energies
+    atomic_energies_dict = {}
+    for head_config in head_configs:
+        if head_config.atomic_energies_dict is None or len(head_config.atomic_energies_dict) == 0:
+            assert head_config.E0s is not None, "Atomic energies must be provided"
+            if all(check_path_ase_read(f) for f in head_config.train_file) and head_config.E0s.lower() != "foundation":
+                atomic_energies_dict[head_config.head_name] = get_atomic_energies(
+                    head_config.E0s, head_config.collections.train, head_config.z_table
+                )
+            elif head_config.E0s.lower() == "foundation":
+                assert args.foundation_model is not None
+                z_table_foundation = AtomicNumberTable(
+                    [int(z) for z in model_foundation.atomic_numbers]
+                )
+                foundation_atomic_energies = model_foundation.atomic_energies_fn.atomic_energies
+                if foundation_atomic_energies.ndim > 1:
+                    foundation_atomic_energies = foundation_atomic_energies.squeeze()
+                    if foundation_atomic_energies.ndim == 2:
+                        foundation_atomic_energies = foundation_atomic_energies[0]
+                        logging.info("Foundation model has multiple heads, using the first head as foundation E0s.")
+                atomic_energies_dict[head_config.head_name] = {
+                    z: foundation_atomic_energies[
+                        z_table_foundation.z_to_index(z)
+                    ].item()
+                    for z in z_table.zs
+                }
+            else:
+                atomic_energies_dict[head_config.head_name] = get_atomic_energies(head_config.E0s, None, head_config.z_table)
+        else:
+            atomic_energies_dict[head_config.head_name] = head_config.atomic_energies_dict
+
+    # Atomic energies for multiheads finetuning
+    if args.multiheads_finetuning:
+        assert (
+            model_foundation is not None
+        ), "Model foundation must be provided for multiheads finetuning"
+        z_table_foundation = AtomicNumberTable(
+            [int(z) for z in model_foundation.atomic_numbers]
+        )
+        foundation_atomic_energies = model_foundation.atomic_energies_fn.atomic_energies
+        if foundation_atomic_energies.ndim > 1:
+            foundation_atomic_energies = foundation_atomic_energies.squeeze()
+            if foundation_atomic_energies.ndim == 2:
+                foundation_atomic_energies = foundation_atomic_energies[0]
+                logging.info("Foundation model has multiple heads, using the first head as foundation E0s.")
+        atomic_energies_dict["pt_head"] = {
+            z: foundation_atomic_energies[
+                z_table_foundation.z_to_index(z)
+            ].item()
+            for z in z_table.zs
+        }
+    heads = sorted(heads, key=lambda x: -1000 if x == "pt_head" else 0)
+    # Padding atomic energies if keeping all elements of the foundation model
+    if args.foundation_model_elements and model_foundation:
+        atomic_energies_dict_padded = {}
+        for head_name, head_energies in atomic_energies_dict.items():
+            energy_head_padded = {}
+            for z in z_table.zs:
+                energy_head_padded[z] = head_energies.get(z, 0.0)
+            atomic_energies_dict_padded[head_name] = energy_head_padded
+        atomic_energies_dict = atomic_energies_dict_padded
+
+    if args.model == "AtomicDipolesMACE":
+        atomic_energies = None
+        dipole_only = True
+        args.compute_dipole = True
+        args.compute_energy = False
+        args.compute_forces = False
+        args.compute_virials = False
+        args.compute_stress = False
+    else:
+        dipole_only = False
+        if args.model == "EnergyDipolesMACE":
+            args.compute_dipole = True
+            args.compute_energy = True
+            args.compute_forces = True
+            args.compute_virials = False
+            args.compute_stress = False
+        else:
+            args.compute_energy = True
+            args.compute_dipole = False
+        # atomic_energies: np.ndarray = np.array(
+        #     [atomic_energies_dict[z] for z in z_table.zs]
+        # )
+        atomic_energies = dict_to_array(atomic_energies_dict, heads)
+        for head_config in head_configs:
+            try:
+                logging.info(f"Atomic Energies used (z: eV) for head {head_config.head_name}: " + "{" + ", ".join([f"{z}: {atomic_energies_dict[head_config.head_name][z]}" for z in head_config.z_table.zs]) + "}")
+            except KeyError as e:
+                raise KeyError(f"Atomic number {e} not found in atomic_energies_dict for head {head_config.head_name}, add E0s for this atomic number") from e
+
+    # Load datasets for each head, supporting multiple files per head
+    valid_sets = {head: [] for head in heads}
+    train_sets = {head: [] for head in heads}
+
+    for head_config in head_configs:
+        train_datasets = []
+
+        logging.info(f"Processing datasets for head '{head_config.head_name}'")
+        ase_files = [f for f in head_config.train_file if check_path_ase_read(f)]
+        non_ase_files = [f for f in head_config.train_file if not check_path_ase_read(f)]
+
+        if ase_files:
+            #dataset = load_dataset_for_path(
+            #file_path=ase_files,
+            #r_max=args.r_max,
+            #z_table=z_table,
+            #head_config=head_config,
+            #heads=heads,
+            #collection=head_config.collections.train,
+            #)
+            # ✨ 使用新的 TextDataset 包装内存中已划分好的数据 (collections.train)
+            logging.info(f"Wraping training data into TextDataset for head '{head_config.head_name}'")
+            dataset = TextDataset(
+                configurations=head_config.collections.train,  # 直接传入预处理好的 Configuration 列表
+                r_max=args.r_max,
+                z_table=z_table,
+                head=head_config.head_name,
+                heads=heads
+            )
+            train_datasets.append(dataset)
+            logging.debug(f"Successfully loaded dataset from ASE files: {ase_files}")
+
+        for file in non_ase_files:
+            dataset = load_dataset_for_path(
+            file_path=file,
+            r_max=args.r_max,
+            z_table=z_table,
+            head_config=head_config,
+            heads=heads,
+            )
+            train_datasets.append(dataset)
+            logging.debug(f"Successfully loaded dataset from non-ASE file: {file}")
+
+        if not train_datasets:
+            raise ValueError(f"No valid training datasets found for head {head_config.head_name}")
+
+        train_sets[head_config.head_name] = combine_datasets(train_datasets, head_config.head_name)
+
+        if head_config.valid_file:
+            valid_datasets = []
+
+            valid_ase_files = [f for f in head_config.valid_file if check_path_ase_read(f)]
+            valid_non_ase_files = [f for f in head_config.valid_file if not check_path_ase_read(f)]
+
+            if valid_ase_files:
+                #valid_dataset = load_dataset_for_path(
+                #    file_path=valid_ase_files,
+                #    r_max=args.r_max,
+                #    z_table=z_table,
+                #    head_config=head_config,
+                #    heads=heads,
+                #    collection=head_config.collections.valid,
+                #)
+                #valid_datasets.append(valid_dataset)
+                #logging.debug(f"Successfully loaded validation dataset from ASE files: {valid_ase_files}")
+                # ✨ 使用新的 TextDataset 包装内存中已划分好的验证数据
+                logging.info(f"Wraping validation data into TextDataset for head '{head_config.head_name}'")
+                valid_dataset = TextDataset(
+                    configurations=head_config.collections.valid, # 直接传入预处理好的 Configuration 列表
+                    r_max=args.r_max,
+                    z_table=z_table,
+                    head=head_config.head_name,
+                    heads=heads
+                )
+                valid_datasets.append(valid_dataset)
+                logging.debug(f"Successfully loaded validation dataset from ASE files: {valid_ase_files}")
+            for valid_file in valid_non_ase_files:
+                valid_dataset = load_dataset_for_path(
+                file_path=valid_file,
+                r_max=args.r_max,
+                z_table=z_table,
+                head_config=head_config,
+                heads=heads,
+            )
+                valid_datasets.append(valid_dataset)
+                logging.debug(f"Successfully loaded validation dataset from {valid_file}")
+
+            # Combine validation datasets
+            if valid_datasets:
+                valid_sets[head_config.head_name] = combine_datasets(valid_datasets, f"{head_config.head_name}_valid")
+                logging.info(f"Combined validation datasets for {head_config.head_name}")
+
+        # If no valid file is provided but collection exist, use the validation set from the collection
+        if head_config.valid_file is None and head_config.collections.valid:
+            valid_sets[head_config.head_name] = [
+                AtomicData.from_config(
+                    config, z_table=z_table, cutoff=args.r_max, heads=heads
+                )
+                for config in head_config.collections.valid
+            ]
+        if not valid_sets[head_config.head_name]:
+            raise ValueError(f"No valid datasets found for head {head_config.head_name}, please provide a valid_file or a valid_fraction")
+
+        # Create data loader for this head
+        if isinstance(train_sets[head_config.head_name], list):
+            dataset_size = len(train_sets[head_config.head_name])
+        else:
+            dataset_size = len(train_sets[head_config.head_name])
+        logging.info(f"Head '{head_config.head_name}' training dataset size: {dataset_size}")
+
+        train_sampler_head = None
+        if args.distributed:
+            train_sampler_head = torch.utils.data.distributed.DistributedSampler(
+                train_sets[head_config.head_name],
+                num_replicas=world_size,
+                rank=rank,
+                shuffle=True,
+                drop_last=(not args.lbfgs),
+                seed=args.seed,
+            )
+        train_loader_head = torch_geometric.dataloader.DataLoader(
+            dataset=train_sets[head_config.head_name],
+            batch_size=args.batch_size,
+            sampler=train_sampler_head,
+            shuffle=(train_sampler_head is None),
+            drop_last=(train_sampler_head is None and not args.lbfgs),
+            pin_memory=args.pin_memory,
+            num_workers=args.num_workers,
+            generator=torch.Generator().manual_seed(args.seed),
+        )
+        head_config.train_loader = train_loader_head
+
+    # concatenate all the trainsets
+    train_set = ConcatDataset([train_sets[head] for head in heads])
+    train_sampler, valid_sampler = None, None
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(
+            train_set,
+            num_replicas=world_size,
+            rank=rank,
+            shuffle=True,
+            drop_last=(not args.lbfgs),
+            seed=args.seed,
+        )
+        valid_samplers = {}
+        for head, valid_set in valid_sets.items():
+            valid_sampler = torch.utils.data.distributed.DistributedSampler(
+                valid_set,
+                num_replicas=world_size,
+                rank=rank,
+                shuffle=True,
+                drop_last=True,
+                seed=args.seed,
+            )
+            valid_samplers[head] = valid_sampler
+    train_loader = torch_geometric.dataloader.DataLoader(
+        dataset=train_set,
+        batch_size=args.batch_size,
+        sampler=train_sampler,
+        shuffle=(train_sampler is None),
+        drop_last=(train_sampler is None and not args.lbfgs),
+        pin_memory=args.pin_memory,
+        num_workers=args.num_workers,
+        generator=torch.Generator().manual_seed(args.seed),
+    )
+    valid_loaders = {heads[i]: None for i in range(len(heads))}
+    if not isinstance(valid_sets, dict):
+        valid_sets = {"Default": valid_sets}
+    for head, valid_set in valid_sets.items():
+        valid_loaders[head] = torch_geometric.dataloader.DataLoader(
+            dataset=valid_set,
+            batch_size=args.valid_batch_size,
+            sampler=valid_samplers[head] if args.distributed else None,
+            shuffle=False,
+            drop_last=False,
+            pin_memory=args.pin_memory,
+            num_workers=args.num_workers,
+            generator=torch.Generator().manual_seed(args.seed),
+        )
+
+    loss_fn = get_loss_fn(args, dipole_only, args.compute_dipole)
+    args.avg_num_neighbors = get_avg_num_neighbors(head_configs, args, train_loader, device)
+
+    # Model
+    model, output_args = configure_model(args, train_loader, atomic_energies, model_foundation, heads, z_table, head_configs)
+    model.to(device)
+
+    logging.debug(model)
+    logging.info(f"Total number of parameters: {tools.count_parameters(model)}")
+    logging.info("")
+    logging.info("===========OPTIMIZER INFORMATION===========")
+    logging.info(f"Using {args.optimizer.upper()} as parameter optimizer")
+    logging.info(f"Batch size: {args.batch_size}")
+    if args.ema:
+        logging.info(f"Using Exponential Moving Average with decay: {args.ema_decay}")
+    logging.info(
+        f"Number of gradient updates: {int(args.max_num_epochs*len(train_set)/args.batch_size)}"
+    )
+    logging.info(f"Learning rate: {args.lr}, weight decay: {args.weight_decay}")
+    logging.info(loss_fn)
+
+    # Cueq
+    if args.enable_cueq:
+        logging.info("Converting model to CUEQ for accelerated training")
+        assert model.__class__.__name__ in ["MACE", "ScaleShiftMACE"]
+        model = run_e3nn_to_cueq(deepcopy(model), device=device)
+    # Optimizer
+    param_options = get_params_options(args, model)
+    optimizer: torch.optim.Optimizer
+    optimizer = get_optimizer(args, param_options)
+    if args.device == "xpu":
+        logging.info("Optimzing model and optimzier for XPU")
+        model, optimizer = ipex.optimize(model, optimizer=optimizer)
+    logger = tools.MetricsLogger(
+        directory=args.plot_dir, tag=tag + "_train"
+    )  # pylint: disable=E1123
+
+    lr_scheduler = LRScheduler(optimizer, args)
+
+    swa: Optional[tools.SWAContainer] = None
+    swas = [False]
+    if args.swa:
+        swa, swas = get_swa(args, model, optimizer, swas, dipole_only)
+
+    checkpoint_handler = tools.CheckpointHandler(
+        directory=args.checkpoints_dir,
+        tag=tag,
+        keep=args.keep_checkpoints,
+        swa_start=args.start_swa,
+    )
+
+    start_epoch = 0
+    restart_lbfgs = False
+    opt_start_epoch = None
+    if args.restart_latest:
+        try:
+            opt_start_epoch = checkpoint_handler.load_latest(
+                state=tools.CheckpointState(model, optimizer, lr_scheduler),
+                swa=True,
+                device=device,
+            )
+        except Exception:  # pylint: disable=W0703
+            try:
+                opt_start_epoch = checkpoint_handler.load_latest(
+                    state=tools.CheckpointState(model, optimizer, lr_scheduler),
+                    swa=False,
+                    device=device,
+                )
+            except Exception: # pylint: disable=W0703
+                restart_lbfgs = True
+        if opt_start_epoch is not None:
+            start_epoch = opt_start_epoch
+
+    ema: Optional[ExponentialMovingAverage] = None
+    if args.ema:
+        ema = ExponentialMovingAverage(model.parameters(), decay=args.ema_decay)
+    else:
+        for group in optimizer.param_groups:
+            group["lr"] = args.lr
+
+    if args.lbfgs:
+        logging.info("Switching optimizer to LBFGS")
+        optimizer = LBFGS(model.parameters(),
+                          history_size=200,
+                          max_iter=20,
+                          line_search_fn="strong_wolfe")
+        if restart_lbfgs:
+            opt_start_epoch = checkpoint_handler.load_latest(
+                state=tools.CheckpointState(model, optimizer, lr_scheduler),
+                swa=False,
+                device=device,
+            )
+            if opt_start_epoch is not None:
+                start_epoch = opt_start_epoch
+
+    if args.wandb:
+        setup_wandb(args)
+    if args.distributed:
+        distributed_model = DDP(model, device_ids=[local_rank])
+    else:
+        distributed_model = None
+
+
+    train_valid_data_loader = {}
+    for head_config in head_configs:
+        data_loader_name = "train_" + head_config.head_name
+        train_valid_data_loader[data_loader_name] = head_config.train_loader
+    for head, valid_loader in valid_loaders.items():
+        data_load_name = "valid_" + head
+        train_valid_data_loader[data_load_name] = valid_loader
+
+    if args.plot and args.plot_frequency > 0:
+        try:
+            plotter = TrainingPlotter(
+                results_dir=logger.path,
+                heads=heads,
+                table_type=args.error_table,
+                train_valid_data=train_valid_data_loader,
+                test_data={},
+                output_args=output_args,
+                device=device,
+                plot_frequency=args.plot_frequency,
+                distributed=args.distributed,
+                swa_start=swa.start if swa else None
+                )
+        except Exception as e:  # pylint: disable=W0718
+            logging.debug(f"Creating Plotter failed: {e}")
+    else:
+        plotter = None
+
+    if args.dry_run:
+        logging.info("DRY RUN mode enabled. Stopping now.")
+        return
+
+
+    tools.train(
+        model=model,
+        loss_fn=loss_fn,
+        train_loader=train_loader,
+        valid_loaders=valid_loaders,
+        optimizer=optimizer,
+        lr_scheduler=lr_scheduler,
+        checkpoint_handler=checkpoint_handler,
+        eval_interval=args.eval_interval,
+        start_epoch=start_epoch,
+        max_num_epochs=args.max_num_epochs,
+        logger=logger,
+        patience=args.patience,
+        save_all_checkpoints=args.save_all_checkpoints,
+        output_args=output_args,
+        device=device,
+        swa=swa,
+        ema=ema,
+        max_grad_norm=args.clip_grad,
+        log_errors=args.error_table,
+        log_wandb=args.wandb,
+        distributed=args.distributed,
+        distributed_model=distributed_model,
+        plotter=plotter,
+        train_sampler=train_sampler,
+        rank=rank,
+    )
+
+    logging.info("")
+    logging.info("===========RESULTS===========")
+
+    train_valid_data_loader = {}
+    for head_config in head_configs:
+        data_loader_name = "train_" + head_config.head_name
+        train_valid_data_loader[data_loader_name] = head_config.train_loader
+    for head, valid_loader in valid_loaders.items():
+        data_load_name = "valid_" + head
+        train_valid_data_loader[data_load_name] = valid_loader
+    test_sets = {}
+    stop_first_test = False
+    test_data_loader = {}
+    if all(
+        head_config.test_file == head_configs[0].test_file
+        for head_config in head_configs
+    ) and head_configs[0].test_file is not None:
+        stop_first_test = True
+    if all(
+        head_config.test_dir == head_configs[0].test_dir
+        for head_config in head_configs
+    ) and head_configs[0].test_dir is not None:
+        stop_first_test = True
+    for head_config in head_configs:
+        if all(check_path_ase_read(f) for f in head_config.train_file):
+            for name, subset in head_config.collections.tests:
+                test_sets[name] = [
+                    AtomicData.from_config(
+                        config, z_table=z_table, cutoff=args.r_max, heads=heads
+                    )
+                    for config in subset
+                ]
+        if head_config.test_dir is not None:
+            if not args.multi_processed_test:
+                test_files = get_files_with_suffix(head_config.test_dir, "_test.h5")
+                for test_file in test_files:
+                    name = os.path.splitext(os.path.basename(test_file))[0]
+                    test_sets[name] = HDF5Dataset(
+                        test_file, r_max=args.r_max, z_table=z_table, heads=heads, head=head_config.head_name
+                    )
+            else:
+                test_folders = glob(head_config.test_dir + "/*")
+                for folder in test_folders:
+                    name = os.path.splitext(os.path.basename(test_file))[0]
+                    test_sets[name] = dataset_from_sharded_hdf5(
+                        folder, r_max=args.r_max, z_table=z_table, heads=heads, head=head_config.head_name
+                    )
+        for test_name, test_set in test_sets.items():
+            test_sampler = None
+            if args.distributed:
+                test_sampler = torch.utils.data.distributed.DistributedSampler(
+                    test_set,
+                    num_replicas=world_size,
+                    rank=rank,
+                    shuffle=True,
+                    drop_last=True,
+                    seed=args.seed,
+                )
+            try:
+                drop_last = test_set.drop_last
+            except AttributeError as e:  # pylint: disable=W0612
+                drop_last = False
+            test_loader = torch_geometric.dataloader.DataLoader(
+                test_set,
+                batch_size=args.valid_batch_size,
+                sampler=test_sampler,
+                shuffle=(test_sampler is None),
+                drop_last=drop_last,
+                num_workers=args.num_workers,
+                pin_memory=args.pin_memory,
+            )
+            test_data_loader[test_name] = test_loader
+        if stop_first_test:
+            break
+
+    for swa_eval in swas:
+        epoch = checkpoint_handler.load_latest(
+            state=tools.CheckpointState(model, optimizer, lr_scheduler),
+            swa=swa_eval,
+            device=device,
+        )
+        model.to(device)
+        if args.distributed:
+            distributed_model = DDP(model, device_ids=[local_rank], broadcast_buffers=False)
+        model_to_evaluate = model if not args.distributed else distributed_model
+        if swa_eval:
+            logging.info(f"Loaded Stage two model from epoch {epoch} for evaluation")
+        else:
+            logging.info(f"Loaded Stage one model from epoch {epoch} for evaluation")
+
+        if rank == 0:
+            # Save entire model
+            if swa_eval:
+                model_path = Path(args.checkpoints_dir) / (tag + "_stagetwo.model")
+            else:
+                model_path = Path(args.checkpoints_dir) / (tag + ".model")
+            logging.info(f"Saving model to {model_path}")
+            model_to_save = deepcopy(model)
+            if args.enable_cueq:
+                print("RUNING CUEQ TO E3NN")
+                print("swa_eval", swa_eval)
+                model_to_save = run_cueq_to_e3nn(deepcopy(model), device=device)
+            if args.save_cpu:
+                model_to_save = model_to_save.to("cpu")
+            torch.save(model_to_save, model_path)
+            extra_files = {
+                "commit.txt": commit.encode("utf-8") if commit is not None else b"",
+                "config.yaml": json.dumps(
+                    convert_to_json_format(extract_config_mace_model(model))
+                ),
+            }
+            if swa_eval:
+                torch.save(
+                    model_to_save, Path(args.model_dir) / (args.name + "_stagetwo.model")
+                )
+                try:
+                    path_complied = Path(args.model_dir) / (
+                        args.name + "_stagetwo_compiled.model"
+                    )
+                    logging.info(f"Compiling model, saving metadata {path_complied}")
+                    model_compiled = jit.compile(deepcopy(model_to_save))
+                    torch.jit.save(
+                        model_compiled,
+                        path_complied,
+                        _extra_files=extra_files,
+                    )
+                except Exception as e:  # pylint: disable=W0718
+                    pass
+            else:
+                torch.save(model_to_save, Path(args.model_dir) / (args.name + ".model"))
+                try:
+                    path_complied = Path(args.model_dir) / (
+                        args.name + "_compiled.model"
+                    )
+                    logging.info(f"Compiling model, saving metadata to {path_complied}")
+                    model_compiled = jit.compile(deepcopy(model_to_save))
+                    torch.jit.save(
+                        model_compiled,
+                        path_complied,
+                        _extra_files=extra_files,
+                    )
+                except Exception as e:  # pylint: disable=W0718
+                    pass
+
+        logging.info("Computing metrics for training, validation, and test sets")
+        for param in model.parameters():
+            param.requires_grad = False
+        skip_heads = args.skip_evaluate_heads.split(",") if args.skip_evaluate_heads else []
+        if skip_heads:
+            logging.info(f"Skipping evaluation for heads: {skip_heads}")
+        table_train_valid = create_error_table(
+            table_type=args.error_table,
+            all_data_loaders=train_valid_data_loader,
+            model=model_to_evaluate,
+            loss_fn=loss_fn,
+            output_args=output_args,
+            log_wandb=args.wandb,
+            device=device,
+            distributed=args.distributed,
+            skip_heads=skip_heads,
+        )
+        logging.info("Error-table on TRAIN and VALID:\n" + str(table_train_valid))
+
+        if test_data_loader:
+            table_test = create_error_table(
+                table_type=args.error_table,
+                all_data_loaders=test_data_loader,
+                model=model_to_evaluate,
+                loss_fn=loss_fn,
+                output_args=output_args,
+                log_wandb=args.wandb,
+                device=device,
+                distributed=args.distributed,
+            )
+            logging.info("Error-table on TEST:\n" + str(table_test))
+        if args.plot:
+            try:
+                plotter = TrainingPlotter(
+                    results_dir=logger.path,
+                    heads=heads,
+                    table_type=args.error_table,
+                    train_valid_data=train_valid_data_loader,
+                    test_data=test_data_loader,
+                    output_args=output_args,
+                    device=device,
+                    plot_frequency=args.plot_frequency,
+                    distributed=args.distributed,
+                    swa_start=swa.start if swa else None
+                )
+                plotter.plot(epoch, model_to_evaluate, rank)
+            except Exception as e:  # pylint: disable=W0718
+                logging.debug(f"Plotting failed: {e}")
+
+        if args.distributed:
+            torch.distributed.barrier()
+
+    logging.info("Done")
+    if args.distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()