DeepCompile: A Compiler-Driven Approach to Optimizing Distributed Deep Learning Training

The rapid growth of deep learning models has increased the demand for efficient distributed training strategies. Fully sharded approaches like ZeRO-3 and FSDP partition model parameters across GPUs and apply optimizations such as prefetching and unsharding to reduce communication overhead. However, these systems lack fine-grained control over memory and communication scheduling, making it difficult to balance computation--communication overlap with memory requirements. Coordinating multiple optimizations such as prefetching and unsharding is also difficult, since their effects on memory usage can influence each other. To tackle these challenges, we propose DeepCompile, a compiler-based optimization framework for distributed training. DeepCompile transforms user-defined models into computation graphs and applies a series of profiling-guided optimization passes, each modifying the graph based on profiling information such as execution time and memory usage. This design allows each pass to flexibly insert, reorder, or remove operations such as all-gather and memory allocation, improving communication--computation overlap and reducing memory pressure. Each pass can access updated profiling feedback from earlier passes, enabling coordinated optimizations. We further enhance DeepCompile by three additional optimizations: proactive prefetching, selective unsharding, and adaptive offloading. Our evaluation shows that DeepCompile achieves up to 1.28$\times$ and 1.54$\times$ speedups over ZeRO-3 and FSDP baselines, respectively, and up to a 7.01$\times$ throughput increase in settings with limited GPU resources using offloading.