1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 48,
6  "content": "def fused_mlp_forward(x, w_fc, w_proj):\n    \"\"\"Memory-efficient fused MLP with gated ReLU²-SiLU activation.\n\n    Novel contribution: split the hidden projection into two equal chunks and use\n    a multiplicative gate act(a) * sigmoid(b) where act(a)=relu(a)^2. This keeps\n    the original 4x expansion/parameter count fixed while improving expressivity\n    over plain GELU, and the custom autograd recomputes lightweight activations\n    in backward to avoid saving the full post-activation tensor.\n    \"\"\"\n\n    class _GatedReLUSqMLP(torch.autograd.Function):\n        @staticmethod\n        def forward(ctx, x, w_fc, w_proj):\n            h = x @ w_fc.t()\n            h1, h2 = h.chunk(2, dim=-1)\n            a = F.relu(h1)\n            gate = torch.sigmoid(h2)\n            act = (a * a) * gate\n            out = act @ w_proj[:, :act.shape[-1]].t() + act @ w_proj[:, act.shape[-1]:].t()\n            ctx.save_for_backward(x, w_fc, w_proj, h1, h2)\n            return out\n\n        @staticmethod\n        def backward(ctx, grad_output):\n            x, w_fc, w_proj, h1, h2 = ctx.saved_tensors\n            dtype = grad_output.dtype\n            mid = h1.shape[-1]\n            w_proj_sum = (w_proj[:, :mid] + w_proj[:, mid:]).to(dtype)\n\n            relu_h1 = F.relu(h1)\n            gate = torch.sigmoid(h2)\n            act = (relu_h1 * relu_h1) * gate\n\n            d_act = grad_output @ w_proj_sum.to(dtype)\n            d_h1 = 2.0 * relu_h1.to(dtype) * gate.to(dtype) * d_act\n            d_h2 = (relu_h1 * relu_h1).to(dtype) * gate.to(dtype) * (1.0 - gate.to(dtype)) * d_act\n            d_h = torch.cat([d_h1, d_h2], dim=-1)\n\n            grad_x = d_h @ w_fc.to(dtype)\n            grad_w_fc = d_h.t() @ x.to(dtype)\n            grad_w_proj_half = grad_output.t() @ act.to(dtype)\n            grad_w_proj = torch.cat([grad_w_proj_half, grad_w_proj_half], dim=1)\n            return grad_x, grad_w_fc, grad_w_proj\n\n    return _GatedReLUSqMLP.apply(x, w_fc, w_proj)"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 77,
6  "content": "def fused_mlp_forward(x, w_fc, w_proj):\n    \"\"\"Fused MLP with variance-preserving GEGLU-style ReLU² gate.\n\n    We reinterpret the fixed 4x hidden width as two 2x branches:\n      value = relu(a)^2 / sqrt(2)\n      gate  = gelu_tanh(b)\n      hidden = value * gate\n    and project with the corresponding summed halves of w_proj.\n\n    This is an algorithmic change to the MLP parameterization (gated multiplicative\n    mixing under fixed parameter count) and uses custom autograd with saved\n    preactivations only.\n    \"\"\"\n\n    class _GatedReLUSqGELUMLP(torch.autograd.Function):\n        @staticmethod\n        def forward(ctx, x, w_fc, w_proj):\n            h = x @ w_fc.t()\n            a, b = h.chunk(2, dim=-1)\n            ra = F.relu(a)\n            value = (ra * ra) * (2.0 ** -0.5)\n            c = 0.7978845608028654\n            inner = c * (b.float() + 0.044715 * b.float() * b.float() * b.float())\n            gate = 0.5 * b.float() * (1.0 + torch.tanh(inner))\n            gate = gate.to(h.dtype)\n            act = value * gate\n            mid = act.shape[-1]\n            out = act @ (w_proj[:, :mid] + w_proj[:, mid:]).t()\n            ctx.save_for_backward(x, w_fc, w_proj, a, b)\n            return out\n\n        @staticmethod\n        def backward(ctx, grad_output):\n            x, w_fc, w_proj, a, b = ctx.saved_tensors\n            dtype = grad_output.dtype\n            mid = a.shape[-1]\n            w_proj_sum = (w_proj[:, :mid] + w_proj[:, mid:]).to(dtype)\n\n            ra = F.relu(a)\n            value = (ra * ra) * (2.0 ** -0.5)\n            b_f = b.float()\n            c = 0.7978845608028654\n            inner = c * (b_f + 0.044715 * b_f * b_f * b_f)\n            tanh_inner = torch.tanh(inner)\n            gate = (0.5 * b_f * (1.0 + tanh_inner)).to(dtype)\n            sech2 = 1.0 - tanh_inner * tanh_inner\n            d_inner = c * (1.0 + 3.0 * 0.044715 * b_f * b_f)\n            gelu_grad = (0.5 * (1.0 + tanh_inner) + 0.5 * b_f * sech2 * d_inner).to(dtype)\n\n            act = value.to(dtype) * gate\n            d_act = grad_output @ w_proj_sum\n            d_a = (2.0 ** 0.5) * ra.to(dtype) * gate * d_act\n            d_b = value.to(dtype) * gelu_grad * d_act\n            d_h = torch.cat([d_a, d_b], dim=-1)\n\n            grad_x = d_h @ w_fc.to(dtype)\n            grad_w_fc = d_h.t() @ x.to(dtype)\n            grad_w_proj_half = grad_output.t() @ act\n            grad_w_proj = torch.cat([grad_w_proj_half, grad_w_proj_half], dim=1)\n            return grad_x, grad_w_fc, grad_w_proj\n\n    return _GatedReLUSqGELUMLP.apply(x, w_fc, w_proj)"
7}

1{
2  "n": 2
3}