Search papers, labs, and topics across Lattice.
The paper introduces Entrain, a distributed training framework designed to mitigate workload imbalance in multimodal LLM training caused by data heterogeneity and variability across modalities and samples. Entrain uses a static model-parallel configuration optimized for macroscopic batches, coupled with a hierarchical microbatch assignment algorithm to stabilize variability at the microbatch level. Experiments demonstrate that Entrain reduces workload variability by up to 10.6x and improves end-to-end training throughput by up to 1.40x compared to existing methods.
Static model parallelism can outperform dynamic approaches for load balancing in distributed multimodal LLM training by shifting the profiling paradigm to macroscopic batches.
Multimodal LLM datasets are inherently heterogeneous, with significant data variability. Although each modality exhibits independent variability, sample-level entanglement makes it difficult to balance workloads across both modalities and batches. We present Entrain, a distributed MLLM training framework that addresses both heterogeneity and variability in multimodal training workloads. Entrain challenges the intuition that dynamic data variability requires dynamic model parallelism by shifting the profiling paradigm from micro-level samples to macroscopic batches. We prove that a single, static model-parallel configuration suffices for optimal load balancing under this paradigm. At the microscopic scale, Entrain introduces a hierarchical microbatch assignment algorithm that defers excess workload within each iteration to stabilize variability across microbatches. Evaluations show that Entrain reduces workload variability across microbatches by up to 10.6$\times$, improving end-to-end training throughput by up to 1.40$\times$ over existing baselines.
