Communication Efficient Distributed Training with Distributed Lion

1. Incompatible with all reduce.

Our current algorithm indeed requires a customized all_reduce, but we believe the code should be relatively simple to apply to various real-world applications. Additionally, we are exploring ways to optimize the communication process for low-bit information. You are right, due to the limited operator support in practical all reduce such as PyTorch, which is ring-based reduce, it makes the low-bit information hard to broadcast and maintain its information when broadcast through the whole ring. This issue widely exists in all our baselines. Our current practical solution is packing and encoding our binary information into a low-bit tensor and then using an int8 tensor as a container to broadcast through the ring all reduce. This can still compress the information from 2-4x depending on the number of workers, we show the empirical result under different works below together with question 2. For sure, we are looking forward to some core updates that can be made by the NCCL library or pytorch team to support more low-bit friendly operators or datatypes like fp4, and fp8 during communication, which could greatly help with low-bit communication algorithms.

2. Computation overhead.

Please check the common response.

3. The difference between Distributed Lion and SIGNUM-like algorithm.

The main differences are two-fold:

• Lion has an additional reweighting before the sign(), as in sign( $\beta_1 m_t + (1 - \beta_1) g_t$) and the weight decay term, which are crucial for performance improvement. The unique form of Lion presents the special $\ell_\infty$ constraint optimization view of the algorithm.

• Usually, when applying SIGNUM to distributed training, we aggregate the gradient and then apply the sign operation. In distributed Lion, we transmit the signed results and their majority vote, not the gradient. The weight decay, as a result, is applied locally in distributed Lion.

1. Comparison to quantization methods.

The actual update on each worker is actually not the gradient, but rather the Lion’s update (the sign() plus weight decay). To our knowledge, the quantization methods often quantize the gradients before feeding the quantized gradient to the optimizer. In our case, the sign() operation is applied to the reweighted momentum, which is computed using exact gradients. Moreover, unlike quantization methods, the sign() function in distributed Lion is a feature, instead of any approximation.

Given the limited time for rebuttal, we are not able to run QSGD and SignSGD distributedly. However, note that the Ternary gradient baseline is a specific quantization method, where the gradient is quantized to a ternary. Moreover, we provide a comparison against a stronger baseline than SignSGD (Sign SGD momentum, a.k.a., SIGNUM), and also the SGD momentum (as requested by the reviewer) in the following:

From the result, we see that the distributed Lion still results in better performance.

2. Sensitivity to hyperparameters.

Empirically we observe that the performance of distributed Lion is not sensitive to hyperparameters. In our experiments, we did not tune the betas, and chose the default betas suggested in the Lion paper, and only tuned the learning rate and weight decay, which are the two hyperparameters shared by all optimizers. According to our findings in Table 4, the best hyperparameters for distributed Lion are the same as in the Global Lion. Therefore, we expect that the default hyperparameters for Lion will always be a good configuration for distributed Lion.

3. Reproducibility.

The actual implementation requires further internal review to be made publicly available. But we provide an implementation of our distributed Lion optimizer in this link.

4. Performance on CIFAR-10.

Note that as mentioned in Line 224, we are applying all optimizers to ViT models with 6 layers and 8 heads, as it is easier and faster to train. We chose ViT models as they are broadly used in practice. As a result, given the limited size of the ViT model, a performance with nearly 90% accuracy is already pretty high. The overall architecture is similar to the ViT small from this codebase [https://github.com/kentaroy47/vision-transformers-cifar10/tree/main] (note that it is ViT small, not the ViT small (timm transfer)). The training result from their repo is around 80% accuracy.

5. Convergence Curve.

We provide the training curves of Distributed Lion, Global Lion, and Global AdamW on ImageNet training in the uploaded PDF from the common response. From the curve, we can tell that Distributed Lion performs slightly better throughout the training.

6. Wall-clock time comparison.

We refer the reviewer to the common response for this question.

We thank the reviewer for your followups :)

1. Why SIGNUM outperforms SGD Mom. This is in fact expected. So the sign() operator is not merely compressing the momentum information, it also normalizes the update. Take the Adam optimizer as an example, at the first step, Adam's update recovers the normalized gradient descent $g / ||g||$, and at each step, Adam's update is also roughly doing that. It is known that these kinds of normalized gradient descent perform well in practice. One reason is that it makes each entry of the parameter update at a similar pace (therefore similar to a second-order method), the other reason is that it can avoid saddle points (as the update norm is relatively constant). SIGNUM (and similarly Lion) also mimics what Adam does by having the sign operation (their update norm is a constant). Based on my personal experience, SGD Momentum works better for convolution-based networks. But for Transformer models, often Adam-like optimizers work better.

2. Why the Optimizer State Communication (ms) is half less? As we mentioned in the common response, currently during our implementation, we are using the int8 all reduce instead of binary all reduce. This is because in practice all-reduce in PyTorch does not support low-bit all-reduce in a ring-based fashion, which is the default way to perform distributed training. Our current practical solution is packing and encoding our binary information into a low-bit tensor and then using an int8 tensor as a container to broadcast through the ring all reduce. This can still achieve a 2-4x compression, depending on the number of workers.

On the other hand, we are looking forward to some core updates from the NCCL library or pytorch team to support more low-bit friendly operators or datatypes like fp4, and fp8 during communication, which could greatly help with low-bit communication algorithms. As a result, our current work mainly focuses on demonstrating the theoretical convergence and empirical performance of distributed Lion algorithms, rather than the practical aspects of implementation.

1. Wall-clock time reduction.

We refer the reviewer to the common response for this question.

2. Compatibility with ZeRO3 data parallelism.

Although large-scale parallelism techniques such as ZeRO3 and FSDP require additional inter-node gather operations that cannot be accelerated by our algorithm, we can still optimize intra-node optimizer state synchronization. In the common large-scale training scenario where the model size increases and intra-node latency becomes significant, our optimizer can reduce communication time by around 50%. While our algorithm does not reduce inter-node gather time in this case, it is compatible with other gradient-size compression algorithms like Galore[1]. In optimal cases, combining these algorithms can lead to even greater reductions in communication time.

3. Difference between majority vote and averaging.

Based on our empirical observations, we do not observe too much difference between the two schemes and one could potentially perform better than the other on different datasets. The two reduction schemes are indeed similar, as the majority vote method can be seen as applying an actual sign() activation on top of the averaging.

---

[1] Zhao et al, GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection.

1. i.i.d assumption.

Indeed, currently, we assume data are i.i.d (the dataset on each worker is pre-sharded before training). We leave it as a future work to show the convergence of distributed Lion under a non-i.i.d setting.

2. Wall-clock time comparison and communication reduction.

Please check the common response.

3. Whether the method can be applied to federated learning?

Yes, in principle we expect the method to work in a federated fashion. However, when distributed Lion is applied in a local-SGD fashion, the local update will involve multiple steps of the Lion update, so the aggregated server update will be quantized instead of being binary.