| |
| |
| |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch import Tensor |
|
|
|
|
| def grab_first_if_tuple(x): |
| if x.__class__.__name__ == "tuple": |
| return x[0] |
| return x |
|
|
|
|
| class RMSNorm(torch.nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| self.eps = config.eps |
| self.hidden_size = config.hidden_size |
| self.scale = torch.nn.Parameter(torch.ones(self.hidden_size)) |
| self.register_parameter("scale", self.scale) |
| self.use_flash_rmsnorm = config.get("use_flash_rmsnorm", False) |
| if self.use_flash_rmsnorm: |
| from flash_attn.ops.rms_norm import rms_norm as rmsnorm_func |
| self.rmsnorm_func = rmsnorm_func |
|
|
| def forward(self, x): |
| if self.use_flash_rmsnorm: |
| return self.rmsnorm_func(x, self.scale, self.eps) |
| y = x / (x.norm(2, dim=-1, keepdim=True) * self.hidden_size ** (-1.0 / 2) + self.eps) |
| return self.scale * y |
|
|
|
|
| class ParallelGatedMLP(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| multiple_of = config.get("inner_size_multiple_of", 64) |
| self.act_type = config.get("mlp_activation", "silu") |
| if self.act_type == "gelu": |
| self.act = F.gelu |
| elif self.act_type == "silu": |
| self.act = F.silu |
| else: |
| raise NotImplementedError(f"Unknown mlp_activation: {self.act_type}") |
|
|
| self.multiple_of = multiple_of * config.model_parallel_size |
| inner_size = int(2 * config.hidden_size * 4 / 3) |
| inner_size = self.multiple_of * ((inner_size + self.multiple_of - 1) // self.multiple_of) |
| if config.get("inner_mlp_size", None) is not None: |
| inner_size = config.inner_mlp_size |
|
|
| self.l1 = nn.Linear(config.hidden_size, inner_size, bias=False) |
| self.l2 = nn.Linear(config.hidden_size, inner_size, bias=False) |
| self.l3 = nn.Linear(inner_size, config.hidden_size, bias=False) |
|
|
| def forward(self, z): |
| z1, z2 = self.l1(z), self.l2(z) |
| z1, z2 = grab_first_if_tuple(z1), grab_first_if_tuple(z2) |
| y = self.l3(self.act(z1) * z2) |
| return grab_first_if_tuple(y) |
|
|
|
|
| class VocabParallelEmbedding(nn.Embedding): |
| """Single-process variant of the original VocabParallelEmbedding. |
| |
| The original supports tensor-parallel embedding sharding. We keep the |
| naming so existing checkpoints load directly, but drop all distributed |
| paths since this minimal port runs on a single device. |
| """ |
|
|
| def __init__(self, config): |
| vocab_size = config.vocab_size |
| padding_idx = config.get("padding_idx", None) |
| super().__init__(vocab_size, embedding_dim=config.hidden_size, padding_idx=padding_idx) |
|
|
| def embed(self, x: Tensor) -> Tensor: |
| return self.forward(x) |
|
|
| def unembed(self, u: Tensor) -> Tensor: |
| return u @ self.weight.T |
|
|
