FlowTrain: Decoupled Multimodal Training Framework Delivers 1.7x Throughput and 50%+ MFU

Background and Challenges

Multimodal large models are becoming critical for intelligent retail, content understanding, customer service, and industrial inspection. Unlike text-only LLMs, multimodal training processes heterogeneous inputs (images, text, video) through a vision encoder, projection layer, and LLM backbone. These modules differ significantly in compute patterns, parallelism strategies, and resource demands, creating efficiency bottlenecks.

In traditional frameworks, samples are organized in fixed batches. The encoder must finish feature extraction before passing embeddings to the backbone. Large variations in resolution, image count, and text length cause per-microbatch processing time fluctuations, leading to three core problems:

Batch-level synchronization wait: Encoder and backbone are bound by fixed batch boundaries; the faster side idles waiting for the slower side, wasting compute.

Unified parallel strategy mismatch: Encoders suit lightweight DP/TP combinations, while LLM backbones typically require DP/TP/PP composite parallelism. A single strategy for all modules causes resource misallocation.

Variable-length sample load imbalance: Padding waste, attention compute differences, and pipeline stage imbalance further reduce cluster MFU.

FlowTrain: Flow-Based Decoupled Training Paradigm

To address these challenges, JD Retail AI Infra proposes FlowTrain, reconstructing multimodal training from a synchronous pipeline into a producer-consumer data flow. The framework comprises three core innovations.

Unified Memory Pool: Temporal Decoupling via Global Address Space

The Unified Memory Pool is the foundation for temporal decoupling. It manages cross-node HBM and DRAM via a 64-bit global virtual address space, placing encoder-produced multimodal embeddings and text tokens into a globally addressable buffer. The backbone continuously accesses ready data via metadata and global addresses without waiting for a full batch of encoder output. The domestic super-node's high-bandwidth interconnect and cross-device access enable transparent embedding flow and continuous backbone consumption, absorbing encoder-side runtime fluctuations without propagating them to the backbone stage.

Heterogeneous Parallel Allocator: Throughput Matching for Heterogeneous Parallelism

Encoders and LLM backbones have distinct compute characteristics. Encoders favor DP/TP; backbones need DP/TP/PP. A uniform strategy causes resource mismatch. FlowTrain's Heterogeneous Parallel Allocator reframes resource allocation as matching producer and consumer rates rather than relying on heuristic tuning. It first builds a throughput proxy model via small-scale profiling, then uses Mixed-Integer Linear Programming (MILP) within a given resource budget to select differentiated parallel configurations for encoder and backbone. The optimization objective prioritizes maximizing global throughput while further reducing producer-consumer mismatch. This mechanism incorporates cluster topology, communication costs, operator execution efficiency, and parallel strategy combinations to more accurately model real training throughput.

Dynamic Packing Scheduler: Runtime Cost-Aware Microbatch Construction

Multimodal samples are inherently variable-length: text length, image count, and visual token count differ. Even with similar total token counts, the longest subsequence significantly impacts attention compute. FlowTrain's Dynamic Packing Scheduler moves beyond static bucketing by token count or sample size. After encoding, it dynamically constructs microbatches based on the LLM side's actual compute cost. The scheduler uses a segment-tree-based best-fit packing to reduce padding waste, then buckets microbatches by real compute cost, and applies zig-zag dispatch to balance both pipeline-parallel (PP) and data-parallel (DP) dimensions, effectively handling compute fluctuations and load imbalance.

Engineering Implementation on MindSpeed MM

FlowTrain is implemented atop MindSpeed MM, which provides data engineering, model components, parallel training, and operator optimization for multimodal training. Additionally, FlowTrain leverages MC2 communication-computation overlap, embedding parallel communication tasks into compute windows to reduce critical-path communication exposure. Under related engineering optimizations, overall performance improves by approximately 10%.

Experimental Results

End-to-End Training Throughput Improvement

FlowTrain is evaluated on three model scales (VLM-6B, VLM-32B, VLM-72B) against Megatron-LM and DistTrain baselines. FlowTrain consistently achieves the highest throughput. On InfoVQA and an industrial dataset, FlowTrain delivers up to 1.7x and 1.4x throughput gains respectively. The performance gap widens with model scale: on InfoVQA, FlowTrain's throughput advantage over Megatron-LM grows from 1.3x at VLM-6B to 1.7x at VLM-72B. This is because baselines are constrained by batch-level synchronization barriers, lowering hardware utilization, whereas FlowTrain enables decoupled execution and coordinates heterogeneous modules via shared memory abstraction.

Significant Compute Efficiency Gains (MFU)

Under the same environment, FlowTrain substantially outperforms Megatron-LM in MFU across all configurations. On VLM-72B training, FlowTrain achieves 52.9% MFU on InfoVQA and 50.0% on the industrial dataset, while Megatron-LM reaches only 31.2% and 36.9% due to its monolithic design. DistTrain's decoupled orchestration mitigates some parallelism mismatch but its rigid pipeline execution keeps MFU below FlowTrain. FlowTrain allows the backbone to proceed without waiting for the full encoder batch and consumes balanced microbatches, boosting resource utilization.

Strong Horizontal Scalability

Scalability is tested by training VLM-32B on InfoVQA across varying cluster sizes, with pure LLM training as an upper-bound reference. FlowTrain maintains over 53% MFU at all scales, delivering 1.1-1.4x throughput improvement, demonstrating robustness. The average MFU gap to pure LLM training is 5.0 percentage points for FlowTrain, versus 18.7 for Megatron-LM and 12.1 for DistTrain, confirming that FlowTrain significantly narrows the efficiency gap between multimodal and single-modality training.

Real-World Deployment on JD E-Commerce Models

FlowTrain is deployed on JD's JoyAI-OxygenVLM 32B Dense and 35B-A3B MoE models. Training throughput increases by 1.7x and 1.4x respectively. General capabilities match same-size open-source models, while on the OxyEcomBench retail evaluation benchmark the models achieve industry SOTA. The framework supports 7+ retail scenarios including AI image search, product material quality inspection, outfit recommendation, item-description mismatch detection, AI hallucination detection, product violation patrol, and risk merchant patrol. AI image search shows a confident +2.24% UCVR improvement over the online model, with multiple core metrics showing confident positive gains.

Conclusion and Outlook

JD has heavily invested in building a domestic AI ecosystem. This work, a collaboration between JD Retail AI Infra and domestic chip teams, combines high-speed interconnect hardware, MemFabric memory pooling, and the MindSpeed MM software stack with multimodal training characteristics to produce a system-level solution that has been successfully applied in JD's e-commerce business, yielding significant efficiency gains.

The team observes that large model training is shifting from single-model optimization to system-level co-optimization. As model scale, data modality, and business complexity grow, training systems must deliver higher compute utilization, stronger heterogeneous parallelism, and more flexible data scheduling. JD Retail AI Infra will continue advancing training frameworks, scheduling algorithms, and system optimizations for real business scenarios, scaling FlowTrain to larger models, more complex data distributions, and additional industrial applications, while sharing practices with the developer community to foster an efficient, open multimodal training ecosystem.