| # ============================================================ | |
| # 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 核心数 | |