Rethinking Machine Learning Collective Communication as a Multi-Commodity Flow Problem

TL;DR

This paper introduces TECCL, a new method for ML collective communication scheduling that models the problem as a multi-commodity flow, leading to faster and more efficient schedules on large GPU topologies.

Contribution

It redefines ML collective communication scheduling as a multi-commodity flow problem and proposes TECCL, which outperforms existing methods in speed and efficiency on large-scale GPU topologies.

Findings

TECCL produces faster schedules for ML collectives.

TECCL reduces communication bytes compared to previous methods.

TECCL scales better on larger GPU topologies.

Abstract

We show communication schedulers' recent work proposed for ML collectives does not scale to the increasing problem sizes that arise from training larger models. These works also often produce suboptimal schedules. We make a connection with similar problems in traffic engineering and propose a new method, TECCL, that finds better quality schedules (e.g., finishes collectives faster and/or while sending fewer bytes) and does so more quickly on larger topologies. We present results on many different GPU topologies that show substantial improvement over the state-of-the-art.