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