Distributed Data Parallel Training on Modal - Multi GPU

This project demonstrates PyTorch Distributed Data Parallel (DDP) training running on Modal's cloud GPUs, including NanoGPT speedrun implementations.

This repository is part of First Break AI — a free, open cohort for learning training, inference, and AI product building. The roadmap, checklist, and community live on that site.

Prerequisites

1. Install the Modal CLI:

   pip install modal

2. Authenticate with Modal:

   modal token new

3. Create secrets for speedruns (Projects 2 and 3):

   modal secret create wandb-secret WANDB_API_KEY=your_api_key_here
   modal secret create HF_TOKEN HF_TOKEN=your_huggingface_token_here

   Get your W&B key from wandb.ai/authorize.
   HF_TOKEN is only needed if you still have to download FineWeb into the fineweb-data volume.

4. For H100 jobs (Tyler + Keller speedruns), add a payment method in Modal Settings → Usage & Billing. Credits alone are not enough for H100.

---

Project 1: Basic DDP Training

Run basic DDP training on 2x L40S GPUs:

modal run modal_ddp.py::train_single_node

This will:

• Launch a container with 2x L40S GPUs on Modal
• Run ddp.py using torchrun for distributed training
• Save profiler traces to a persistent Modal Volume
• Log training metrics to Weights & Biases

---

Project 2: NanoGPT Speedrun (Tyler Romero)

Progressive GPT-2 training optimization following Tyler Romero's worklog. Each --step copies that step’s train_gpt2.py and runs its run.sh on Modal.

Modal entrypoint: nanogpt-speedrun/src/runfiles/modal_runner.py

Steps Overview

Times below are Tyler’s original 2×RTX 4090 wall-clock. This repo’s Modal runner uses 8×H100, so wall-clock is much shorter; the optimization sequence is the same.

Step	Folder	Original 2×4090 time	Key Changes
1	01-Initialbaseline	~8 hrs	Base GPT-2 from llm.c
2	02-ArchitecturalChanges	~7.5 hrs	RoPE, ReLU², trapezoidal LR
3	03-MuonOptimizer	~4.5 hrs	Muon optimizer
4	04-DataLoadingTwerks	~4.3 hrs	Micro-batch loading
5	05-LogitSoftCappingat30	~4 hrs	Logit soft-capping
6	06-LongerSequenceLength	~2.5 hrs	FlexAttention, 32K seq

Before you run

Shared with Project 3 (modded-nanogpt):

Need	Name
Volume (data)	fineweb-data
Volume (traces)	ddp-traces
Secret	wandb-secret (WANDB_API_KEY)
Secret	HF_TOKEN (for download if data is missing)
GPUs	H100:8

Modal requires a payment method on file to use H100s (credits alone are not enough). Add one under Settings → Usage & Billing.

If fineweb-data already has fineweb_train_*.bin and fineweb_val_000000.bin, skip download_data.

Running Tyler's Speedrun on Modal

# Must run from nanogpt-speedrun/ (image syncs this directory)
cd nanogpt-speedrun

# Only if fineweb-data is empty (~900M tokens, a few minutes)
modal run src/runfiles/modal_runner.py::download_data

# Step 1 — Initial baseline (01-Initialbaseline)
# Target: FineWeb val_loss <= 3.28 (or finish 24576 steps)
modal run src/runfiles/modal_runner.py::train --step 1

# Optional notes for the W&B run
modal run src/runfiles/modal_runner.py::train --step 1 --notes "first attempt"

# Later optimization steps
modal run src/runfiles/modal_runner.py::train --step 2
modal run src/runfiles/modal_runner.py::train --step 3
modal run src/runfiles/modal_runner.py::train --step 4
modal run src/runfiles/modal_runner.py::train --step 5
modal run src/runfiles/modal_runner.py::train --step 6

cd ..

What train --step 1 does:

1. Mounts fineweb-data at src/data/fineweb10B and ddp-traces at /traces
2. Copies src/runfiles/01-Initialbaseline/train_gpt2.py → src/train_gpt2.py
3. Runs 01-Initialbaseline/run.sh with torchrun on 8 GPUs (grad_accum_steps=1 so global batch stays 262144)
4. Logs to W&B project tyler-nanogpt-run

Watch for val_loss every 128 steps (not every train_loss line). Walkthrough: tutorial/tyler-speedrun/Whatisbaseline.md.

Wandb Project: tyler-nanogpt-run
GPUs: H100:8
App name: tyler-nanogpt-speedrun

---

Project 3: Modded-NanoGPT (Keller Jordan)

World-record NanoGPT speedrun by Keller Jordan. Trains GPT-2 to 3.28 val loss in under 100 seconds on 8x H100.

Running Keller's Speedrun

# Navigate to modded-nanogpt folder first (required for Dockerfile)
cd modded-nanogpt

# First time: download data
modal run modal_modded_nanogpt.py::download_data

# Run the speedrun (requires 8x H100)
modal run modal_modded_nanogpt.py::train

# Optional: capture PyTorch profiler traces to ddp-traces + W&B Artifacts (adds overhead)
modal run modal_modded_nanogpt.py::train --profiler

# Return to root folder
cd ..

Note: torch.compile adds ~7 minutes latency on first run.

Wandb Project: modded-nanogpt-run

Metrics vs traces

Output	Default train	train --profiler
W&B Charts (val_loss, train_time_ms, …)	Yes	Yes
W&B Artifacts (profiler-traces)	No	Yes (if traces written)
Modal volume ddp-traces (.pt.trace.json)	No	Yes

download_data does not write traces; only training with the profiler flag does.

---

Viewing Training Logs in Wandb

After training, view your logs at wandb.ai:

• ddp-training - Basic DDP training
• tyler-nanogpt-run - Tyler Romero's speedrun steps
• modded-nanogpt-run - Keller Jordan's world-record speedrun

Modded-nanogpt: Charts tab shows scalars from every run. The Artifacts tab gets profiler-traces only when you run modal run modal_modded_nanogpt.py::train --profiler (or set ENABLE_PROFILER=1 in the training env).

Downloading Profiler Traces

After a profiler-enabled modded-nanogpt run (or Project 1 DDP training), download traces from Modal:

# List available traces
modal volume ls ddp-traces

# Download all traces to a local directory
modal volume get ddp-traces / ./local_traces

Project 1 (modal_ddp.py) always profiles. Keller's speedrun profiles only with --profiler.

Viewing Traces in Perfetto

1. Open ui.perfetto.dev in your browser
2. Click "Open trace file" or drag and drop
3. Select the .pt.trace.json file from your downloaded traces
4. Explore the timeline to analyze:
   • Forward/backward pass timing
   • Data movement overhead
   • Gradient synchronization
   • Optimizer step timing

---

Project Structure

├── ddp.py                          # Basic DDP training script
├── modal_ddp.py                    # Modal wrapper for DDP
├── training_utils/                 # Shared utilities
│   ├── __init__.py
│   ├── memory.py
│   ├── trun.py
│   └── utils.py
├── nanogpt-speedrun/               # Tyler Romero's speedrun
│   ├── src/runfiles/
│   │   ├── modal_runner.py         # Modal wrapper
│   │   ├── 01-Initialbaseline/
│   │   ├── 02-ArchitecturalChanges/
│   │   ├── 03-MuonOptimizer/
│   │   ├── 04-DataLoadingTwerks/
│   │   ├── 05-LogitSoftCappingat30/
│   │   └── 06-LongerSequenceLength/
│   └── pyproject.toml
├── modded-nanogpt/                 # Keller Jordan's speedrun
│   ├── modal_modded_nanogpt.py     # Modal wrapper
│   ├── train_gpt.py
│   ├── Dockerfile
│   └── requirements.txt
└── readme.md

GPU Configuration

Project	Default GPU	Cost/hr
DDP Training	L40S:2	~$2.60
Tyler's Speedrun	H100:8	~$32/hr
Keller's Speedrun	H100:8	~$32

Edit the Modal wrapper files to change GPU type/count.

---

How Modal Runs These Projects

All three projects run on Modal's cloud infrastructure. When you run modal run <file>::<function>, Modal:

1. Parses your Python file
2. Builds (or reuses cached) container image
3. Spins up container with GPUs attached
4. Mounts volumes and injects secrets
5. Runs your function
6. Tears down container when done

┌─────────────────────────────────────────────────────────────┐
│  Your laptop                                                │
│  $ modal run modal_modded_nanogpt.py::train                 │
└─────────────────────────────────────────────────────────────┘
                            │
                            ▼
┌─────────────────────────────────────────────────────────────┐
│  Modal Cloud                                                │
│                                                             │
│  1. Check if image is cached                                │
│     ├── Cached? → Use existing image                        │
│     └── Not cached? → Build new image                       │
│                                                             │
│  2. Spin up container with:                                 │
│     - GPUs attached (L40S, H100, etc.)                      │
│     - Volume mounts (fineweb-data, ddp-traces)              │
│     - Secrets injected (WANDB_API_KEY)                      │
│                                                             │
│  3. Run function: train()                                   │
│     └── subprocess: torchrun train_gpt.py                   │
│                                                             │
│  4. Container stops when function returns                   │
│     - Volumes persisted (data survives)                     │
│     - Container deleted (ephemeral)                         │
└─────────────────────────────────────────────────────────────┘

Container Image Methods

Modal supports two ways to define container images:

Method	Description	When to Use
Image Builder	Chain of Modal commands	Simple pip dependencies
Dockerfile	Standard Docker file	Complex builds, custom OS setup

How Each Project Builds Its Image

Project	Method	Why
Basic DDP	Image Builder	Simple - just pip packages
Tyler's Speedrun	Image Builder	Uses uv package manager
Keller's Modded	Dockerfile	Complex - custom Python 3.12, CUDA 12.6, Flash Attention 3

Basic DDP (modal_ddp.py):

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install("torch~=2.5.1", "accelerate", ...)
)

Tyler's Speedrun (modal_runner.py):

image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("curl")
    .run_commands("curl -LsSf https://astral.sh/uv/install.sh | sh")
    .add_local_dir(".", remote_path="/root/nanogpt-speedrun", copy=True)
    .run_commands("uv sync --all-extras")
)

Keller's Modded (modal_modded_nanogpt.py):

image = modal.Image.from_dockerfile("Dockerfile")

Shared Resources

All projects share these Modal resources:

Resource	Type	Purpose
fineweb-data	Volume	Pre-tokenized training data (1.9 GB)
ddp-traces	Volume	Profiler traces
wandb-secret	Secret	Weights & Biases API key

Data downloaded once is available to all projects.