A Parallel Alternative for Energy-Efficient Neural Network Training and Inferencing

TL;DR

The paper introduces phantom parallelism (PP) as an energy-efficient alternative to tensor parallelism (TP) for training and inference in large neural networks, demonstrated on feedforward networks (FFNs) as a proof of principle. PP inserts phantom layers with ghost neurons to compress local activations before inter-rank communication, reducing both computation and communication costs. The authors derive forward and backward PP operators, implement them as custom autograd primitives, and show through end-to-end experiments up to 256 GPUs that PP can achieve about a 50% reduction in training energy compared to TP, with memory and runtime benefits particularly at scale. They also show that smaller PP models on fewer GPUs can reach the same loss as larger TP models on more GPUs, suggesting substantial practical energy savings, while acknowledging the need to generalize to transformers and other architectures in future work.

Abstract

Energy efficiency of training and inferencing with large neural network models is a critical challenge facing the future of sustainable large-scale machine learning workloads. This paper introduces an alternative strategy, called phantom parallelism, to minimize the net energy consumption of traditional tensor (model) parallelism, the most energy-inefficient component of large neural network training. The approach is presented in the context of feed-forward network architectures as a preliminary, but comprehensive, proof-of-principle study of the proposed methodology. We derive new forward and backward propagation operators for phantom parallelism, implement them as custom autograd operations within an end-to-end phantom parallel training pipeline and compare its parallel performance and energy-efficiency against those of conventional tensor parallel training pipelines. Formal analyses that predict lower bandwidth and FLOP counts are presented with supporting empirical results on up to 256 GPUs that corroborate these gains. Experiments are shown to deliver approximately 50% reduction in the energy consumed to train FFNs using the proposed phantom parallel approach when compared with conventional tensor parallel methods. Additionally, the proposed approach is shown to train smaller phantom models to the same model loss on smaller GPU counts as larger tensor parallel models on larger GPU counts offering the possibility for even greater energy savings.

Figures (7)

• Figure 1: Parallel partitioning of a fully connected single layer FFN ($p=2$).
• Figure 2: A 3-way TP partition of a FFN with two layers. Mono-colored block computations do not depend on remote information while bi-colored block computations depend on information in corresponding colored ranks. The colored lines indicate inter-process communications.
• Figure 3: Illustration of a phantom layer of width $k$ inserted between the local $n/p$-width portions of the input and output layers on rank $i$ in a single layer of an FNN of width $n$ partitioned across $p$ processes.
• Figure 4: PP components and operations between layer $\ell-1$ and layer $\ell+1$ with $p=3$. Each thick solid-bordered box represents the weight matrix for a specific partition. Dashed boxes denote individual processors, encapsulating their associated weights and intermediate variables. Colored thick lines indicate inter-processor communication.
• Figure 5: Comparison of TP and PP performance across model sizes and GPU counts. The phantom layer sizes (numbers of ghost neurons), $k$, is indicated in (b) and (c) for PP.