Automated Graph Partitiong for DL Workload Analysis

Abstract

AI models are rapidly growing in size (hitting trillion parameter mark). We may require days and weeks to train such large models efficiently. Performance of AI workloads at such scale depends on the choice of parallelization techniques. One solution is Distributed learning that uses two primary approaches to split the workloads among the available GPUs, ie Model Parallel and Data Parallel. There are many techniques available like Gpipe, PipeDream that uses combination of distributed learning techniques, but these techniques have some limitations. Architects need to study the structure of that DL workload to split it. This may lead to sub-optimal partition of workload where most of system resource are employed to communicate the intermediate results and rather than computation. In such scenario selection of the optimal values of MP and DP to get the best split is arduous as exercising all combinations to arrive at best configuration is difficult. Automated Graph Partitioning eliminates the need of studying individual workloads and figuring out which layers are required to be split and what is the optimal split values for the given workload. This thesis aims to resolve the above defined problem by computing the most optimal MP and DP value for any given workload and also provide the optimal way to split that workload