ZeRO One: Sharding the Optimizer

Reviewing DDP

Before we get started, let's quickly summarize DDP. In DDP, on each GPU, we have a copy of our model weights, the gradients of the weights, and the optimizer state for each weight. Once we do a forward pass on each GPU, we calculate the gradients in the backward pass on each GPU and use an all_reduce operation to sync gradients across GPUs. Finally we step the optimizer on each GPU to update the model weights with the gradients we just synchronized.

If you're thinking that we're storing a lot of duplicate parameters and doing a lot of duplicate work, you're right, we are! Fortunately we can shard some of this data across GPUs to reduce memory.

Training Memory Requirements

We can calculate the memory requirements of our model by multiplying the number of parameters by the number of bytes per parameter. If we're using full precision FP32 training our parameters and gradients require 4 bytes each.

If we are using the Adam optimizer we have to store the averaged momentem and the variance of the gradients, which means storing another 8 bytes (4 bytes each) for every parameter. In total we have to store 16 bytes per parameter.

We'll follow the notation from the paper, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, and use $ \Psi $ to represent the model size in terms of number of parameters.

For full FP32 training our total memory requirements for the weights, gradients, and optimizer states are:

$$ \begin{aligned} m_{params} &= 4 \cdot \Psi \newline m_{grad} &= 4 \cdot \Psi \newline m_{opt} &= (4 + 4) \cdot \Psi \newline \newline m_{total} &= m_{params} + m_{grad} + m_{opt} \newline & = 16 \cdot \Psi \end{aligned} $$

The paper uses $K$ to denote the memory multiplier of the optimizer states, and the memory required to store the optimizer state is $K\Psi$ bytes. In full FP32 training Adam has $K = 8$. In total we require

$$ 4 \Psi + 4 \Psi + K \Psi = 16 \Psi $$

In mixed precision training we use BF16 to store the parameters and gradient at 2 bytes each, but now we need to keep an additional copy of the parameters in FP32 for the optimizer on top of the normal 8 bytes per paramter required for the momentum and variance used by Adam. In total we need 12 bytes per parameter for the optimizer, or $K = 12$, which is more than the 8 bytes per parameter required for FP32 training.

$$ 2 \Psi + 2 \Psi + K \Psi = 16 \Psi $$

Even though it does not save memory mixed precision training is generally faster. If we decided to accumulate our gradients in FP32 to improve stability we would need an additional 4 bytes per parameter for a total of $20 \Psi$.

For mixed precision training our total memory requirements for the weights, gradients, and optimizer states are:

$$ \begin{aligned} m_{params} &= 2 \cdot \Psi \newline m_{grad} &= 2 \cdot \Psi \newline m_{params \_ fp32} &= 4 \cdot \Psi \newline m_{opt} &= (4 + 4) \cdot \Psi \newline \end{aligned} $$

The Adam optimizer takes half the total memory in FP32 training, and 75-80% in mixed precision training.

For FP32 training of a 1 billion parameter model we would need 16GB, so for our 320M parameter model we estimate

$$ 16 GB \cdot .32 = 5.12 GB $$

will be required for FP32 training, though in practice it will require a bit more for launching CUDA kernels or for temporary tensors like activations. We'll use the PyTorch profiler later to profile the memory and look in more detail at how memory is allocated during training.

The Basic Training Loop

As we noted earlier, our optimizer parameters are duplicated which means its possible to shard them, and because it accounts for a large portion of the memory consumption the reduction will be the most effective.

In a typical training loop, stepping the optimizer to update the model weights usually happens as the last step in the loop. By this point we've already made a prediction on a unique batch of data on each GPU, calculated the loss, and computed the gradients in the backwards pass. Because we are using replicas of the model on each GPU we have to sync our gradients to benefit from distributed training. Finally we would step the optimizer, but how can we do it in the distributed world where our optimizer is sharded?

When we shard the optimizer we divide its parameters to store a portion locally on each GPU. In order to update the model weights on every GPU we have to either, communicate the local optimizer's parameters to every other GPU before updating the weights on every GPU, or we can update the weights corresponding to our optimizer's local parameters and share those weights from the GPU where they were updated to all other GPUs.

Sharding Strategy

The former approach would result in duplicating the computation of updating the model weights on each GPU. The latter approach means we do the computation once per GPU. In both approaches the number of times we communicate is the same, but the former approach requires sending more data, as we saw previously the optimizer requires more memory than the weights!

Zero Redundancy Optimizer

In the paper ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, the authors made similar observations and introduced two model sharding strategies to mitigate the redundancies. The first strategy, ZeRO One, focuses on sharding the optimizer state.

The authors note that we can group the optimizer states into $N_d$ equal partitions, where $N_d$ denotes the degree of data parallelism (in our case $N_d = 4$) so that the $i$th GPU only updates the optimizer states corresponding to the $i$th partition. With this change each GPU now only needs to store and update $\frac{1} {N_d}$ of the total optimizer states, and update the $\frac{1} {N_d}$ parameters corresponding to those optimizer states.

Our equation for memory requirements, taking the mixed-precision training example, was $4 \Psi + K \Psi$ but now it becomes:

$$ 4 \Psi + \frac {K \Psi} {N_d} $$

We can see that when $N_d$ is larger than $K$, the memory requirements reduce to a negligible amount. This is a 75% reduction when $K = 12$.

This is represented in "Figure 1" of the paper. Here $P_{os}$ represents parameter optimizer sharding, considered the first stage of sharding for ZeRO, a.k.a ZeRO One.

Sharing Weight Updates

How do we correctly update the local parameters and share those updates across GPUS?

On each GPU, after updating only the weights corresponding to the local optimizer's parameters, the model's weights across GPUs are no longer in sync. Each model instance has a distinct slice of parameters with up to date weights. Since we already synced the gradients after the backwards pass the weight updates are based on the same gradients across GPUs.

Now that each GPU has $\frac{1} {N_d}$ of parameters updated, the other parameters need to be synced with the modal that has the updated weights every layer. We can use broadcast to share the updated weights from the GPU which contains the update to every other GPU.

!uv pip install --system datasets "git+https://github.com/muellerzr/nbdistributed"

%load_ext nbdistributed

%dist_init --num-processes 4

Starting 4 distributed workers...
✓ Successfully started 4 workers
Available commands:
  %%distributed - Execute code on all ranks (explicit)
  %%rank [0,n] - Execute code on specific ranks
  %sync - Synchronize all ranks
  %dist_status - Show worker status
  %dist_mode - Toggle automatic distributed mode
  %dist_shutdown - Shutdown workers

🚀 Distributed mode active: All cells will now execute on workers automatically!
   Magic commands (%, %%) will still execute locally as normal.

🐍 Below are auto-imported and special variables auto-generated into the namespace to use
  `torch`
  `dist`: `torch.distributed` import alias
  `rank` (`int`): The local rank
  `world_size` (`int`): The global world size
  `gpu_id` (`int`): The specific GPU ID assigned to this worker
  `device` (`torch.device`): The current PyTorch device object (e.g. `cuda:1`)

import random
import numpy as np
import torch

# Enable TF32 for matrix multiplications
torch.backends.cuda.matmul.allow_tf32 = True
# Enable TF32 for cuDNN (convolution operations)
torch.backends.cudnn.allow_tf32 = True

common_seed = 314

def set_seed(seed: int = common_seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL_NAME = "HuggingFaceTB/SmolLM2-360M-Instruct"

def get_smol_model():
    return AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME, num_labels=2, torch_dtype="bfloat16"
    )

def get_smol_tokenizer():
    return AutoTokenizer.from_pretrained(MODEL_NAME)

from datasets import load_dataset

def get_dataset():
    dataset = load_dataset("glue", "mrpc")
    tokenizer = get_smol_tokenizer()

    def tokenize_func(examples):
        return tokenizer(
            examples["sentence1"],
            examples["sentence2"],
            max_length=None,
            truncation=True,
            padding=True,
        )

    dataset = dataset.map(
        tokenize_func, batched=True, remove_columns=["idx", "sentence1", "sentence2"]
    )
    dataset = dataset.rename_columns({"label": "labels"})
    return dataset

tokenizer = get_smol_tokenizer()
dataset = get_dataset()["train"]

from torch.utils.data import DataLoader

train_ds = dataset.shuffle(seed=common_seed)

def collate_func(batch):
    return tokenizer.pad(
        batch,
        padding="longest",
        max_length=None,
        pad_to_multiple_of=8,
        return_tensors="pt",
    )

def get_dataloader(ds, batch_size = 16, seed = common_seed):
    return DataLoader(
        ds,
        batch_size=batch_size,
        collate_fn=collate_func,
        drop_last=True,
        shuffle=True
    )

from hta.trace_analysis import TraceAnalysis
import json

def get_trace_dicts(trace_dir: str):
    analyzer = TraceAnalysis(trace_dir = trace_dir)

    time_spent_df = analyzer.get_temporal_breakdown(visualize=True)
    overlap_df = analyzer.get_comm_comp_overlap(visualize=True)
    kernel_type_metrics_df, _ = analyzer.get_gpu_kernel_breakdown(visualize=True)

def get_memory_timeline(mem_file: str):
    mem_analyzer = MemoryAnalysis(mem_path)
    mem_timeline = mem_analyzer._process_raw_events()
    mem_timeline.plot_memory_timeline()

def get_profiler(file_prefix: str, steps: int = 4):
    worker_name=f"gpu-rank-0{dist.get_rank()}"
    tensorboard_trace_handler = torch.profiler.tensorboard_trace_handler(
        f"{file_prefix}",
        worker_name=worker_name,
        use_gzip=True,
    )

    def trace_handler(prof: torch.profiler.profile):
        tensorboard_trace_handler(prof)
        prof.export_memory_timeline(f"{file_prefix}-{worker_name}-memory-timeline.raw.json.gz")

    return torch.profiler.profile(
        activities=[
            # profile activity on the CPU and GPU
            torch.profiler.ProfilerActivity.CPU,
            torch.profiler.ProfilerActivity.CUDA,
        ],
        # Setup the profiler schedule
        schedule=torch.profiler.schedule(wait=0, warmup=0, active=steps + 1),
        # Save the trace files in this directory when it is ready
        on_trace_ready=trace_handler,
        profile_memory=True,
        record_shapes=True,
        with_stack=True,
    )

ws = dist.get_world_size()
rank = dist.get_rank()

train_ds = dataset.shuffle(seed=common_seed)

# To ensure equal distribution of the dataset, we calculate the dataset length per
# rank/GPU by dividing total dataset length by world size
ds_length = len(train_ds)
ds_length_per_gpu = ds_length // ws

# We calculate the start and end index of the data shard
# Remember the value of `rank` will be unique on each GPU
start = rank * ds_length_per_gpu
end = start + ds_length_per_gpu if rank != ws - 1 else ds_length

train_shard = train_ds.select(list(range(start, end)))

Initialization

First let's walk through sharding the optimizer step by step. Since we also need to sync the gradients for ZeRO One, we'll reuse the SimpleDDP class from my previous post on DDP as a starting point.

class SimpleDDP:
    def __init__(self, model: torch.nn.Module):
        self.model = model

        # Verify our model parameters are in sync at initialization as we did before
        for p in model.parameters():
            rank0_param = p.data.clone()
            dist.broadcast(rank0_param, src=0)

            if not torch.equal(p.data, rank0_param):
                raise ValueError(
                    "Model parameters are not synced across ranks during __init__.",
                    "Use set_seed"
                )

    def sync_gradients(self):
        for p in self.model.parameters():
            if p.grad is not None:
                # Sum the gradient across all GPUs and take the average
                dist.all_reduce(p.grad, op=dist.ReduceOp.SUM)
                # Divide by the total number of GPUs to obtain the average
                p.grad /= dist.get_world_size()

    def __call__(self, *args, **kwargs):
        return self.model(*args, **kwargs)

Lets setup our model.

set_seed()

train_dataloader = get_dataloader(train_shard)

model = get_smol_model()
model.to(device)
ddp_model = SimpleDDP(model) # Wrap the model in our DDP class

Next we create the optimizer, we'll use Adam.

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

Sharding Parameters

Now we need to get a list of the parameters in the optimizer. In PyTorch, optimizer parameters are stored in parameter groups.

og_param_groups = optimizer.param_groups # Keep a copy of the original parameter groups
params = [p for g in og_param_groups for p in g['params']] # Get a list of parameters

Before we divide the parameters we need to calculate the number of params per GPU, including the remainder if they can not be divided exactly evenly.

params_per_rank = len(params) // ws
remainder = len(params) % ws

len(params), params_per_rank, remainder

🔹 Rank 0:
  (72, 3)

🔹 Rank 1:
  (72, 3)

🔹 Rank 2:
  (72, 3)

🔹 Rank 3:
  (72, 3)

Now let's use that to calculate the start and end index of the parameters to assign to each GPU

# Use params_per_rank to calculate the start and
# end indices with an added offset for any remainder
start_idx = rank * params_per_rank + min(rank, remainder)
end_idx = start_idx + params_per_rank + (1 if rank < remainder else 0)

start_idx, end_idx

🔹 Rank 0:
  (0, 73)

🔹 Rank 1:
  (73, 146)

🔹 Rank 2:
  (146, 219)

🔹 Rank 3:
  (219, 291)

Using the indices we just calculated, we can now get the local parameters on each GPU.

local_param_indices = list(range(start_idx, end_idx))
# Get a set of parameters if they are in the range of local indices
local_params = set(params[i] for i in local_param_indices)

Now we overwrite the optimizer's param groups, keeping only the local parameters. The other parameters will be dereferenced and the memory representing those parameters will be cleaned up, achieving a memory reduction on each GPU!

for g in optimizer.param_groups:
    g['params'] = [p for p in g['params'] if p in local_params]

Making a Prediction

Now let's train on a batch of data to verify our sharded optimizer is working.

batch = next(iter(train_dataloader))
batch = {k: v.to(device) for k, v in batch.items()}

output = model(**batch) # Make a prediction on one batch of data
output.loss.backward()  # Backpropagate to calculate the gradients

ddp_model.sync_gradients() # Sync the gradients

After our gradients have synchronized we need to step the optimizer. Since we've overwritten the parameters to remove the non-local ones, it will only update the model parameters which correspond to the optimizer's local parameters.

optimizer.step()

Now on each model/GPU the weights corresponding to the optimizer's local parameters are updated, but the remaining non-local parameters are out of sync. We need to broadcast the updated local parameters to all other GPUs so that every copy of our model has the current weights before the next step in our training loop.

But first let's show that our model's parameters are out of sync. The code below will check if the parameters are equal, we should expect it to raise an exception.

# We check every parameter of the model
for p in model.parameters():
    # First we allocate some space for receiving the broadcasted parameter.
    rank0_param = p.data.clone()
    # Broadcast the `rank0_param` values from the rank 0 GPU into `rank0_param` on other
    # GPUs, they are the same shape.
    dist.broadcast(rank0_param, src=0)

    # Raise an error if they ar enot equal
    if not torch.equal(p.data, rank0_param):
        raise ValueError("Model parameters are not in sync")

Synchronizing Parameters

Let's synchronize our parameters by iterating through the original full list of parameters and broadcasting local or receiving a broadcast of non-local parameters.

for i, p in enumerate(params):
    # Compute the owner rank
    if i < (params_per_rank + 1) * remainder:
        owner_rank = i // (params_per_rank + 1)
    else:
        owner_rank = (i - remainder) // params_per_rank

    # The owner rank will broadcast the data will other GPUs will receive it
    dist.broadcast(p.data, src=owner_rank)

Finally we can zero the gradients.

optimizer.zero_grad()

We've stepped through one batch in our training loop in which we:

1. Updated the model's parameters locally
2. Broadcasted the updated local parameters to all other GPUs
3. Received broadcast of all non-local parameters from their owner GPUs

Let's confirm that our model weights are in sync using the same code as before.

# We check every parameter of the model
for p in model.parameters():
    # First we allocate some space for receiving the broadcasted parameter.
    rank0_param = p.data.clone()
    # Broadcast the `rank0_param` values from the rank 0 GPU into `rank0_param` on other
    # GPUs, they are the same shape.
    dist.broadcast(rank0_param, src=0)

    # Raise an error if they ar enot equal
    if not torch.equal(p.data, rank0_param):
        raise ValueError("Model parameters are not in sync")

This time we don't see an error which means it worked!

Free Memory

del model.model, model
del optimizer, params, local_params, local_param_indices
torch.cuda.empty_cache()

Distributed Data Parallelism

Before we implement and profile ZeRO One we should establish a baseline to which we can compare our implementation. For our baseline we'll use DDP without ZeRO One; we initialize our model and optimizer below.

set_seed()

train_dataloader = get_dataloader(train_shard)

model = get_smol_model()
model.to(device)
model.train()
ddp_model = SimpleDDP(model)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

max_steps = 4

Then we run a training loop with the PyTorch profiler.

with get_profiler("/workspace/zero1/ddp-baseline", max_steps) as prof:
    for step, batch in enumerate(train_dataloader):
        if step >= max_steps:
            break

        batch = {k: v.to(device) for k, v in batch.items()}

        output = ddp_model(**batch)
        output.loss.backward()

        ddp_model.sync_gradients()

        optimizer.step()
        optimizer.zero_grad()

        prof.step()

For more details on using the profiler, you can refer to the documentation. I also cover a bit of how to get started with the PyTorch profiler in my previous post on DDP.

Profiling

We'll examine the profiler trace and use the HolisticTraceAnalysis library to get more insights on how efficiently we are utilizing our GPUs.

To make it easier to ingest for loading in the browser, I've downsampled the memory timeline data by a factor of ~10.

Looking at the data for the "GPU Memory Timeline" chart, we can see the total memory peaks at around ~5.76 GB and optimizer state consuming ~1.39 GB. Since we'll be using 4 GPUs we expect this number to theoretically come down to ~.35 GB, however we'll see that in practice this will be a bit higher.

The PyTorch memory profiler can only track memory allocated through PyTorch, this means memory allocated directly from CUDA APIs won't be visible to it. When NCCL allocates memory for communication buffers through CUDA APIs, the PyTorch memory profiler can't properly categorize that memory allocation and marks it as "unknown". For more information see the note in the PyTorch documentation

If we look at the "Temporal Breakdown Across Ranks" chart we can see that we are at about 16% non-compute time, and the "Computation-Communication Overlap" chart shows that there is 0% overlap.

For convenience (and fun!) I've also embedded a trace viewer into this post; clicking the "Show trace explore" heading below will reveal it.

Sharding the Optimizer State

Now that we've established a baseline, let's pull the code from where we sharded the optimizer into a class we can reuse.

class ZeROOne:
    def __init__(self, optimizer):
        self.optimizer = optimizer

        # Keep track of the original parameter groups
        self.og_param_groups = optimizer.param_groups
        # Flatten the list of parameters for splitting
        self.params = [
            p for group in self.og_param_groups for p in group['params']
        ]

        world_size = dist.get_world_size()
        rank = dist.get_rank()

        # Calculate which parmeters to keep locally
        self.params_per_rank = len(self.params) // world_size
        self.remainder = len(self.params) % world_size

        start_idx = rank * self.params_per_rank + min(rank, self.remainder)
        end_idx = start_idx + self.params_per_rank + (1 if rank < self.remainder else 0)

        local_param_indices = list(range(start_idx, end_idx))
        local_params = set(self.params[i] for i in local_param_indices)

        self._shard_optimizer_params(local_params)

    def _shard_optimizer_params(self, local_params):
        # Update the optimizer's param groups to only include local params
        for g in self.optimizer.param_groups:
            g['params'] = [p for p in g['params'] if p in local_params]

    def step(self):
        # Step the local parameters
        self.optimizer.step()
        # Zero the gradients of the local parameters that have been stepped
        self.optimizer.zero_grad()

        for i, p in enumerate(self.params):
            # Compute the owner rank
            if i < (self.params_per_rank + 1) * self.remainder:
                owner_rank = i // (self.params_per_rank + 1)
            else:
                owner_rank = (i - self.remainder) // self.params_per_rank

            # Broadcast the data from the owner rank to all other ranks
            dist.broadcast(p.data, src=owner_rank)

    def zero_grad(self):
        # We zero out every parameter's gradient since optimizer.zero_grad
        # will only zero the gradients of locally tracked parameters
        for p in self.params:
            p.grad = None

Training

Let's run and profile a training loop using the ZeROOne class we just defined.

set_seed()

train_dataloader = get_dataloader(train_shard)

model = get_smol_model()
model.to(device)
model.train()
ddp_model = SimpleDDP(model)

opt = torch.optim.Adam(model.parameters(), lr=1e-3)

max_steps = 4

We'll wrap the Adam optimizer with our ZeROOne class so that it gets sharded across GPU ranks.

optimizer = ZeROOne(opt)

Each GPU now has one shard of our optimizer state! Now we run a training loop and profile it with the PyTorch profiler.

with get_profiler("/workspace/zero1/zero1-ddp", max_steps) as prof:
    for step, batch in enumerate(train_dataloader):
        if step >= max_steps:
            break

        batch = {k: v.to(device) for k, v in batch.items()}

        output = ddp_model(**batch)
        output.loss.backward()

        # Wait for all async gradient sync operations to complete
        ddp_model.sync_gradients()

        optimizer.step()
        optimizer.zero_grad()

        prof.step()

Profiling

This time we see a our optimizer is now using .484 GB of memory, down from 1.39 GB without ZeRO One, a 65% reduction! We achive an overall reduction in memory usage of ~12.24%, using a peak of 5.05 GB.

Unfortunately, the "Temporal Breakdown Across Ranks" chart shows that we have about 22% non-compute time and the "Computation-Communication Overlap" chart shows that there is 0% overlap.

Asynchronous Communication

Let's improve the computation-communication overlap by making our gradient synchronization use non-blocking communication. We will run our all_reduce asynchronously!

Below is a diagram showing how our initial synchronous version works.

All Params: [────Backward────]──[AllReduce]──[Step]──[Broadcast]

We can note that once a gradient has been computed in the backwards pass we can immediately begin to synchronize it across GPU ranks, and since the gradient has already been computed the update is independent of other parameters' gradients. After making these changes our operations should look like this:

Param 4:  [Backward]──[AllReduce]────────────────────────────────────────[Step]──[Broadcast]

Param 3:              [Backward]──[AllReduce]────────────────────────────[Step]──[Broadcast]

Param 2:                          [Backward]──[AllReduce]────────────────[Step]──[Broadcast]

Param 1:                                      [Backward]──[AllReduce]────[Step]──[Broadcast]

We'll reuse the asynchronous version of DDP from my previous post. The main change from the synchronous version is that we pass async_op=True to all_reduce, this allows us to register a callback to run after the asynchronous operation completes. In this callback we'll complete the process of averaging the gradient by dividing by the total number of GPUs, which is four in our case.

class SimpleDDPAsync(SimpleDDP):
    def __init__(self, model: torch.nn.Module):
        super().__init__(model)
        self.pending_futures = []  # Store futures to wait on
        self.register_backward_hook()

    def _sync_gradient(self, param):
        if param.grad is None:
            return

        # Sum the gradient across all GPUs and take the average
        work = dist.all_reduce(param.grad, op=dist.ReduceOp.SUM, async_op=True)

        def after_sync_cb(fut):
            # Average in-place once reduction finished
            ws = dist.get_world_size()
            param.grad.div_(ws)
            return None

        # Get the future and chain the callback
        future = work.get_future().then(after_sync_cb)

        # Store the future so we can wait on it later
        self.pending_futures.append(future)

    def register_backward_hook(self):
        self.sync_hooks = []

        for p in self.model.parameters():
            if p.requires_grad:
                h = p.register_post_accumulate_grad_hook(self._sync_gradient)
                self.sync_hooks.append(h)

    def wait_for_gradients(self):
        # Wait for all pending async gradient operations to be enqueued
        # to the CUDA stream
        for future in self.pending_futures:
            future.wait()
        self.pending_futures.clear()

Training

We set up and run our training run as we did before.

set_seed()

train_dataloader = get_dataloader(train_shard)

model = get_smol_model()
model.to(device)
model.train()

opt = torch.optim.Adam(model.parameters(), lr=1e-3)

optimizer = ZeROOne(opt)
ddp_model = SimpleDDPAsync(model)

max_steps = 4

with get_profiler("/workspace/zero1/zero1-ddp-async", max_steps) as prof:
    for (i, batch) in enumerate(train_dataloader):
        if i >= max_steps:
            break
        batch = {k: v.to(device) for k, v in batch.items()}

        output = ddp_model(**batch)
        output.loss.backward()

        # Wait for all async gradient sync operations to complete
        ddp_model.wait_for_gradients()
        torch.cuda.synchronize()

        optimizer.step()
        optimizer.zero_grad()

        prof.step()

Profiling

We see in our "GPU Memory Timeline" that peak memory is similar to what we saw in when we used synchronous DDP. However, in the "Temporal Breakdown Across Ranks" chart, we see that the non-compute time has dropped to 13% from 22%, a ~41% improvement! Additionally the "Computation-Communication Overlap" chart shows that we now have some overlap and are at nearly 77% overlap.

Bucketing

Instead of synchronizing every paramter with a separate all_reduce call, we can improve even further by bucketing our gradients and instead synchronizing when our bucket is full. Again, we use a modified version of DDP with bucketed gradient syncs from my previous post on DDP. See that post for a detailed explanation on how the bucketing is implemented.

class SimpleDDPBucket(SimpleDDP):
    def __init__(self, model: torch.nn.Module, bucket_cap_mb=25):
        super().__init__(model)
        self.bucket_cap_bytes = bucket_cap_mb * 1024 * 1024
        self._bucket = []
        self._bucket_size = 0
        # List of pending Futures to wait on
        self._pending_futures = []

        self._register_hooks()

    def _register_hooks(self):
        self._bucket_hooks = []
        # Register hooks to collect gradients and synchronize in buckets
        for p in self.model.parameters():
            if p.requires_grad:
                h = p.register_post_accumulate_grad_hook(self._bucket_hook)
                self._bucket_hooks.append(h)

    def _bucket_hook(self, param):
        if param.grad is not None:
            # Append the param to the bucket so it will be synced
            self._bucket.append(param)
            # Increase the bucket size counter
            self._bucket_size += param.grad.numel() * param.grad.element_size()
            # Sync if bucket full
            if self._bucket_size >= self.bucket_cap_bytes:
                self._sync_bucket()

    def _sync_bucket(self):
        # Synchronize all gradients in the bucket using a coalesced reduce
        grads = [p.grad for p in self._bucket]
        if not grads:
            return

        bucket_params = list(self._bucket)  # snapshot before clearing
        # Clear bucket state early so hooks can refill while async op in flight
        self._bucket = []
        self._bucket_size = 0

        # Coalesce gradients into a single contiguous tensor
        flat_grads = torch.cat([g.view(-1) for g in grads])

        # Launch async all-reduce on the coalesced tensor
        work = dist.all_reduce(flat_grads, op=dist.ReduceOp.SUM, async_op=True)

        # Register callback to run after gradient sync
        def _after_sync_cb(fut):
            # Average and copy back once reduction finished
            flat_grads.div_(dist.get_world_size())

            # Copy back gradients
            offset = 0
            for p in bucket_params:
                numel = p.grad.numel()
                p.grad = flat_grads[offset:offset + numel].view_as(p.grad)
                offset += numel

            return None

        # Get the future, chain the callback, and store the resulting future
        future = work.get_future().then(_after_sync_cb)
        self._pending_futures.append(future)

    def flush_bucket(self):
        # Flush any remaining gradients in the bucket
        if self._bucket:
            self._sync_bucket()

    def wait_for_gradients(self):
        # Wait for all pending async gradient operations to be enqueued
        # to the CUDA stream
        for future in self._pending_futures:
            future.wait()
        self._pending_futures.clear()

Training

Now we setup and run our training loop again.

set_seed()

train_dataloader = get_dataloader(train_shard)

model = get_smol_model()
model.to(device)
model.train()

opt = torch.optim.Adam(model.parameters(), lr=1e-3)

optimizer = ZeROOne(opt)
ddp_model = SimpleDDPBucket(model)

max_steps = 4

with get_profiler("/workspace/zero1/zero1-ddp-bucket", max_steps) as prof:
    for (i, batch) in enumerate(train_dataloader):
        if i >= max_steps:
            break
        batch = {k: v.to(device) for k, v in batch.items()}

        output = ddp_model(**batch)
        output.loss.backward()

        # Flush any remaining gradients in the bucket
        ddp_model.flush_bucket()

        # Wait for all async gradient sync operations to complete
        ddp_model.wait_for_gradients()
        torch.cuda.synchronize()

        optimizer.step()
        optimizer.zero_grad()

        prof.step()

Profiling

Our results show below that implementing bucketing slightly reduced non-compute time, but did not yield much improvement. This is likely due to the the lack of overlap during broadcasts, which occur after all gradients have synced and we've stepped the optimizer. Let's try to improve this.

Overlapping Parameter Synchronizations

We'll make one last optimization, however this time our change well center on the optimizer and paramter broadcasts rather than the gradient all_reduces as we did througout this post.

Similarly to how we exploited gradient synchronizations being independent of each other to immediately trigger the all_reduces after gradient computation, we can observe that updating a parameter with a gradient is independent of other parameter updates. This means that as soon as a gradient is synchronized, we can immediately update the parameter of that gradient and broadcast the update. Visualized, this would look like:

Bucket 1:  [Backward]──[AllReduce]──[Step]──[Broadcast]

Bucket 2:              [Backward]──[AllReduce]──[Step]──[Broadcast]

Bucket 3:                          [Backward]──[AllReduce]──[Step]──[Broadcast]

Bucket 4:                                      [Backward]──[AllReduce]──[Step]──[Broadcast]

Calling .step() on the PyTorch optimizer will update every one of the paramters the optimizer is tracking. So instead of sharding the optimizer by removing different parameters from them on each GPU, we can instantiate one optimizer per parameter. While there is some additional overhead for each optimizer Python object, we'll see that this overhead is minimal, especially when compared to the optimizer state of the parameter (its averaged momentum and variance).

Now when we call .step() on an optimizer it will only update the one paramter it is responsible for, which means we can do it immediately after the gradient has synchronized. Then we can immediately overlap and asynchronously fire the broadcast call to synchronize the updated parameters across to GPU ranks.

class ShardedOptimizer:
    def __init__(self, optimizer_cls, model, **optimizer_kwargs):
        self.optimizer_cls = optimizer_cls
        self.optimizer_kwargs = optimizer_kwargs

        # Flatten all parameters
        self.params = list(model.parameters())

        world_size = dist.get_world_size()
        rank = dist.get_rank()

        # Calculate parameter partitioning
        self.params_per_rank = len(self.params) // world_size
        self.remainder = len(self.params) % world_size

        # Create optimizer per parameter and track ownership
        self.param_to_optimizer = {}
        self.param_to_owner_rank = {}

        for i, p in enumerate(self.params):
            if not p.requires_grad:
                continue

            # Compute owner rank for this parameter
            if i < (self.params_per_rank + 1) * self.remainder:
                owner_rank = i // (self.params_per_rank + 1)
            else:
                owner_rank = (i - self.remainder) // self.params_per_rank

            self.param_to_owner_rank[p] = owner_rank

            # Create optimizer only for locally owned parameters
            if owner_rank == rank:
                self.param_to_optimizer[p] = optimizer_cls([p], **optimizer_kwargs)

    def step_parameter(self, param):
        current_rank = dist.get_rank()
        owner_rank = self.param_to_owner_rank[param]

        if current_rank == owner_rank:
            optimizer = self.param_to_optimizer[param]
            optimizer.step()

        # Clear gradient immediately to free memory
        param.grad = None

        # Broadcast updated parameter from owner rank
        owner_rank = self.param_to_owner_rank[param]
        return dist.broadcast(param.data, src=owner_rank, async_op=True)

class SimpleDDPBucketAsync(SimpleDDPBucket):
    def __init__(self, model: torch.nn.Module, zero1_optimizer, bucket_cap_mb=25):
        super().__init__(model, bucket_cap_mb)
        self.zero1_optimizer = zero1_optimizer

        # Separate tracking of gradient and broadcast futures
        self._pending_gradient_futures = []
        self._pending_broadcast_futures = []

    def _sync_bucket(self):
        grads = [p.grad for p in self._bucket]
        if not grads:
            return

        bucket_params = list(self._bucket)  # snapshot before clearing
        # Clear bucket state early so hooks can refill while async op in flight
        self._bucket = []
        self._bucket_size = 0

        # Coalesce gradients into a single contiguous tensor
        flat_grads = torch.cat([g.view(-1) for g in grads])

        # Launch async all-reduce on the coalesced tensor
        work = dist.all_reduce(flat_grads, op=dist.ReduceOp.SUM, async_op=True)

        # Register callback to run after gradient sync
        def _after_sync_cb(fut):
            # Average gradients
            flat_grads.div_(dist.get_world_size())

            # Process each parameter: copy back gradient, step optimizer, broadcast
            offset = 0
            for p in bucket_params:
                # Copy back gradient
                numel = p.grad.numel()
                p.grad.data = flat_grads[offset:offset + numel].view_as(p.grad)
                offset += numel

                work = self.zero1_optimizer.step_parameter(p)
                self._pending_broadcast_futures.append(work.get_future())

            return None

        # Get the future, chain the callback, and store the resulting future
        future = work.get_future().then(_after_sync_cb)
        self._pending_gradient_futures.append(future)

    def wait_for_gradients(self):
        for future in self._pending_gradient_futures:
            future.wait()
        self._pending_gradient_futures.clear()

    def wait_for_broadcasts(self):
        # Wait for all pending async parameter operations to be enqueued
        # to the CUDA stream
        for future in self._pending_broadcast_futures:
            future.wait()
        self._pending_broadcast_futures.clear()

Training

We set up our training loop as we did before, this time with some slight adjustments for our new ShardedOptimizer and SimpleDDPBucketAsync classes.

set_seed()

train_dataloader = get_dataloader(train_shard)

model = get_smol_model()
model.to(device)
model.train()

# Create async ZeRO-1 optimizer
zero1_opt = ShardedOptimizer(torch.optim.Adam, model, lr=1e-3)
ddp_model = SimpleDDPBucketAsync(model, zero1_optimizer=zero1_opt)

max_steps = 4

with get_profiler("/workspace/zero1/zero1-ddp-bucket-opt", max_steps) as prof:
    for (i, batch) in enumerate(train_dataloader):
        if i >= max_steps:
            break
        batch = {k: v.to(device) for k, v in batch.items()}

        output = ddp_model(**batch)
        output.loss.backward()

        # Flush any remaining gradients in the bucket
        ddp_model.flush_bucket()

        # Wait for all async gradient sync operations to be enqueued to the CUDA stream
        ddp_model.wait_for_gradients()

        # Wait for all parameter broadcasts to be enqueued to the CUDA stream
        ddp_model.wait_for_broadcasts()

        # Ensure GPU operations complete
        torch.cuda.synchronize()

        prof.step()

Profiling

Our "GPU Memory Timeline" shows that we are using a little bit more memory than before, and now we have a reduction of 10% from our baseline instead of 12% we saw without the immediate optimizer steps.

We can see an improvement in our "Temporal Breakdown Across Ranks" and "Computation-Communication Overlap" charts, we are now at 7% non-compute time and 87% overlap! That is a 70% reduction in non-compute time from our initial ZeRO One implementation!

As we saw, these improvements comes by trading off a slight increase in memory. This tradeoff might be worthwhile since if we want to maximize the utilization of the GPUs we are paying for!