scripts/demo/README.md

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. 运行训练

# 方式一：直接运行（交互式或在已分配的 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. 复制一个最接近的配置文件：

   cp configs/DMC.yaml configs/my_experiment.yaml
   

2. 编辑 YAML 中的参数（只需改参数值，不用碰任何 shell 脚本）

3. 运行：

   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/                # 训练输出（自动创建）