# MACE Training Demo

## Quick Start

### 1. Select a Configuration

The `configs/` directory provides several predefined experiment configurations:

| Configuration | Dataset | GPUs | Description |
| --- | --- | --- | --- |
| `DMC.yaml` | DMC solvent XTB | 1 | Simplest introductory example |
| `water_1dcu.yaml` | Water | 1 | Single-GPU training |
| `water_4dcu.yaml` | Water | 4 | Four-GPU distributed training |
| `water_8dcu.yaml` | Water | 8 | Eight-GPU distributed training |
| `ani1x_8dcu.yaml` | ANI-1x | 8 | Distributed training |
| `nanotube_l0_8dcu.yaml` | Carbon nanotube | 8 | `max_L=0` |
| `nanotube_l2_8dcu.yaml` | Carbon nanotube | 8 | `max_L=2` |
| `nanotube_l2_16dcu.yaml` | Carbon nanotube | 2x8 | Multi-node distributed training |

### 2. Run Training

```bash
# Option 1: Run directly (interactively or on an allocated SLURM node)
bash run.sh --config configs/DMC.yaml

# Option 2: Submit a SLURM job
bash run.sh --config configs/DMC.yaml --submit

# Option 3: Preview the command without running it
bash run.sh --config configs/DMC.yaml --dry-run
```

### 3. View the Outputs

Training outputs are automatically saved to `outputs/{experiment_name}_{timestamp}/` and include:

- Model checkpoints
- Training logs
- A snapshot of the configuration used for the run (`config.yaml`)

## Create a Custom Experiment

1. Copy the closest configuration:

   ```bash
   cp configs/DMC.yaml configs/my_experiment.yaml
   ```

2. Edit the parameters in the YAML file. Only parameter values need to change; no shell scripts need to be modified.

3. Run the experiment:

   ```bash
   bash run.sh --config configs/my_experiment.yaml
   ```

## YAML Configuration Fields

### `train_args` - Training Arguments

Every field maps directly to a `train.py` command-line argument. A Boolean value of `true` becomes a flag (for example, `swa: true` becomes `--swa`), while `false` is omitted.

Common arguments:

| Argument | Description | Example |
| --- | --- | --- |
| `model` | Model type | `MACE` |
| `r_max` | Cutoff radius (Å) | `4.0` - `6.0` |
| `num_channels` | Number of channels | `64`, `256` |
| `max_L` | Maximum angular-momentum quantum number | `0`, `2` |
| `batch_size` | Training batch size | `2` - `128` |
| `E0s` | Atomic reference energies | `average`, `isolated`, or an explicit dictionary |
| `swa` | Enable stochastic weight averaging | `true` |
| `ema` | Enable exponential moving average | `true` |
| `distributed` | Enable distributed training | `true` (added automatically for multiple GPUs) |

### `launch` - Launch Configuration

| Argument | Description | Launch method |
| --- | --- | --- |
| `num_nodes: 1, num_gpus: 1` | Single GPU | `python train.py` |
| `num_nodes: 1, num_gpus: N` | Multiple GPUs on one node | `torchrun --nproc_per_node=N` |
| `num_nodes: M, num_gpus: N` | Multiple nodes | `srun` (requires `--submit`) |

### `env` - Environment Configuration

| Argument | Description |
| --- | --- |
| `conda_env` | Conda environment name |
| `modules` | List of modules to load |

### `slurm` - SLURM Job Configuration

| Argument | Description |
| --- | --- |
| `partition` | SLURM partition |
| `time` | Job time limit |
| `cpus_per_task` | Number of CPU cores |

### `nccl` - Multi-Node Communication Configuration (Optional)

| Argument | Description |
| --- | --- |
| `socket_ifname` | InfiniBand interface name |
| `ib_hca` | IB HCA device name |
| `proto` | NCCL protocol |

## Directory Structure

```text
demo/
  run.sh                  # Unified entry-point script
  _parse_config.py        # Configuration parser (internal use)
  README.md               # This file
  configs/                # Experiment configurations
  templates/              # Script templates
    env_setup.sh          # Environment initialization
    preflight_check.sh    # Pre-training checks
    slurm_header.template # SLURM header template
  outputs/                # Training outputs (created automatically)
```