PyLO: Towards Accessible Learned Optimizers in PyTorch

TL;DR

PyLO introduces a user-friendly PyTorch library for learned optimizers, enabling large-scale pre-training and performance improvements in real-world machine learning tasks.

Contribution

The paper presents PyLO, a PyTorch-based library with CUDA-accelerated learned optimizers, making advanced optimization techniques accessible and practical for large-scale applications.

Findings

CUDA implementations significantly increase training throughput

Learned optimizers benefit from combining with existing tools

PyLO enables large-scale pre-training with learned optimizers

Abstract

Learned optimizers have been an active research topic over the past decade, with increasing progress toward practical, general-purpose optimizers that can serve as drop-in replacements for widely used methods like Adam. However, recent advances such as VeLO, which was meta-trained for 4000 TPU-months, remain largely inaccessible to the broader community, in part due to their reliance on JAX and the absence of user-friendly packages for independently using the optimizers after meta-training. To address this gap, we introduce PyLO, a PyTorch-based library that brings learned optimizers to the remaining ~70% of machine learning community via the familiar torch.optim.Optimizer interface. Unlike prior work focused on limited-scale academic tasks, our emphasis is on applying learned optimization to real-world large-scale pre-training tasks. Our systems contribution includes CUDA-accelerated implementations of the small fc lopt(Metz et al., 2022a) and VeLO(Metz et al., 2022b) learned optimizers, achieving substantial performance gains, with training throughput on ViT-B/16 (batch size 32) increasing from 39.36 and 49.73 to 205.59 and 191.18 samples per second, respectively. PyLO has the versatility that allows us to easily combine learned optimizers with existing optimization tools such as learning rate schedules and weight decay. When doing so, we discover that learned optimizers can substantially benefit from it. Our code is available at https://github.com/Belilovsky-Lab/pylo