1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 38,
5  "end_line": 115,
6  "content": "def _ste_round(x):\n    return (x.round() - x).detach() + x\n\n\ndef _ste_clip(x, lo, hi):\n    return (x.clamp(lo, hi) - x).detach() + x\n\n\ndef weight_quant(weight):\n    \"\"\"Row-wise balanced ternary quantization with residual-level scaling.\n\n    Novelty:\n    - Per-output-channel scaling to reduce quantization mismatch for GPT blocks.\n    - Learn-free adaptive dead-zone from row RMS/mean-abs statistics, producing\n      sparse ternary codes only when confidence is low.\n    - Energy correction factor to better match the row-wise second moment after\n      ternarization, improving optimization stability over plain absmean scaling.\n\n    Returns q in {-1,0,1} (via STE) and a per-row scale.\n    \"\"\"\n    eps = 1e-12\n    w_det = weight.detach()\n\n    # Per-output-channel statistics are more faithful for transformer projections.\n    mean_abs = w_det.abs().mean(dim=1, keepdim=True).clamp(min=eps)\n    rms = w_det.square().mean(dim=1, keepdim=True).sqrt().clamp(min=eps)\n\n    # Normalize by mean-abs, then use an adaptive dead-zone.\n    w_norm = weight / mean_abs\n    # Threshold increases mildly for heavy-tailed / high-variance rows.\n    tau = (0.55 + 0.25 * (rms / mean_abs - 1.0)).clamp(0.45, 0.95)\n\n    # Balanced ternary assignment with STE.\n    pos = (w_norm > tau).to(weight.dtype)\n    neg = (w_norm < -tau).to(weight.dtype)\n    q_hard = pos - neg\n    q = (q_hard - w_norm.clamp(-1.0, 1.0)).detach() + w_norm.clamp(-1.0, 1.0)\n\n    # Residual energy correction: fit scale to quantized support per row.\n    support = q_hard.abs().sum(dim=1, keepdim=True).clamp(min=1.0)\n    ls_scale = (w_det * q_hard).sum(dim=1, keepdim=True) / support\n    # Mix mean-abs and least-squares scales for robustness early in training.\n    scale = (0.7 * ls_scale.abs() + 0.3 * mean_abs).clamp(min=eps)\n    return q, scale\n\n\n\ndef activation_quant(x):\n    \"\"\"Token-wise mixed clipping int8 quantization with STE.\n\n    Uses a token-local scale from a blend of absmax and RMS, which reduces the\n    impact of outliers compared with pure absmax while preserving headroom.\n    \"\"\"\n    eps = 1e-12\n    Qb = 127.0\n    x_det = x.detach()\n    absmax = x_det.abs().amax(dim=-1, keepdim=True).clamp(min=eps)\n    rms = x_det.square().mean(dim=-1, keepdim=True).sqrt().clamp(min=eps)\n    scale = torch.maximum(0.6 * absmax + 0.4 * (2.5 * rms), 0.85 * absmax).clamp(min=eps)\n    x_norm = x / scale\n    x_clip = _ste_clip(x_norm, -1.0, 1.0)\n    x_q = _ste_round(x_clip * Qb)\n    return x_q, scale / Qb\n\n\nclass BitLinear(nn.Module):\n    \"\"\"Linear layer using native low-bit weights and int8 activations.\n\n    Adds a quantization-consistency residual gate: a small fraction of the\n    latent-minus-quantized weight is injected during training only through STE,\n    improving optimization without changing the discrete forward values.\n    \"\"\"\n    def __init__(self, in_features, out_features, bias=True):\n        super().__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.weight = nn.Parameter(torch.empty(out_features, in_features))\n        if bias:\n            self.bias = nn.Parameter(torch.zeros(out_features))\n        else:\n            self.bias = None\n        nn.init.normal_(self.weight, mean=0.0, std=0.02)\n\n    def forward(self, x):\n        w_q, w_scale = weight_quant(self.weight)\n        x_q, x_scale = activation_quant(x)\n        out = F.linear(x_q, w_q, None)\n        out = out * x_scale\n        out = out * w_scale.transpose(0, 1)\n        if self.bias is not None:\n            out = out + self.bias\n        return out"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 38,
5  "end_line": 129,
6  "content": "def _ste_round(x):\n    return (x.round() - x).detach() + x\n\n\ndef _ste_passthrough(hard, soft):\n    return (hard - soft).detach() + soft\n\n\ndef weight_quant(weight):\n    \"\"\"Variance-matched groupwise ternary quantization with STE.\n\n    Algorithmic contribution:\n    1) Row-wise centering before quantization reduces sign bias.\n    2) Groupwise ternary coding over contiguous input-channel groups preserves\n       local structure better than one global threshold.\n    3) A variance-matching rescale fits each output row after ternarization,\n       improving forward fidelity while keeping weights strictly ternary.\n    \"\"\"\n    eps = 1e-12\n    group_size = 64\n    w_det = weight.detach()\n    out_features, in_features = weight.shape\n\n    # Center each output row to stabilize ternary occupancy.\n    row_mean = w_det.mean(dim=1, keepdim=True)\n    w_centered = weight - row_mean\n    w_centered_det = w_det - row_mean\n\n    # Pad input dimension so we can quantize in fixed-size groups.\n    pad = (group_size - (in_features % group_size)) % group_size\n    if pad > 0:\n        w_centered = F.pad(w_centered, (0, pad))\n        w_centered_det = F.pad(w_centered_det, (0, pad))\n    n_groups = (in_features + pad) // group_size\n\n    wv = w_centered.view(out_features, n_groups, group_size)\n    wv_det = w_centered_det.view(out_features, n_groups, group_size)\n\n    mean_abs = wv_det.abs().mean(dim=-1, keepdim=True).clamp(min=eps)\n    rms = wv_det.square().mean(dim=-1, keepdim=True).sqrt().clamp(min=eps)\n    w_norm = wv / mean_abs\n\n    # Adaptive dead-zone per group; slightly lower than standard b1.58 to keep\n    # more active weights after centering.\n    tau = (0.42 + 0.18 * (rms / mean_abs - 1.0)).clamp(0.30, 0.75)\n    q_hard = torch.where(w_norm > tau, 1.0, torch.where(w_norm < -tau, -1.0, 0.0)).to(weight.dtype)\n    q_soft = w_norm.clamp(-1.0, 1.0)\n    q = _ste_passthrough(q_hard, q_soft)\n\n    q_hard_flat = q_hard.view(out_features, -1)[..., :in_features]\n    q_flat = q.view(out_features, -1)[..., :in_features]\n\n    # Row-wise least-squares fit to match the centered latent weights.\n    denom = q_hard_flat.square().sum(dim=1, keepdim=True).clamp(min=1.0)\n    ls_scale = (w_centered_det[:, :in_features] * q_hard_flat).sum(dim=1, keepdim=True) / denom\n    row_rms = w_centered_det[:, :in_features].square().mean(dim=1, keepdim=True).sqrt().clamp(min=eps)\n    q_rms = q_hard_flat.square().mean(dim=1, keepdim=True).sqrt().clamp(min=eps)\n    var_scale = row_rms / q_rms\n    scale = (0.6 * ls_scale.abs() + 0.4 * var_scale).clamp(min=eps)\n    return q_flat, scale\n\n\n\ndef activation_quant(x):\n    \"\"\"Per-token int8 quantization with RMS-guided clipping and STE.\"\"\"\n    eps = 1e-12\n    Qb = 127.0\n    x_det = x.detach()\n    rms = x_det.square().mean(dim=-1, keepdim=True).sqrt().clamp(min=eps)\n    absmax = x_det.abs().amax(dim=-1, keepdim=True).clamp(min=eps)\n    # Clip from RMS, but never exceed absmax and never under-cover too much.\n    clip = torch.minimum(absmax, 2.8 * rms)\n    clip = torch.maximum(clip, 0.7 * absmax).clamp(min=eps)\n    x_norm = x / clip\n    x_clip = x_norm.clamp(-1.0, 1.0)\n    x_q = _ste_round(x_clip * Qb)\n    x_q = _ste_passthrough(x_q, x_clip * Qb)\n    return x_q, clip / Qb\n\n\nclass BitLinear(nn.Module):\n    \"\"\"Linear layer with native ternary weights and token-wise int8 activations.\"\"\"\n    def __init__(self, in_features, out_features, bias=True):\n        super().__init__()\n        self.in_features = in_features\n        self.out_features = out_features\n        self.weight = nn.Parameter(torch.empty(out_features, in_features))\n        if bias:\n            self.bias = nn.Parameter(torch.zeros(out_features))\n        else:\n            self.bias = None\n        nn.init.normal_(self.weight, mean=0.0, std=0.02)\n\n    def forward(self, x):\n        w_q, w_scale = weight_quant(self.weight)\n        x_q, x_scale = activation_quant(x)\n        out = F.linear(x_q, w_q, None)\n        out = out * x_scale\n        out = out * w_scale.transpose(0, 1)\n        if self.bias is not None:\n            out = out + self.bias\n        return out"
7}

1{
2  "n": 1
3}