File size: 2,976 Bytes
bc88a9c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | # Copyright (c) Together
# Apache 2.0 - Author: Michael Poli
# Adapted for the minimal Evo1 HF port.
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
|