File size: 2,940 Bytes
f15d29e | 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 | #!/bin/bash
# MatterGen Slurm training entry. Users may edit the defaults below directly.
# The recommended interface is: cd demo && bash run.sh --config configs/train_8dcu.yaml --submit
# Direct interface: cd examples/matchem/mattergen && sbatch submit_train.sh
#SBATCH --job-name=mattergen_train
#SBATCH --partition=hx1hdexclu12
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gres=dcu:8
#SBATCH --cpus-per-task=64
#SBATCH --time=48:00:00
#SBATCH --output=slurm_%j.out
#SBATCH --error=slurm_%j.err
set -euo pipefail
SCRIPT_DIR="${SCRIPT_DIR:-${SLURM_SUBMIT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}}"
if [[ ! -f "$SCRIPT_DIR/train.py" ]]; then
echo "ERROR: MatterGen root not found: $SCRIPT_DIR" >&2
echo "Run this script from examples/matchem/mattergen or use demo/run.sh." >&2
exit 1
fi
NUM_DEVICES="${NUM_DEVICES:-8}"
TRAIN_BATCH_SIZE="${TRAIN_BATCH_SIZE:-4}"
VAL_BATCH_SIZE="${VAL_BATCH_SIZE:-4}"
ACCUMULATE_GRAD_BATCHES="${ACCUMULATE_GRAD_BATCHES:-16}"
SMOKE_TEST="${SMOKE_TEST:-0}"
TRAIN_ARGS=(
data_module=mp_20
"trainer.devices=$NUM_DEVICES"
trainer.num_nodes=1
"trainer.accumulate_grad_batches=$ACCUMULATE_GRAD_BATCHES"
"data_module.batch_size.train=$TRAIN_BATCH_SIZE"
"data_module.batch_size.val=$VAL_BATCH_SIZE"
"data_module.batch_size.test=$VAL_BATCH_SIZE"
data_module.num_workers.train=2
data_module.num_workers.val=2
data_module.num_workers.test=2
'~trainer.logger'
)
if [[ -n "${CONFIG_PATH:-}" ]]; then
mapfile -t TRAIN_ARGS < <(python3 - "$CONFIG_PATH" <<'PY'
import os
import sys
import yaml
config = yaml.safe_load(open(sys.argv[1]))
if config.get("task", "train") != "train":
raise SystemExit("submit_train.sh only accepts task: train")
for key, value in config.get("args", {}).items():
if isinstance(value, bool):
if value:
print(key)
elif value is not None:
print(f"{key}={os.path.expandvars(str(value))}")
PY
)
fi
if [[ "$SMOKE_TEST" == "1" ]]; then
TRAIN_ARGS+=(
trainer.max_epochs=1
trainer.check_val_every_n_epoch=1
+trainer.limit_train_batches=1
+trainer.limit_val_batches=1
+trainer.num_sanity_val_steps=0
)
fi
for arg in "${TRAIN_ARGS[@]}"; do
case "$arg" in
trainer.devices=*) NUM_DEVICES="${arg#*=}" ;;
data_module.batch_size.train=*) TRAIN_BATCH_SIZE="${arg#*=}" ;;
data_module.batch_size.val=*) VAL_BATCH_SIZE="${arg#*=}" ;;
trainer.accumulate_grad_batches=*) ACCUMULATE_GRAD_BATCHES="${arg#*=}" ;;
esac
done
echo "MatterGen DDP training"
echo " devices: $NUM_DEVICES"
echo " train batch per device: $TRAIN_BATCH_SIZE"
echo " validation batch per device: $VAL_BATCH_SIZE"
echo " gradient accumulation: $ACCUMULATE_GRAD_BATCHES"
echo " effective global batch: $((NUM_DEVICES * TRAIN_BATCH_SIZE * ACCUMULATE_GRAD_BATCHES))"
cd "$SCRIPT_DIR"
srun python train.py "${TRAIN_ARGS[@]}"
|