Fix decoder architecture to match v2 training

app.py

@@ -54,16 +54,19 @@ class GcodeDecoder(nn.Module):
         self.token_embed = nn.Embedding(config.vocab_size, config.hidden_size)
         self.pos_embed = nn.Embedding(config.max_seq_len, config.hidden_size)
         
-        decoder_layer = nn.TransformerDecoderLayer(
-            d_model=config.hidden_size,
-            nhead=config.num_heads,
-            dim_feedforward=config.hidden_size * 4,
-            dropout=config.dropout,
-            activation='gelu',
-            batch_first=True,
-            norm_first=True,
-        )
-        self.decoder = nn.TransformerDecoder(decoder_layer, config.num_layers)
+        # Individual layers (matches v2 training architecture)
+        self.layers = nn.ModuleList([
+            nn.TransformerDecoderLayer(
+                d_model=config.hidden_size,
+                nhead=config.num_heads,
+                dim_feedforward=config.hidden_size * 4,
+                dropout=config.dropout,
+                activation='gelu',
+                batch_first=True,
+                norm_first=True,
+            )
+            for _ in range(config.num_layers)
+        ])
         
         self.ln_f = nn.LayerNorm(config.hidden_size)
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
@@ -82,7 +85,9 @@ class GcodeDecoder(nn.Module):
         
         causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device)
         
-        x = self.decoder(x, memory, tgt_mask=causal_mask)
+        for layer in self.layers:
+            x = layer(x, memory, tgt_mask=causal_mask)
+        
         x = self.ln_f(x)
         return self.lm_head(x)