optzero: Custom ZeRO Distributed Training Framework

optzero is a custom PyTorch optimizer and engine implementing ZeRO-1 (Optimizer State Partitioning) and ZeRO-2 (Gradient Partitioning) from scratch. It shards optimizer states and gradients across distributed GPUs to prevent Out-Of-Memory (OOM) errors when training large models.

Built entirely using torch.distributed and raw PyTorch autograd hooks. No HuggingFace Trainer, no PyTorch Lightning, no PyTorch FSDP.

Architecture

The engine intercepts PyTorch's autograd engine using post-accumulate gradient hooks. Network communication is overlapped with gradient computation using asynchronous collective operations.

graph TB
    Model[Full Model on All GPUs] --> Partition[Parameter Partitioner]
    Partition -->|Rank 0 Shard| Opt0[ZeroAdam Rank 0]
    Partition -->|Rank 1 Shard| Opt1[ZeroAdam Rank 1]

    Backward[Backward Pass] -->|Async Gradient Hook| AllReduce[dist.all_reduce async_op]
    AllReduce --> Wait[work.wait in step]
    Wait -->|ZeRO-2: Free Non-local Gradients| Opt0
    Wait -->|ZeRO-2: Free Non-local Gradients| Opt1

    Opt0 -->|Update Shard 0| Broadcast[dist.broadcast]
    Opt1 -->|Update Shard 1| Broadcast
    Broadcast --> Synced[Synced Weights]

Loading

Features

• ZeRO-1 & ZeRO-2 Partitioning: Each GPU allocates optimizer states (Adam momentum/variance) and gradients only for its specific parameter shard, reducing VRAM usage by 1/N.
• Custom AdamW (ZeroAdam): A from-scratch implementation of the AdamW algorithm using raw tensor operations (exp_avg, addcdiv_).
• Asynchronous Communication: Uses async_op=True with dist.all_reduce and PyTorch 2.0+ register_post_accumulate_grad_hook to overlap network I/O with gradient computation, preventing autograd deadlocks.
• Distributed Checkpointing: Fault-tolerant save_checkpoint and load_checkpoint methods that save the model centrally on Rank 0 and the sharded optimizer states per-rank.
• Real Model Training: Proven to train a ResNet18 model on the CIFAR-10 dataset across multiple processes.
• Strict Typing: 100% mypy --strict compliance and ruff linting.

Tech Stack

• Python 3.11+
• PyTorch (Autograd, torch.distributed, torch.optim)
• Torchvision (ResNet18, CIFAR-10)
• Gloo / NCCL (Collective communication backends)

Repository Structure

optzero/
├── src/
│   └── optzero/
│       ├── core/
│       │   ├── engine.py           # ZeroEngine: async hooks, ZeRO-2, checkpointing
│       │   ├── optimizer.py        # ZeroAdam: custom sharded AdamW
│       │   └── partition.py        # Parameter slicing logic
│       └── __init__.py
├── scripts/
│   ├── test_dist_train.py          # Multi-process spawn test (Checkpointing)
│   └── train_real.py               # ResNet18 on CIFAR-10 training script
├── tests/
│   └── unit/                       # Optimizer and partition logic tests
└── pyproject.toml

How to Run

1. Simulated Distributed Training (CPU/Gloo)

Tests the core engine, ZeRO-2 partitioning, and distributed checkpointing. Spins up 2 processes using torch.multiprocessing.

python scripts/test_dist_train.py

Expected output: Loss drops, and Weights match after load: True is printed.

2. Real Model Training (ResNet18 on CIFAR-10)

Trains a real CNN model. Automatically falls back to mp.spawn if torchrun is not used, ensuring cross-platform compatibility.

python scripts/train_real.py

Expected output: Step-by-step loss reduction for ResNet18.

(To run on a real multi-GPU Linux cluster, use: torchrun --nproc_per_node=2 scripts/train_real.py)

Design Decisions

Why manual parameter partitioning instead of PyTorch FSDP? PyTorch FSDP is a black box. Building ZeRO-1/2 from scratch proves an understanding of the exact memory bottlenecks in distributed training and the collective communication patterns required to solve them.

Why asynchronous all_reduce? Blocking collective operations inside PyTorch's autograd engine cause deadlocks because threads can acquire locks in different orders across ranks. async_op=True launches the network request and returns immediately, letting autograd finish its work before we wait() for the network sync.

Why register_post_accumulate_grad_hook? Standard DDP wraps the entire model and reduces gradients in a single blocking call. Per-parameter hooks allow gradients to be reduced as soon as their layer's backward pass finishes, maximizing overlap between compute and network I/O.

Local Development

python -m venv venv
source venv/bin/activate  # Windows: venv\Scripts\activate
pip install -e ".[dev]"

# Run strict quality gates
ruff check .
mypy src/
pytest tests/