Spaces:
Paused
Paused
File size: 3,403 Bytes
6f75f79 06d2e6b 6f75f79 06d2e6b 6f75f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Baselines to compare the counter synapse against — in plain PyTorch.
Right now: a BitNet-b1.58-style ternary QAT linear (full FP32 master + Adam, ternary in the
forward) which isolates "cost of the 6-bit counter optimizer" vs "cost of ternarization".
The README roadmap lists the memory-efficient-training baselines worth adding next (8-bit
Adam / GaLore / LoMo), so the memory claim is measured against real competitors, not only
FP32+Adam.
"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
__all__ = ["TernaryQATLinear", "make_linear"]
class TernaryQATLinear(nn.Module):
"""FP32 master weight, per-row absmean ternary in forward, straight-through gradient to
the master; Adam updates the master. Same ternary inference as the counter layer, but a
full FP32 optimizer -- so (counter vs this) isolates the 6-bit optimizer cost."""
def __init__(self, fin: int, fout: int, init_gain: float = 1.0) -> None:
super().__init__()
self.weight = nn.Parameter(torch.empty(fout, fin))
nn.init.normal_(self.weight, mean=0.0, std=init_gain * math.sqrt(1.0 / fin))
def forward(self, x: torch.Tensor) -> torch.Tensor:
w = self.weight
scale = w.abs().mean(dim=1, keepdim=True).clamp_min(1e-5)
wq = (w / scale).round().clamp_(-1, 1) * scale
wq = w + (wq - w).detach() # straight-through estimator
return F.linear(x, wq)
def make_linear(kind: str, fin: int, fout: int, init_gain: float = 1.0, **counter_kw):
"""Factory used by the GPT harness.
kind in {dense, qat, counter, counter_rms, counter_packed, counter_triton, slowfast, group}.
'slowfast' (M3) and 'group' (M2) are counter-linear drop-ins: SlowFastCounterLinear folds a
low-rank fp residual into the base every K steps (cuts the base correlation 8-32x); the extra
rank/merge_every knobs ride in counter_kw. GroupCounterLinear is a 2:4 structured-sparse counter
linear (group/keep/hysteresis in counter_kw). Both replace any nn.Linear in the block."""
from .counter import CompactCounterLinear, RMSCounterLinear
from .packed import PackedRMSCounterLinear
if kind == "counter":
return CompactCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
if kind == "counter_rms":
return RMSCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
if kind == "counter_packed":
return PackedRMSCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
if kind == "counter_triton":
from .triton_counter import TritonCounterLinear
return TritonCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
if kind == "slowfast": # M3
from .slowfast import SlowFastCounterLinear
return SlowFastCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
if kind == "group": # M2
from .group_counter import GroupCounterLinear
return GroupCounterLinear(fin, fout, init_gain=init_gain, **counter_kw)
if kind == "qat":
return TernaryQATLinear(fin, fout, init_gain)
if kind == "dense":
lin = nn.Linear(fin, fout, bias=False)
nn.init.normal_(lin.weight, mean=0.0, std=init_gain * math.sqrt(1.0 / fin))
return lin
raise ValueError(f"unknown linear kind: {kind!r}")
|