DOCSEARCHbyAlgolia
Submit Docs

ASTRA-sim Documentation

repository·master·Indexed 20 days ago

https://github.com/astra-sim/astra-sim

ASTRA-sim is a distributed AI system simulator that models the full stack from workload scheduling to hardware architecture, including compute, memory, and network. It supports multi-fidelity simulation using analytical models, ns-3, and HTsim network backends. The simulator accepts MLCommons Chakra Execution Traces as workload inputs and allows for custom collective implementations via the CollectiveAPI.

Tokens
4.2K
Snippets
7
Records
20
Agent score
71%

What's inside ASTRA-sim

  1. Overview of ASTRA-sim

    master

    ASTRA-sim is a distributed AI system simulator designed to model the end-to-end software and hardware stack of modern AI systems. It covers workload scheduling, collective communication algorithms, and hardware architectures (including compute, memory, and network).

    Key features include:

    • Multi-fidelity simulation: Supports end-to-end modeling for designing and deploying next-generation distributed AI systems.
    • Plug-and-play APIs: Allows users to integrate external open or proprietary components to model specific parts of the AI system.
    • Workload Inputs: Accepts MLCommons Chakra Execution Traces as inputs at the workload layer.

  3. Select a network backend for simulation

    master

    ASTRA-sim supports multiple network backend input types depending on the level of fidelity required:

    • analytical: Uses analytical network models.
    • ns3: Uses the ns-3 network simulator backend.
    • htsim: Uses the HTSim network simulator backend.

  4. Configure ASTRA-sim system layers

    master

    ASTRA-sim allows you to define how collectives are implemented at the system layer using different configuration files:

    • native_collectives: Uses ASTRA-sim's built-in native collective algorithm implementations.
    • custom_collectives: Uses custom collective implementations via ASTRA-sim's CollectiveAPI.

  6. Explore ASTRA-sim example input files

    master

    The examples/ directory provides sample input files for workloads, system configurations, and network backends to help you get started with ASTRA-sim simulations.

    Workload Examples

    • microbenchmarks: Contains simple Chakra ET (Execution Trace) files for 1MB collective communication microbenchmarks. These are available for 4, 8, and 16 NPUs and cover the following operations:
      • All-Reduce
      • All-Gather
      • Reduce-Scatter
      • All-to-All

    Note: To generate synthetic ETs for realistic workloads (including compute and communication), use STG. For real-system ET collection, refer to the Chakra wiki.

  7. Access ASTRA-sim documentation and resources

    master

    For detailed usage instructions and technical insights, use the following resources:

  8. Configure the logical topology using LOGICAL_TOPOLOGY

    master

    The LOGICAL_TOPOLOGY environment variable is used in build.sh to define how many NPUs are used and their logical arrangement. This value determines the physical_dims vector.

    • A single value like {64} allocates the first 64 NPUs in a 1D topology.
    • Multiple values like {8,8} allocates the same number of NPUs (64) but in a 2D topology.

    Note: The logical topology allows you to use a subset of a larger physical cluster. For example, if your TOPOLOGY_FILE defines 128 nodes but your LOGICAL_TOPOLOGY is {64}, only the first 64 nodes will be allocated to the workload.

    # Example usage concept
export LOGICAL_TOPOLOGY='{8,8}'
./build/astra_ns3/build.sh

  9. Run sample ASTRA-sim simulations using scripts

    master

    You can execute sample simulations by running the provided .sh scripts in the examples/ directory. The scripts are organized by network backend:

    • analytical: Scripts for analytical network backends. Includes two modes:
      • congestion_unaware
      • congestion_aware
    • ns3: Scripts for the ns-3 network backend.
    • htsim: Scripts for the HTsim network backend.

  10. Configure the network input file for ns3 simulation

    master

    To run the ns3 simulation, you must provide a network input file using the NETWORK environment variable when executing the build script build/astra_ns3/build.sh.

    Inside the network input file (e.g., mix/config.txt), you must define the TOPOLOGY_FILE parameter. This parameter points to a file describing the physical topology, including link connections, switches, and nodes.

    # Example of setting the environment variable before running the build script
export NETWORK='mix/config.txt'
./build/astra_ns3/build.sh

  11. Simulation Lifecycle and Execution Flow

    master

    The congestion-aware analytical simulation follows a specific lifecycle to ensure all components are correctly initialized before execution:

    1. Initialization: Initialize the AstraSim::LoggerFactory using the provided logging configurations.
    2. Event Queue Setup: Create a shared EventQueue and register it with the Topology via Topology::set_event_queue(event_queue).
    3. Topology Generation: Parse the network configuration using NetworkParser and construct the topology using construct_topology(network_parser).
    4. Network API Configuration: Set the global EventQueue and Topology for the network API using CongestionAwareNetworkApi::set_event_queue(event_queue) and CongestionAwareNetworkApi::set_topology(topology).
    5. Resource Allocation: Instantiate AnalyticalRemoteMemory and create a Sys object for each NPU. Each Sys object is initialized with its own CongestionAwareNetworkApi instance.
    6. Workload Triggering: Start the simulation by calling .workload->fire() on every Sys instance.
    7. Simulation Loop: Execute the simulation by repeatedly calling event_queue->proceed() until event_queue->finished() returns true.
    8. Cleanup: Delete Sys instances and call AstraSim::LoggerFactory::shutdown().

  12. How the ASTRA-sim and NS3 interaction works

    master

    The interaction between the ASTRA-sim System layer and the NS3 simulator (Network layer) follows an event-driven pattern:

    1. Issuing Events: The System layer issues send/receive events via send_flow.
    2. Simulation: The System layer waits while NS3 simulates the conclusion of these events.
    3. Completion: When NS3 completes an event (e.g., a node finishes sending or receiving), it calls qp_finish.
    4. Callbacks: qp_finish performs lookups in internal maps to trigger the appropriate MsgEvent callback handlers, signaling the System layer that the event has finished.

    This mechanism allows the System layer to synchronize its collective communication phases with the network simulation.