PaperlayerView exact revision on arXiv

Source-linked AI summary

TACCL: Guiding Collective Algorithm Synthesis using Communication Sketches

Aashaka Shah, Vijay Chidambaram, Meghan Cowan, Saeed Maleki, Madan Musuvathi, Todd Mytkowicz, Jacob Nelson, Olli Saarikivi, Rachee Singh

arXiv:2111.04867v4cs.DCcs.LG

TL;DR

Large distributed ML workloads need efficient collective communication, but synthesizing algorithms for heterogeneous multi-node GPU topologies is difficult and computationally expensive. TACCL combines communication sketches with a scalable synthesis procedure, producing collectives that outperform NCCL by up to 6.7× and accelerate end-to-end training by 11%–2.4×.

  • Problem

    Efficient collective communication is needed because inter-GPU communication can leave GPUs idle, while exhaustive synthesis over heterogeneous multi-node topologies is computationally prohibitive.

  • Method

    TACCL uses communication sketches to constrain synthesis and combines relaxed routing, heuristic ordering, and exact scheduling in a scalable encoding.

  • Results

    Up to 6.7× faster than NCCL, TACCL synthesized efficient algorithms across ALLGATHER, ALLTOALL, and ALLREDUCE and improved end-to-end training by 11%–2.4%.

  • Takeaways & Limitations

    TACCL provides topology- and input-size-aware collective algorithms for multi-node distributed ML training and inference.

  • Takeaways & Limitations

    TACCL remains based on an NP-hard synthesis problem whose complexity grows rapidly with scale, and its logical topology does not capture connection-count effects on switch congestion.

Abstract

from arXiv · show

Machine learning models are increasingly being trained across multiple GPUs and servers. In this setting, data is transferred between GPUs using communication collectives such as AlltoAll and AllReduce, which can become a significant bottleneck in training large models. Thus, it is important to use efficient algorithms for collective communication. We develop TACCL, a tool that enables algorithm designers to guide a synthesizer into automatically generating algorithms for a given hardware configuration and communication collective. TACCL uses a novel communication sketch abstraction to get crucial information from the designer to significantly reduce the search space and guide the synthesizer towards better algorithms. TACCL also uses a novel encoding of the problem that allows it to scale beyond single-node topologies. We use TACCL to synthesize algorithms for three collectives and two hardware topologies: DGX-2 and NDv2. We demonstrate that the algorithms synthesized by TACCL outperform the Nvidia Collective Communication Library (NCCL) by up to 6.7x. We also show that TACCL can speed up end-to-end training of Transformer-XL and BERT models by 11%--2.3x for different batch sizes.

1 Introduction

TACCL addresses the difficulty of designing efficient collective communication algorithms for heterogeneous, multi-node GPU systems by combining communication sketches with scalable synthesis. It generates algorithms that outperform NCCL across several collectives and improve end-to-end training performance.

  • Motivation: Growing model sizes require distributed training across multiple GPUs and servers, making inter-GPU communication increasingly important.Large models use data, model, pipeline, and expert parallelism across multi-GPU systems.
  • Challenges: Collective algorithms must balance latency and bandwidth while accounting for heterogeneous GPU connectivity and topology-specific constraints.The search space over routing and scheduling is computationally prohibitive, especially beyond single-node systems.
  • Approach: Communication sketches let designers provide high-level algorithmic structure while TACCL synthesizes routing and scheduling details.The abstraction reduces the search space analogously to program sketching, where a solver fills in unspecified details.
  • Approach: TACCL uses a novel MILP encoding that first determines routing in a bandwidth-relaxed problem before applying ordering heuristics and subsequent scheduling.This approach builds on solver-based synthesis while targeting larger cluster topologies.
  • Results: 6.7× faster than NCCL, TACCL’s synthesized collectives improve communication performance across DGX-2 and NDv2 systems.Reported gains include up to 25% for large-input ALLGATHER and up to 66% for ALLTOALL.
  • Results: 11%–2× faster end-to-end training was achieved for Transformer-XL, alongside a 17% speed-up for a mixture-of-experts model.The NCCL replacement required only a two-line PyTorch code change.

2 Background and Motivation

Distributed ML workloads rely on collectives whose algorithms must efficiently move or combine data across heterogeneous GPU topologies. TACCL reduces the otherwise intractable synthesis search space through lightweight designer guidance and a staged synthesis procedure.

  • Collective communication: ALLGATHER, ALLTOALL, and ALLREDUCE support common distributed ML communication patterns such as gathering data, shuffling chunks, and aggregating values.Different parallelism strategies select different collectives for gradients, expert routing, and embedding lookups.
  • Existing approaches: NCCL relies largely on predefined, topology-specific templates, including rings, pairwise transfers, and hardcoded algorithm choices.These templates may be topology-agnostic or based on previously profiled latency and bandwidth assumptions.
  • Synthesis scalability: SCCL failed to synthesize nearly all tested multi-node collective-topology combinations within 24 hours.The only successful case was a latency-optimal ALLGATHER algorithm for two Azure NDv2 nodes.
  • Search-space reduction: Communication sketches provide low-effort high-level constraints that discard large portions of the routing and scheduling search space.They avoid requiring designers to handcraft complete algorithms while making synthesis more tractable.
  • Synthesis technique: TACCL first solves a relaxed routing problem, then heuristically orders chunks on shared links according to bandwidth constraints.This staged procedure makes the NP-hard synthesis problem more tractable even after sketch-based reduction.

3 Communication Sketches

Communication sketches encode designer knowledge about topology, connectivity, congestion, and symmetry without requiring complete algorithm implementations. TACCL uses these constraints to guide synthesis while retaining freedom to choose detailed routes and schedules.

  • Sketch design: A communication sketch asks designers for high-level routing guidance while leaving tedious scheduling details to the synthesizer.Routing is emphasized because its global properties are easier to express and scheduling depends on route intersections.
  • Logical topology: The logical topology selects a subset of the physical topology and can exclude links that designers prefer the algorithm to avoid.It preserves the physical cost model while abstracting away NICs and switches into GPU-to-GPU links.
  • Topology-specific guidance: For NDv2 systems, sketches can favor NVLink within a node while retaining carefully selected PCIe paths for distributed communication.PCIe selection must account for host-memory transfers and oversubscribed links.
  • Limitation: TACCL does not model the bandwidth drop caused by increasing switch connections, although switch-hyperedges partially control that connectivity.The limitation is most relevant when connection count affects queuing latency.
  • Switch-hyperedges: Switch-hyperedges let TACCL control how many logical connections use a switch through maximize, minimize, or free-choice policies.These policies adapt connectivity to congestion and input-size conditions.
  • Symmetry: Designers can enforce symmetry by supplying topology-and-collective automorphisms that require corresponding chunk transfers to follow the same pattern.This supports structures such as rings and hierarchical symmetry across machines.

4 Physical Topologies of GPU systems

TACCL profiles heterogeneous GPU interconnects and uses those measurements to infer physical topologies and α-β communication costs. This addresses undocumented or obscured connectivity, especially in Azure NDv2 systems.

  • Motivation: TACCL includes a physical topology profiler because heterogeneous link characteristics and even cloud topologies may be sparsely documented.The profiler measures performance characteristics needed before users can sketch algorithms for a target system.
  • Model Scope: TACCL adopts the α-β model for tractability, while communication sketches expose controls for concurrency and congestion that the model omits.The model simplifies encoding, but topology-aware sketch parameters provide additional control over excessive concurrency.
  • Interconnects: Modern GPU systems combine PCIe, NVLink, and InfiniBand interconnects with different bandwidth and connectivity properties.PCIe provides limited shared bandwidth, while NVLink offers dedicated intra-node GPU connections.
  • α-β Cost Model: TACCL estimates α and β by timing sequential and concurrent peer-to-peer transfers, then solving for latency and inverse bandwidth.Sequential transfers estimate n·(α+β·s), while concurrent transfers attribute time to α+n·β·s.
  • α-β Cost Model: 17% faster transfer time is expected when two 32 KB IB chunks between Azure NDv2 nodes are sent together as one 64 KB chunk.TACCL uses these α-β values to decide whether and when to send chunks together, while accounting for forwarding delays.
  • Topology Inference: TACCL infers NDv2 PCIe topology using latency and bandwidth probes that reveal CPU proximity, shared PCIe switches, and NIC connectivity.The profiler also selects an NVLink topology automorphism so the NIC is consistently placed near GPU 0.

5 TACCL Synthesizer

TACCL’s synthesizer partitions collective data into chunks and generates schedules through a staged MILP-based process. Communication sketches and user parameters guide routing, ordering, contiguity, and topology-specific connectivity choices.

  • Problem Formulation: TACCL partitions each collective’s data into C equal chunks and schedules their transfers so every chunk reaches its required destination GPUs.The synthesizer decides chunk transfer schedules across every network link from a communication sketch and collective specification.
  • Problem Formulation: The MILP uses timing, transmission, bandwidth, and correctness constraints to minimize the maximum arrival time of required chunks.Its variables include chunk availability times, link-transfer indicators, and send times.
  • Contiguity: Sending chunks contiguously pays one α latency cost instead of n×α, but separate sends can preserve ordering flexibility for dependent transfers.TACCL applies this merging feature only to high-latency InfiniBand transfers, not lower-latency NVLinks.
  • Scalability: O(C^2) ordering decisions per link make direct MILP synthesis infeasible as topology size and chunk counts grow.Chunk counts grow linearly or quadratically for collectives such as ALLGATHER and ALLTOALL, increasing solver time and memory requirements.
  • Synthesis Process: TACCL reduces per-link binary variables from O(C^2) to O(C) by solving routing first, heuristically ordering chunks, and then optimizing contiguity.The routing stage determines paths without fixing chunk ordering, after which follow-up stages address ordering and latency costs.
  • Synthesizer Hyperparameters: The uc-max policy favors many switch connections for small data, whereas uc-min favors fewer connections when large-data congestion matters.These policies let the synthesizer trade off concurrency and congestion in switch-hyperedges.
  • Combining Collectives: Combining collectives such as REDUCESCATTER and ALLREDUCE are synthesized by adapting non-combining collective synthesis with reductions.REDUCESCATTER can be viewed as an inverse of ALLGATHER, but direct send inversion is insufficient because reductions cannot follow arbitrary simultaneous sends.

6 Backend

The TACCL backend lowers synthesized abstract algorithms into executable TACCL-EF programs and runs them through an NCCL-based runtime. It maps chunk operations to buffers, threadblocks, dependencies, and parallel instances.

  • Runtime: TACCL runtime extends NCCL with an interpreter that executes an entire synthesized communication algorithm in a single kernel launch.This avoids the multiple kernel launches required by a straightforward point-to-point implementation.
  • Executable Format: TACCL-EF represents collective algorithms over input, output, and scratch buffers divided into equal-sized chunks.Each algorithm step operates on chunk references in these buffers.
  • Execution Model: Threadblocks execute sequential instructions consisting of sends, receives with optional reduction, and local copies.For implementation simplicity, each threadblock communicates with at most one GPU in each direction.
  • Integration: TACCL runtime remains backward compatible with NCCL’s API, so integrating it into frameworks such as PyTorch requires a one-line library substitution.This enables dynamic selection of generated collective algorithms through torch.distributed.
  • Lowering: Lowering transforms abstract chunk sends into concrete buffer-index send and receive operations, then records their execution dependencies.The backend also allocates scratch buffers and groups operations into executable threadblocks while preserving abstract order.
  • Parallel Instances: TACCL generates n parallel algorithm instances by subdividing chunks into subchunks, allowing one-link bandwidth saturation beyond a single threadblock.The duplicated instruction groups and threadblocks execute concurrently along the parent chunks’ paths.

7 Evaluation

TACCL is evaluated on three collectives across DGX-2 and NDv2 multi-node configurations, with synthesized algorithms generally outperforming NCCL and improving end-to-end training throughput. The evaluation also examines sketch choices, synthesis time, and scalability trade-offs.

  • Evaluation setup: TACCL evaluates ALLGATHER, ALLTOALL, and ALLREDUCE on two-node and four-node DGX-2 and NDv2 configurations, comparing algorithm bandwidth against NCCL.The evaluation uses up to 32 GPUs and measures algorithm bandwidth as input buffer size divided by execution time.
  • ALLTOALL: 53%−66% faster than NCCL, TACCL’s NDv2 ALLTOALL algorithms improve performance for 16MB−1GB buffers, while other sketches target smaller buffers.On DGX-2, TACCL is up to 15% faster for buffers of at least 2MB and up to 55% faster for 1KB−16KB buffers.
  • ALLREDUCE: 6.4× faster than NCCL, TACCL’s DGX-2 ALLREDUCE improves performance for 1KB−4MB buffers, while it is at most 9% slower for buffers of 512MB and larger.On NDv2, TACCL is up to 28% faster for buffers up to 1MB and 28%−2.7× faster for larger buffers.
  • End-to-end training: 11%−1.94× faster for Transformer-XL and 12%−2.36× faster for BERT, TACCL improves end-to-end training throughput on two nodes.On four nodes, the corresponding speedups are 2%−1.44× and 7%−1.74%.
  • End-to-end training: 17% higher end-to-end throughput, TACCL improves an internal mixture-of-experts workload using synthesized ALLTOALL and ALLREDUCE algorithms.The workload uses approximately 6MB ALLTOALL and 256MB ALLREDUCE operations on two NDv2 nodes.

8 Related Work

Prior work optimizes collective communication for specific or homogeneous topologies, but existing synthesis methods have limited scale or constrained formulations. TACCL instead combines fine-grained synthesis with communication sketches to search larger, heterogeneous spaces.

  • MPI collectives: Prior MPI collective work targets efficient implementations for homogeneous interconnects such as mesh, hypercube, and fat-tree networks.Some prior methods combine bandwidth- and latency-optimal algorithms based on input sizes, but only for mesh networks.
  • GPU collectives: NCCL uses topology and buffer size to choose Ring or Tree algorithms but is agnostic to exact link performance, making it slower than TACCL in the paper’s comparisons.NCCL is optimized for NVLINK and Infiniband interconnects.
  • Topology-specialized synthesis: SCCL, Blink, and Plink specialize algorithms for underlying topologies, but prior synthesis approaches remain restricted in scale or formulation.SCCL synthesizes a Pareto-optimal ALLGATHER for a DGX1 node but is limited to single-node topologies.
  • Hierarchical approaches: Hierarchical approaches select known decompositions, whereas TACCL uses communication sketches with fine-grained algorithm synthesis for larger networks.The paper identifies combining these complementary approaches as future work.
  • Program sketching: Program sketching provides the conceptual basis for using partial specifications to search large spaces of communication algorithms.TACCL applies sketching to collective algorithm synthesis.
  • Network-flow formulations: Network-flow techniques cannot directly generate collective algorithms because collective communication can violate flow conservation when non-source GPUs replicate chunks across links.TACCL instead uses communication sketches and an encoding relaxation technique.

9 Conclusion and Future Work

TACCL combines communication sketches with a new problem formulation to synthesize efficient collective algorithms across multi-node and varied topologies. Its synthesized algorithms improve communication and end-to-end training performance, while sketch exploration remains an open direction.

  • Scalability: 1.7× higher algorithm bandwidth than NCCL was achieved for ALLGATHER across different data sizes on 8 Azure NDv2 nodes.
  • Generality across different topologies: TACCL synthesized ALLGATHER for a 2D 6×8 Torus, but communication-sketch exploration may be more limited there than on hierarchical topologies.
  • Exploring communication sketches: Different communication sketches can optimize different input-size ranges, motivating automated controllers for exploring sketch choices.
  • TACCL uses communication sketches and a novel problem formulation to generate algorithms for ALLGATHER, ALLTOALL, and ALLREDUCE.
  • 6.7× faster than NCCL and 11%−2.4× faster end-to-end training are reported for TACCL-generated algorithms.

A Communication Sketch Input

A communication sketch is a user-provided, high-level specification that guides TACCL’s synthesis search. It describes topology and communication policies, symmetry, chunking, and input-size information.

  • A communication sketch guides TACCL by specifying high-level structure while the synthesizer fills in remaining routing and scheduling details.
  • The sketch includes a logical topology, switch-hyperedge strategy, symmetry information, input size, and other hyperparameters.
  • Intra-node policy: The example sketch uses an intra-node NVSwitch policy with a uc-min switch-hyperedge strategy.
  • Inter-node policy: For inter-node transfers, the example uses relay routing and fixes which GPUs send to and receive data across nodes.
  • Symmetry and chunking: The sketch specifies rotational symmetry offsets and partitions each GPU’s input into 2 independently routable chunks.

B TACCL Synthesizer in Detail

TACCL’s synthesizer models collective communication over a logical topology derived from the hardware and communication sketch. Its formulation represents chunks, ranks, links, switched connections, and transfer costs.

  • TACCL defines chunks, GPU ranks, preconditions, postconditions, and logical links for the collective being synthesized.
  • Switch-hyperedge source and destination sets represent the links grouped at switched communication points.
  • The link cost combines alpha latency and beta bandwidth terms, with lat(r1,r2) representing total single-chunk transfer cost.

B.1 Routing

TACCL separates synthesis into routing, heuristic ordering, and contiguity and exact scheduling. The stages progressively determine paths, transfer order, and bandwidth-respecting contiguous transfers while minimizing collective time.

  • Routing: The routing stage chooses chunk paths and minimizes the time needed to reach the collective postcondition.
  • Routing: Routing initializes precondition chunks at time zero and enforces that a chunk is sent only after arriving at its source rank.
  • Routing: The routing encoding uses relaxed bandwidth constraints, so its generated solution may send multiple chunks simultaneously over one link.
  • Routing: For inter-node chunks, routing requires at least one inter-node link, while switch-hyperedge constraints encode whether to maximize or minimize utilized links.
  • Heuristic ordering: Heuristic ordering derives chunk paths from routing and selects transfers using shortest-path and longest-remaining-path heuristics.
  • Contiguity and exact scheduling: Contiguity and exact scheduling groups chunks to trade lower link latency against reduced pipelining, then applies strict link and switch bandwidth constraints.

C Standalone Experiments on Four Azure NDv2 Nodes

On four-node Azure NDv2, TACCL synthesized collective algorithms using the ndv2-sk-1 communication sketch and evaluated the best result across buffer sizes. The synthesized algorithms outperform NCCL across ALLGATHER, ALLTOALL, and ALLREDUCE, with gains varying by collective and buffer size.

  • TACCL synthesized ALLGATHER, ALLTOALL, and ALLREDUCE algorithms using the ndv2-sk-1 communication sketch, lowered with either 1 or 8 instances.The evaluation plots the best of the two resulting algorithms over different buffer sizes.
  • 10%−2.2× faster than NCCL, TACCL’s ALLGATHER algorithms outperform the baseline across all buffer sizes.
  • Up to 46% faster than NCCL for buffer sizes greater than 1MB, TACCL’s ALLTOALL algorithms benefit most from the ndv2-sk-1 sketch at large sizes.
  • Up to 34% faster than NCCL for small buffers and 1.9×−2.1× faster for larger buffers, TACCL’s ALLREDUCE algorithms improve across both size regimes.
Loading 2111.04867v4…