"""GradientAggregator — averages worker gradients and steps the global model. This is the "software interconnect". Two modes: * **sync** — buffer ``world_size`` gradient sets computed against the *current* global version, average them element-wise, run one optimizer step. The averaged update is exactly what you'd get from one large batch = sum of the workers' micro-batches. This is gradient accumulation, but over HTTP instead of NVLink. * **async** — apply every gradient as soon as it lands, scaled down by how stale it is (``1 / (1 + age)``). Higher throughput, noisier convergence, no barrier. All model mutation happens under a single lock, so concurrent uploads from many workers are serialized and can never corrupt the weights. """ from __future__ import annotations import threading from dataclasses import dataclass from typing import Dict, List, Optional, Tuple import torch from server.state import GlobalState class StaleGradientError(Exception): """Raised when an upload is too far behind the current global version.""" def __init__(self, worker_version: int, global_version: int, tolerance: int): self.worker_version = worker_version self.global_version = global_version self.tolerance = tolerance super().__init__( f"gradient at v{worker_version} is stale " f"(global v{global_version}, tolerance {tolerance})" ) class InvalidGradientError(Exception): """Raised when an upload's tensor keys don't match the model parameters.""" @dataclass class SubmitResult: accepted: bool applied: bool # True iff an optimizer step was taken version: int # global version after this submission step: int pending: int # gradients buffered toward the next step (sync mode) last_loss: float class GradientAggregator: def __init__(self, state: GlobalState): self.state = state self.cfg = state.swarm_cfg self._lock = threading.Lock() # sync-mode buffer of gradient dicts awaiting averaging. self._buffer: List[Dict[str, torch.Tensor]] = [] self._buffer_losses: List[float] = [] self._expected_keys = set(state.param_names) # ---- public API --------------------------------------------------------- def submit( self, grads: Dict[str, torch.Tensor], worker_version: int, loss: Optional[float] = None, ) -> SubmitResult: self._validate_keys(grads) with self._lock: age = self.state.version - worker_version if age < 0 or age > self.cfg.staleness_tolerance: raise StaleGradientError( worker_version, self.state.version, self.cfg.staleness_tolerance ) if self.cfg.mode == "async": return self._submit_async(grads, age, loss) return self._submit_sync(grads, loss) def status(self) -> dict: with self._lock: return { "mode": self.cfg.mode, "version": self.state.version, "step": self.state.step, "pending": len(self._buffer), "world_size": self.cfg.world_size, "last_loss": self.state.last_loss, "num_params": self.state.model.num_params(), } # ---- sync --------------------------------------------------------------- def _submit_sync( self, grads: Dict[str, torch.Tensor], loss: Optional[float] ) -> SubmitResult: self._buffer.append(grads) if loss is not None: self._buffer_losses.append(loss) if len(self._buffer) < self.cfg.world_size: return SubmitResult( accepted=True, applied=False, version=self.state.version, step=self.state.step, pending=len(self._buffer), last_loss=self.state.last_loss, ) averaged = self._average(self._buffer) mean_loss = ( sum(self._buffer_losses) / len(self._buffer_losses) if self._buffer_losses else None ) self._apply_and_step(averaged, mean_loss) self._buffer.clear() self._buffer_losses.clear() return SubmitResult( accepted=True, applied=True, version=self.state.version, step=self.state.step, pending=0, last_loss=self.state.last_loss, ) # ---- async -------------------------------------------------------------- def _submit_async( self, grads: Dict[str, torch.Tensor], age: int, loss: Optional[float] ) -> SubmitResult: scale = 1.0 / (1.0 + max(age, 0)) scaled = {k: v * scale for k, v in grads.items()} self._apply_and_step(scaled, loss) return SubmitResult( accepted=True, applied=True, version=self.state.version, step=self.state.step, pending=0, last_loss=self.state.last_loss, ) # ---- shared step -------------------------------------------------------- def _apply_and_step( self, grads: Dict[str, torch.Tensor], loss: Optional[float] ) -> None: """Load grads onto params, clip, optimizer.step(), bump version. (lock held)""" model = self.state.model opt = self.state.optimizer opt.zero_grad(set_to_none=True) for name, p in model.named_parameters(): g = grads.get(name) p.grad = g.to(p.device, p.dtype) if g is not None else None if self.cfg.grad_clip and self.cfg.grad_clip > 0: torch.nn.utils.clip_grad_norm_(model.parameters(), self.cfg.grad_clip) opt.step() opt.zero_grad(set_to_none=True) self.state.version += 1 self.state.step += 1 if loss is not None: self.state.last_loss = float(loss) self._maybe_checkpoint() def _maybe_checkpoint(self) -> None: every = self.cfg.checkpoint_every if every and self.state.step % every == 0 and self.cfg.checkpoint_path: try: self.state.save_checkpoint(self.cfg.checkpoint_path) except OSError: # Persistence is best-effort; never let a disk hiccup kill training. pass # ---- helpers ------------------------------------------------------------ @staticmethod def _average(buffer: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: n = len(buffer) keys = buffer[0].keys() out: Dict[str, torch.Tensor] = {} for k in keys: acc = buffer[0][k].clone() for d in buffer[1:]: acc += d[k] out[k] = acc / n return out def _validate_keys(self, grads: Dict[str, torch.Tensor]) -> None: keys = set(grads.keys()) if keys != self._expected_keys: missing = self._expected_keys - keys extra = keys - self._expected_keys raise InvalidGradientError( f"gradient keys mismatch (missing={sorted(missing)[:3]}, " f"extra={sorted(extra)[:3]})" )