[Paper Review] Scaling Distributed Machine Learning with In-Network Aggregation

Amedeo Sapio, Marco Canini+8 more|arXiv (Cornell University)|Feb 22, 2019

Stochastic Gradient Optimization Techniques73 references119 citations

TL;DR

SwitchML introduces in-network aggregation using programmable switches to sum worker updates, reducing data traffic and achieving up to 5.5x end-to-end speedups for distributed training.

ABSTRACT

Training machine learning models in parallel is an increasingly important workload. We accelerate distributed parallel training by designing a communication primitive that uses a programmable switch dataplane to execute a key step of the training process. Our approach, SwitchML, reduces the volume of exchanged data by aggregating the model updates from multiple workers in the network. We co-design the switch processing with the end-host protocols and ML frameworks to provide an efficient solution that speeds up training by up to 5.5$ imes$ for a number of real-world benchmark models.

Motivation & Objective

• Motivate the shift of training bottlenecks from computation to communication in large-scale distributed ML.
• Propose a practical in-network aggregation primitive implemented in programmable switches.
• Co-design switch processing with end-host protocols and ML frameworks to accelerate SGD-based training.
• Demonstrate end-to-end training speedups and compare against traditional all-reduce and parameter-server approaches.

Proposed method

• Decompose model updates into chunks that fit switch pipeline limits for streaming aggregation.
• Implement a pool-based, in-switch integer aggregation to sum updates, with floating-point values mapped to fixed-point via adaptive scaling.
• Coordinate end-host aggregation protocols to drive synchronized, slot-based utilization of switch aggregators.
• Provide fault-tolerant mechanisms for packet loss using lightweight host-driven retransmission and switch-side shadow copies.
• Integrate SwitchML with existing ML frameworks (e.g., PyTorch, TensorFlow) and Horovod for practical deployment.

Experimental results

Research questions

• RQ1 Can in-network aggregation inside programmable switches reduce communication overhead for synchronized distributed SGD?
• RQ2 What are the design trade-offs (precision, memory, reliability) to implement gradient aggregation in switch dataplanes?
• RQ3 How does SwitchML compare to traditional all-reduce and parameter-server approaches under varying network speeds and model sizes?
• RQ4 What is the end-to-end training performance impact when integrating SwitchML into real ML pipelines?

Key findings

• SwitchML achieves up to 5.5x end-to-end training speedups for real-world DNN benchmarks.
• In-network aggregation consistently outperforms ring-all-reduce across tested workloads, with up to 80% faster at 10 Gbps and up to 67% faster at 100 Gbps for network-bound models in synthetic benchmarks.
• Compared to NCCL with RDMA, SwitchML-based aggregation is up to 2.9x faster, and up to 9.1x faster than NCCL with TCP in the reported microbenchmark.
• Gradient compression (Top-k, QSGD) is not universally beneficial; in many cases, the compression overhead offsets data reduction, and SwitchML often dominates when hardware and workload conditions limit compression gains.
• The design leverages a pool-based streaming aggregation and a two-pool shadow-copy mechanism to tolerate packet loss with minimal switch complexity.