Akash Mehra

2papers

2 Papers

93.4LGMay 26Code
Heterogeneous Parallelism for Multimodal Large Language Model Training

Yashaswi Karnati, Kamran Jafari, Akash Mehra et al.

Foundation model training is becoming multimodal, from post-training pipelines to large-scale pretraining. As modality coverage broadens, context windows grow, and encoder LLM scales diverge, a single LLM-centric TP/CP/PP/DP/EP layout increasingly limits throughput. This coupling forces encoders to inherit LLM-driven sharding and placement choices that can add communication, limit encoder parallelism, or constrain the LLM schedule; the mismatch is most pronounced at long contexts, where LLM context parallelism is needed for the fused multimodal sequence but encoder inputs remain bounded. We present heterogeneous parallelism for multimodal large language model training, an abstraction that lets modules in one end-to-end graph use independent layouts and rank placements, supporting colocated execution on shared GPUs and non-colocated execution on disjoint rank sets. The key challenge is preserving boundary tensor semantics across independent layouts: forward activations must be materialized for the destination layout, while backward gradients must be routed back to the source layout. We address this with boundary communicators that implement forward and backward layout transforms, plus scheduling extensions for both placement modes. We evaluate optimized homogeneous, colocated heterogeneous, and non-colocated heterogeneous configurations across multimodal workloads and GPU scales to characterize when added layout and placement freedom exposes a better operating point. Across this sweep, colocated heterogeneity improves TFLOPS/GPU by up to 49.3%, while non-colocated heterogeneity improves aggregate token throughput by up to 13.0% and TFLOPS/GPU by up to 9.6%. We validate loss convergence parity against homogeneous baselines and release the system as an open-source Megatron-LM extension.

HCApr 6, 2020
Leveraging GANs to Improve Continuous Path Keyboard Input Models

Akash Mehra, Jerome R. Bellegarda, Ojas Bapat et al.

Continuous path keyboard input has higher inherent ambiguity than standard tapping, because the path trace may exhibit not only local overshoots/undershoots (as in tapping) but also, depending on the user, substantial mid-path excursions. Deploying a robust solution thus requires a large amount of high-quality training data, which is difficult to collect/annotate. In this work, we address this challenge by using GANs to augment our training corpus with user-realistic synthetic data. Experiments show that, even though GAN-generated data does not capture all the characteristics of real user data, it still provides a substantial boost in accuracy at a 5:1 GAN-to-real ratio. GANs therefore inject more robustness in the model through greatly increased word coverage and path diversity.