1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 38,
5  "end_line": 115,
6  "content": "def weight_quant(weight):\n    \"\"\"Ternary quantization with per-channel absmean scaling.\n    \n    For each output channel, compute absmean as scale, normalize weights,\n    then round to nearest integer in {-1, 0, +1} using STE.\n    \"\"\"\n    # weight shape: [out_features, in_features]\n    # compute per-channel absmean: shape [out_features]\n    scale = weight.detach().abs().mean(dim=1, keepdim=True).clamp(min=1e-12)\n    w_normed = weight / scale\n    # STE round: forward round, backward identity\n    w_q = w_normed.clamp(-1, 1)\n    w_q = (w_q.round() - w_q).detach() + w_q\n    # rescale\n    w_q = w_q * scale\n    return w_q, scale.squeeze(1)\n\n\ndef activation_quant(x):\n    \"\"\"Absmax 8-bit activation quantization with STE.\n    \n    Quantizes activations to 127 levels (int8 range) using per-tensor\n    absmax scaling, following the BitNet b1.58 paper.\n    \"\"\"\n    Qb = 127  # int8 range\n    scale = x.detach().abs().max().clamp(min=1e-12)\n    x_normed = x / scale\n    x_q = (x_normed * Qb).round().clamp(-Qb, Qb)\n    # STE: forward uses quantized, backward passes through\n    x_q = (x_q - x_normed * Qb).detach() + x_normed * Qb\n    return x_q, scale / Qb\n\n\nclass BitLinear(nn.Module):\n    \"\"\"BitNet b1.58 linear layer with ternary {-1, 0, +1} weights and per-channel scaling.\n    \n    During both training and eval: weights are ternarized via per-channel absmean\n    + round-clip, activations are quantized to int8 range. Output is\n    rescaled by weight_scale * activation_scale.\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        # w_scale is per-channel, shape [out_features]; we need to reshape for broadcasting\n        # out shape: (..., out_features)\n        # w_scale shape: (out_features,) -> (1, 1, ..., out_features) for broadcasting\n        # x_scale is scalar\n        out = out * (w_scale.view(1, -1) * x_scale)\n        if self.bias is not None:\n            out = out + self.bias\n        return out"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 38,
5  "end_line": 115,
6  "content": "def weight_quant(weight):\n    \"\"\"Ternary quantization: {-1, 0, +1} via absmean with STE.\n\n    Forward: normalize by absmean, round-then-clip to {-1, 0, +1}\n    Backward: STE (gradient passes through rounding as identity)\n    \"\"\"\n    scale = weight.detach().abs().mean().clamp(min=1e-12)\n    w_normed = weight / scale\n    # STE round: (round(x) - x).detach() + x\n    w_q = w_normed.clamp(-1, 1)\n    w_q = (w_q.round() - w_q).detach() + w_q\n    return w_q, scale\n\n\ndef activation_quant(x):\n    \"\"\"Absmax 8-bit activation quantization with STE.\n\n    Quantizes activations to 127 levels (int8 range) using per-tensor\n    absmax scaling, following the BitNet b1.58 paper.\n    \"\"\"\n    Qb = 127  # int8 range\n    scale = x.detach().abs().max().clamp(min=1e-12)\n    x_normed = x / scale\n    x_q = (x_normed * Qb).round().clamp(-Qb, Qb)\n    # STE: forward uses quantized, backward passes through\n    x_q = (x_q - x_normed * Qb).detach() + x_normed * Qb\n    return x_q, scale / Qb\n\n\nclass BitLinear(nn.Module):\n    \"\"\"BitNet b1.58 linear layer with ternary {-1, 0, +1} weights.\n\n    During both training and eval: weights are ternarized via absmean\n    + round-clip, activations are quantized to int8 range. Output is\n    rescaled by weight_scale * activation_scale.\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 * (w_scale * x_scale)\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": 92,
6  "content": "class TernaryQuantize(torch.autograd.Function):\n    \"\"\"Ternary quantization with per-channel scaling and gradient scaling.\n    \n    Forward:\n        scale = mean(|weight|, dim=1)\n        w_norm = weight / scale\n        w_clip = clamp(w_norm, -1, 1)\n        w_round = round(w_clip)\n        w_q = w_round * scale\n        \n    Backward (STE):\n        gradient w.r.t weight = gradient w.r.t w_q * (1 / scale) * mask\n        where mask = (w_norm.abs() <= 1) (optional clipping gradient)\n        We implement STE: gradient passes through round as identity.\n        Also scale gradient by 1/scale to counteract scaling effect.\n    \"\"\"\n    @staticmethod\n    def forward(ctx, weight):\n        # weight shape: [out_features, in_features]\n        scale = weight.abs().mean(dim=1, keepdim=True).clamp(min=1e-12)\n        w_norm = weight / scale\n        w_clip = w_norm.clamp(-1, 1)\n        w_round = w_clip.round()\n        w_q = w_round * scale\n        ctx.save_for_backward(scale, w_norm)\n        return w_q, scale.squeeze(1)\n    \n    @staticmethod\n    def backward(ctx, grad_w_q, grad_scale):\n        scale, w_norm = ctx.saved_tensors\n        # STE: gradient w.r.t weight = grad_w_q * (w_norm.abs() <= 1) * (1/scale)\n        # where clipping gradient mask optional; we can ignore for simplicity.\n        # Use STE: treat round as identity, clip as identity.\n        # Scale gradient by 1/scale as per BitNet gradient scaling.\n        grad_weight = grad_w_q / scale\n        # grad_weight shape matches weight\n        # grad_scale is not needed (scale is computed from weight, not trainable)\n        return grad_weight, None\n\n\ndef weight_quant(weight):\n    \"\"\"Ternary quantization using custom autograd Function.\"\"\"\n    return TernaryQuantize.apply(weight)\n\n\ndef activation_quant(x):\n    \"\"\"Absmax 8-bit activation quantization with STE.\n\n    Quantizes activations to 127 levels (int8 range) using per-tensor\n    absmax scaling, following the BitNet b1.58 paper.\n    \"\"\"\n    Qb = 127  # int8 range\n    scale = x.detach().abs().max().clamp(min=1e-12)\n    x_normed = x / scale\n    x_q = (x_normed * Qb).round().clamp(-Qb, Qb)\n    # STE: forward uses quantized, backward passes through\n    x_q = (x_q - x_normed * Qb).detach() + x_normed * Qb\n    return x_q, scale / Qb\n\n\nclass BitLinear(nn.Module):\n    \"\"\"BitNet b1.58 linear layer with ternary {-1, 0, +1} weights and per-channel scaling.\n\n    During both training and eval: weights are ternarized via per-channel absmean\n    + round-clip, activations are quantized to int8 range. Output is\n    rescaled by weight_scale * activation_scale.\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        # w_scale is per-channel, shape [out_features]; we need to reshape for broadcasting\n        # out shape: (..., out_features)\n        # w_scale shape: (out_features,) -> (1, 1, ..., out_features) for broadcasting\n        # x_scale is scalar\n        out = out * (w_scale.view(1, -1) * x_scale)\n        if self.bias is not None:\n            out = out + self.bias\n        return out"
7}