BigMac – Xiaohongshu's Open-Source Framework for Pipeline-Parallel Training of Multimodal Large Models

1. What is BigMac

BigMac is an open-source pipeline-parallel training framework for multimodal large models developed by the dots infra team at Xiaohongshu. It aims to address the fundamental contradiction between computational efficiency and GPU memory consumption during the training of multimodal large models (MLLMs). The framework introduces a dependency-safe nested pipeline paradigm, using a mature large language model (LLM) pipeline as a stable backbone. It embeds forward and backward computations of encoders and generators without disrupting the execution order, thereby achieving a 1.08 to 1.9 times increase in training speed while maintaining stable GPU memory usage. BigMac has been deployed as a core component in Xiaohongshu's multimodal training system, proving its effectiveness and reliability in real-world scenarios.

Image source: Official article
Image source: official article

Technical Positioning and Domain: BigMac belongs to the domain of large-scale distributed training frameworks, focusing specifically on pipeline-parallel training for multimodal large models. Unlike traditional frameworks such as Megatron-LM, BigMac is optimized for the complex dependencies between encoders, language models, and generators in multimodal models. By leveraging nested pipeline technology, it breaks the Pareto trade-off between computational efficiency and GPU memory consumption, offering a new, efficient, and memory-controlled solution for multimodal training.

Development Background: This framework was developed by the dots infra team at Xiaohongshu, driven by the practical needs of internal multimodal training. When training large-scale multimodal models (such as vision-language models), the team found that traditional pipeline-parallel methods introduced significant pipeline bubbles and memory pressure once encoder and generator computations were incorporated. Thus, BigMac was developed to address this pain point. Its design philosophy reflects the industry's pursuit of training efficiency and resource utilization.

Core Value: The core value of BigMac lies in breaking the Pareto frontier between computational efficiency and GPU memory consumption. In traditional approaches, improving computational efficiency often comes at the cost of increased memory usage (e.g., reducing pipeline bubbles but retaining more activations), and vice versa. BigMac achieves dependency-safe nested pipeline scheduling, maintaining the efficient execution of the LLM pipeline while reducing the activations of encoders and generators to O(1) levels. This allows for increased training speed with controllable memory usage, supporting larger batch sizes and enhancing the scalability of model training.

Technical Features: The technical features of BigMac include dependency-safe nested pipeline scheduling, memory-bounded modal activations, decoupled global scheduling and execution, pipeline-transparent interfaces, and a schedule-aware toolchain. Among these, the dependency-safe nested pipeline is the core innovation, inserting modal computations into the "gaps" of the LLM pipeline by analyzing data dependencies between modules. This ensures the correctness of execution order while avoiding bubbles across modules. Additionally, the decoupled design of global scheduling and execution allows adjustments to scheduling strategies without affecting the execution layer, thereby improving system flexibility.

2. Key Features

• Dependency-safe nested pipeline: Using the LLM pipeline as a stable backbone, the computation of encoders and generators is inserted during "gaps" where input is ready and does not disrupt the execution order of the LLM, by analyzing data dependencies between modules. This design avoids cross-module bubbles caused by treating modal modules as independent stages, while maintaining the continuity of LLM execution, thereby improving overall training throughput.
• Memory-bounded modal activation: The forward activation memory usage of encoders and generators is bounded at O(1) level, meaning it does not grow linearly with sequence length or batch size. This ensures that during training, the memory consumption of the modal computation remains constant, while the activation behavior of the LLM remains unchanged, enabling support for larger batch sizes without causing memory overflow (OOM), significantly improving memory utilization.
• Global scheduling and execution decoupling: The Scheduler generates a global operator table that covers all ranks, microbatches, and modules. This table explicitly describes the scheduling strategy for the entire training process. The Executor is responsible for interpreting the corresponding operator sequence locally and dispatching it to the LLM backend (Megatron Core) or modal backend (data parallelism/FSDP). This decoupling allows scheduling strategies to be adjusted without modifying the execution layer code, facilitating rapid iteration.
• Pipeline-transparent interface: Algorithm engineers only need to describe the input and output boundaries of each module through a PP-transparent interface. BigMac automatically handles pipeline stage partitioning, activation handoff, gradient handoff, and inter-device communication. The model code remains modular and does not need to be aware of the underlying pipeline runtime, reducing development complexity and enabling seamless scaling of single-card validated models to pipeline parallelism.
• Schedule-aware toolchain: Built-in profiler, simulator, and visualization tools support trace exports at the per-operator level and are compatible with the Perfetto UI. Engineers can use these tools to simulate the performance of different parallel strategies before training, and during training, identify pipeline bubbles, slow operators, and communication waits, thereby accelerating the performance optimization process.
• Flexible parallel strategy configuration: Supports configuration of pipeline parallelism (PP) degree, virtual pipeline parallelism (VPP) degree, microbatch count, and placement strategies for modules on devices. Users can flexibly adjust these settings based on specific model structures, cluster topologies, and memory budgets to find the optimal training configuration.

3. How to Use

1. Environment Requirements and Prerequisites: BigMac is built on PyTorch and Megatron Core, so you need to install PyTorch 2.0 or higher, along with the Megatron Core library. It is recommended to use a Linux operating system with NVIDIA GPU (such as A100 or H100), and have CUDA and NCCL communication libraries properly configured. Additionally, you need to prepare the definition of the multimodal model, including the module code for the encoder, LLM, and generator, and ensure that these modules can correctly execute during forward and backward propagation.

2. Define Module Boundaries: Write the encoder_forward, llm_forward, and generator_forward functions in a modular fashion. Each function must clearly declare the input and output tensor shapes and data types. For example, the encoder may take an image tensor and output visual features, the LLM may take text embeddings and output hidden states, and the generator may generate the final output from the LLM's output. These boundary definitions will serve as the foundation for BigMac's scheduler to analyze dependencies.

3. Integrate with BigMac Interface: Register the above modules into the BigMac framework using the PP-transparent interface. Specific operations include calling the bigmac.register_module() function, passing in the module instance, input specification (input_spec), and output specification (output_spec). After registration, BigMac will automatically handle pipeline stage partitioning, activation handoff, and gradient handoff setup, without requiring manual intervention.

4. Configure Parallelism Strategy: Set the pipeline parallelism (PP) degree, virtual pipeline parallelism (VPP) degree, microbatch count, and placement strategy of modules on devices. These parameters can be passed through a configuration file (such as YAML) or a configuration dictionary within the code. BigMac provides default configurations, but it is recommended to adjust them based on the specific model size and cluster topology to achieve optimal performance.

5. Run Training and Performance Diagnosis: Start the Scheduler to generate a global operator plan, then execute it on each rank via the Executor. During training, you can use the built-in profiler to collect the execution time of each operator, export the rank-aligned timeline to the Perfetto UI for visual analysis. You can also use the simulator to model the performance of different parallelism strategies before training, helping you select the optimal configuration and avoid performance loss during actual training.

4. Pros and Cons Analysis

5. Comparative Analysis with Similar Tools

Selection Recommendations: For teams focused on training large-scale multimodal models (such as vision-language models or image generation models), BigMac offers the best scheduling efficiency and memory control, especially suitable for dual-path scenarios involving understanding and generation. Its pipeline-transparent interface and built-in toolchain significantly reduce development and tuning costs. If a team is already using the Megatron-LM ecosystem and primarily conducts pure text or simple multimodal training, Megatron-LM remains a mature and stable choice, though manual handling of pipeline integration for modal modules is required. For teams aiming to reduce manual tuning by leveraging automatic parallelism, ColossalAI provides flexible parallel combinations, but additional configuration is needed in multimodal scenarios to achieve similar efficiency. DeepSpeed is suitable for teams that are already deeply integrated with DeepSpeed and require ZeRO optimization for large-scale training, though extra configuration is needed for pipeline-related aspects.

6. Editor's Summary

BigMac introduces a dependency-safe nested pipeline paradigm for stream parallel scheduling, representing an innovative technical breakthrough. By embedding encoder and generator computations into the "gaps" within the LLM pipeline, it effectively addresses the trade-off between computational efficiency and GPU memory usage in traditional stream parallelism, enabling multi-modal large model training to maintain high throughput while controlling memory overhead. From a technological innovation perspective, BigMac's decoupled global scheduling and execution design, bounded memory activation mechanism, and pipeline-agnostic interface all demonstrate deep consideration in the design of distributed training systems. Its practical value has already been validated in Xiaohongshu's internal production environment, with training speed improvements ranging from 1.08 to 1.9 times, highlighting its effectiveness in real-world scenarios. BigMac is primarily aimed at algorithm engineers and distributed systems engineers working on multi-modal large model research and development, especially suitable for teams requiring efficient training of large-scale multi-modal models. Currently, the maturity of its community ecosystem and documentation are key factors influencing widespread adoption, and future investments in these areas will determine its rate of adoption. Overall, BigMac is a technically solid open-source project with practical application value, deserving attention and experimentation in the field of multi-modal training.

7. Application Scenarios

• Multimodal Understanding Model Pre-training: For end-to-end pre-training of MLLMs such as image-text understanding or video understanding, efficient processing of encoder collaboration with LLMs like ViT and Whisper is required. BigMac ensures that encoder computations do not interfere with the execution order of the LLM through secure nested pipelines, while reducing encoder activation memory usage to O(1), resulting in a 1.08 to 1.9 times increase in pre-training speed, stable memory consumption, and support for larger batch data processing.
• Multimodal Generation Model Training: Training models that include both understanding and generation paths (such as image generation and speech synthesis) requires coordination of bidirectional dependencies between LLMs and generators like MMDiT and LDM. BigMac's global scheduler can handle such complex dependencies, seamlessly inserting computations when the generator is ready, avoiding bubbles, while maintaining constant activation memory usage for the generator, thereby improving the training efficiency of generation models.
• Large-scale Batch Training: In memory-constrained clusters, the goal is often to achieve a larger global batch size. Traditional computation-efficient designs often lead to OOM (out of memory) errors due to activation accumulation. BigMac's memory-bounded activation mechanism keeps the memory usage of both encoders and generators constant, enabling support for larger batch sizes, increasing training throughput, and accelerating model convergence.
• Training Performance Optimization: Engineering teams can use BigMac's built-in schedule-aware profiler and simulator to simulate the performance of different parallel strategies before training, and to identify pipeline bubbles and slow operators during training. This toolchain allows for optimization of parallel configurations and module placement strategies, thereby improving training efficiency. It supports analysis at the per-operator level, making performance tuning more precise.

8. FAQ

Q: What is the main difference between BigMac and Megatron-LM?
A: BigMac is specifically optimized for multi-modal model training, embedding encoder and generator computations into the LLM pipeline via a secure nested pipeline, thereby avoiding pipeline bubbles and memory overhead caused by introducing modal modules as independent stages in traditional Megatron-LM. Additionally, BigMac provides a more user-friendly PP-transparent interface and an integrated schedule-aware toolchain, reducing the barrier to entry.

Q: Which hardware platforms does BigMac support?
A: Currently, BigMac is based on Megatron Core and primarily supports NVIDIA GPU clusters. High-end GPUs such as A100 or H100 are recommended. Official support for other hardware platforms (e.g., AMD GPU, Huawei Ascend) has not been announced yet, but future expansion is possible.

Q: Does BigMac support single-card training?
A: BigMac is mainly designed for pipeline parallel training, but it can be run on a single card by setting PP=1. However, the scheduling optimization effects are limited in this case. Algorithm teams can validate model architectures in single-card or data-parallel environments, and then seamlessly scale up to pipeline parallelism using BigMac's pipeline-transparent interface without rewriting code.

Q: How can I get BigMac and start using it?
A: The source code of BigMac has been open-sourced on GitHub (https://github.com/Dots-Infra/BigMac). Users can follow the instructions in the README to install dependencies and run examples. Technical papers and project documentation are also available. It is recommended that users first read the technical paper to understand the design principles, and then refer to the official examples for practical implementation.

Q: How much performance improvement can BigMac achieve?
A: According to official data, BigMac can achieve a 1.08 to 1.9 times increase in training speed on typical multi-modal training tasks, while maintaining stable memory usage. The exact improvement depends on the model architecture, hardware configuration, and parallel strategy. Users are advised to use the built-in simulator to test performance on their own setup.

Q: Does BigMac support custom operators or modules?
A: Yes, BigMac's scheduler supports any module as long as it defines clear input and output boundaries. Users can register custom encoder or generator operators via the PP-transparent interface, and BigMac will automatically handle dependency analysis and scheduling.

9. Project Links

• Project Website: https://dots-infra.github.io/BigMac/
• GitHub Repository: https://github.com/Dots-Infra/BigMac/
• arXiv Technical Paper: https://arxiv.org/pdf/2605.25451