Spaces:
Sleeping
Sleeping
| """Gradient compression — the bandwidth kernel. | |
| On a real GPU cluster the interconnect is NVLink at ~600 GB/s. Here it is HTTP | |
| over the public internet at maybe 10 MB/s, so **bytes on the wire are the | |
| bottleneck**, not FLOPs. A worker that spends 200 ms computing a gradient and | |
| 8 s uploading it is 97% idle. | |
| The fix is standard in distributed training but rarely built from scratch: | |
| * **Top-K sparsification** — only the ``ratio`` fraction of gradient entries with | |
| the largest magnitude are transmitted (as index/value pairs). At ratio=0.01 | |
| that is a ~50x reduction in payload. | |
| * **Error feedback** — the entries we *dropped* are not discarded; they are kept | |
| in a local residual buffer and added to the next step's gradient. A small | |
| consistent gradient direction therefore accumulates until it becomes large | |
| enough to be sent. Without this, sparsification biases training and | |
| convergence stalls; with it, Top-K is near-lossless in practice. | |
| (Seide et al. 2014; Lin et al., "Deep Gradient Compression", 2018.) | |
| * **int8 quantization** — the surviving values are optionally quantized to 8-bit | |
| with a per-tensor scale, for another ~4x on top. | |
| The server reconstructs a dense gradient before aggregating, so compression is | |
| entirely transparent to :class:`~server.aggregator.GradientAggregator`. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from typing import Dict, Mapping, Optional, Sequence, Tuple | |
| import torch | |
| # Wire-format suffixes. safetensors payloads are flat {name: tensor} maps, so a | |
| # compressed tensor is split across a few synthetic keys. | |
| IDX = "::idx" | |
| VAL = "::val" | |
| SCALE = "::scale" | |
| # Values for the X-Compression header. | |
| NONE = "none" | |
| TOPK = "topk" | |
| TOPK_INT8 = "topk-int8" | |
| MODES = (NONE, TOPK, TOPK_INT8) | |
| def _numel(shape: Sequence[int]) -> int: | |
| return int(math.prod(shape)) if len(shape) else 1 | |
| def topk_count(numel: int, ratio: float) -> int: | |
| """How many entries survive Top-K for a tensor of ``numel`` elements.""" | |
| if ratio >= 1.0: | |
| return numel | |
| return max(1, min(numel, int(numel * ratio))) | |
| def compress( | |
| tensors: Mapping[str, torch.Tensor], | |
| mode: str = TOPK, | |
| ratio: float = 0.01, | |
| residual: Optional[Dict[str, torch.Tensor]] = None, | |
| ) -> Dict[str, torch.Tensor]: | |
| """Compress ``{name: dense_gradient}`` into a sparse wire payload. | |
| ``residual`` (error feedback) is read *and updated in place* when provided: | |
| dropped mass is carried into the next call so no gradient signal is lost, | |
| only delayed. | |
| """ | |
| if mode == NONE: | |
| return dict(tensors) | |
| if mode not in MODES: | |
| raise ValueError(f"unknown compression mode {mode!r}; expected one of {MODES}") | |
| quantize = mode == TOPK_INT8 | |
| out: Dict[str, torch.Tensor] = {} | |
| for name, grad in tensors.items(): | |
| g = grad.detach().to(torch.float32) | |
| if residual is not None: | |
| prev = residual.get(name) | |
| if prev is not None: | |
| g = g + prev | |
| flat = g.reshape(-1) | |
| k = topk_count(flat.numel(), ratio) | |
| # Select by magnitude; sorted=False is cheaper and order is irrelevant. | |
| _, idx = torch.topk(flat.abs(), k, sorted=False) | |
| vals = flat[idx] | |
| if residual is not None: | |
| # Everything we did NOT send stays behind for the next step. | |
| carry = flat.clone() | |
| carry[idx] = 0.0 | |
| residual[name] = carry.view_as(grad) | |
| if quantize: | |
| peak = vals.abs().max() | |
| scale = (peak / 127.0) if peak > 0 else torch.tensor(1.0) | |
| out[name + VAL] = torch.round(vals / scale).clamp_(-127, 127).to(torch.int8) | |
| out[name + SCALE] = scale.reshape(1).to(torch.float32) | |
| else: | |
| out[name + VAL] = vals | |
| out[name + IDX] = idx.to(torch.int64) | |
| return out | |
| def decompress( | |
| payload: Mapping[str, torch.Tensor], | |
| shapes: Mapping[str, torch.Size], | |
| mode: str = TOPK, | |
| ) -> Dict[str, torch.Tensor]: | |
| """Rebuild dense gradients from a sparse payload. | |
| Shapes come from the server's own model, so they never travel on the wire. | |
| """ | |
| if mode == NONE: | |
| return dict(payload) | |
| if mode not in MODES: | |
| raise ValueError(f"unknown compression mode {mode!r}; expected one of {MODES}") | |
| out: Dict[str, torch.Tensor] = {} | |
| for name, shape in shapes.items(): | |
| try: | |
| idx = payload[name + IDX] | |
| val = payload[name + VAL] | |
| except KeyError as e: # pragma: no cover - surfaced as a 400 by the app | |
| raise KeyError(f"compressed payload missing entry for {name!r}: {e}") from e | |
| scale = payload.get(name + SCALE) | |
| if scale is not None: | |
| val = val.to(torch.float32) * float(scale.reshape(-1)[0]) | |
| flat = torch.zeros(_numel(shape), dtype=torch.float32) | |
| flat[idx.to(torch.int64)] = val.to(torch.float32) | |
| out[name] = flat.view(shape) | |
| return out | |
| def payload_bytes(payload: Mapping[str, torch.Tensor]) -> int: | |
| """Total tensor bytes in a payload (excludes safetensors' small header).""" | |
| return sum(t.numel() * t.element_size() for t in payload.values()) | |
| def compression_report( | |
| dense: Mapping[str, torch.Tensor], wire: Mapping[str, torch.Tensor] | |
| ) -> Tuple[int, int, float]: | |
| """(dense_bytes, wire_bytes, ratio) — ratio is how many times smaller the wire is.""" | |
| d = payload_bytes(dense) | |
| w = payload_bytes(wire) | |
| return d, w, (d / w if w else 1.0) | |