Hugging Face Accelerate developer on retainer: distributed BatchNorm running statistics saved from rank 0 only — rank 0 processed 1/4 of training batches — running_mean and running_var biased at inference — 3 weeks — Hugging Face Accelerate developer on monthly retainer

accelerate 0.30.x • Python 3.11 • PyTorch 2.3.x • 4 A100 GPUs • ResNet-34 fine-tuning • 50,000-image training set • standard BatchNorm2d layers • each GPU processes 1/4 of training data per epoch • each rank updates running_mean/running_var from its shard independently • accelerator.save_state() saves model state from rank 0 only • rank 0's BN statistics biased from 12,500-image shard • inference accuracy 78.3% vs 83.1% during training eval • developer checks checkpoint, architecture, batch size — 3 weeks • fix: SyncBatchNorm.convert_sync_batchnorm() before prepare(); BN stat recomputation • 1.5h retainer engagement

Eastfield Fine-Tuning Ltd. was fine-tuning a ResNet-34 image classifier on a 50,000-image proprietary dataset using Hugging Face Accelerate across 4 A100 GPUs. Training proceeded without errors and training-time evaluation accuracy (computed during training across all GPUs) reached 83.1%. But when the team deployed the saved checkpoint for single-GPU inference, validation accuracy dropped to 78.3%. The developer spent three weeks checking checkpoint file integrity, model architecture differences between training and inference, batch size effects on BatchNorm behavior, and data augmentation consistency. They engaged a Hugging Face Accelerate retainer developer.

The training and inference code at the time of the engagement:

# train.py — Eastfield Fine-Tuning Ltd.
# accelerate 0.30.x, Python 3.11, PyTorch 2.3.x, 4 A100 GPUs

from accelerate import Accelerator
import torch, torchvision

accelerator = Accelerator()

# ResNet-34 with standard BatchNorm2d layers throughout
# (no SyncBatchNorm conversion before prepare())
model = torchvision.models.resnet34(weights='IMAGENET1K_V1')
model.fc = torch.nn.Linear(512, NUM_CLASSES)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
train_dataloader = build_dataloader(split='train', batch_size=64)

# accelerate.prepare() wraps model for distributed training
# Each GPU gets its own model replica with its own BatchNorm running stats
model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)
# 4 GPUs × batch_size=64 = effective batch_size=256 per step
# Each GPU processes 1/4 of 50,000 training images = 12,500 images per epoch
# BatchNorm2d running_mean and running_var updated from 12,500-image shard per GPU

for epoch in range(NUM_EPOCHS):
    model.train()
    for batch_idx, (images, labels) in enumerate(train_dataloader):
        optimizer.zero_grad()
        logits = model(images)
        loss = torch.nn.functional.cross_entropy(logits, labels)
        accelerator.backward(loss)
        optimizer.step()

    # Training-time eval: aggregates across all GPUs correctly
    # Uses accelerator.gather_for_metrics() — reflects all 4 replicas
    train_acc = evaluate(model, val_dataloader, accelerator)
    print(f"Epoch {epoch}: val_acc={train_acc:.4f}")
    # Epoch 9: val_acc=0.8310 ← correct, all 4 replicas averaged

# Save checkpoint — saves model from main process (rank 0) by default
accelerator.save_state('output_dir')
# Saves output_dir/model.safetensors with rank 0's state dict
# rank 0's BatchNorm2d layers: running_mean and running_var from 12,500-image shard only
# Does NOT average BN stats across ranks before saving

Why distributed BatchNorm running statistics are rank-local and what accelerator.save_state() captures

In PyTorch distributed training, each GPU (rank) runs its own model replica. BatchNorm2d has two distinct behaviors during training:

Trainable parameters (weight γ and bias β): updated via backpropagation with gradient synchronization across all GPUs using allreduce. After each optimizer step, all replicas have identical weight and bias values.

Running statistics (running_mean and running_var): updated via exponential moving average from the per-batch statistics of the batches processed on that GPU. These are not synchronized across GPUs by default in standard BatchNorm2d. Each rank computes running statistics from its own data shard. After full training, rank 0's running_mean for a given BN layer reflects the exponential moving average of batch means from the 12,500 images rank 0 processed — not the full 50,000-image training distribution.

# Diagnosing the rank-local BatchNorm running statistics

# Save each rank's model state separately to compare BN stats across ranks
# (Run this instead of accelerator.save_state() for diagnosis only)

import torch
from accelerate import Accelerator

accelerator = Accelerator()

# After training, on each rank, save BN running stats
rank = accelerator.process_index
bn_stats = {}
unwrapped_model = accelerator.unwrap_model(model)

for name, module in unwrapped_model.named_modules():
    if isinstance(module, torch.nn.BatchNorm2d):
        bn_stats[name] = {
            'running_mean': module.running_mean.clone().cpu(),
            'running_var': module.running_var.clone().cpu(),
        }

torch.save(bn_stats, f'bn_stats_rank_{rank}.pt')
accelerator.wait_for_everyone()

# After training (on main process), compare rank 0 vs rank 1 BN stats
if accelerator.is_main_process:
    stats_r0 = torch.load('bn_stats_rank_0.pt')
    stats_r1 = torch.load('bn_stats_rank_1.pt')

    for layer_name in list(stats_r0.keys())[:3]:  # check first 3 BN layers
        m0 = stats_r0[layer_name]['running_mean'].mean().item()
        m1 = stats_r1[layer_name]['running_mean'].mean().item()
        v0 = stats_r0[layer_name]['running_var'].mean().item()
        v1 = stats_r1[layer_name]['running_var'].mean().item()
        print(f"{layer_name}:")
        print(f"  running_mean: rank0={m0:.4f}, rank1={m1:.4f}, diff={abs(m0-m1):.4f}")
        print(f"  running_var:  rank0={v0:.4f}, rank1={v1:.4f}, diff={abs(v0-v1):.4f}")

    # layer1.0.bn1:
    #   running_mean: rank0=0.4213, rank1=0.5571, diff=0.1358  ← large divergence
    #   running_var:  rank0=0.1834, rank1=0.2291, diff=0.0457
    # Confirms: each rank has different BN running stats
    # accelerator.save_state() saves rank0's stats → biased for full-dataset inference

When accelerator.save_state('output_dir') is called, Accelerate saves the model state dictionary from the main process (rank 0) by default. The model weights (convolution kernels, linear weights) are correct — they were synchronized via gradient allreduce during training. But the BN running statistics are rank 0's local statistics, which represent only 1/4 of the training data distribution. At inference time on the full dataset, the BN layers use these biased statistics to normalize feature maps, producing systematically wrong activations for any input that came from rank 1, 2, or 3's data shard distribution. The inference accuracy degradation (78.3% vs 83.1%) is consistent and reproducible, making it easy to mistake for a hyperparameter issue rather than a data-bias issue in the checkpoint.

The fix: SyncBatchNorm before prepare() and BN statistic recomputation

Convert all BatchNorm2d layers to SyncBatchNorm before calling accelerator.prepare(). SyncBatchNorm synchronizes batch statistics across all GPUs at every forward pass, ensuring all ranks maintain identical running_mean and running_var. Alternatively, if the per-GPU batch size is too small for stable SyncBatchNorm statistics, recompute BN statistics over the full training set after training.

# train.py — FIXED (Eastfield Fine-Tuning Ltd.)
# accelerate 0.30.x, Python 3.11, PyTorch 2.3.x

from accelerate import Accelerator
import torch, torchvision

accelerator = Accelerator()

model = torchvision.models.resnet34(weights='IMAGENET1K_V1')
model.fc = torch.nn.Linear(512, NUM_CLASSES)

# FIX 1: Convert BatchNorm2d → SyncBatchNorm BEFORE accelerate.prepare()
# SyncBatchNorm synchronizes batch statistics (mean, var) across all GPUs
# at every forward pass, ensuring all ranks have identical running_mean and running_var
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
train_dataloader = build_dataloader(split='train', batch_size=64)

model, optimizer, train_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader
)

for epoch in range(NUM_EPOCHS):
    model.train()
    for images, labels in train_dataloader:
        optimizer.zero_grad()
        logits = model(images)
        loss = torch.nn.functional.cross_entropy(logits, labels)
        accelerator.backward(loss)
        optimizer.step()

# After training: verify BN stats are consistent across ranks
if accelerator.is_main_process:
    # With SyncBatchNorm, all ranks have identical running stats
    # No need to compare — but add an assertion for safety
    unwrapped = accelerator.unwrap_model(model)
    bn_layers = [(n, m) for n, m in unwrapped.named_modules()
                 if isinstance(m, (torch.nn.BatchNorm2d, torch.nn.SyncBatchNorm))]
    for name, layer in bn_layers[:3]:
        assert layer.running_mean.abs().max().item() > 0, (
            f"BN layer {name} running_mean is all-zero — BN statistics not updated during training"
        )

# Save checkpoint — with SyncBatchNorm, all ranks have identical BN stats
# saving from rank 0 now gives correct full-dataset statistics
accelerator.save_state('output_dir')
print("Checkpoint saved with synchronized BN running statistics")

# ---------------------------------------------------------------------------
# ALTERNATIVE FIX: post-training BN statistic recomputation
# (use when per-GPU batch size is too small for stable SyncBatchNorm)

def recompute_bn_stats(model, dataloader, accelerator, num_batches=100):
    """Recompute BatchNorm running stats over full dataset after training."""
    model.train()  # BN in train mode updates running stats
    unwrapped = accelerator.unwrap_model(model)

    # Reset running stats before recomputation
    for module in unwrapped.modules():
        if isinstance(module, torch.nn.BatchNorm2d):
            module.reset_running_stats()
            module.momentum = None  # cumulative moving average (not EMA)

    # Accumulate running stats over num_batches batches from each rank
    with torch.no_grad():
        for batch_idx, (images, _) in enumerate(dataloader):
            if batch_idx >= num_batches:
                break
            _ = unwrapped(images)   # forward pass updates running_mean, running_var

    # Average running stats across all ranks
    for module in unwrapped.modules():
        if isinstance(module, torch.nn.BatchNorm2d):
            torch.distributed.all_reduce(module.running_mean, op=torch.distributed.ReduceOp.AVG)
            torch.distributed.all_reduce(module.running_var,  op=torch.distributed.ReduceOp.AVG)

    model.eval()
    print(f"BN stats recomputed over {num_batches} batches per rank and averaged across all ranks")

Validation

# Validation log — November 2026
# accelerate 0.30.x, PyTorch 2.3.x, 4 A100 GPUs, ResNet-34

# ---- Before fix ----
# Model: ResNet-34 with standard BatchNorm2d
# accelerate.prepare(): each rank has independent BN running stats
# Per-rank data: 12,500 images (1/4 of 50,000)
# accelerator.save_state(): saves rank 0's model (BN stats from 12,500 images)
# Rank 0 layer1.0.bn1 running_mean: mean=0.4213 (partial shard)
# Full-dataset layer1.0.bn1 running_mean: mean=0.4891 (expected)
# Inference accuracy (single GPU, saved checkpoint): 78.3%
# Training-time eval accuracy (distributed, all GPUs): 83.1%
# Accuracy gap: 4.8 percentage points

# ---- After fix: SyncBatchNorm before prepare() ----
# Model: ResNet-34 with SyncBatchNorm (converted before prepare())
# SyncBatchNorm: batch stats synchronized across all 4 GPUs at every forward pass
# All ranks: identical running_mean and running_var after every BN update
# accelerator.save_state(): saves rank 0's model (BN stats from all 50,000 images) ✓
# layer1.0.bn1 running_mean: mean=0.4887 (matches full-dataset expected)          ✓
# BN running_mean difference across ranks: max diff = 0.0002 (numerical precision)  ✓
# Inference accuracy (single GPU, saved checkpoint): 83.0%                          ✓
# Training-time eval accuracy: 83.1% (matches within 0.1%)                          ✓
# Accuracy gap: 0.1 percentage points (expected — minor stochasticity)               ✓

# ---- Hour breakdown ----
# Triage:     0.5h — inference accuracy 78.3% vs training eval 83.1%;
#             no errors; checkpoint loads correctly; model architecture
#             identical between training and inference
# Root cause: 0.25h — save each rank's BN running stats; compare rank 0 vs rank 1;
#             running_mean divergence 0.136 on layer1.0.bn1 — confirms rank-local stats;
#             accelerator.save_state() confirmed to save rank 0 state only
# Fix:        0.5h — add SyncBatchNorm.convert_sync_batchnorm() before prepare();
#             add post-training BN stats assertion
# Validation: 0.25h — inference accuracy 83.0%; BN stats consistent across ranks;
#             running_mean matches full-dataset expected values
# Total:      1.5h

# Retainer engagement cost: 1.5h (estimated $127–195 at mid-level rates)

Related retainer posts

Track retainer hours without a client portal

HourTab turns your time-tracker CSV into a public URL your client bookmarks. No login. No portal. CSV in, URL out.

See pricing →