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

MAKE IMPACT

Masahiro Tanaka , Canbing Li , Umesh Chand , Zafar Ali , Haiying Shen , Olatunji Ruwase · YOU? · · 2025 · Open Access · · DOI: https://doi.org/10.48550/arxiv.2504.09983

The increasing scale of deep learning models has led to the development of various parallelization strategies for distributed training across accelerators. For example, fully sharded approaches like DeepSpeed ZeRO-3 and FSDP partition the parameters of each layer across multiple GPUs and gather them through communication when needed. These methods rely on optimizations such as prefetching, which initiates communication early to overlap it with computation and reduce communication overhead, and unsharding, which retains as many parameters in their unsharded form as possible to reduce communication volume. Although the timing of prefetching should be adjusted in response to dynamic memory usage during execution, these systems lack the flexibility to control it, which limits the benefits of prefetching. Moreover, they cannot anticipate how memory usage will change after prefetching is applied, making it difficult to combine it effectively with other optimizations such as unsharding. We present DeepCompile, which compiles user-defined models into computation graphs and applies a sequence of profiling-guided optimization passes for distributed training. Taking dynamic memory usage into account, these passes flexibly insert, reorder, or remove operations to improve communication-computation overlap, reduce memory pressure, and coordinate multiple optimizations in a unified manner. To evaluate the effectiveness of this design, we implemented a fully sharded approach like ZeRO-3 and FSDP on top of DeepCompile, along with three optimizations: proactive prefetching, selective unsharding, and adaptive offloading. We evaluate DeepCompile on the training of Llama 3 70B and Mixtral 8x7B MoE models. DeepCompile achieves up to 1.28x and 1.54x performance improvements over ZeRO-3 and FSDP baselines, respectively, and up to a 7.01x throughput increase with limited GPU resources, using offloading.

Type

preprint

Language

en

Landing Page

http://arxiv.org/abs/2504.09983

PDF

https://arxiv.org/pdf/2504.09983

OA Status

green

OpenAlex ID

https://openalex.org/W4415158855

Raw OpenAlex JSON

OpenAlex ID

https://openalex.org/W4415158855

Canonical identifier for this work in OpenAlex

DOI

https://doi.org/10.48550/arxiv.2504.09983

Digital Object Identifier

Title

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

Work title

Type

preprint

OpenAlex work type

Language

en

Primary language

Publication year

2025

Year of publication

Publication date

2025-04-14

Full publication date if available

Authors

Masahiro Tanaka, Canbing Li, Umesh Chand, Zafar Ali, Haiying Shen, Olatunji Ruwase

List of authors in order

Landing page

https://arxiv.org/abs/2504.09983

Publisher landing page

PDF URL

https://arxiv.org/pdf/2504.09983

Direct link to full text PDF

Open access

Yes

Whether a free full text is available

OA status

green

Open access status per OpenAlex

OA URL

https://arxiv.org/pdf/2504.09983

Direct OA link when available

Cited by

0

Total citation count in OpenAlex

Full payload