メニュー
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.
デフォルトでは、最初にソースを確認する Prompt が選択されています。直接コマンドに切り替えるか、ローカルコピーをダウンロードすることもできます。
インストールを決める前に、SKILL.md と SkillsMP に表示されている付属ファイルをお読みください。
Codex または Claude でインストール この Prompt をコピーして Codex、Claude、または他のアシスタントに貼り付けると、Skill ページを確認してインストールできます。
直接コマンドでは確認用 Prompt が省略されます。実行前にソースを確認してください。
npx skills add https://github.com/ADu2021/skillXiv --skill noloco-low-communication-trainingコマンドは1行のまま表示されます。コピー前に横へスクロールして全体を確認してください。
ローカルで確認しますか？SkillsMP が現在取得できるファイルをダウンロードできます。
SOC 職業分類に基づく
SKILL.md を表示中
| 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 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.
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):
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
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
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)
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
Configuration Parameters:
Synchronization Frequency:
Network Characteristics:
When to Use NoLoCo:
Scaling Considerations: