Mixtures of Subspaces for Bandwidth Efficient Context Parallel Training
For researchers training large language models with long contexts in bandwidth-constrained or decentralized settings, this method dramatically reduces communication costs without sacrificing performance.
The paper tackles communication overhead in context-parallel training of large language models over low-bandwidth networks, achieving over 95% compression with no convergence loss, enabling decentralized training with context lengths exceeding 100K tokens on 300Mbps networks to match centralized speed on 100Gbps interconnects.
Pretraining language models with extended context windows enhances their ability to leverage rich information during generation. Existing methods split input sequences into chunks, broadcast them across multiple devices, and compute attention block by block which incurs significant communication overhead. While feasible in high-speed clusters, these methods are impractical for decentralized training over low-bandwidth connections. We propose a compression method for communication-efficient context parallelism in decentralized settings, achieving a remarkable compression rate of over 95\% with negligible overhead and no loss in convergence. Our key insight is to exploit the intrinsic low-rank structure of activation outputs by dynamically constraining them to learned mixtures of subspaces via efficient reparameterizations. We demonstrate scaling billion-parameter decentralized models to context lengths exceeding 100K tokens on networks as slow as 300Mbps, matching the wall-clock convergence speed of centralized models on 100Gbps interconnects.