Balancing Pipeline Parallelism with Vocabulary Parallelism

Balancing Pipeline Parallelism with Vocabulary Parallelism

Speakers: Yeung Man Tsung, Penghui Qi, Min Lin, Xinyi Wan

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This talk, presented by Yeung Man Tsung from the National University of Singapore, details a novel approach called Vocabulary Parallelism (VP). The research, conducted during an internship at ByteDance's CAI lab, addresses a critical bottleneck in scaling large language models (LLMs) with pipeline parallelism (PP): the disproportionate computational and memory burden imposed by the input and output embedding layers. While traditional pipeline parallelism effectively distributes transformer layers across multiple devices, the embedding layers, particularly the output embedding layer, often become a significant performance impediment, especially with the burgeoning vocabulary sizes seen in modern multi-language models.

The core problem stems from the fact that the first input embedding layer and the last output embedding layer are typically handled by the initial and final pipeline devices, respectively. As vocabulary sizes grow, the compute and memory requirements of these layers can exceed those of an entire transformer layer, leading to severe workload imbalance, idle time (pipeline bubbles), and inefficient resource utilization. Vocabulary Parallelism proposes a method to partition these embedding layers and distribute their computation and memory footprint more evenly across pipeline stages, thereby restoring balance and improving overall training and inference efficiency.

This work is highly significant for the ML/systems community. As LLMs continue to scale in size and complexity, incorporating multi-language support and larger vocabularies, efficient distributed training becomes paramount. Existing solutions for handling embedding layer imbalance often introduce their own drawbacks, such as increased activation memory or inability to fully balance both compute and memory. Vocabulary Parallelism offers a robust, generalizable, and memory-efficient solution that integrates seamlessly with established pipeline parallelism schedules, promising substantial performance gains for the next generation of large-scale AI models.

Background

▶ Watch: Introduction: Imbalanced workload in pipeline parallelism (0:00)

The landscape of deep learning, particularly with the advent of large language models (LLMs), is characterized by an insatiable demand for computational resources. To train and deploy these colossal models, distributed training strategies are indispensable. Among these, pipeline parallelism (PP) has emerged as a widely adopted technique. Pipeline parallelism partitions a deep learning model, typically a sequence of transformer layers, across multiple devices. Each device processes a subset of the layers, and micro-batches of data flow sequentially through the pipeline, overlapping computation and communication to maximize throughput. The 1F1B (one forward, one backward) schedule is a classic example, where a stable phase is achieved with continuous forward and backward passes, minimizing idle time.

However, a fundamental challenge arises when considering the complete architecture of LLMs: the input and output embedding layers. In a standard pipeline parallelism setup, the initial input embedding layer is typically assigned to the first pipeline device, and the final output embedding layer (often referred to as the vocabulary layer or head) is handled by the last pipeline device. While this seems straightforward, it creates a significant imbalance in workload. The problem is exacerbated by two key factors:

1. Imbalanced Workload Distribution: Even in smaller models, assigning the embedding layers exclusively to the first and last stages can create a non-uniform distribution of computation. This leads to persistent "pipeline bubbles" – periods where some devices are idle waiting for others to complete their embedding layer computations – which cannot be alleviated simply by increasing the number of in-flight micro-batches.
2. Growing Vocabulary Sizes: Modern LLMs, especially those designed for multi-language support, feature increasingly vast vocabulary sizes. The talk highlights that for some open-source models, the ratio of FLOPs (floating-point operations) for the output embedding layer compared to a single transformer layer can exceed one. This means the output embedding layer alone can demand more compute than an entire transformer block. Furthermore, both the compute and memory requirements of these layers scale dramatically with vocabulary size, as demonstrated with GML29B model, where they grow "very large." This makes the imbalance not just a minor inefficiency but a severe bottleneck, causing the first and last pipeline stages to become overloaded in terms of both computation and memory.

Prior attempts to address this issue have met with limited success or introduced new problems:

• Layer-level Redistribution: A common naive approach involves reducing the number of transformer layers assigned to the first and last pipeline devices by one. While this attempts to redistribute compute, it struggles to achieve true balance for several reasons. Firstly, the input embedding layer often has less compute but still demands significant memory, making it difficult to balance both aspects simultaneously. Secondly, the slight differences in compute/memory profiles between embedding layers and transformer layers prevent perfect balance.
• Interlace Pipeline: This more sophisticated method alternates between pipeline parallelism for transformer layers and tensor parallelism (TP) for vocabulary layers. The idea is to apply tensor parallelism to the large embedding matrices, distributing their computation across devices. However, a major drawback of the interlace pipeline is its substantial increase in peak activation memory, up to 1.5 times the original. Given that activation memory is already a critical constraint for large models trained with pipeline parallelism, this increase makes the interlace pipeline difficult to adopt in real-world scenarios.

Recognizing these limitations, the talk introduces Vocabulary Parallelism as a targeted and efficient solution to overcome the embedding layer bottleneck, aiming for balanced compute and memory without the prohibitive memory overhead of prior methods.

Key Findings

▶ Watch: Limitations of current naive solutions (2:00)

The central contribution of this work is the proposal of Vocabulary Parallelism (VP), a novel method designed to alleviate the computational and memory imbalance caused by large input and output embedding layers in pipeline parallelism. The key findings and contributions can be summarized as follows:

1. Balanced Workload Distribution: VP's primary objective is to partition the vocabulary layers (both input and output embeddings) and distribute their workload equally across all pipeline devices. This directly addresses the imbalance that arises when these layers are exclusively handled by the first and last stages, thereby eliminating the persistent pipeline bubbles and improving overall resource utilization.
2. Dual Balancing of Compute and Memory: Unlike naive layer redistribution, VP is engineered to balance both the computation and memory requirements of the vocabulary layers. This is crucial as embedding layers, particularly the input embedding, can have disparate compute and memory footprints.
3. Activation Memory Efficiency: A critical design goal for VP was to incorporate vocabulary layer partitioning without significantly increasing activation memory usage. The proposed scheduling mechanism ensures that the increase in in-flight micro-batches is minimal (a constant increase of one), thereby avoiding the prohibitive memory overhead seen in approaches like the interlace pipeline (which increases peak activation memory by 1.5x).
4. Novel Communication Barrier Reduction: For the output embedding layer, a key technical challenge was the multiple communication barriers required for synchronization during partitioning. The authors developed an approach inspired by online softmax, effectively reducing the number of communication barriers from three to just one. This simplification is vital for efficient pipeline scheduling and reduced activation memory.
5. Generalizable Pipeline Scheduling: VP integrates its vocabulary passes into existing pipeline schedules using a building block framework. By modifying the repeating pattern (building block) of schedules like 1F1B or V-half, VP inserts its specialized passes (S for computation, T for weight gradient updates) into the existing offset between forward and backward passes. This demonstrates the generalizability of VP across various pipeline parallelism schedules.
6. Significant Performance Improvements: Empirical validations show substantial performance gains:

• MFU (Model FLOPs Utilization) improved by 5% to 51% over existing approaches, showcasing a much more efficient use of computational resources.
• Improved memory usage in scenarios with large vocabularies due to better balancing across pipeline stages.
• When combined with memory-balanced schedules like V-half, VP achieved full memory balance across all pipeline devices and significantly improved throughput over baselines.

1. Overall Performance Gains Despite Overhead: While partitioning introduces some overhead (e.g., smaller kernel dimensions reducing single-kernel MFU, or non-parallelizable operations), the overall performance gains from balanced workload distribution outweigh these costs, even for models with relatively smaller vocabulary sizes.

In essence, Vocabulary Parallelism provides a practical, efficient, and broadly applicable solution to a long-standing challenge in distributed LLM training, enabling better scaling and higher utilization of hardware resources.

Technical Deep Dive

▶ Watch: Challenges in partitioning output embedding layer (4:00)

The technical core of Vocabulary Parallelism revolves around two primary innovations: efficiently partitioning the large embedding layers and seamlessly integrating these partitioned components into existing pipeline schedules without incurring significant overhead.

Partitioning the Embedding Layers

While the talk states that partitioning the input embedding layer is "relatively trivial" and can be "piggybacked to the output embedding passes," the main technical challenge lies in the output embedding layer. This layer is more complex due to its role in generating logits for the entire vocabulary and the subsequent softmax and gradient computations.

The computation graph of the output embedding layer involves:

1. A matrix multiplication (often referred to as the "unembedding" or "head" layer) that projects the transformer's output to the vocabulary dimension.
2. A softmax operation to convert logits into probabilities.
3. Two backward matrix multiplications to compute gradients for the input to this layer and for the weights of the embedding matrix itself.

A naive approach to parallelizing this layer would split these computations across devices. The challenge is that the softmax operation and the backward passes inherently require global information, leading to synchronization points. Specifically, a naive split would result in three distinct passes (F1, F2, and B) each requiring an all-reduce communication step to synchronize data across devices. Each all-reduce acts as a barrier, forcing all devices to wait, thereby increasing pipeline complexity and introducing bubbles.

Crucially, the number of communication barriers directly impacts activation memory usage. If there are multiple barriers, the forward and backward passes of a single micro-batch must be split further apart, requiring a larger number of in-flight micro-batches to fill the resulting gaps and maintain a bubble-free stable phase. This directly translates to higher activation memory consumption. The goal, therefore, is to minimize these communication barriers.

The proposed solution to reduce communication barriers for the output embedding layer is inspired by online softmax. The key insight is to move the synchronization point out of the computation's critical path. The backward matrix multiplications are linear transformations, meaning the global softmax can be effectively "fixed" or adjusted after initial local computations.

The refined approach proceeds as follows:

1. Local Computation: Each device independently performs its local portion of the matrix multiplication and an initial local softmax calculation. It then proceeds with the local backward matrix multiplications to compute input and weight gradients based on these local (un-synchronized) softmax values.
2. Single Communication Barrier (C1): After these local computations, a single all-reduce communication (C1) is performed. This barrier synchronizes the necessary information (e.g., partial sums for softmax normalization, or aggregated gradients) across all devices.
3. Result Fixing: Post-C1, all ranks have the necessary global information to "fix" their local softmax values and the results of the backward matrix multiplications, obtaining the correct global gradients for the input embedding.
4. Weight Gradient Update (T pass): A separate pass, denoted as T, is then performed to update the weight gradients of the output embedding layer. This T pass is designed to be "out of the computation critical path," meaning it can be scheduled more flexibly and potentially overlapped with other operations using multi-stream execution without blocking the main pipeline flow.

This strategy effectively fuses all the primary vocabulary layer computations into a single pass, referred to as the S pass, followed by the single synchronization barrier C1, and then the non-critical T pass. This reduces the number of communication barriers from three to one, significantly simplifying scheduling and reducing activation memory requirements.

Pipeline Scheduling with the Building Block Framework

Integrating these specialized vocabulary passes (S and T) into existing pipeline parallelism schedules is achieved using the building block framework, a concept from prior work that represents most pipeline schedules as a repeating pattern of operations.

The standard 1F1B schedule, for instance, has a specific building block defining the sequence of forward (F) and backward (B) passes for each micro-batch across pipeline stages. The proposed method modifies this building block:

1. Introduce an Offset: The first step is to strategically increase the offset between the forward (F) and backward (B) passes within the building block. This creates a deliberate "gap" in the schedule.
2. Insert Vocabulary Passes:

• The S pass (fused vocabulary computation) is inserted within the first interval created by this offset.
• The T pass (weight gradient update, out of critical path) is inserted within the interval after the offset.
• This insertion is repeated for every single micro-batch.

A crucial advantage of this approach is that this sequence of operations only increases the number of in-flight micro-batches by one, which is a constant. This minimal increase is key to maintaining low activation memory usage, a significant improvement over methods like the interlace pipeline.

The generalizability of this building block approach is a major strength. The talk illustrates how this modification can be applied not only to 1F1B but also to other complex schedules like V-half, which is part of the "zero bubble family" and aims for uniform activation memory usage. By applying the same steps (increasing offset, inserting S and T passes), Vocabulary Parallelism can be integrated with various existing pipeline schedules, extending its benefits broadly.

In summary, the technical innovations in Vocabulary Parallelism lie in a clever re-architecture of the output embedding layer's computation to minimize synchronization overhead and a flexible, generalizable scheduling mechanism that integrates these optimized passes into standard pipeline parallelism without compromising memory efficiency.

Experimental Setup & Results

▶ Watch: Goals: Reduce barriers, integrate vocabulary passes (5:36)

The evaluation of Vocabulary Parallelism (VP) focused on demonstrating its effectiveness in balancing compute and memory for embedding layers, reducing pipeline bubbles, and improving overall throughput and Model FLOPs Utilization (MFU) in large language model training scenarios.

Model and Scenarios:

The experiments primarily used GML29B as a representative large language model, particularly emphasizing scenarios with growing vocabulary sizes. The analysis of GML29B at varying vocabulary sizes highlighted that compute and memory requirements for embedding layers can become "very large," underscoring the problem VP aims to solve.

Baselines:

The performance of VP was compared against "existing approaches," which implicitly include naive methods like layer-level redistribution that attempt to balance workload by adjusting transformer layer assignments. The limitations of these baselines (inability to balance both compute and memory, persistent bubbles) served as the motivation for VP. The interlace pipeline, while mentioned as an alternative, was implicitly dismissed due to its severe activation memory overhead, making it an impractical baseline for direct comparison in memory-constrained settings.

Metrics:

The key performance indicators used in the evaluation included:

• MFU (Model FLOPs Utilization): A measure of how efficiently the available computational FLOPs are being used by the model. Higher MFU indicates better performance.
• Bubble Rate: The percentage of time pipeline devices are idle, waiting for data or synchronization. Lower bubble rate is better.
• Activation Memory Usage: The total memory consumed by intermediate activations during the forward pass. Minimizing this is crucial for scaling large models.
• Throughput: The number of tokens or samples processed per unit of time.
• Compute and Memory Balance: Qualitative and quantitative assessment of how evenly the compute and memory load of vocabulary layers are distributed across pipeline stages.

Headline Results:

1. MFU Improvement: Vocabulary Parallelism consistently demonstrated significant improvements in MFU, ranging from 5% to 51% over existing approaches. This wide range indicates that VP's benefits are particularly pronounced in scenarios where embedding layer imbalance is most severe (e.g., very large vocabulary sizes or complex embedding layer computations).
2. Improved Memory Usage: For large vocabulary scenarios, VP led to "better balancing among the pipeline stages," which translated into improved overall memory usage. This is a direct consequence of distributing the memory-intensive embedding layers more evenly, preventing single stages from becoming memory bottlenecks.
3. Bubble Rate and Activation Memory Control: A critical finding was that VP successfully maintained a low bubble rate while only increasing the number of in-flight micro-batches by a constant factor of one. This confirms that VP achieves its balancing goals without introducing new significant pipeline bubbles or incurring substantial activation memory overhead, a key design objective.
4. Full Memory Balance with Memory-Balanced Schedules: The talk highlighted a specific experiment where VP was combined with V-half, a pipeline schedule known for its ability to achieve uniform activation memory usage across pipeline devices. When integrated with V-half, Vocabulary Parallelism was able to achieve full memory balance across all stages, encompassing both the transformer layers and the vocabulary layers. This synergistic effect resulted in significantly improved throughput over the baseline, demonstrating that VP can complement and enhance other advanced scheduling techniques.
5. Overhead Analysis and Net Gains: The authors acknowledged that partitioning the vocabulary layer introduces some scaling overheads. These include:

• Reduced single kernel MFU: When a large kernel (like a full embedding matrix multiplication) is split into smaller partitions, each sub-kernel operates on smaller dimensions. This can sometimes lead to less efficient utilization of GPU cores due as the overhead of launching kernels dominates for smaller workloads.
• Extra computation for non-parallelizable operations: Some operations, particularly those involving data loading or certain pre-processing steps for the input embedding layer, might not be perfectly parallelizable across the vocabulary dimension. This means all ranks might have to perform these operations, adding a small amount of redundant computation.

Despite these overheads, the experimental results clearly showed that Vocabulary Parallelism "still brings overall performance gains even for small vocabulary sizes," indicating that the benefits of balanced workload distribution outweigh the overheads across a broad spectrum of model configurations.

In summary, the experimental results robustly validate Vocabulary Parallelism as an effective and efficient solution for addressing embedding layer imbalance in pipeline-parallel LLMs, leading to substantial performance improvements in MFU, throughput, and memory efficiency, especially in the context of growing vocabulary sizes.

Practical Implications

▶ Watch: Integrating into existing pipeline scheduling framework (7:15)

Vocabulary Parallelism (VP) presents significant practical implications for researchers, infrastructure teams, and model builders working with large-scale AI models, particularly large language models (LLMs). Its core benefits address critical pain points in distributed training and deployment, making it a valuable addition to the ML systems toolkit.

For Practitioners and Model Builders:

• Enabling Larger and More Diverse Models: VP directly tackles the bottleneck posed by large vocabulary sizes, which are increasingly common in multi-language LLMs, multimodal models, and models with specialized tokenizers. By efficiently distributing the compute and memory of embedding layers, VP allows model builders to design and train models with much larger vocabularies without incurring severe performance penalties. This facilitates the development of more inclusive and capable AI systems.
• Improved Training Efficiency and Reduced Costs: The demonstrated MFU improvements (5-51%) and increased throughput translate directly into faster training times. For computationally intensive LLM training, this means models can be trained in less time, or larger models can be trained within existing timeframes. This also has a direct impact on operational costs, as efficient hardware utilization reduces the total GPU-hours required.
• Better Resource Utilization: By balancing the workload across all pipeline devices, VP ensures that all GPUs are actively engaged in computation, rather than some sitting idle due to overloaded embedding layers on the first/last stage. This maximizes the return on investment for expensive accelerator hardware.
• Memory Footprint Management: The ability to balance memory usage, especially when combined with memory-balanced schedules like V-half, is crucial for fitting ever-larger models onto available hardware. VP helps prevent single devices from becoming memory bottlenecks, potentially allowing for larger batch sizes or larger models to be trained on the same cluster.
• Generalizability: The use of the building block framework means VP is not tied to a specific pipeline schedule. Practitioners can integrate it with their preferred or existing pipeline parallelism setups (e.g., 1F1B, V-half), offering flexibility and easing adoption.

For Infrastructure and ML Systems Teams:

• Enhanced Pipeline Parallelism Robustness: VP makes pipeline parallelism a more robust and scalable solution for LLMs. It addresses one of the primary limitations of PP in the context of rapidly evolving model architectures and tokenization strategies.
• Simplified System Design: By providing an efficient, in-pipeline solution for embedding layers, VP reduces the need for complex, ad-hoc workarounds or hybrid parallelism strategies that might introduce their own complexities (like the interlace pipeline's memory overhead).
• Predictable Performance: Achieving balanced workloads leads to more predictable and stable performance during distributed training, simplifying capacity planning and resource allocation.
• Potential for Extension: While the talk focuses on decoder-only language models, the speaker explicitly notes that the method can extend to other imbalanced components, such as encoders in vision-language models or multi-token prediction heads. This suggests a broader applicability for balancing non-uniform modules in complex multi-modal or multi-task architectures.

Tradeoffs and Limitations:

• Scaling Overhead: As acknowledged, VP introduces some overheads, such as reduced single-kernel MFU due to smaller partitioned dimensions and redundant computation for non-parallelizable operations. While the overall gains outweigh these costs, systems teams should be aware of these factors, especially when dealing with very small vocabulary sizes where the relative overhead might be higher.
• Implementation Complexity: While the conceptual framework is elegant, implementing the specialized S and T passes and integrating them precisely into existing pipeline schedulers, especially with multi-stream execution for overlapping communication, requires careful engineering within deep learning frameworks.
• Focus on Decoder-Only Models: The current work primarily addresses decoder-only language models. While the principles are extensible, adapting VP to other architectures (e.g., encoder-decoder models, multi-modal models with diverse input/output modalities) would require further research and implementation effort to identify and balance their specific non-uniform components.

In essence, Vocabulary Parallelism offers a practical and powerful solution to a critical scaling challenge in modern AI. By intelligently distributing the load of embedding layers, it paves the way for more efficient, larger, and more capable deep learning models across a wider range of applications.

Key Takeaways

• Vocabulary Parallelism (VP) addresses a critical imbalance: It resolves the computational and memory bottleneck caused by large input and output embedding layers in pipeline-parallel training of LLMs, especially those with growing vocabulary sizes.
• Efficient Embedding Layer Partitioning: VP partitions vocabulary layers and distributes their workload evenly across pipeline devices, balancing both compute and memory requirements.
• Reduced Communication Overhead: A key innovation is the reduction of communication barriers for the output embedding layer from three to one, inspired by online softmax, significantly improving pipeline efficiency and reducing activation memory.
• Memory-Efficient Scheduling: VP integrates its specialized "S" (computation) and "T" (weight gradient update) passes into existing pipeline schedules (e.g., 1F1B, V-half) by modifying the building block, increasing in-flight micro-batches by only a constant one, thus minimizing activation memory overhead.
• Significant Performance Gains: Experiments show MFU improvements of 5-51% and better memory usage, particularly in large vocabulary scenarios. When combined with memory-balanced schedules like V-half, VP achieves full memory balance and substantially higher throughput.
• Generalizable and Practical: The building block framework makes VP applicable to various pipeline schedules, offering a robust and generalizable solution that can extend to other non-uniform modules in complex AI architectures.

About the Speaker(s)

The primary presenter of this work is Yeung Man Tsung, who is affiliated with the National University of Singapore. This research was conducted during an internship at the CAI lab within ByteDance. The co-authors on this paper include Penghui Qi, Min Lin, and Xinyi Wan, indicating a collaborative effort originating from this academic and industrial partnership.