fix weight initialization for residual block

2025-02-03 15:50:50 +05:45 · 2025-02-03 15:50:50 +05:45 · 5f84ee6370
parent 6e9cd28ef9
commit 5f84ee6370
1 changed files with 1 additions and 1 deletions
--- a/model/model.py
+++ b/model/model.py
@ -333,7 +333,7 @@ class Transformer(PreTrainedModel):
        self.apply(self._init_weights)

        for pn, p in self.named_parameters():
-            if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
+            if pn.endswith('w2.weight') or pn.endswith('wo.weight'):
                torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * params.n_layers))

        self.last_loss = None