NoLoCo: No-all-reduce Low Communication Training Method for Large Models
Authors: name: Jari Kolehmainen, name: Nikolay Blagoev, name: John Donaghy, name: Oğuzhan Ersoy, name: Christopher Nies
Organization: Gensyn

Training large language models is generally done via optimization methods on
clusters containing tens of thousands of accelerators, communicating over a
high-bandwidth interconnect. Scaling up these clusters is expensive and can
become impractical, imposing limits on the size of models that can be trained.
Several recent studies have proposed training methods that are less
communication intensive, avoiding the need for a highly connected compute
cluster. These state-of-the-art low communication training methods still employ
a synchronization step for model parameters, which, when performed over all
model replicas, can become costly on a low-bandwidth network.
  In this work, we propose a novel optimization method, NoLoCo, that does not
explicitly synchronize all model parameters during training and, as a result,
does not require any collective communication. NoLoCo implicitly synchronizes
model weights via a novel variant of the Nesterov momentum optimizer by
partially averaging model weights with a randomly selected other one. We
provide both a theoretical convergence analysis for our proposed optimizer as
well as empirical results from language model training.
  We benchmark NoLoCo on a wide range of accelerator counts and model sizes,
between 125M to 6.8B parameters. Our method requires significantly less
communication overhead than fully sharded data parallel training or even widely
used low communication training method, DiLoCo. The synchronization step itself
is estimated to be one magnitude faster than the all-reduce used in DiLoCo for
few hundred accelerators training over the internet. We also do not have any
global blocking communication that reduces accelerator idling time. Compared to
DiLoCo, we also observe up to 4% faster convergence rate with wide range of
model sizes and accelerator counts.