Large Batch Training › Code Implementation
Michael BrenndoerferJanuary 27, 202614 min read
Language AI HandbookAll chapters Part 32 of 67Training Optimization
Large Batch Training › Code ImplementationOne section · Sections in this chapter7 / 12 ↓
Sections in this chapter12
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
We will implement and compare several training scenarios: standard SGD at small batch size, SGD at large batch size without scaling, SGD at large batch size with linear scaling, and LAMB. For tractability, we use a simple convolutional model on CIFAR-10.
Setup and Imports
Data Loading
We use CIFAR-10 with standard preprocessing. To make the experiments run quickly, we use a subset of the training data.
In[4]:
Codeimport numpy as np
import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import Dataset, Subset
torch.manual_seed(42)
np.random.seed(42)
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
)
class HFImageDataset(Dataset):
"""Adapt a Hugging Face image split to the (tensor, label) pairs DataLoader expects.
We source CIFAR-10 from the Hub rather than through
`torchvision.datasets.CIFAR10(download=True)`, whose only mirror is the
original university host — a single server with no CDN that commonly
delivers the 170 MB archive at well under 100 KB/s.
"""
def __init__(self, split, transform):
self.split = split
self.transform = transform
def __len__(self):
return len(self.split)
def __getitem__(self, idx):
record = self.split[idx]
return self.transform(record["img"].convert("RGB")), record["label"]
try:
from datasets import load_dataset
cifar = load_dataset("uoft-cs/cifar10")
full_train = HFImageDataset(cifar["train"], transform)
full_test = HFImageDataset(cifar["test"], transform)
except Exception:
full_train = torchvision.datasets.FakeData(
size=10000, image_size=(3, 32, 32), num_classes=10, transform=transform
)
full_test = torchvision.datasets.FakeData(
size=2000, image_size=(3, 32, 32), num_classes=10, transform=transform
)
# Use a compact subset so the complete experiment remains practical to rerun
train_subset = Subset(full_train, indices=range(1000))
test_subset = Subset(full_test, indices=range(500))Model Definition
In[5]:
Codeimport torch.nn as nn
class SmallCNN(nn.Module):
"""A simple 3-layer CNN for CIFAR-10."""
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(64 * 8 * 8, 256),
nn.ReLU(),
nn.Linear(256, 10),
)
def forward(self, x):
return self.classifier(self.features(x))LAMB Optimizer Implementation
We implement a simplified LAMB to show the core mechanics. Production use would rely on the torch-optimizer or apex packages.
In[6]:
Codeimport torch
import torch.optim as optim
class LAMB(optim.Optimizer):
"""
Layer-wise Adaptive Moments for Batch training.
Simplified implementation following You et al. (2019).
"""
def __init__(
self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.01
):
defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
beta1, beta2 = group["betas"]
for p in group["params"]:
if p.grad is None:
continue
grad = p.grad
state = self.state[p]
if len(state) == 0:
state["step"] = 0
state["exp_avg"] = torch.zeros_like(p)
state["exp_avg_sq"] = torch.zeros_like(p)
exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
state["step"] += 1
t = state["step"]
# Adam moment updates
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
# Bias correction
bias_correction1 = 1 - beta1**t
bias_correction2 = 1 - beta2**t
m_hat = exp_avg / bias_correction1
v_hat = exp_avg_sq / bias_correction2
# Adam update direction
adam_update = m_hat / (v_hat.sqrt() + group["eps"])
# Weight decay in update direction
if group["weight_decay"] != 0:
adam_update.add_(p, alpha=group["weight_decay"])
# LARS trust ratio (layer-wise)
weight_norm = p.norm(p=2)
update_norm = adam_update.norm(p=2)
if weight_norm == 0 or update_norm == 0:
trust_ratio = 1.0
else:
trust_ratio = weight_norm / update_norm
# Apply scaled update
p.add_(adam_update, alpha=-group["lr"] * trust_ratio)
return lossTraining Loop
In[7]:
Codefrom torch.utils.data import DataLoader
def train_one_epoch(model, loader, optimizer, device):
"""Train for one epoch and return average loss."""
model.train()
total_loss = 0.0
criterion = nn.CrossEntropyLoss()
for xb, yb in loader:
xb, yb = xb.to(device), yb.to(device)
optimizer.zero_grad()
loss = criterion(model(xb), yb)
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / len(loader)
def evaluate(model, loader, device):
"""Evaluate model accuracy on a dataset."""
model.eval()
correct = 0
total = 0
with torch.no_grad():
for xb, yb in loader:
xb, yb = xb.to(device), yb.to(device)
preds = model(xb).argmax(dim=1)
correct += (preds == yb).sum().item()
total += len(yb)
return correct / total
def run_experiment(
batch_size,
lr,
optimizer_name="sgd",
epochs=20,
warmup_epochs=0,
device="cpu",
):
"""Run a single training experiment and return loss/accuracy curves."""
torch.manual_seed(42)
model = SmallCNN().to(device)
train_loader = DataLoader(
train_subset, batch_size=batch_size, shuffle=True, num_workers=0
)
test_loader = DataLoader(
test_subset, batch_size=256, shuffle=False, num_workers=0
)
if optimizer_name == "sgd":
optimizer = optim.SGD(
model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4
)
elif optimizer_name == "adam":
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
elif optimizer_name == "lamb":
optimizer = LAMB(model.parameters(), lr=lr, weight_decay=0.01)
train_losses = []
test_accs = []
base_lr = lr / max(warmup_epochs, 1)
for epoch in range(epochs):
# Linear warmup
if warmup_epochs > 0 and epoch < warmup_epochs:
current_lr = base_lr + (lr - base_lr) * epoch / warmup_epochs
for pg in optimizer.param_groups:
pg["lr"] = current_lr
elif warmup_epochs > 0 and epoch == warmup_epochs:
for pg in optimizer.param_groups:
pg["lr"] = lr
epoch_loss = train_one_epoch(model, train_loader, optimizer, device)
acc = evaluate(model, test_loader, device)
train_losses.append(epoch_loss)
test_accs.append(acc)
return train_losses, test_accs
device = "cpu"
epochs = 20Now we run four experiments: small batch baseline, large batch without scaling, large batch with linear scaling, and large batch with LAMB.
In[8]:
Code# Experiment 1: Small batch baseline (B=64, lr=0.05)
base_batch = 64
base_lr_sgd = 0.05
loss_small, acc_small = run_experiment(
batch_size=base_batch,
lr=base_lr_sgd,
optimizer_name="sgd",
epochs=epochs,
warmup_epochs=0,
device=device,
)
# Experiment 2: Large batch without LR scaling (B=512, same lr=0.05)
large_batch = 512
k = large_batch // base_batch # scaling factor = 8
loss_large_no_scale, acc_large_no_scale = run_experiment(
batch_size=large_batch,
lr=base_lr_sgd,
optimizer_name="sgd",
epochs=epochs,
warmup_epochs=0,
device=device,
)
# Experiment 3: Large batch with linear scaling (B=512, lr=0.4=8*0.05)
scaled_lr = base_lr_sgd * k
loss_large_scaled, acc_large_scaled = run_experiment(
batch_size=large_batch,
lr=scaled_lr,
optimizer_name="sgd",
epochs=epochs,
warmup_epochs=3,
device=device,
)
# Experiment 4: Large batch with LAMB
lamb_lr = 0.01
loss_lamb, acc_lamb = run_experiment(
batch_size=large_batch,
lr=lamb_lr,
optimizer_name="lamb",
epochs=epochs,
warmup_epochs=3,
device=device,
)Out[9]:
ConsoleBase batch size: 64, base LR: 0.05 Large batch size: 512, scale factor k: 8 Linearly scaled LR: 0.400 Final test accuracy (small batch SGD): 0.380 Final test accuracy (large batch, unscaled LR): 0.306 Final test accuracy (large batch, scaled LR): 0.208 Final test accuracy (large batch, LAMB): 0.340
The compact experiment illustrates both the motivation for scaling and its limits. The small-batch baseline reaches the highest final accuracy. The unscaled large-batch run trails it, but applying the full 8x learning-rate increase is too aggressive here and performs worse. LAMB is the strongest of the large-batch configurations, although it still does not match the small-batch baseline. This is why the linear rule is a starting point rather than a guarantee: the scaled rate must remain within the stable range of the particular model, data, and training budget.
Training Dynamics Visualization
Out[10]:
Visualization
Simulating the Gradient Variance Reduction
To make the statistical intuition concrete, we can directly compute how gradient variance changes with batch size on a fixed parameter snapshot.
In[11]:
Codetorch.manual_seed(42)
# Use a fixed model and fixed mini dataset for reproducibility
ref_model = SmallCNN()
criterion_ref = nn.CrossEntropyLoss()
# Pre-load a set of samples
sample_loader = DataLoader(
train_subset, batch_size=1, shuffle=True, num_workers=0
)
gradients = []
# Collect per-sample gradients for first conv layer bias
for i, (xb, yb) in enumerate(sample_loader):
if i >= 128:
break
ref_model.zero_grad()
loss_ref = criterion_ref(ref_model(xb), yb)
loss_ref.backward()
grad = ref_model.features[0].bias.grad.detach().clone()
gradients.append(grad)
gradients = torch.stack(gradients) # shape: (200, 32)
# Compute variance of gradient estimates at different batch sizes
batch_sizes_var = [1, 2, 4, 8, 16, 32, 64]
variances = []
for bs in batch_sizes_var:
n_batches = len(gradients) // bs
batch_means = [
gradients[i * bs : (i + 1) * bs].mean(0) for i in range(n_batches)
]
batch_means_t = torch.stack(batch_means)
var = batch_means_t.var(0).mean().item()
variances.append(var)Out[12]:
ConsoleGradient estimate variance by batch size:
Batch Size | Variance | Var / Var(B=1)
----------------------------------------------
1 | 0.000237 | 1.0000
2 | 0.000119 | 0.5007
4 | 0.000061 | 0.2549
8 | 0.000037 | 0.1540
16 | 0.000018 | 0.0758
32 | 0.000009 | 0.0362
64 | 0.000006 | 0.0233The variance decreases inversely with batch size. When we move from batch size 1 to batch size 8, the variance drops by a factor of roughly 8. This is exactly the relationship derived from statistics. The linear scaling rule compensates for this: each individual update becomes more accurate but less frequent, so you can afford a proportionally larger step.
Out[13]:
Visualization
Estimating the Critical Batch Size
The critical batch size can be approximated empirically by measuring how many training examples you need to process to reach a target loss level, at different batch sizes with linearly scaled learning rates.
In[14]:
Codedef compute_examples_to_threshold(losses, batch_size, threshold):
"""Compute cumulative examples processed to reach a loss threshold."""
dataset_size = len(train_subset)
cumulative_examples = 0
for epoch_loss in losses:
cumulative_examples += dataset_size
if epoch_loss <= threshold:
return cumulative_examples
return None
# Define loss threshold we want to reach
target_loss = 1.75
# Run experiments at several batch sizes with linearly scaled LR
exp_batch_sizes = [32, 64, 128, 256, 512, 1024]
base_lr_for_crit = 0.025 # smaller base to allow scaling
examples_to_thresh = {}
all_crit_losses = {}
for bs in exp_batch_sizes:
k_crit = max(bs // 32, 1)
lr_crit = min(base_lr_for_crit * k_crit, 0.4)
warmup_e = min(3, max(1, k_crit // 2))
losses_crit, _ = run_experiment(
batch_size=bs,
lr=lr_crit,
optimizer_name="sgd",
epochs=15,
warmup_epochs=warmup_e,
device=device,
)
all_crit_losses[bs] = losses_crit
n_examples = compute_examples_to_threshold(losses_crit, bs, target_loss)
examples_to_thresh[bs] = n_examplesOut[15]:
ConsoleTarget loss threshold: 1.75
Batch Size | Examples to Threshold
--------------------------------------
32 | 4,000
64 | 5,000
128 | 6,000
256 | Not reached
512 | Not reached
1024 | Not reachedOut[16]:
Visualization
This compact run does not show an initial efficiency gain from larger batches. Instead, examples required rise from batch size 32 through 128, and batches of 256 or more miss the target within the training budget. That outcome places the critical regime below 256 for this particular model, subset, learning-rate rule, and 15-epoch budget. In larger experiments the transition is often flatter, but the diagnostic principle is the same: once larger batches stop reducing the examples needed to reach a target, extra parallelism no longer improves statistical efficiency.
LAMB Trust Ratio Across Layers
A key property of LAMB is that the trust ratio varies significantly across layers. Early layers (especially embeddings and the first convolutional layers) tend to have different weight-to-gradient-norm ratios compared to later layers. We can visualize this directly on our trained model.
In[17]:
Code# Compute trust ratios for each layer of the LAMB-trained model
torch.manual_seed(42)
lamb_inspect_model = SmallCNN()
lamb_inspect_opt = LAMB(
lamb_inspect_model.parameters(), lr=lamb_lr, weight_decay=0.01
)
# Run a single forward/backward pass to compute gradients
inspect_loader = DataLoader(
train_subset, batch_size=large_batch, shuffle=True, num_workers=0
)
xb_inspect, yb_inspect = next(iter(inspect_loader))
criterion_inspect = nn.CrossEntropyLoss()
lamb_inspect_opt.zero_grad()
loss_inspect = criterion_inspect(lamb_inspect_model(xb_inspect), yb_inspect)
loss_inspect.backward()
# Collect trust ratios before the step (simulate first LAMB step)
layer_names = []
weight_norms = []
grad_norms = []
trust_ratios = []
for name, p in lamb_inspect_model.named_parameters():
if p.grad is not None and p.dim() > 1: # only weight tensors (not biases)
wn = p.norm(p=2).item()
gn = p.grad.norm(p=2).item()
if gn > 0:
tr = wn / gn
layer_names.append(name.replace(".weight", ""))
weight_norms.append(wn)
grad_norms.append(gn)
trust_ratios.append(tr)Out[18]:
ConsoleLayer | Weight Norm | Grad Norm | Trust Ratio -------------------------------------------------------------------------- features.0 | 3.3163 | 0.0150 | 221.2744 features.3 | 4.6290 | 0.0565 | 81.9249 classifier.1 | 9.2346 | 0.2364 | 39.0607 classifier.3 | 1.8339 | 0.0603 | 30.3973
The trust ratios vary considerably across layers, confirming that LAMB's layer-wise scaling is doing real work. Layers with large weight norms relative to their gradient norms receive larger effective learning rates, while layers where gradients are proportionally larger receive smaller effective learning rates.
Out[19]:
Visualization
