SuperAPI AI Directory

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

Author: AI小集 · Original source: AI工具集

What is BigMac

BigMac is a pipeline parallel training framework for multimodal large models open-sourced by Xiaohongshu's dots infra team. Targeting the difficulty of balancing computational efficiency and memory usage in MLLM training, it proposes a dependency-safe nested pipeline paradigm. The framework uses a mature LLM pipeline as its backbone and embeds encoder and generator computations without disrupting the execution order, achieving a 1.08–1.9x training speedup while maintaining stable memory usage. It has been deployed in production as a core component of dots multimodal training.

BigMac

BigMac's Main Features

  • Dependency-Safe Nested Pipeline: Uses the LLM pipeline as a stable backbone and embeds encoder / generator computations at positions where dependencies are satisfied, avoiding cross-module bubble interference.
  • Memory-Bounded Modality Activations: Encoder and generator activation memory is reduced to O(1), while LLM activation behavior remains unchanged, supporting larger batch sizes without OOM.
  • Decoupling of Global Scheduling and Execution: The Scheduler generates a global operator table covering all ranks, microbatches, and modules; the Executor interprets the local sequence and dispatches it to the corresponding backends.
  • Pipeline-Transparent Interface: Algorithm engineers only need to describe module input/output boundaries, and BigMac automatically handles schedule, handoff, and cross-device communication, keeping model code modular.
  • Schedule-Aware Toolchain: Built-in profiler, simulator, and visualization tools support per-operator level trace export (compatible with Perfetto UI), enabling quick identification of bubbles, slow operators, and communication waits.

BigMac's Technical Principles

  • Dependency-Safe Nested Pipeline Scheduling: BigMac uses this as the base timeline. The scheduler analyzes the data dependencies of encoder forward / LLM forward / generator forward and their corresponding backward passes, inserting modality computations only into "gaps" where inputs are ready and the LLM execution order will not be disrupted. The LLM continues to advance according to its original schedule, and encoder and generator activations can be released as soon as gradients are ready, with backward executed in a timely manner, breaking the Pareto trade-off of "high computational efficiency necessarily means high memory usage, and high memory efficiency necessarily means more bubbles".
  • Global Operator Table and Backend Decoupling: The system makes the scheduling strategy explicit as a global operator table containing the forward/backward operators of LLM / encoder / generator on all pipeline ranks and module boundary communications. The execution layer only interprets the local operator sequence, LLM operators reuse Megatron Core, and modality operators retain their own data-parallel / FSDP implementations, so schedule adjustments require no rewriting of the model runtime.
Dig into GitHub

Follow on WeChat and reply "Open Source" to join theAI Open Source Project Community Group

How to Use BigMac

  • Define Module Boundaries: Write encoder_forward, llm_forward, and generator_forward functions in a modular way, declaring the inputs and outputs of each module.
  • Integrate with the BigMac Interface: Register modules through the PP-transparent interface, and BigMac automatically handles pipeline stage partitioning, activation handoff, gradient handoff, and cross-rank communication.
  • Configure Parallel Strategy: Set parameters such as PP / VPP, number of microbatches, and module placement strategy.
  • Run Training: Launch the Scheduler to generate a global plan, and the Executor executes the operator sequence on each rank.
  • Performance diagnosis: Use the built-in profiler to collect per-operator time, export a rank-aligned timeline to the Perfetto UI to analyze bubbles and bottlenecks; or use the simulator to quickly iterate on parallel strategies before starting training.

BigMac's core advantages

  • Breaking the Pareto frontier: Achieve both high computational efficiency and low memory usage simultaneously, with no need to choose one over the other.
  • Zero interference in LLM pipelines: Preserve the mature LLM schedule, and fluctuations in encoder/generator latency will not propagate along the pipeline.
  • Scalable memory usage: Encoder/generator activation memory is O(1), supporting production-level large-batch training.
  • Algorithm-friendly: Model code does not need to be aware of the pipeline runtime, and experiments validated on a single GPU can naturally scale to pipeline parallelism.
  • Debuggable and deployable: The global schedule is visible and inspectable, and the accompanying trace/simulation toolchain reduces the cost of performance optimization.

BigMac's 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

Comparison of BigMac with similar competing products

Dimension BigMac Megatron-LM
Scheduling paradigm Dependency-safe nested pipeline, with the LLM as the backbone A single memory-efficient pipeline, with modality modules as independent stages
Computational efficiency High, with no cross-module bubbles Medium, fluctuations in encoder/generator introduce pipeline bubbles
Memory usage Stable, with modality activations O(1) Relatively low, but limited room for memory optimization as modules become coupled
Scalability Supports large-scale training for both understanding and generation paths Bubbles and memory pressure become significant when scaling the generation path
Interface friendliness PP-transparent, with zero intrusion into algorithm code Requires understanding and adapting to the pipeline runtime and stage partitioning
Toolchain Built-in schedule-aware profiler + simulator Relies on traditional nsys / torch trace and manual correlation analysis

BigMac's application scenarios

  • Pretraining multimodal understanding models: End-to-end pretraining of MLLMs such as image-text understanding and video understanding, requiring efficient pipeline coordination between encoders such as ViT / Whisper and the LLM.
  • Training multimodal generative models: Involving both understanding and generation paths (such as image generation and speech synthesis), requiring coordination of bidirectional dependencies between the LLM and generators such as MMDiT / LDM.
  • Large-batch training: Pursuing larger global batch sizes on memory-constrained clusters, avoiding OOM caused by activation accumulation in traditional computationally efficient designs.
  • Rapid experimental iteration: Algorithm teams can validate model structures on a single GPU or in a DP environment, and seamlessly scale to pipeline parallelism through the BigMac interface without rewriting code.
  • Training performance tuning: Engineering teams use the schedule-aware profiler and simulator to locate pipeline bubbles and optimize parallel configurations and module placement strategies.