File size: 3,947 Bytes
0bda409 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 | # 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)
```
|