ResBM: Residual Bottleneck Models for Low-Bandwidth Pipeline Parallelism

TL;DR

ResBM introduces a novel architecture with residual bottleneck modules that enable efficient low-bandwidth pipeline parallelism, achieving high activation compression with minimal loss in training performance.

Contribution

The paper proposes ResBM, a new architecture that facilitates low-bandwidth pipeline parallelism by integrating residual bottleneck modules trained end-to-end.

Findings

ResBMs achieve 128x activation compression.

ResBMs maintain convergence rates similar to standard models.

ResBMs incur minimal memory and compute overhead.

Abstract

Unlocking large-scale low-bandwidth decentralized training has the potential to utilize otherwise untapped compute resources. In centralized settings, large-scale multi-node training is primarily enabled by data and pipeline parallelism, two techniques that require ultra-high-bandwidth communication. While efficient methods now exist for decentralized data parallelism, pipeline parallelism remains the primary challenge. Recent efforts, such as Subspace Models (SM), have claimed up to 100x activation compression but rely on complex constrained optimization and diverge from true end-to-end training. In this paper, we propose a different approach, based on an architecture designed from the ground up to be native to low-bandwidth communication environments while still applicable to any standard transformer-based architecture. We call this architecture the Residual Bottleneck Model or ResBM, it introduces a residual encoder-decoder bottleneck module across pipeline boundaries that can be trained end-to-end as part of the model's parameters while preserving an explicit low-rank identity path. We show that ResBMs achieve state-of-the-art 128x activation compression without significant loss in convergence rates and without significant memory or compute overhead.