NequIP / demo /configs /tutorial_fcu_8dcu.yaml
dadadaxi's picture
Upload folder using huggingface_hub
3e02ab8 verified
Raw History Blame Contribute Delete
4.11 kB
# Eight-DCU DDP training plan for the official NequIP fcu.xyz tutorial.
# DataLoader batch size is per rank: batch_size=5 gives a global batch of 40.
# The learning rate remains conservative at the official 0.01; tune it for
# production based on validation behavior rather than scaling it blindly.
run: [train, test]
cutoff_radius: 5.0
num_layers: 4
l_max: 1
num_features: 32
model_type_names: [C, H, O, Cu]
chemical_species: ${model_type_names}
monitored_metric: val0_epoch/weighted_sum
data:
_target_: onescience.datapipes.materials.nequip.datamodule.ASEDataModule
seed: 456
split_dataset:
file_path: ${oc.env:ONESCIENCE_DATASETS_DIR}/matchem/NequIP/fcu.xyz
train: 0.8
val: 0.1
test: 0.1
transforms:
- _target_: onescience.datapipes.materials.nequip.transforms.ChemicalSpeciesToAtomTypeMapper
model_type_names: ${model_type_names}
- _target_: onescience.datapipes.materials.nequip.transforms.NeighborListTransform
r_max: ${cutoff_radius}
train_dataloader:
_target_: torch.utils.data.DataLoader
batch_size: 5
num_workers: 0
shuffle: true
val_dataloader:
_target_: torch.utils.data.DataLoader
batch_size: 10
num_workers: 0
test_dataloader: ${data.val_dataloader}
stats_manager:
_target_: onescience.datapipes.materials.nequip.CommonDataStatisticsManager
dataloader_kwargs:
batch_size: 10
type_names: ${model_type_names}
trainer:
_target_: lightning.Trainer
accelerator: gpu
devices: 8
num_nodes: 1
strategy:
_target_: lightning.pytorch.strategies.DDPStrategy
enable_checkpointing: true
max_epochs: 1000
max_time: "03:00:00:00"
log_every_n_steps: 20
logger:
_target_: lightning.pytorch.loggers.CSVLogger
save_dir: ${hydra:runtime.output_dir}
name: metrics
version: 0
enable_progress_bar: true
callbacks:
- _target_: onescience.utils.nequip.train.callbacks.PlainTextMetricsLogger
- _target_: lightning.pytorch.callbacks.EarlyStopping
monitor: ${monitored_metric}
min_delta: 1e-3
patience: 20
- _target_: lightning.pytorch.callbacks.ModelCheckpoint
monitor: ${monitored_metric}
dirpath: ${hydra:runtime.output_dir}/checkpoints
filename: best
save_last: true
training_module:
_target_: onescience.utils.nequip.train.EMALightningModule
ema_decay: 0.999
loss:
_target_: onescience.utils.nequip.train.EnergyForceLoss
per_atom_energy: true
coeffs:
total_energy: 1.0
forces: 1.0
val_metrics:
_target_: onescience.utils.nequip.train.EnergyForceMetrics
coeffs:
total_energy_mae: 1.0
forces_mae: 1.0
train_metrics: ${training_module.val_metrics}
test_metrics: ${training_module.val_metrics}
optimizer:
_target_: torch.optim.Adam
lr: 0.01
lr_scheduler:
scheduler:
_target_: torch.optim.lr_scheduler.ReduceLROnPlateau
factor: 0.6
patience: 5
threshold: 0.2
min_lr: 1.0e-6
monitor: ${monitored_metric}
interval: epoch
frequency: 1
model:
_target_: onescience.models.nequip.model.NequIPGNNModel
compile_mode: eager
seed: 456
model_dtype: float32
type_names: ${model_type_names}
r_max: ${cutoff_radius}
num_bessels: 8
bessel_trainable: false
polynomial_cutoff_p: 6
num_layers: ${num_layers}
l_max: ${l_max}
parity: true
num_features: ${num_features}
radial_mlp_depth: 2
radial_mlp_width: 64
avg_num_neighbors: ${training_data_stats:num_neighbors_mean}
per_type_energy_scales: ${training_data_stats:per_type_forces_rms}
per_type_energy_shifts: ${training_data_stats:per_atom_energy_mean}
per_type_energy_scales_trainable: false
per_type_energy_shifts_trainable: false
pair_potential:
_target_: onescience.models.nequip.nn.pair_potential.ZBL
units: metal
chemical_species: ${chemical_species}
name: nequip_fcu_tutorial_8dcu
launch:
mode: auto
num_nodes: 1
num_gpus: 8
slurm:
partition: hx1hdnormal01
nodelist: ""
time: "3-00:00:00"
cpus_per_task: 8
env:
OMP_NUM_THREADS: 8
NCCL_DEBUG: "WARN"