AIONetBench

Distributed AI Networking & Performance Stack

AIONetBench is a full-stack distributed AI benchmarking framework inspired by Meta’s AI Networking group. It simulates large-scale GPU training environments using PyTorch Distributed (DDP/FSDP) with built‑in observability, reliability, and performance tuning.

What It Does

• Distributed Training: GPT‑mini (LLM) and ResNet18 (Vision) benchmarks using DDP and FSDP.
• Fault Tolerance: Simulated slowdowns and crashes with checkpoint recovery.
• Observability: Prometheus metrics + Grafana dashboard for live tokens/sec, images/sec, step time, loss, and restarts.
• Performance Tuning: Automated sweeps of batch size, LR, and communication settings.

🧩 System Architecture

flowchart TD
    subgraph User["🧑‍💻 User / Developer"]
        A1["Run Orchestrator (train.py)"]
        A2["Select Config (YAML)"]
    end

    subgraph Benchmarks["⚙️ Distributed Benchmarks"]
        B1["GPT-Mini (LLM)"]
        B2["ResNet18 (Vision)"]
        B3["Microbenchmarks (NCCL-style)"]
    end

    subgraph Engine["🧮 PyTorch Distributed Engine"]
        E1["DDP / FSDP Backends"]
        E2["Gloo / NCCL Communication"]
        E3["Fault Injection + Checkpoint Recovery"]
    end

    subgraph Metrics["📊 Observability"]
        M1["Prometheus Exporter (Port 8000)"]
        M2["Prometheus Server (Port 9090)"]
        M3["Grafana Dashboard (Port 3000)"]
    end

    subgraph Automation["🤖 Automation & Tuning"]
        T1["Parameter Sweeps (Batch, LR, Comm)"]
        T2["Tuner (tuner.py)"]
        T3["Results → CSV + Grafana"]
    end

    A1 --> A2 --> B1 & B2 & B3
    B1 & B2 & B3 --> E1 --> E2 --> E3 --> M1 --> M2 --> M3
    E3 --> T1 --> T2 --> T3

Loading

Run It

# Setup
make build

# LLM Benchmark
python orchestrator/train.py --config configs/ddp_gptmini.yaml

# Vision Benchmark
python orchestrator/train.py --config configs/ddp_vision.yaml

# Tuner
make tune_gpt

Dashboard

Prometheus scrapes metrics at localhost:8000, Prometheus server runs on 9090, and Grafana at 3000.
Import observe/grafana/aionet_dashboard.json into Grafana to view real‑time metrics.