| 284 | self.reset_parameters() |
| 285 | |
| 286 | def reset_parameters(self): |
| 287 | std = 1.0 / math.sqrt(self.dim) |
| 288 | nn.init.trunc_normal_(self.ffnn.ffnn.gateup_proj.weight, std=std, a=-3 * std, b=3 * std) |
| 289 | nn.init.trunc_normal_(self.attn.attn.proj_qkv.weight, std=std, a=-3 * std, b=3 * std) |
| 290 | nn.init.trunc_normal_(self.attn.attn.attn_out.weight, std=std, a=-3 * std, b=3 * std) |
| 291 | |
| 292 | xstd = 1.0 / math.sqrt(self.expand_dim) |
| 293 | nn.init.trunc_normal_(self.ffnn.ffnn.down_proj.weight, std=xstd, a=-3 * xstd, b=3 * xstd) |
| 294 | |
| 295 | def forward(self, x: Tensor, kv: Optional[Tensor] = None) -> Tensor: |
| 296 | """ |