FlexiTrain

Abstract: Modern machine learning workloads involve training massive models with billions of parameters across tens of thousands of devices. Failures are inevitable at this scale due to the sheer number of devices and duration of training. Existing methods handle failures by periodically checkpointing the model state at various point throughout training. However, this approach to fault tolerance incurs additional overhead that can significantly reduce training throughput. We present FlexiTrain, a checkpoint-free framework for fault-tolerant, distributed model training. FlexiTrain gracefully handles machine failures by discarding the failed machines, redistributing their training data, and rescaling gradient updates among the remaining nodes. We find that FlexiTrain can improve training throughput by 20% with no observable loss on training quality. We empirically validate our approach by training a vision model and a large language model with simulated faults.

Paper