Scale distributed LLM training without all-reduce synchronization using dynamic pipeline routing and modified Nesterov momentum, achieving 4% faster convergence than DiLoCo with exponentially lower communication.
Instrucciones de origen · Vista previa de solo lectura
name
noloco-low-communication-training
title
NoLoCo: No-all-reduce Low Communication Training Method for Large Models
version
0.0.2
engine
skillxiv-v0.0.2-claude-opus-4.6
license
MIT
url
https://arxiv.org/abs/2506.10911
keywords
["distributed training","low communication","synchronization","large language models","optimization"]
description
Scale distributed LLM training without all-reduce synchronization using dynamic pipeline routing and modified Nesterov momentum, achieving 4% faster convergence than DiLoCo with exponentially lower communication.
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
classDynamicPipelineRouter:
"""
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):
replicas = []
_ (.num_workers):
replica = copy.deepcopy(.model)
replicas.append(replica)
replicas
():
selected_worker = random.randint(, .num_workers - )
torch.no_grad():
layer_idx == :
activation = .worker_models[selected_worker].layers[](x)
:
activation = x
l (layer_idx):
activation = .worker_models[selected_worker].layers[l](activation)
activation = .worker_models[selected_worker].layers[layer_idx](activation)
activation, selected_worker
():
activations = []
routes = []
step (num_inner_steps):
current = x
layer_idx (.num_layers):
activation, worker_id = .forward_with_routing(current, layer_idx)
current = activation
routes.append(worker_id)
activations.append(current)
activations, routes
"""Replicate model across all workers."""
for
in
range
self
self
return
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
0
self
1
# Get activation from selected worker
with
if
0
# First layer processes raw input
self
0
else
# Later layers: forward through prior layers on same worker
# then apply current layer
for
in
range
self
# Apply current layer
self
return
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.
"""
for
in
range
for
in
range
self
self
return
Step 2: Modified Nesterov Momentum Optimizer
import torch
from torch.optim.optimizer import Optimizer
classModifiedNesterovMomentum(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)
defstep(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 = Noneif closure isnotNone:
loss = closure()
for group inself.param_groups:
weight_decay = group['weight_decay']
momentum = group['momentum']
lr = group['lr']
divergence_penalty = group['divergence_penalty']
for p_idx, p inenumerate(group['params']):
if p.grad isNone:
continue
d_p = p.grad.data
# Standard L2 weight decayif weight_decay != 0:
d_p = d_p.add(p.data, alpha=weight_decay)
# Initialize momentum state
param_state = self.state[p]
if'momentum_buffer'notin 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 divergenceif local_group_weights isnotNoneandlen(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
classNoLoCoTrainer:
"""
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 workersself.replicas = [copy.deepcopy(model) for _ inrange(num_workers)]
# Optimizer configured for local group updatesself.optimizer = ModifiedNesterovMomentum(
self.model.parameters(),
lr=0.7,
momentum=0.5,
divergence_penalty=1.0
)
# Pipeline router for implicit mixingself.router = DynamicPipelineRouter(num_workers, len(model.layers), model)
# Tracking for diagnosticsself.loss_history = []
deftrain_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 inrange(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 workerfor worker_id, replica inenumerate(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 isnotNone
)
# OUTER LOOP: Pair-wise synchronization with weight averagingfor outer_step inrange(outer_steps):
# Collect weights from all replicas
all_weights = [
[p.data.clone() for p in replica.parameters()]
for replica inself.replicas
]
# Apply optimizer step with local group averaging# Select pairs for synchronizationfor i inrange(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 stepself.optimizer.step(
closure=lambda: self._compute_loss(input_ids, labels),
local_group_weights=group_weights
)
# Sync replicas with updated master modelfor replica inself.replicas:
replica.load_state_dict(self.model.state_dict())
# Compute loss for loggingwith 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
classConvergenceAnalyzer:
"""
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 = []
defmeasure_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 inself.trainer.replicas[0].parameters()]
divergence = 0.0for replica inself.trainer.replicas[1:]:
replica_weights = [p.data.clone() for p in replica.parameters()]
for ref_w, rep_w inzip(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
defrecommend_sync_frequency(self, max_divergence_threshold=0.1):
"""
Recommend synchronization frequency based on divergence trajectory.
"""iflen(self.divergence_scores) < 10:
return50# Default: sync every 50 steps# Compute divergence growth rate
recent_divergences = self.divergence_scores[-10:]
growth_rate = (recent_divergences[-1] - recent_divergences[0]) / 10if growth_rate > max_divergence_threshold:
# Divergence growing too fast: sync more frequentlyreturn25elif growth_rate < 0.001:
# Divergence stable: can sync less frequentlyreturn100else:
# Normal: default frequencyreturn50
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
