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