PaperlayerView exact revision on arXiv

Source-linked AI summary

ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale

William Won, Taekyung Heo, Saeed Rashidi, Srinivas Sridharan, Sudarshan Srinivasan, Tushar Krishna

arXiv:2303.14006v1cs.DCcs.LG

TL;DR

Large-model training increasingly requires distributed platforms whose parallelization, networking, and memory choices create a complex co-design space. This paper extends ASTRA-sim with graph-based arbitrary parallelism, parameterizable multidimensional topology modeling, and richer memory support, then uses case studies to examine emerging systems and reports speedups including 4.6x over Zero-Infinity.

  • Problem

    Rapidly growing model and data requirements make distributed training necessary, while evolving parallelization, network, and memory designs create a complex co-design space.

  • Method

    The paper extends ASTRA-sim with graph-based arbitrary parallelism, parameterizable multidimensional topology generation with analytical estimates, and local, remote, and in-network memory modeling.

  • Results

    Case studies show that conventional systems can match wafer-scale performance with suitable scheduling and parallelization, while selected disaggregated-memory configurations provide 4.6x speedup over Zero-Infinity.

  • Takeaways & Limitations

    ASTRA-sim 2.0 enables system designers to explore distributed-training software and hardware co-designs for emerging platforms at scale.

Abstract

from arXiv · show

As deep learning models and input data are scaling at an unprecedented rate, it is inevitable to move towards distributed training platforms to fit the model and increase training throughput. State-of-the-art approaches and techniques, such as wafer-scale nodes, multi-dimensional network topologies, disaggregated memory systems, and parallelization strategies, have been actively adopted by emerging distributed training systems. This results in a complex SW/HW co-design stack of distributed training, necessitating a modeling/simulation infrastructure for design-space exploration. In this paper, we extend the open-source ASTRA-sim infrastructure and endow it with the capabilities to model state-of-the-art and emerging distributed training models and platforms. More specifically, (i) we enable ASTRA-sim to support arbitrary model parallelization strategies via a graph-based training-loop implementation, (ii) we implement a parameterizable multi-dimensional heterogeneous topology generation infrastructure with analytical performance estimates enabling simulating target systems at scale, and (iii) we enhance the memory system modeling to support accurate modeling of in-network collective communication and disaggregated memory systems. With such capabilities, we run comprehensive case studies targeting emerging distributed models and platforms. This infrastructure lets system designers swiftly traverse the complex co-design stack and give meaningful insights when designing and deploying distributed training platforms at scale.

I. INTRODUCTION

Rapidly growing model computation and memory needs make distributed training necessary, but its interdependent software and hardware choices require co-design support. ASTRA-sim 2.0 extends the infrastructure for arbitrary parallelism, hierarchical networks, and local or remote memory, then applies it to emerging platforms.

  • Trillion-parameter models require tens of terabytes of memory and zeta floating-point operations, far exceeding single-NPU capacity.
  • Distributed training is necessary because DNN resource requirements are increasing beyond individual processing-unit capabilities.
  • Designing distributed training systems requires co-design across parallelization, NPU, network, and scheduling choices because these choices are interdependent.
  • Emerging software uses diverse parallelism strategies, including 3D parallelism, FSDP, ZeRO, and expert parallelism, beyond the original limited set.
  • ASTRA-sim 2.0 adds arbitrary parallelism, hierarchical network, and memory-model support, including local HBM and networked remote memories.
  • Case studies report that suitable scheduling and parallelization can let conventional systems match wafer-scale performance, while wafer-scale systems achieve up to 2.51× better collective time when scaled.
  • In disaggregated-memory MoE training, communication time dominates, and selected configurations hide communication to provide 4.6x speedup over Zero-Infinity.

II. BACKGROUND

Distributed training organizes computation, memory, and communication across devices through parallelization strategies and training loops. Collective operations and topology-aware hierarchical algorithms coordinate the resulting data movement across multidimensional networks.

  • Distributed training fundamentals: Synchronous distributed training commonly uses All-Reduce to synchronize distributed information across NPUs.
  • Distributed training fundamentals: Data, model, and pipeline parallelism determine how parameters, inputs, and computation are distributed across NPUs.
  • Distributed training fundamentals: A training loop specifies the ordering of communication and computation required to execute distributed training.
  • Collective communication: Collective communication patterns include Reduce-Scatter, All-Gather, All-Reduce, and All-to-All.
  • Hierarchical collectives: Topology-aware collective algorithms must account for multidimensional physical networks because basic logical topologies may mismatch the underlying topology.
  • Hierarchical collectives: Hierarchical All-Reduce runs Reduce-Scatter across dimensions in ascending order, followed by All-Gather in descending order.

C. ASTRA-sim

ASTRA-sim models the distributed-training software and hardware search space through workload, system, and network abstractions. The section motivates extensions because the original framework cannot represent arbitrary parallelisms or the broader emerging platform space.

  • ASTRA-sim overview: ASTRA-sim represents workload descriptions, collective algorithms, operation scheduling, compute-communication overlap, NPU compute inputs, and simulated communication.
  • ASTRA-sim overview: Its workload layer describes target DNN models, parallelization strategies, and training loops, while the system layer schedules computation and communication.
  • Extension requirements: ASTRA-sim therefore requires support for arbitrary parallelisms, multidimensional hierarchical networks, and memory systems.
  • Extension requirements: The original framework cannot support complex strategies such as pipeline and 3D parallelism because it assumes synchronized NPU operations and couples strategies to separate frontend training loops.

B. Ability to Model Multi-dimensional Networks

Distributed training is often communication-bound, motivating heterogeneous multidimensional networks and topology-aware modeling. The surrounding design space also includes local and disaggregated memory, in-switch collectives, and mechanisms for faster large-scale simulation.

  • Network motivation: Large-scale distributed training is usually communication-bound, so systems combine multiple networking technologies to improve performance.
  • Network motivation: Examples include DGX-A100’s two-dimensional NVLink-plus-InfiniBand or Ethernet topology and TPUv4’s three-dimensional torus.
  • Network representation: ASTRA-sim’s predefined topology support limits exploration because new topologies require implementing both Garnet networks and topology-aware collective algorithms.
  • Network representation: A more powerful backend should represent arbitrary multidimensional topologies systematically for rapid design-space exploration.
  • Memory and collectives: Efficient memory systems matter because DNN parameters must be loaded and stored, while limited GPU capacity motivates local HBM and disaggregated-memory modeling.
  • Memory and collectives: In-switch collective communication is proposed as a design-space option for reducing distributed-training communication time.

IV. EXTENSIONS TO ASTRA-SIM

ASTRA-sim is extended with a graph-based execution engine and a generic representation for arbitrary hierarchical multi-dimensional network topologies. The extensions use standardized execution traces and topology-aware collective algorithms to model diverse distributed training systems.

  • Graph-based Execution Engine: The graph-based execution engine decouples parallelization strategies from the frontend and represents model execution as input graphs.Graphs can be generated from frameworks including PyTorch, TensorFlow, and FlexFlow.
  • Graph-based Execution Engine: ASTRA-sim ET standardizes execution traces with compute, memory, and communication nodes carrying simulation metadata.A converter translates other execution-trace formats into ASTRA-sim ET; PyTorch and FlexFlow are currently supported.
  • Multi-dimensional Network Representation: The topology taxonomy builds multi-dimensional networks by hierarchically stacking Ring, FullyConnected, and Switch building blocks.These blocks provide corresponding topology-aware collective communication algorithms.
  • Multi-dimensional Network Representation: Arbitrary topology dimensions and block order can be represented, including 2D tori, switched hierarchies, DragonFly, and 3D tori.The notation can represent 4D, 5D, and higher-dimensional networks and maps examples to distributed training platforms.
  • Multi-dimensional Network Representation: Topology-aware multi-rail collectives run the basic collective algorithm iteratively across each network dimension.The chosen building blocks have known congestion-free collective algorithms, simplifying hierarchical collective design.

C. Analytical Network Backend

The analytical network backend estimates communication delay with an equation-based model instead of packet-level cycle-accurate simulation. It closely matches measured All-Reduce performance while enabling much faster simulation of large systems.

  • Motivation: Analytical modeling supports first-order exploration of topology shape and bandwidth at the scale of thousands of NPUs.Cycle-level Garnet simulation is too slow for state-of-the-art and futuristic systems at this scale.
  • Analytical Model: The analytical backend estimates communication delay using link latency, hop count, message size, and link bandwidth rather than simulating packet-level network behavior.This approach is intended for large, bandwidth-bound collectives and can be extended with additional effects.
  • Limitations: The analytical backend has limitations for networks with non-trivial congestion or link oversubscription.First-order congestion modeling is identified as future work.
  • Validation: The analytical backend achieved 5% mean simulation error across 64 MB–1.5 GB All-Reduce configurations on real systems.Validation used 4- and 16-GPU NVIDIA V100 systems with Ring topology and 150 GB/s NVLink.
  • Speedup: 756× speedup reduced a 1MB All-Reduce simulation on a 64-NPU 3D Torus from 21.42 minutes to 1.70 seconds.The backend also simulated a 4K-NPU 3D Torus in 3.14 seconds.

D. Memory Models

ASTRA-sim adds memory models spanning local memory, remote disaggregated memory, and in-switch collective communication. These models estimate tensor-transfer cycles using configurable memory and network parameters, including pipelined hierarchical transfers.

  • Memory API: The memory API models local, remote, and in-switch memory systems using tensor location, size, bandwidth, and system-design inputs.It returns the cycles required to load or store a tensor.
  • Local Memory Model: The local memory model combines memory access latency with tensor size divided by memory bandwidth.Latency and bandwidth are system parameters, while payload size is encoded in the execution trace.
  • Remote Memory Model: The remote memory model estimates transfer time for hierarchical disaggregated memory using remote-memory design parameters and pipelined data movement.Remote memory groups form a shared pool connected through out-node switches.
  • Remote Memory Model: Remote-memory transfer time is determined by the critical path, with each pipeline stage set by the maximum data-transfer time among its links.The model partitions traffic across remote memory groups, out-node switches, and chunks.
  • In-switch Collective Communication: In-switch collective communication gathers parameters during loading and shards them during storing, changing the transferred data size on each link.The model supports All-Gather and Reduce-Scatter operations and uses pipelined parameter transfers.

V. CASE STUDIES

The case studies evaluate ASTRA-sim’s extended capabilities using a fixed NPU compute model based on measured A100 GPU performance.

  • Experimental Setup: All experiments assume an NPU compute power of 234 TFLOPS based on measurements of an A100 GPU.

A. Conventional System vs Wafer-scale System

The study compares conventional hierarchical systems with wafer-scale systems across scheduling and scaling. Appropriate collective scheduling can make conventional systems competitive, while wafer-scale systems gain up to 2.51× when scaling if on-wafer bandwidth is sufficient.

  • System Comparison: The experiments model 512-NPU conventional and wafer-scale systems using multidimensional hierarchical and low-dimensional high-bandwidth on-wafer networks.Wafer-scale proxies include 1D topologies with 300, 500, and 600 GB/s on-wafer bandwidth and a 2D topology with 250 × 250 GB/s bandwidth.
  • Impact of Scheduling: Multi-dimensional topologies can suffer pipelining bubbles and unbalanced bandwidth, while the 1D W-1D topology achieves the best overall performance.
  • Impact of Scheduling: Themis scheduling heavily benefits W-2D, Conv-3D, and Conv-4D by balancing collective load across dimensions; W-1D gains no benefit because it has one dimension.For single All-Reduce and DLRM, conventional systems with Themis match corresponding wafer-scale systems with equivalent bandwidth per NPU.
  • Impact of Scheduling: Conventional systems with Themis can match wafer-scale performance for single All-Reduce and DLRM when their bandwidth per NPU is equivalent.
  • Impact of Scaling: 2.51× speedup is achieved by wafer-scale scaling over scale-out when sufficient on-wafer bandwidth shifts communication load away from other dimensions.Once the on-wafer dimension becomes the bottleneck, collective time begins to increase.

B. Comparing Disaggregated Memory Systems

The study compares ZeRO-Infinity with HierMem for trillion-parameter MoE training and finds nearly equivalent performance, while exploring bandwidth configurations to improve HierMem.

  • System comparison: ZeRO-Infinity and HierMem are compared because current model sizes exceed available GPU memory capacity.ZeRO-Infinity uses local HBM, CPU memory, and NVMe, whereas HierMem supports arbitrary remote memory groups.
  • Evaluation method: The trillion-parameter MoE training task is evaluated using compute, exposed local-memory, exposed remote-memory, exposed communication, and exposed idle time.These components form the execution-time breakdown used to compare the systems.
  • Results: 0.1%: ZeRO-Infinity performs better than HierMem overall, with both systems showing similar performance because they use almost equivalent resources.HierMem’s small performance drop comes from additional data-transfer stages through multi-level switches.
  • Design-space exploration: HierMem’s design space is explored by varying in-node pooled-fabric bandwidth and remote-memory-group bandwidth because exposed communication is the bottleneck.The sweep varies pooled-fabric bandwidth from 256GB/s to 2048GB/s and remote-memory-group bandwidth from 100GB/s to 500GB/s.

VI. RELATED WORK

Prior simulators address general-purpose distributed workloads or HPC communication, while this work extends ASTRA-sim toward arbitrary DNN traces on multidimensional topologies and disaggregated memory systems.

  • Existing simulators: General-purpose distributed-system simulators face a trade-off among simulation accuracy, simulation speed, and engineering effort.HPC communication simulators such as LogGOPSim and SMPI address communication-performance modeling.
  • ASTRA-sim extension: The work extends ASTRA-sim to capture arbitrary parallelization strategies, training loops, multidimensional network topologies, and complex memory systems.It uses the framework for end-to-end, full-stack co-design-space exploration of distributed training.
  • Practical role: The resulting infrastructure helps system designers navigate distributed-training design spaces and obtain first-order insights for building training platforms at scale.The stated benefit is faster exploration of the complex co-design stack.
Loading 2303.14006v1…