EPAS: Efficient Training with Progressive Activation Sharing

TL;DR

EPAS introduces a progressive activation sharing method that enhances training and inference efficiency in transformer models by reducing redundant computations, leading to significant throughput improvements without sacrificing model performance.

Contribution

This paper proposes a novel progressive activation sharing technique for transformers, enabling variable sharing regions and improving training and inference efficiency.

Findings

Up to 11.1% increase in training throughput.

Up to 29% increase in inference throughput.

10% accuracy improvement in continual pretraining.

Abstract

We present a novel method for Efficient training with Progressive Activation Sharing (EPAS). This method bridges progressive training paradigm with the phenomenon of redundant QK (or KV ) activations across deeper layers of transformers. EPAS gradually grows a sharing region during training by switching decoder layers to activation sharing mode. This results in throughput increase due to reduced compute. To utilize deeper layer redundancy, the sharing region starts from the deep end of the model and grows towards the shallow end. The EPAS trained models allow for variable region lengths of activation sharing for different compute budgets during inference. Empirical evaluations with QK activation sharing in LLaMA models ranging from 125M to 7B parameters show up to an 11.1% improvement in training throughput and up to a 29% improvement in inference throughput while maintaining similar loss curve to the baseline models. Furthermore, applying EPAS in continual pretraining to transform TinyLLaMA into an attention-sharing model yields up to a 10% improvement in average accuracy over state-of-the-art methods, emphasizing the significance of progressive training in cross layer activation sharing models.