Source-linked AI summary
Scaling Distributed Machine Learning with In-Network Aggregation
Amedeo Sapio, Marco Canini, Chen-Yu Ho, Jacob Nelson, Panos Kalnis, Changhoon Kim, Arvind Krishnamurthy, Masoud Moshref, Dan R. K. Ports, Peter Richtárik
TL;DR
Distributed ML training increasingly faces communication bottlenecks from synchronizing large model updates across workers. SwitchML addresses this by aggregating updates inside programmable switches, co-designed with end-host transport and ML frameworks. Across real-world DNN benchmarks, it reports training-throughput speedups of up to 5.5×.
Problem
Distributed training requires frequent synchronization of large model updates, while accelerator performance has increasingly outpaced network performance.
Method
SwitchML uses programmable-switch in-network aggregation, co-designed with end-host transport and ML frameworks, to synchronize workers’ updates.
Results
Up to 5.5× training-throughput speedup is reported across eight real-world DNN benchmarks, with SwitchML generally outperforming state-of-the-art collective libraries.
Takeaways & Limitations
In-network aggregation can reduce communication overhead and accelerate distributed DNN training at single-rack scale.
Takeaways & Limitations
The prototype targets single-rack scale, and its performance depends on neural-network architecture and physical network speed.
Abstract
from arXiv · showhide
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$\times$ for a number of real-world benchmark models.
1 Introduction
Distributed training is increasingly constrained by network communication, motivating SwitchML’s programmable-switch aggregation of model updates. The system co-designs switch processing, end-host transport, and ML frameworks to reduce synchronization traffic and accelerate training.
- Motivation: Distributed training is increasingly network-bound as accelerator performance has improved faster than cloud network performance.Synchronous model updates alternate with computation, making communication a growing performance constraint.
- Contribution: SwitchML aggregates workers’ model updates in programmable switches to reduce synchronization traffic and improve throughput, latency, and training time.The primitive executes aggregation in the network rather than exchanging all updates among workers.
- Design: SwitchML co-designs in-switch processing with an end-host transport layer and ML frameworks to perform aggregation at line rate.Parameter updates are divided into chunks processed through the switch pipeline.
- Design: Adaptive scaling converts floating-point updates to fixed-point values with negligible approximation loss.This addresses the integer-only arithmetic supported by programmable switch dataplanes.
- Deployment: SwitchML integrates with PyTorch and TensorFlow and initially targets rack-scale training with a centrally aggregating switch.Commercial switches can service up to 64 nodes at 100 Gbps or 256 nodes at 25 Gbps.
- Results: 5.5× is the reported maximum end-to-end training-performance improvement for popular DNN models.The communication microbenchmark is up to 2.9× faster than NCCL with RDMA and 9.1× faster than NCCL with TCP.
2 Network bottlenecks in ML training
Data-parallel synchronous SGD repeatedly aggregates large worker updates, creating communication bursts that can stall training. Accelerator improvements and limited communication overlap increasingly make networking the bottleneck, especially as models grow.
- Distributed SGD: Data-parallel SGD partitions data across workers, which independently compute updates that are aggregated and added to the model each iteration.Training repeats these mini-batch updates across multiple epochs until validation error is acceptable.
- Communication Cost: Model updates can reach hundreds of megabytes or gigabytes, and aggregation is required after every iteration.Increasing mini-batch size can hurt convergence, limiting an obvious way to reduce synchronization frequency.
- Communication Cost: Training alternates computation with intense communication bursts, and workers stall until model-update synchronization completes.Both parameter-server and all-reduce approaches incur this synchronization traffic.
- Network Bottleneck: 10× and 20× accelerator improvements for floating-point and mixed-precision calculations outpaced the 10× Ethernet improvement standardized over 8 years.This disparity shifts the distributed-training bottleneck from computation toward communication.
- Network Bottleneck: Communication overlap depends on DNN structure and is marginal for models with large initial layers.The study profiles eight DNNs on eight P100 workers at 10 Gbps and 100 Gbps to quantify communication and overlap.
- Network Bottleneck: α < 4 for half the workloads, indicating that network performance can constrain systems using faster GPUs with 100 Gbps networking.The scaling factor marks the hypothetical GPU speed at which communication saturates 100 Gbps once communication begins.
3 In-network aggregation
In-network aggregation sums model updates inside the network and distributes only the result, minimizing communication cost and latency compared with conventional exchanges. SwitchML evaluates this approach against ring all-reduce and gradient compression, showing performance benefits alongside model-quality trade-offs for aggressive compression.
- Communication advantages: 2|U| bytes is the minimum communication cost per worker for in-network aggregation, compared with 4|U| n−1 n for bandwidth-optimal all-reduce.The parameter-server approach can match 2|U| bytes but requires greater resource cost, potentially doubling machines and network bandwidth.
- Communication advantages: Sub-RTT latency follows because in-network aggregation avoids end-host processing required for aggregation.The contrasted all-reduce and parameter-server approaches cannot achieve this latency.
- Resource trade-offs: When parameter-server nodes are co-located with workers, effective bandwidth per node is halved, doubling latency.This is a resource and latency cost associated with matching in-network aggregation’s communication volume.
- Comparison with ring all-reduce: Practical SwitchML speedups exceed the synthetic analysis because real RAR implementations do not achieve theoretically optimal bandwidth utilization.System overheads and difficulty exploiting all available bandwidth contribute to the gap.
- Evaluation scope: The profiling environment uses two-generation-old NVIDIA P100 accelerators, while faster GPUs increase the relative impact of communication overheads.This bounds the interpretation of the profiling setup but motivates evaluation with faster GPUs.
- Comparison with gradient compression: INA outperforms QSGD across all workloads at both 64 and 256 levels, while Top-k underperforms INA at 10% compression.Top-k can outperform INA at 1%, but aggressive compression may require more iterations or harm accuracy.
- Comparison with gradient compression: Top-k at 1% compression outperforms INA but can fail to converge or reduce model quality, as illustrated by NCF’s 95.8% versus 93.6% hit rate at 20 epochs.The cited NCF comparison is without compression versus Top-k at 10%; Top-k at 1% fails to converge.
4 Design
SwitchML partitions aggregation across switches and hosts to overcome programmable-switch limits in computation, storage, numeric representation, and reliability. It streams integer aggregation through reusable slots, handles packet loss, and preserves training accuracy while reducing communication overhead.
- Combined switch-host architecture: SwitchML places integer aggregation in the programmable switch while end hosts manage reliability and complex computations.
- Pool-based streaming aggregation: Pool-based streaming aggregation processes limited vector chunks through reusable switch slots instead of storing an entire model update.Each slot aggregates k integers, allowing packet-level processing under switch memory and parsing constraints.
- Fault-tolerant protocols: SwitchML uses lightweight packet-loss recovery and traditional mechanisms for worker or network failures.The protocol must prevent lost updates from being discarded and retransmissions from applying updates twice.
- Switch-side aggregation protocol: The switch adds contributions from all workers and multicasts the completed aggregate, while end hosts perform the final division required for averaging.Addition is used because it is commutative and associative, whereas the switch cannot efficiently perform division.
- Dealing with packet loss: Shadow copies and seen bitmasks increase switch memory usage, although the required slot count is reported to remain below switch capacity.
- Quantized integer-based aggregation: Adaptive scaling converts floating-point gradients to 32-bit fixed-point values for switch aggregation with negligible approximation loss.The fixed-point representation is used only for gradient aggregation; weights and activations remain in 32-bit floating point.
- Dealing with floating-point numbers: SwitchML trains ResNet110 on CIFAR10 to similar accuracy and in a similar number of iterations as an unquantized network.With 8 workers and 64,000 steps, its final test accuracy is about 91–93%, similar to TensorFlow under the same setup.
5 Implementation
SwitchML is implemented as a collective library integrated with PyTorch and TensorFlow, with worker-side packet processing and P4 switch logic. The implementation supports multiple packet sizes, pipelines, and an RDMA-based path for higher throughput.
- Framework integration: SwitchML integrates with PyTorch DistributedDataParallel and TensorFlow via Horovod, using C++ worker and P4 switch components.The worker and switch components contain approximately 3,100 and 3,700 lines of code, respectively.
- Switch implementation: The P4 program distributes aggregation across ingress stages and implements flow control, retransmission, exponent calculation, and result replication.A single pipeline processes 64 elements per packet, while all four pipelines support 256-element packets.
- Packet sizing: Larger packets improve bandwidth efficiency by offsetting network framing overhead, motivating 64-element and 256-element configurations.The 256-element configuration recirculates packets through all four switch pipelines.
- RDMA support: A subset of RDMA is implemented in the switch so NICs can break large messages into packets and offload host packet processing.
- RDMA support: RDMA Write Immediate messages move data between the switch and GPUs, while client CPUs handle protocol operations and the switch writes aggregated results to destination buffers.
6 Evaluation
SwitchML is evaluated against collective libraries and parameter-server-like approaches using tensor microbenchmarks and eight DNN training benchmarks. It approaches line-rate aggregation and improves training performance most for network-bound models, while benefits vary by architecture and transport.
- Experimental setup: The evaluation uses tensor microbenchmarks and eight DNN benchmarks across PyTorch and TensorFlow, with NCCL over TCP or RDMA as the default baseline.Experiments use eight-worker configurations on 10 and 100 Gbps networks, with measurements of aggregation and training throughput.
- Tensor aggregation microbenchmarks: SwitchML’s RDMA variant with 256-value packets operates within 2% of maximum achievable goodput and maintains predictable ATE/s as workers increase.Using 64-value packets reduces performance, highlighting the importance of the multi-pipeline design.
- Tensor aggregation microbenchmarks: SwitchML outperforms parameter-server-like alternatives: Dedicated PS falls short despite twice the machines and network capacity, while Colocated PS reaches only half its performance.The comparison indicates lower performance for host-based aggregation under these configurations.
- SwitchML improves training speed: SwitchML accelerates batch processing especially for larger DNNs, reaching at most 2.1× over NCCL-RDMA and up to one order of magnitude over NCCL-TCP.Measured speedups can exceed emulated communication predictions because NCCL’s RAR implementation is not theoretically maximally efficient.
- SwitchML improves training speed: Benefits are substantial for some real-world DNNs but modest or absent for models with small tensors or models that are not network-bound at 100 Gbps.BERT and NCF have lower speedups, while UGATIT, SSD, and ResNet are not network-bound in this setting.
- SwitchML improves training speed: End-to-end training speedups over NCCL range from 1.13–2.27× with RDMA and 2.05–5.55× with TCP for four network-bottlenecked models.These measurements use the DPDK implementation with 256-value packets and do not directly correspond to the batch-processing results.
- Overheads: At loss probabilities of 0.1% or higher, SwitchML completes tensor aggregation significantly faster than Gloo or NCCL, while 0.01% loss minimally affects TAT.At 1% loss, unevenly affected slots cause slowdown because work-stealing is not implemented.
- Overheads: Tensor scaling and type-conversion overheads are negligible, and using float16 doubles performance.The implementation uses x86 SSE/AVX instructions and GPU offload.
7 Extensions
SwitchML can extend beyond a single rack through hierarchical composition of switch instances and can support multiple pipelines and concurrent jobs. Its current scope assumes synchronous SGD, unencrypted traffic, and dedicated networks without explicit congestion control.
- Scaling beyond a rack: Multiple racks can be supported by hierarchically composing top-of-rack aggregation switches, although the authors lack a sufficiently large testbed to evaluate this design.Each rack aggregates worker updates before forwarding them upward.
- Scaling beyond a rack: With p-port pipelines, a switch acts as u = ⌈d/p⌉ virtual switches, each aggregating traffic from a subset of downstream ports.For the described switches, p = 16.
- Scaling beyond a rack: Hierarchical composition is bandwidth-optimal for all-to-all communication, reducing bandwidth cost proportionality from n workers to u upstream ports and enabling p:1 oversubscription.This supports large clusters with relatively shallow hierarchies.
- Scaling beyond a rack: Bitmap and shadow-copy mechanisms allow packet-loss recovery to work across multiple racks by forwarding updated values toward upper-layer switches.Retransmissions are recognized by switches that already processed the original packet.
- Congestion control: SwitchML has no explicit congestion-control algorithm; self-clocking provides rudimentary control that is sufficient for dedicated networks but may require a general congestion-control scheme elsewhere.A congested worker can reduce the sending rate of all workers through self-clocking.
- Multi-job (tenancy): Each concurrent job requires a separate aggregator pool, while one reduction uses less than 10% of switch capabilities.Multiple independent pipelines can provide additional resources for multi-job operation.
- Scope: The described deployment does not accommodate encrypted traffic and targets synchronous SGD for reproducible cluster training.These are explicit scope choices of the presented system.
8 Related work
SwitchML builds on prior in-network aggregation and programmable-data-plane research while targeting complete, high-throughput ML aggregation on Ethernet. Its distinctions include programmable hardware, packet recovery for unreliable Ethernet, and support for large models under constrained switch memory.
- In-network aggregation: Prior systems placed application-specific aggregation in switch-attached middleboxes, direct-connect servers, or rack-scale parameter servers for partition-aggregate and MapReduce workloads.These approaches establish related uses of network-side aggregation outside SwitchML’s target design.
- In-network aggregation: DAIET is the closest prior work, but it presents no complete design for ML applications and only a proof-of-concept on a P4 emulator.The authors state that its viability on a real switch is unclear.
- In-network aggregation: SHARP uses fixed-function FPUs in InfiniBand, whereas SwitchML uses programmable hardware on standard Ethernet with a new packet-recovery protocol.SHARP can leverage link-layer flow control and lossless guarantees; SwitchML uses unreliable connections.
- In-network aggregation: Programmability lets operators repurpose aggregation ALUs for traditional networking and evolve protocols without replacing fixed-function hardware.The paper cites experimentation with new floating-point representations and sparse-vector protocols.
- In-network aggregation: iSwitch supports native floating point at lower bandwidth and stores an entire gradient vector, limiting scalability for large DNN models.SwitchML instead targets large models and high throughput under limited on-chip memory.
- Accelerating DNN training: The broader literature includes data and model parallelism, combinations of both, and automated methods for improving DNN training performance.These approaches address training acceleration through system or algorithmic changes rather than SwitchML’s specific aggregation primitive.
9 Conclusion
SwitchML reduces communication overhead by aggregating model updates in the network and reports up to 5.5× training-throughput speedup on eight real-world DNN benchmarks. Its recovery protocol handles packet loss through retransmissions and shadow-copy state.
- Conclusion: SwitchML minimizes communication overhead by synchronizing distributed workers’ model updates through in-network aggregation at single-rack scale.The evaluation covers eight real-world DNN benchmarks on 10 and 100 Gbps GPU clusters.
- Fault tolerance: The packet-loss example follows workers sending slot updates, retransmitting lost packets, and the switch completing aggregation once the missing contribution arrives.The switch then creates a shadow copy and sends response packets to workers.
- Fault tolerance: A lost downstream response causes the affected worker to retransmit its update, after which the switch safely reuses the slot for the next offset.The slot roles flip once the result is confirmed received by every worker.
B Implementation details
SwitchML’s implementation co-designs switch processing, host transport, RDMA variants, and ML-framework integration to aggregate distributed updates efficiently under programmable-switch constraints.
- Switch component: The switch design maximizes computation per packet while managing tight pipeline-stage and memory constraints.Execution resources, especially pipeline stages, are tighter constraints than storage in the SwitchML data plane.
- Switch component: SwitchML confines aggregation to the ingress pipeline, then duplicates completed results for multicast.Using both ingress and egress pipelines required recirculation, added dependencies, and limited packet processing.
- ML framework integration: The implementation exposes a synchronous all-reduce interface while virtually treating ordered tensor reductions as one reduction.Results are steered back to the requesting thread, and multiple threads may call all-reduce when tensor order matches across workers.
- Host processing: A single CPU core sustains reduction at line rate on 10 Gbps, while higher rates use multiple cores, NIC queues, and disjoint slot sets.Flow Director uniformly distributes incoming traffic across receive queues, and each core processes batches independently.
- RDMA implementation: SwitchML RDMA uses UC and adapts packet handling, sequencing, retransmission, and GPU-memory writes for collective communication.The RDMA variant sends multi-slot messages, allocates queue pairs per in-flight message, and does not require lossless Ethernet.
C Tuning the pool size
SwitchML tunes pool size to sustain line rate without violating slot-use correctness, selecting practical power-of-two sizes for different link rates.
- Sizing criterion: Pool size s should make s·b match the bandwidth-delay product to sustain line-rate transmission.The delay includes end-to-end and end-host processing time; larger initial windows unnecessarily increase worker queuing.
- Sizing criterion: Correctness requires that two in-flight packets from one worker never use the same slot, yielding s = ⌈BDP/b⌉ as the optimal size.This lower bound simultaneously supports line rate and prevents same-worker slot reuse across phases.
- Evaluation context: The evaluation compares model and task metrics across the DDL benchmarks, including compression-related convergence outcomes.Table 4 notes that NCF at Top-1% did not converge and that BERT reports the median of six fine-tuning runs.
- Practical configuration: SwitchML rounds the formula up to the next power of two for batched DPDK operations and uses pool sizes 128 at 10 Gbps and 512 at 100 Gbps.These choices occupy 256 KB and 1 MB of switch register space, respectively.
D Compression affects convergence
Compression can alter convergence behavior and model quality in task-dependent ways, so loss traces alone do not reliably predict accuracy.
- Convergence behavior: Loss and accuracy do not necessarily correlate well under the evaluated compression schemes.For SSD, methods have similar loss traces but an obvious accuracy gap.
- Convergence behavior: NCF with Top-k compression at 1% does not converge, although its accuracy can still increase.The result illustrates that convergence traces and final accuracy may diverge under aggressive compression.
E Model quantization
SwitchML quantizes floating-point updates into fixed-point integers for switch aggregation, using adaptive scaling to control overflow and approximation error. The formal analysis bounds this error and shows it is negligible in typical applications.
- Motivation: Programmable Ethernet switches lack native floating-point operations, motivating integer-based aggregation and conversion strategies.Restricted 16-bit floating-point conversion is possible on Tofino using lookup tables and ALUs, but conversion placement remains a design choice.
- Related work: Existing compression work studies how quantization, dithering, and sparsification affect iterations needed to meet convergence criteria.These methods provide context for evaluating the implications of SwitchML’s quantization strategy.
- Quantization mechanism: Workers divide updates into packet-sized blocks, scale each block adaptively, round entries to integers without overflow, and send them for switch aggregation.The aggregated integer update is scaled back on workers before updating the model.
- Example: Using f = 100 in the example produces an aggregated integer of 579 and the exact model update 5.79 after rescaling.Integer-only switch aggregation can therefore reproduce the unconverted update in this example.
- Example: Using f = 10 instead produces update 5.8 rather than 5.79, incurring error 0.01.The example demonstrates the rounding trade-off from choosing a smaller scaling factor.
- Quantization mechanism: Scaling factor f must be large enough to reduce rounding error but bounded so individual and aggregated integers do not overflow.The strategy chooses f per block while enforcing representability constraints for both quantized entries and their sum.
- Formal analysis: Theorem 1 bounds aggregation error, with the bound improving as f increases and the worker count n decreases.The analysis assumes bounded model parameters and uses this condition to control quantization and overflow.
- Formal analysis: Under bounded parameters, Theorem 2 establishes no overflow, and choosing f = 2^31−n/(nB) bounds the exact-update difference by n^2B/2^31−n.For typical applications where n^2B ≪ 2^31, the introduced error is negligible.
F Encrypted traffic
SwitchML does not address encrypted datacenter traffic, but the paper outlines possible alternatives because in-switch aggregation uses simple integer summation.
- Encryption of datacenter traffic is explicitly out of scope for SwitchML.The authors note that encryption is commonly performed at the NIC level.
- One proposed alternative is in-line decryption and re-encryption at switches using hardware accelerators, although this is likely costly.
- Partially homomorphic cryptosystems could let switches aggregate encrypted updates because E(x)·E(y) = E(x+y).The paper identifies integer summation as the simple switch operation compatible with this approach.
- The encrypted-data approach would require customizing end-host encryption and remains future work to validate.