scripts/demo/README.md

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

# 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:

   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 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

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)