Distributed DP-SGD Training¶

This notebook demonstrates how to run DP-SGD training across multiple processes using PyTorch's Distributed Data Parallel (DDP). The key differences from single-device training are:

1. Each rank computes clipped gradients on its local data shard.
2. Clipped gradients are summed across ranks with sum_gradients().
3. All ranks add the same noise (using the same RNG key).
4. The dataset is sharded with local_shard() and each rank gets a per-rank sampling key via fold_in().

Prerequisites: DP-SGD Training, familiarity with torchrun.

Components exercised: sum_gradients, local_shard, PoissonSampler, gaussian_noise with shared RNG key.

Note: This notebook is meant to be saved as a .py script and launched with torchrun --nproc_per_node=N. The cells below show the complete code with explanations.

Setup and initialization¶

Initialize the distributed process group and assign each rank a device.

In [ ]:

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset, default_collate

import opake.accounting as acc
import opake.dpsgd.accounting as dpsgd_acc
from opake.dpsgd.clipping import clipped_grad
from opake.dpsgd.noise import gaussian_noise
from opake.dpsgd.sampling import PoissonSampler
from opake.functional import empty_collate, make_functional
from opake.distributed.gradients import sum_gradients_
from opake.random import key, fold_in
from opake.distributed import local_shard

# Initialize distributed (set by torchrun)
dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo")
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu")
if device.type == "cuda":
    torch.cuda.set_device(device)

print(f"Rank {rank}/{world_size} on {device}")

import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset, default_collate import opake.accounting as acc import opake.dpsgd.accounting as dpsgd_acc from opake.dpsgd.clipping import clipped_grad from opake.dpsgd.noise import gaussian_noise from opake.dpsgd.sampling import PoissonSampler from opake.functional import empty_collate, make_functional from opake.distributed.gradients import sum_gradients_ from opake.random import key, fold_in from opake.distributed import local_shard # Initialize distributed (set by torchrun) dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") rank = dist.get_rank() world_size = dist.get_world_size() device = torch.device(f"cuda:{rank}" if torch.cuda.is_available() else "cpu") if device.type == "cuda": torch.cuda.set_device(device) print(f"Rank {rank}/{world_size} on {device}")

Model and data¶

Every rank loads the same model. The dataset is explicitly sharded using local_shard(), so each rank trains on a disjoint partition.

In [ ]:

torch.manual_seed(42)

# Synthetic dataset (same on all ranks)
dataset_size = 10_000
X = torch.randn(dataset_size, 20)
y = (X[:, 0] + X[:, 1] > 0).float()
dataset = TensorDataset(X, y)


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(20, 64)
        self.fc2 = nn.Linear(64, 1)

    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x))).squeeze(-1)


model = MLP().to(device)
fmodel, params = make_functional(model)

torch.manual_seed(42) # Synthetic dataset (same on all ranks) dataset_size = 10_000 X = torch.randn(dataset_size, 20) y = (X[:, 0] + X[:, 1] > 0).float() dataset = TensorDataset(X, y) class MLP(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(20, 64) self.fc2 = nn.Linear(64, 1) def forward(self, x): return self.fc2(F.relu(self.fc1(x))).squeeze(-1) model = MLP().to(device) fmodel, params = make_functional(model)

DP components¶

Calibrate noise for the global training run. The sample rate is based on the full dataset size (not per-rank).

Pass the same key(seed) on every rank so all devices generate identical noise.

In [ ]:

BATCH_SIZE = 64
NUM_STEPS = 500
TARGET_EPS = 3.0
DELTA = 1e-5
LR = 0.1

sample_rate = BATCH_SIZE / dataset_size

# Calibrate noise multiplier
result = acc.calibrate(
    acc.epsilon_budget(TARGET_EPS, delta=DELTA),
    lambda nm: dpsgd_acc.poisson(dpsgd_acc.gaussian(nm), sample_rate) * NUM_STEPS,
    param_min=0.1,
    param_max=10.0,
)
noise_multiplier = result.param

if rank == 0:
    print(f"Noise multiplier: {noise_multiplier:.4f}")


# Loss function
def loss_fn(params, x, y):
    logit = fmodel(params, x.unsqueeze(0)).squeeze()
    return F.binary_cross_entropy_with_logits(logit, y)


# Clipped gradient function
grad_fn, clip_state = clipped_grad(
    loss_fn,
    argnums=0,
    batch_argnums=(1, 2),
    clipping_norm=1.0,
    normalize_by=BATCH_SIZE,
)

# Noise function — same key on all ranks => identical noise
noise_fn, noise_state = gaussian_noise(
    noise_multiplier=noise_multiplier,
    key=key(42),
)

# Shard dataset and create per-rank sampler
shard = local_shard(dataset, rank=rank, world_size=world_size)
sampler = PoissonSampler(
    shard,
    sample_rate=sample_rate,
    n_steps=NUM_STEPS,
    key=fold_in(key(0), rank),
)
collate = empty_collate(default_collate)
collate([shard[0]])
loader = DataLoader(shard, batch_sampler=sampler, collate_fn=collate)

BATCH_SIZE = 64 NUM_STEPS = 500 TARGET_EPS = 3.0 DELTA = 1e-5 LR = 0.1 sample_rate = BATCH_SIZE / dataset_size # Calibrate noise multiplier result = acc.calibrate( acc.epsilon_budget(TARGET_EPS, delta=DELTA), lambda nm: dpsgd_acc.poisson(dpsgd_acc.gaussian(nm), sample_rate) * NUM_STEPS, param_min=0.1, param_max=10.0, ) noise_multiplier = result.param if rank == 0: print(f"Noise multiplier: {noise_multiplier:.4f}") # Loss function def loss_fn(params, x, y): logit = fmodel(params, x.unsqueeze(0)).squeeze() return F.binary_cross_entropy_with_logits(logit, y) # Clipped gradient function grad_fn, clip_state = clipped_grad( loss_fn, argnums=0, batch_argnums=(1, 2), clipping_norm=1.0, normalize_by=BATCH_SIZE, ) # Noise function — same key on all ranks => identical noise noise_fn, noise_state = gaussian_noise( noise_multiplier=noise_multiplier, key=key(42), ) # Shard dataset and create per-rank sampler shard = local_shard(dataset, rank=rank, world_size=world_size) sampler = PoissonSampler( shard, sample_rate=sample_rate, n_steps=NUM_STEPS, key=fold_in(key(0), rank), ) collate = empty_collate(default_collate) collate([shard[0]]) loader = DataLoader(shard, batch_sampler=sampler, collate_fn=collate)

Training loop¶

The key distributed operation is sum_gradients(), which all-reduces the clipped gradients across ranks. After summing, all ranks add the same noise (from the shared key).

In [ ]:

for step, (xb, yb) in enumerate(loader):
    if step >= NUM_STEPS:
        break

    xb, yb = xb.to(device), yb.to(device)

    # 1. Per-example clipped gradients (local data only)
    grads, clip_state = grad_fn(params, xb, yb, state=clip_state)

    # 2. Sum clipped gradients across all ranks (in-place)
    sum_gradients_(grads)

    # 3. Add noise (identical on all ranks); noisy_grads is a NoisedPytree
    noisy_grads, noise_state = noise_fn(grads, noise_state)

    # 4. SGD update — unwrap .pytree for the raw tensor leaves
    params = tuple(p - LR * g for p, g in zip(params, noisy_grads.pytree))

    if rank == 0 and (step + 1) % 100 == 0:
        with torch.no_grad():
            loss = loss_fn(params, xb[0], yb[0]).item() if len(xb) else float("nan")
        eps = (
            dpsgd_acc.poisson(dpsgd_acc.gaussian(noise_multiplier), sample_rate) * (step + 1)
        ).epsilon_at(DELTA)
        print(f"Step {step + 1:4d}  loss={loss:.4f}  epsilon={eps:.2f}")

for step, (xb, yb) in enumerate(loader): if step >= NUM_STEPS: break xb, yb = xb.to(device), yb.to(device) # 1. Per-example clipped gradients (local data only) grads, clip_state = grad_fn(params, xb, yb, state=clip_state) # 2. Sum clipped gradients across all ranks (in-place) sum_gradients_(grads) # 3. Add noise (identical on all ranks); noisy_grads is a NoisedPytree noisy_grads, noise_state = noise_fn(grads, noise_state) # 4. SGD update — unwrap .pytree for the raw tensor leaves params = tuple(p - LR * g for p, g in zip(params, noisy_grads.pytree)) if rank == 0 and (step + 1) % 100 == 0: with torch.no_grad(): loss = loss_fn(params, xb[0], yb[0]).item() if len(xb) else float("nan") eps = ( dpsgd_acc.poisson(dpsgd_acc.gaussian(noise_multiplier), sample_rate) * (step + 1) ).epsilon_at(DELTA) print(f"Step {step + 1:4d} loss={loss:.4f} epsilon={eps:.2f}")

Cleanup¶

In [ ]:

if rank == 0:
    eps_final = (
        dpsgd_acc.poisson(dpsgd_acc.gaussian(noise_multiplier), sample_rate) * NUM_STEPS
    ).epsilon_at(DELTA)
    print(f"\nFinal epsilon: {eps_final:.2f}")

dist.destroy_process_group()

if rank == 0: eps_final = ( dpsgd_acc.poisson(dpsgd_acc.gaussian(noise_multiplier), sample_rate) * NUM_STEPS ).epsilon_at(DELTA) print(f"\nFinal epsilon: {eps_final:.2f}") dist.destroy_process_group()

Summary¶

Distributed DP-SGD with Opake requires three changes from single-device training:

Step	Single device	Distributed
Gradient aggregation	(not needed)	sum_gradients_(grads)
Noise synchronization	automatic	Pass same key(seed) on all ranks
Data sampling	PoissonSampler(dataset, ...)	local_shard() + PoissonSampler(shard, ..., key=fold_in(key, rank))

Launch with:

torchrun --nproc_per_node=4 distributed_dp_training.py

For adaptive clipping in distributed mode, use sync() from opake.distributed after each step to keep clip norms consistent across ranks. See Distributed Training in the user guide.