Add LBFGS optimizer as an option

[vue1999]

Adds LBFGS optimizer support.

Key details

• Supports distributed training across multiple GPUs when using the —distributed flag
• Handles datasets larger than available memory through batch processing
• Uses strong Wolfe conditions for line search
• modified dataloader to avoid dropping samples when using LBFGS

How to use

• Add the —lbfgs flag
• Use the largest batch_size that fits into memory (for optimal performance)
• Disable EMA

Advantages

• More deterministic training than first order optimizers
• Less sensitivity to random seeds
• Batch size independent convergence path
• Fewer hyperparameters to tune
• Possibly faster and better convergence

Notes

• First-order optimizer hyperparameters (weight decay, learning rate, etc.) are ignored when using LBFGS
• EMA should be disabled as it's redundant with line search
• Unlike first-order optimizers where batch size affects the optimization dynamics, in LBFGS the batch_size parameter is purely computational:
  • Each optimization step requires gradients computed over the entire training set (at several points in the model space)
  • The full-dataset gradient computation is split into batches for memory efficiency and performance (by parallelising across GPUs)
  • The optimizer only sees the accumulated full-dataset gradients, regardless of how they were computed
• While it can be used immediately from the start of training, it is likely that at the beginning ADAM is significantly faster because the precision of LBFGS is not that beneficial initially.
• In the examples we looked at, we first trained the models using ADAM then switched to LBFGS. (This can be done by restarting training from a checkpoint and using the —lbfgs flag.)
• It doesn’t work with multi-head training

Examples

3BPA

SPICE (H,C atoms only subs

et)

[ilyes319]

hey @vue1999, that looks great! is it ready to merge?

[vue1999]

Yes. We didn’t do exhaustive testing, but it works with the multihead options too.

[ilyes319]

would be worth adding some tests for it, inspired with the run train one (test_run_train.py). just a simple training run would be fine.