Cohere Labs - Tycho van der Ouderaa and Matt Beton, Exo Labs

In this talk, we introduce KPOP, a novel preconditioned optimizer developed at Exo Labs that performs adaptive optimization in the Kronecker-factored eigenbasis (KFE). By leveraging efficient Kronecker-factored curvature approximations (KFAC), KPOP brings the benefits of second-order information to modern large-scale training while maintaining scalability and simplicity. Despite a higher memory footprint than AdamW, KPOP consistently outperforms it, both in training iterations as well as wall-clock time, and achieves performance close to state-of-the-art optimizers on the competitive NanoGPT benchmark on 8x NVIDIA H100s set-ups. We also present TopKPOP, a memory- and bandwidth-efficient variant that retains only the top eigenvectors of the KFE. This enables fine-grained control over resource usage and makes the optimizer well-suited for adapting the training algorithm to specific hardware constraints of the training environment. Combined with a hybrid CPU-GPU implementation that fully exploits Apple’s unified memory architecture, these optimizers enable effective training of large language models on consumer hardware. We are the first to successfully demonstrate large-scale distributed training of LLMs on clusters of Apple devices, spanning 2 to 16 nodes.

Add event to calendar

Apple Google Office 365 Outlook Outlook.com Yahoo