GPU Graph Training Tuner

Hardware-aware search tool for tuning graph neural network training workloads on modern GPUs.

This repo is built to signal the kind of work NVIDIA DevTech teams care about:

• graph and deep learning workloads instead of generic app scaffolding
• performance reasoning across HBM bandwidth, tensor throughput, launch overhead, and activation memory
• architecture-aware tradeoffs between precision, kernel fusion, CUDA graphs, checkpointing, and batch construction
• a clean path from cost model to optional PyTorch validation

Problem

Graph neural network training often stops scaling for reasons that are not obvious from model code alone:

• irregular memory access makes scatter-reduce kernels bandwidth-bound
• neighbor sampling and fanout shape the subgraph size more than the nominal batch size does
• mixed precision improves tensor core throughput but can shift the bottleneck back to memory traffic
• checkpointing lowers memory pressure while increasing total step time
• CUDA graphs help only when launch overhead is materially relevant

GPU Graph Training Tuner turns those interactions into a reproducible search problem.

What It Does

Given a scenario JSON, the tuner:

• models a graph workload and a target GPU architecture
• searches over micro-batch size, accumulation steps, fanout, precision, kernel fusion, CUDA graphs, and checkpointing
• estimates step time, samples per second, memory footprint, and energy
• returns the top-ranked configurations for throughput, latency, efficiency, or a balanced objective
• optionally runs a synthetic PyTorch benchmark for the selected config

Quickstart

1. Create a virtual environment

python -m venv .venv
# Windows
.venv\Scripts\activate
# macOS/Linux
source .venv/bin/activate

2. Install the project

pip install -e .[dev]

Install the optional benchmark dependency only if you want the PyTorch validation path:

pip install -e .[dev,bench]

3. Search the sample workload

python -m gpu_graph_training_tuner search scenarios/h100_fraud_graphsage.json --top-k 3

Example Output

rank | micro-batch | accum | eff batch | fanout | precision | fused | graphs | ckpt | step ms | samples/s | memory GB | watts | score
-------------------------------------------------------------------------------------------------------------------------------------
   1 |        8192 |     4 |     32768 |      8 |      bf16 |   yes |    yes |   no |   20.10 | 1630495.8 |      7.00 | 336.5 | 1626.85
   2 |        8192 |     4 |     32768 |      8 |      bf16 |   yes |     no |   no |   20.26 | 1617521.5 |      7.00 | 335.8 | 1613.85
   3 |        8192 |     2 |     16384 |      8 |      bf16 |   yes |    yes |   no |   10.14 | 1615368.0 |      7.00 | 421.4 | 1612.89

Benchmark the Top Configuration

python -m gpu_graph_training_tuner benchmark scenarios/h100_fraud_graphsage.json --rank 1 --device cuda

If PyTorch is not installed, the command exits cleanly and tells you how to enable it.

Architecture

Scenario JSON
    |
    v
ScenarioSpec -> Device profile + graph workload + search space + constraints
    |
    v
GraphTrainingTuner
    |- neighbor-sampling expansion model
    |- step-time cost model
    |- memory model for parameters, optimizer state, activations, and edge buffers
    |- objective scoring for throughput / latency / efficiency / balanced
    |
    v
PlanResult
    |- ranked configs
    |- time breakdown
    |- memory breakdown
    |- rationale notes
    |
    v
Optional PyTorch synthetic benchmark

Detailed design notes live in docs/ARCHITECTURE.md.

Scenario Schema

Each scenario contains:

• device: memory capacity, HBM bandwidth, peak tensor throughput, power, launch overhead
• workload: graph size, layer dimensions, depth, overlap assumptions, optimizer multiplier
• search_space: the knobs the tuner is allowed to search
• constraints: optional memory and step-time caps

See scenarios/h100_fraud_graphsage.json for a complete example.

Project Structure

GPU Graph Training Tuner/
|-- docs/
|-- scenarios/
|-- src/gpu_graph_training_tuner/
|   |-- benchmark.py
|   |-- cli.py
|   |-- models.py
|   |-- optimizer.py
|   |-- report.py
|   `-- scenario.py
|-- tests/
|-- pyproject.toml
`-- README.md

Why This Repo Is Stronger Than a Generic ML Portfolio Project

• It demonstrates GPU systems reasoning instead of only model API usage.
• It gives you a concrete interview story about precision policy, kernel fusion, memory pressure, and launch amortization.
• It maps cleanly to AI Developer Technology work: customer workload analysis, performance tuning, and feedback into software or hardware decisions.

Roadmap

• Add multi-GPU partitioning and NVLink-aware communication costs
• Plug in measured profiler traces to calibrate the cost model
• Add GAT and heterogeneous graph templates
• Export HTML reports for experiment reviews

License

MIT License. See LICENSE.