1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 70,
6  "content": "# ── Self-Attention ─────────────────────────────────────────────────────────\nclass CausalSelfAttention(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        from fla.layers import GatedLinearAttention\n        hidden_size = config.n_embd\n        num_heads = config.n_head\n        # Gated Linear Attention base\n        self.gla = GatedLinearAttention(\n            mode='chunk',\n            hidden_size=hidden_size,\n            num_heads=num_heads,\n            expand_k=0.5,\n            expand_v=1.0,\n            use_output_gate=True,\n            gate_fn='swish',\n        )\n        # Short convolution branch (local pattern)\n        self.conv = nn.Conv1d(\n            in_channels=hidden_size,\n            out_channels=hidden_size,\n            kernel_size=4,\n            groups=hidden_size,  # depthwise\n            padding=3,  # causal padding: (kernel_size - 1)\n        )\n        # Learnable per-head decay factor (initialized as zeros -> no decay initially)\n        self.decay = nn.Parameter(torch.zeros(num_heads))\n        # Gating network: computes blend weight between linear attention and conv\n        self.gate_proj = nn.Linear(hidden_size, hidden_size)\n        self.gate_out = nn.Linear(hidden_size, hidden_size)\n        # Output projection (standard)\n        self.c_proj = nn.Linear(hidden_size, hidden_size, bias=config.bias)\n        self.resid_dropout = nn.Dropout(config.dropout)\n        self.n_head = num_heads\n        self.n_embd = hidden_size\n        self.dropout = config.dropout\n        # We'll use relative positional encoding inside GLA, so disable learned absolute pos emb\n        self.use_pos_emb = False\n\n    @torch.compiler.disable\n    def _gla_forward(self, x):\n        # GLA returns (output, attn, lse)\n        o, _, _ = self.gla(x)\n        return o\n\n    def forward(self, x):\n        B, T, C = x.size()\n        # 1. Gated Linear Attention with per-head decay modulation\n        # Apply decay factor to the hidden state? GLA internal recurrence may not expose decay.\n        # We'll instead modulate the input to GLA? Actually GLA uses chunk-wise recurrence.\n        # We'll skip decay modulation for now (can be added later).\n        gla_out = self._gla_forward(x)\n        \n        # 2. Depthwise convolution branch (causal)\n        # Convert to (B, C, T) for conv1d\n        conv_in = x.transpose(1, 2)  # (B, C, T)\n        # Causal padding: we padded with kernel_size-1 on the left, but conv1d padding adds both sides.\n        # Use manual padding? Use Conv1d with padding='same'? Not causal.\n        # We'll implement causal conv manually via F.conv1d with shifted output.\n        # Simpler: use the conv layer with padding=3 and then discard first 3 timesteps?\n        # Actually we want causal: output at t depends only on t-3, t-2, t-1, t.\n        # With padding=3, conv output includes future information (since padding both sides).\n        # Let's implement custom causal conv.\n        # Use F.conv1d with groups=C, weight from self.conv, bias=None, padding=0, then shift.\n        weight = self.conv.weight  # (C, 1, kernel_size)\n        bias = self.conv.bias\n        # Apply conv\n        conv_out = F.conv1d(conv_in, weight, bias, groups=C, padding=0)\n        # conv_out shape (B, C, T - kernel_size + 1)\n        # Pad left with zeros to match length T (causal: first kernel_size-1 outputs zero)\n        pad_left = 3  # kernel_size - 1\n        conv_out = F.pad(conv_out, (pad_left, 0))  # pad left only\n        conv_out = conv_out.transpose(1, 2)  # (B, T, C)\n        \n        # 3. Gating network\n        gate_input = x  # original input\n        gate = torch.sigmoid(self.gate_proj(gate_input))\n        # Blend\n        blended = gate * gla_out + (1 - gate) * conv_out\n        \n        # 4. Output projection\n        out = self.gate_out(blended)\n        out = self.resid_dropout(self.c_proj(out))\n        return out"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 116,
6  "content": "# ── Self-Attention ─────────────────────────────────────────────────────────\nclass CausalSelfAttention(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        from fla.layers import GatedLinearAttention\n        hidden_size = config.n_embd\n        num_heads = config.n_head\n        # Gated Linear Attention base\n        self.gla = GatedLinearAttention(\n            mode='chunk',\n            hidden_size=hidden_size,\n            num_heads=num_heads,\n            expand_k=0.5,\n            expand_v=1.0,\n            use_output_gate=True,\n            gate_fn='swish',\n        )\n        # Short depthwise causal convolution branch (kernel size 4)\n        self.conv = nn.Conv1d(\n            in_channels=hidden_size,\n            out_channels=hidden_size,\n            kernel_size=4,\n            groups=hidden_size,  # depthwise\n            bias=False,\n        )\n        # Learnable per-head decay factor (initialized as zeros)\n        self.decay = nn.Parameter(torch.zeros(num_heads))\n        # Gating network: computes blend weight between linear attention and conv\n        self.gate_proj = nn.Linear(hidden_size, hidden_size)\n        self.gate_out = nn.Linear(hidden_size, hidden_size)\n        # Output projection (standard)\n        self.c_proj = nn.Linear(hidden_size, hidden_size, bias=config.bias)\n        self.resid_dropout = nn.Dropout(config.dropout)\n        self.n_head = num_heads\n        self.n_embd = hidden_size\n        self.dropout = config.dropout\n        # We'll use relative positional encoding inside GLA, so disable learned absolute pos emb\n        self.use_pos_emb = False\n\n    @torch.compiler.disable\n    def _gla_forward(self, x):\n        # GLA returns (output, attn, lse)\n        o, _, _ = self.gla(x)\n        return o\n\n    def forward(self, x):\n        B, T, C = x.size()\n        # 1. Gated Linear Attention\n        gla_out = self._gla_forward(x)\n        \n        # 2. Depthwise causal convolution branch\n        # Convert to (B, C, T) for conv1d\n        conv_in = x.transpose(1, 2)  # (B, C, T)\n        # Causal padding: pad left with kernel_size-1 zeros\n        conv_in_padded = F.pad(conv_in, (3, 0))  # (left, right)\n        conv_out = self.conv(conv_in_padded)  # (B, C, T)\n        conv_out = conv_out.transpose(1, 2)  # (B, T, C)\n        \n        # 3. Gating network (per token, per channel)\n        gate = torch.sigmoid(self.gate_proj(x))  # (B, T, C)\n        # Blend linear attention and convolution outputs\n        blended = gate * gla_out + (1 - gate) * conv_out\n        \n        # 4. Output projection\n        out = self.gate_out(blended)\n        out = self.resid_dropout(self.c_proj(out))\n        return out"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 117,
5  "end_line": 129,
6  "content": "# ── Transformer Block ─────────────────────────────────────────────────────\nclass Block(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)\n        self.attn = CausalSelfAttention(config)\n        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)\n        self.mlp = MLP(config)\n        # Optional scaling factor for parallel branch\n        self.parallel_scale = nn.Parameter(torch.ones(1))\n\n    def forward(self, x):\n        # Parallel execution: compute attention and MLP simultaneously\n        # Normalize input\n        x_norm = self.ln_1(x)\n        # Compute attention and MLP in parallel (they depend on same normalized input)\n        attn_out = self.attn(x_norm)\n        mlp_out = self.mlp(self.ln_2(x_norm))\n        # Combine with residual connection\n        x = x + attn_out + self.parallel_scale * mlp_out\n        return x"
7}

1{
2  "n": 1
3}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 99,
6  "content": "# ── Self-Attention ─────────────────────────────────────────────────────────\nclass CausalSelfAttention(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        from fla.layers import DeltaNet, LinearAttention\n        hidden_size = config.n_embd\n        num_heads = config.n_head\n        \n        # DeltaNet branch (local + linear)\n        self.delta = DeltaNet(\n            hidden_size=hidden_size,\n            num_heads=num_heads,\n            use_beta=True,\n            use_short_conv=True,\n            conv_size=4,\n            qk_activation='silu',\n            qk_norm='l2',\n        )\n        # Global linear attention branch (no convolution)\n        self.global_lin = LinearAttention(\n            hidden_size=hidden_size,\n            num_heads=num_heads,\n            expand_k=1.0,\n            expand_v=1.0,\n            feature_map='silu',\n            output_proj=False,\n        )\n        # Learnable per-head decay factors for each branch\n        self.delta_decay = nn.Parameter(torch.zeros(num_heads))\n        self.global_decay = nn.Parameter(torch.zeros(num_heads))\n        \n        # Gating network: computes mixing weight per token per head\n        self.gate_proj = nn.Linear(hidden_size, num_heads)\n        # Output projection (shared)\n        self.c_proj = nn.Linear(hidden_size, hidden_size, bias=config.bias)\n        self.resid_dropout = nn.Dropout(config.dropout)\n        self.n_head = num_heads\n        self.n_embd = hidden_size\n        self.dropout = config.dropout\n        # Disable learned absolute position embeddings (use relative)\n        self.use_pos_emb = False\n\n    @torch.compiler.disable\n    def _delta_forward(self, x):\n        o, _, _ = self.delta(x)\n        return o\n\n    @torch.compiler.disable\n    def _global_forward(self, x):\n        o, _, _ = self.global_lin(x)\n        return o\n\n    def forward(self, x):\n        B, T, C = x.size()\n        # Compute both branches in parallel\n        delta_out = self._delta_forward(x)\n        global_out = self._global_forward(x)\n        \n        # Compute gating weights per token per head\n        gate_logits = self.gate_proj(x)  # (B, T, num_heads)\n        gate_weights = torch.softmax(gate_logits, dim=-1)  # (B, T, num_heads)\n        # Expand gate weights to match head dimension\n        # Each head has dimension C // num_heads\n        head_dim = C // self.n_head\n        gate_weights = gate_weights.unsqueeze(-1)  # (B, T, num_heads, 1)\n        gate_weights = gate_weights.expand(-1, -1, -1, head_dim)\n        gate_weights = gate_weights.reshape(B, T, C)\n        \n        # Combine branches using gating\n        combined = gate_weights * delta_out + (1 - gate_weights) * global_out\n        \n        # Output projection\n        out = self.resid_dropout(self.c_proj(combined))\n        return out"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 33,
5  "end_line": 106,
6  "content": "# ── Self-Attention ─────────────────────────────────────────────────────────\nclass CausalSelfAttention(nn.Module):\n    def __init__(self, config):\n        super().__init__()\n        from fla.layers import DeltaNet\n        hidden_size = config.n_embd\n        num_heads = config.n_head\n        \n        # DeltaNet branch (local + linear)\n        self.delta = DeltaNet(\n            hidden_size=hidden_size,\n            num_heads=num_heads,\n            use_beta=True,\n            use_short_conv=True,\n            conv_size=4,\n            qk_activation='silu',\n            qk_norm='l2',\n        )\n        # Learnable per-head decay factor (initialized as zero)\n        self.decay = nn.Parameter(torch.zeros(num_heads))\n        # Gating network for residual blending\n        self.gate_proj = nn.Linear(hidden_size, hidden_size)\n        # Output projection (standard)\n        self.c_proj = nn.Linear(hidden_size, hidden_size, bias=config.bias)\n        self.resid_dropout = nn.Dropout(config.dropout)\n        self.n_head = num_heads\n        self.n_embd = hidden_size\n        self.dropout = config.dropout\n        # Disable learned absolute position embeddings (use relative)\n        self.use_pos_emb = False\n\n    @torch.compiler.disable\n    def _delta_forward(self, x):\n        o, _, _ = self.delta(x)\n        return o\n\n    def forward(self, x):\n        B, T, C = x.size()\n        # DeltaNet forward\n        delta_out = self._delta_forward(x)\n        \n        # Apply per-head decay (scale each head's output)\n        # Reshape delta_out to (B, T, num_heads, head_dim)\n        head_dim = C // self.n_head\n        delta_out = delta_out.view(B, T, self.n_head, head_dim)\n        decay = torch.sigmoid(self.decay)  # (num_heads)\n        delta_out = delta_out * decay.view(1, 1, self.n_head, 1)\n        delta_out = delta_out.view(B, T, C)\n        \n        # Compute gating weights per token per channel\n        gate = torch.sigmoid(self.gate_proj(x))  # (B, T, C)\n        # Blend delta output with skip connection (input)\n        blended = gate * delta_out + (1 - gate) * x\n        \n        # Output projection\n        out = self.resid_dropout(self.c_proj(blended))\n        return out"
7}