adu2021 NoLoCo: No-all-reduce Low Communication Training Method
Core Concept
NoLoCo eliminates explicit all-to-all synchronization in distributed model training through implicit weight convergence via dynamic pipeline routing and modified Nesterov momentum. Instead of global synchronization (all-reduce), the system synchronizes only pairs of accelerators, enabling up to 4% faster convergence than existing low-communication methods while requiring exponentially less communication overhead.
Architecture Overview
- No Collective Communication: Eliminates all-reduce by synchronizing only pairs of accelerators via implicit weight averaging
- Dynamic Pipeline Routing: Inner optimizer steps randomly route inputs through replica stages, implicitly mixing weights without explicit synchronization
- Modified Nesterov Momentum: Specialized momentum update includes local weight averaging term to prevent divergence across distributed instances
- Theoretical Analysis: Convergence proofs under quadratic loss assumptions
- Scalability: Tested from 125M to 6.8B parameter models; outperforms FSDP and DiLoCo
Implementation
Step 1: Dynamic Pipeline Routing
import torch
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
import random
class DynamicPipelineRouter:
"""
Routes samples through randomly selected replica stages during inner optimizer steps.
Enables implicit weight mixing without explicit synchronization.
"""
def __init__(self, num_workers, num_layers, model):
self.num_workers = num_workers
self.num_layers = num_layers
self.model = model
self.worker_models = self._replicate_across_workers()
def _replicate_across_workers(self):
"""Replicate model across all workers."""
replicas = []
for _ in range(self.num_workers):
replica = copy.deepcopy(self.model)
replicas.append(replica)
return replicas
def forward_with_routing(self, x, layer_idx):
"""
Route input through randomly selected worker for current layer.
Probability: each worker equally likely to process layer.
"""
# Select random worker for this layer
selected_worker = random.randint(0, self.num_workers - 1)
# Get activation from selected worker
with torch.no_grad():
if layer_idx == 0:
# First layer processes raw input
activation = self.worker_models[selected_worker].layers[0](x)
else:
# Later layers: forward through prior layers on same worker
# then apply current layer
activation = x
for l in range(layer_idx):
activation = self.worker_models[selected_worker].layers[l](activation)
# Apply current layer
activation = self.worker_models[selected_worker].layers[layer_idx](activation)
return activation, selected_worker
def full_forward_with_dynamic_routing(self, x, num_inner_steps=10):
"""
Complete forward pass with dynamic routing across inner optimizer steps.
Each step can route through different workers.
"""
activations = []
routes = []
for step in range(num_inner_steps):
current = x
for layer_idx in range(self.num_layers):
activation, worker_id = self.forward_with_routing(current, layer_idx)
current = activation
routes.append(worker_id)
activations.append(current)
return activations, routes
Step 2: Modified Nesterov Momentum Optimizer
import torch
from torch.optim.optimizer import Optimizer
class ModifiedNesterovMomentum(Optimizer):
"""
Nesterov momentum with local weight averaging term.
Prevents model weight divergence across distributed replicas.
Update rule:
δₜ,ᵢ = αδₜ₋₁,ᵢ - (β/n)(∑ⱼΔₜ,ⱼ) - γ(ϕₜ,ᵢ - (1/n)∑ⱼϕₜ,ⱼ)
Terms:
- δₜ,ᵢ: momentum term for worker i at timestep t
- Δₜ,ⱼ: gradient for worker j
- ϕₜ,ᵢ: model weights for worker i
- (1/n)∑ⱼϕₜ,ⱼ: average model weights across group
"""
def __init__(self, params, lr=1e-3, momentum=0.5, weight_decay=0.0,
group_size=2, divergence_penalty=0.1):
defaults = dict(
lr=lr,
momentum=momentum,
weight_decay=weight_decay,
group_size=group_size,
divergence_penalty=divergence_penalty
)
super().__init__(params, defaults)
def step(self, closure=None, local_group_weights=None):
"""
Perform single optimizer step with local group weight averaging.
Args:
closure: Optional closure to recompute loss
local_group_weights: List of weight tensors from group members
"""
loss = None
if closure is not None:
loss = closure()
for group in self.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
lr = group['lr']
divergence_penalty = group['divergence_penalty']
for p_idx, p in enumerate(group['params']):
if p.grad is None:
continue
d_p = p.grad.data
# Standard L2 weight decay
if weight_decay != 0:
d_p = d_p.add(p.data, alpha=weight_decay)
# Initialize momentum state
param_state = self.state[p]
if 'momentum_buffer' not in param_state:
buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
else:
buf = param_state['momentum_buffer']
# Nesterov momentum update
buf.mul_(momentum).add_(d_p)
# Local weight averaging penalty: prevent divergence
if local_group_weights is not None and len(local_group_weights) > 0:
# Compute average weight in group
avg_weight = torch.zeros_like(p.data)
for group_weight in local_group_weights:
if p_idx < len(group_weight):
avg_weight += group_weight[p_idx]
avg_weight /= len(local_group_weights)
# Divergence penalty term: γ(ϕₜ,ᵢ - (1/n)∑ⱼϕₜ,ⱼ)
divergence = p.data - avg_weight
buf.add_(divergence, alpha=divergence_penalty)
# Apply Nesterov momentum step
p.data.add_(buf, alpha=-lr)
return loss
Step 3: NoLoCo Training Loop
class NoLoCoTrainer:
"""
Orchestrates distributed training with NoLoCo optimization.
Manages implicit synchronization via dynamic routing and local weight averaging.
"""
def __init__(self, model, num_workers=2, group_size=2):
self.model = model
self.num_workers = num_workers
self.group_size = group_size
# Create replicas across workers
self.replicas = [copy.deepcopy(model) for _ in range(num_workers)]
# Optimizer configured for local group updates
self.optimizer = ModifiedNesterovMomentum(
self.model.parameters(),
lr=0.7,
momentum=0.5,
divergence_penalty=1.0
)
# Pipeline router for implicit mixing
self.router = DynamicPipelineRouter(num_workers, len(model.layers), model)
# Tracking for diagnostics
self.loss_history = []
def train_step(self, batch, inner_steps=50, outer_steps=1):
"""
Single training iteration: inner loop (local updates) + outer loop (synchronization).
Inner loop: Use dynamic routing to implicitly mix weights
Outer loop: Explicit pair-wise synchronization via modified momentum
"""
input_ids, labels = batch
# INNER LOOP: Local gradient accumulation with dynamic routing
accumulated_gradients = [0.0] * len(self.replicas)
for inner_step in range(inner_steps):
# Dynamic routing: forward pass uses different workers for each layer
activations, routes = self.router.full_forward_with_dynamic_routing(
input_ids,
num_inner_steps=1
)
# Backward pass on each worker
for worker_id, replica in enumerate(self.replicas):
output = replica(input_ids)
loss = torch.nn.functional.cross_entropy(output, labels)
loss.backward()
# Accumulate gradients
accumulated_gradients[worker_id] += sum(
p.grad.data.sum().item() for p in replica.parameters()
if p.grad is not None
)
# OUTER LOOP: Pair-wise synchronization with weight averaging
for outer_step in range(outer_steps):
# Collect weights from all replicas
all_weights = [
[p.data.clone() for p in replica.parameters()]
for replica in self.replicas
]
# Apply optimizer step with local group averaging
# Select pairs for synchronization
for i in range(0, self.num_workers, self.group_size):
group_indices = list(range(i, min(i + self.group_size, self.num_workers)))
group_weights = [all_weights[j] for j in group_indices]
# Apply NoLoCo optimizer step
self.optimizer.step(
closure=lambda: self._compute_loss(input_ids, labels),
local_group_weights=group_weights
)
# Sync replicas with updated master model
for replica in self.replicas:
replica.load_state_dict(self.model.state_dict())
# Compute loss for logging
with torch.no_grad():
output = self.model(input_ids)
loss = torch.nn.functional.cross_entropy(output, labels)
self.loss_history.append(loss.item())
return loss.item()
def _compute_loss(self, input_ids, labels):
"""Compute loss on current model state."""
output = self.model(input_ids)
return torch.nn.functional.cross_entropy(output, labels)
Step 4: Convergence Analysis and Synchronization Frequency
class ConvergenceAnalyzer:
"""
Analyzes convergence properties and optimal synchronization frequency.
Provides guidance on outer loop frequency vs communication cost.
"""
def __init__(self, trainer):
self.trainer = trainer
self.divergence_scores = []
def measure_weight_divergence(self):
"""
Measure how much weights diverge across replicas without synchronization.
Higher divergence = more frequent synchronization needed.
"""
# Reference model
reference_weights = [p.data.clone() for p in self.trainer.replicas[0].parameters()]
divergence = 0.0
for replica in self.trainer.replicas[1:]:
replica_weights = [p.data.clone() for p in replica.parameters()]
for ref_w, rep_w in zip(reference_weights, replica_weights):
divergence += torch.norm(ref_w - rep_w).item()
avg_divergence = divergence / len(self.trainer.replicas)
self.divergence_scores.append(avg_divergence)
return avg_divergence
def recommend_sync_frequency(self, max_divergence_threshold=0.1):
"""
Recommend synchronization frequency based on divergence trajectory.
"""
if len(self.divergence_scores) < 10:
return 50 # Default: sync every 50 steps
# Compute divergence growth rate
recent_divergences = self.divergence_scores[-10:]
growth_rate = (recent_divergences[-1] - recent_divergences[0]) / 10
if growth_rate > max_divergence_threshold:
# Divergence growing too fast: sync more frequently
return 25
elif growth_rate < 0.001:
# Divergence stable: can sync less frequently
return 100
else:
# Normal: default frequency
return 50
Practical Guidance
Configuration Parameters:
- Group size: n=2 (pairs) provides best balance between communication and convergence
- Outer learning rate: β=0.7 for both NoLoCo and DiLoCo
- Momentum: α=0.5 (NoLoCo) vs α=0.3 (DiLoCo)
- Divergence penalty: γ=1.0 prevents weight divergence effectively
Synchronization Frequency:
- Inner steps: 50 (local updates before synchronization)
- Outer steps: 1 (single pair-wise sync per outer iteration)
- NoLoCo syncs exponentially less than standard all-reduce
- Theoretical speedup: log₂(n) compared to tree all-reduce
Network Characteristics:
- High-latency networks benefit most from NoLoCo (100-1000ms latency reduction)
- Bandwidth-limited scenarios: communication reduction 10-50x vs FSDP
- Low-latency networks (NVLink, InfiniBand): benefits modest but still meaningful
When to Use NoLoCo:
- Geo-distributed training (edge → cloud)
- Bandwidth-constrained environments (commodity networks)
- Very large model training (8B+ parameters)
- Scenarios where communication is bottleneck
Scaling Considerations:
- Tested up to 6.8B parameter models
- Convergence maintained with proper weight averaging
- Training time: 60% improvement (communication reduced from 8 hours to 3 hours for 1T token training)
Reference
- All-reduce operations: Global synchronization across all workers; bandwidth bottleneck
- Nesterov momentum: Momentum method with lookahead gradient computation
- Weight divergence: Risk in low-communication training; prevented via local averaging
- Pipeline parallelism: Routing enables pipelined execution with implicit mixing