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