training and inference optimization frameworks

Existing large-scale model training systems struggle to flexibly compose diverse parallelization strategies, often relying on manual expert tuning and lacking generality. This work proposes a programmable distributed training system that enables users to declaratively specify composite parallelism strategies—such as data, pipeline, and expert parallelism—through model annotations and scheduling directives. These specifications are compiled via a unified intermediate representation (IR) into device-level execution plans, fully decoupling strategy definition from runtime execution over a global compute-communication DAG. The system is the first to support automatic compilation of user-defined composite strategies, matching the performance of established approaches like ZeRO while significantly improving both performance and memory efficiency in complex scenarios such as DeepSeek-V3’s DualPipe.

Static parallelization strategies in large-scale distributed neural network training suffer from poor resource adaptability, leading to efficiency bottlenecks. To address this, we systematically evaluate the performance boundaries of data parallelism, model parallelism, and hybrid parallelism, and propose a dynamic, topology- and resource-aware scheduling algorithm. This algorithm enables online switching of parallelization strategies within a hybrid parallel framework during training, jointly optimizing communication overhead, computational load balance, and memory constraints. On the CIFAR-100 image classification benchmark, hybrid parallelism achieves a 3.2× speedup over single-GPU training with no accuracy degradation; integrating our adaptive scheduler further improves end-to-end training efficiency by 18%. To the best of our knowledge, this work is the first to incorporate dynamic strategy switching into the hybrid parallel training pipeline, establishing a scalable new paradigm for efficient large-model training in heterogeneous resource environments.

To address the high communication overhead and limited compute-communication overlap in tensor model parallelism (TMP) for large language models, this paper proposes Oases, a dependency-aware, fine-grained TMP auto-optimization framework. Our method introduces: (1) operator-level fine-grained training scheduling; (2) communication-computation overlap–aware modeling and planning; (3) the first automated TMP strategy search algorithm supporting dependency constraints; and (4) end-to-end runtime optimization for TMP. Evaluated across multiple models (e.g., LLaMA, BERT) and hardware platforms (A100/H100), Oases achieves 1.01–1.48× speedup over state-of-the-art approaches and up to 1.9× improvement over Megatron-LM. These results demonstrate significant gains in training efficiency for large-scale foundation models.

Modern large-scale distributed training faces sharply diminishing returns in hardware scaling: as GPU counts reach thousands, communication overhead dominates performance bottlenecks, rendering conventional parallelism strategies—data, tensor, and pipeline parallelism—suboptimal. Method: Leveraging real-world LLM training workloads, this project establishes an empirical analytical framework spanning diverse model scales, hardware configurations, and parallelization strategies. It quantifies the nonlinear relationship between accelerator count and performance gain, precisely identifying critical inflection points across model, data, and compute scaling dimensions. Contribution/Results: We discover that low-communication “suboptimal” strategies become optimal at extreme scale; we empirically determine hardware selection criteria, cluster topology requirements, and optimal parallelism combinations for training billion-parameter models. Our findings provide actionable, deployment-ready optimization guidelines for trillion-parameter LLM training infrastructures.

In deep learning training, repetitive hyperparameter tuning and neural architecture search cause frequent redundant data loading, leading to CPU bottlenecks, low GPU utilization, and resource waste. To address this, we propose the first cross-process shared data loading framework leveraging GPU-to-GPU direct interconnects (NVLink/PCIe RDMA), enabling multi-task reuse of a single preprocessing pipeline. Our approach is zero-intrusive—requiring no model code modification—and hardware- and framework-agnostic. By employing memory-shared mapping and a lightweight coordination protocol, it eliminates duplicate preprocessing operations. Experiments demonstrate up to 100% higher training throughput, a 50% reduction in cloud instance CPU resource demand, and corresponding 50% cost savings. Our method consistently outperforms state-of-the-art systems including CoorDL and Joader.

Tensor accelerators (e.g., NVIDIA Tensor Cores) are increasingly prevalent in CPUs and GPUs, yet their programmability remains limited: existing kernel libraries target only traditional ML and scientific computing workloads, failing to support non-ML linear matrix transform workloads such as image processing. This paper proposes an equivalence-saturation–based flexible tensor instruction selection mechanism, enabling general-purpose, schedulable compilation for tensor hardware. Integrating with the Halide domain-specific language and compiler, our approach retains full compatibility with existing scheduling primitives while substantially broadening the programmability of tensor accelerators. Evaluated on an NVIDIA RTX 4070, our framework achieves a 6.1× speedup on image processing pipelines—including downsampling—demonstrating, for the first time, systematic performance acceleration of tensor hardware in non-ML domains.

This work addresses the substantial accelerator memory consumption of model parameters, gradients, and optimizer states in standard mixed-precision training, which hinders the scalability of large models. The authors propose a memory-efficient training method that significantly reduces quantization error in 8-bit optimizer states through compact master weight partitioning and a novel compression-expansion function. By integrating 16-bit gradients, an improved weight splitting strategy, and a gradient checkpointing mechanism, the approach remains compatible with mainstream optimizers such as SGD, AdamW, and Lion. The method reduces AdamW’s per-parameter memory footprint from 16 bytes to 7 bytes (or 5 bytes when gradients are released) and halves model checkpoint size, achieving lossless training quality across multiple vision and language benchmark tasks.

This work addresses the GPU memory bottleneck in Transformer models caused by high parameter and activation memory demands during training and inference. The authors propose a novel parallelism strategy that integrates tensor parallelism (TP) and sequence parallelism (SP) along the same device axis, enabling each device to simultaneously shard both model weights and input sequences. By leveraging broadcast-based weight sharding with key-value exchange in attention layers and ring-based weight passing with local accumulation in gated MLPs, the method achieves dual compression of both parameter and activation memory. This approach significantly reduces per-device memory consumption, outperforming conventional TP, SP, and their hybrid variants. It demonstrates superior hardware adaptability and scaling efficiency under long-context and memory-constrained settings, while seamlessly integrating with pipeline and expert parallelism.

Existing context parallelism approaches suffer from load imbalance, redundant communication, and suboptimal computational efficiency. This work proposes a communication-efficient and load-balanced context parallel training framework featuring three key innovations: a Whole-Doc document-level sharding strategy integrated with a Per-Doc hybrid sharding mechanism, a shard-aware communication protocol that effectively eliminates redundant transmission of KV caches, and a heuristic search-based near-optimal sharding planning algorithm. Experimental results demonstrate that the proposed method achieves up to a 1.63× speedup over state-of-the-art context parallelism techniques across multiple datasets.

This work addresses the lack of transparent, scalable, and deeply PyTorch-integrated open-source tools for post-training large language models, which hinders research iteration and deployment efficiency. We propose a native PyTorch-based, modular post-training framework centered on the principle of “hackability,” offering composable model builders, training recipes, and a distributed training stack that support diverse fine-tuning strategies and hardware configurations. While maintaining high performance and memory efficiency, the framework significantly enhances code transparency and research flexibility. Empirical evaluations demonstrate that it matches or even surpasses mainstream tools such as Axolotl and Unsloth across multiple post-training scenarios, thereby facilitating efficient and reproducible scientific exploration.