Tula: Optimizing Time, Cost, and Generalization in Distributed Large-Batch Training

TL;DR

Tula is an online service that optimizes large-batch distributed training by balancing time, cost, and model quality, using performance prediction to find the optimal batch-size for various models and resources.

Contribution

It introduces a novel method combining system modeling and statistical prediction to automatically optimize large-batch training, addressing efficiency and generalization issues.

Findings

Predicts training time and cost with 7.5-14% error

Achieves up to 20x speedup in training

Improves test accuracy by 9% on average

Abstract

Distributed training increases the number of batches processed per iteration either by scaling-out (adding more nodes) or scaling-up (increasing the batch-size). However, the largest configuration does not necessarily yield the best performance. Horizontal scaling introduces additional communication overhead, while vertical scaling is constrained by computation cost and device memory limits. Thus, simply increasing the batch-size leads to diminishing returns: training time and cost decrease initially but eventually plateaus, creating a knee-point in the time/cost versus batch-size pareto curve. The optimal batch-size therefore depends on the underlying model, data and available compute resources. Large batches also suffer from worse model quality due to the well-known generalization gap. In this paper, we present Tula, an online service that automatically optimizes time, cost, and convergence quality for large-batch training of convolutional models. It combines parallel-systems modeling with statistical performance prediction to identify the optimal batch-size. Tula predicts training time and cost within 7.5-14% error across multiple models, and achieves up to 20x overall speedup and improves test accuracy by 9% on average over standard large-batch training on various vision tasks, thus successfully mitigating the generalization gap and accelerating training at the same time.