Training Techniques

Extension packages can implement custom training techniques through two main approaches.

Callbacks

Use PyTorch Lightning callbacks for smaller training customizations that don’t alter the core training loop.

As an example, loss coefficients control how different loss terms (energy, forces, stress) are weighted during training - you might want to emphasize forces more heavily in later training epochs. Example callbacks for this featurej include LossCoefficientScheduler and LinearLossCoefficientScheduler.

Training Modules

Training modules refer to PyTorch Lightning’s LightningModule. Subclass NequIPLightningModule for training techniques that require control over weight storage, updates, and manipulation.

Examples include EMALightningModule (exponential moving average), ConFIGLightningModule (multitask training), ScheduleFreeLightningModule (Facebook’s Schedule-Free optimizers), and composed modules like EMAConFIGLightningModule.

Composition Guidelines

When composing multiple LightningModule subclasses, follow these rules to avoid the “deadly diamond of death” multiple inheritance pattern:

• Use unique, distinguishable names for class attributes to prevent overwrites during inheritance

• Be very careful of inheritance order when creating “diamond” subclasses

• Use assertions to verify the composed module behaves as intended