| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from transformers import PreTrainedModel |
|
|
| from stu import STU |
| from modules import Attention |
| from utils import get_spectral_filters, nearest_power_of_two |
| from flash_stu.config import FlashSTUConfig |
|
|
| try: |
| from flashfftconv import FlashFFTConv |
| flash_fft_available = True |
| except ImportError as e: |
| print(f"Unable to import FlashFFTConv: {e}. Falling back to PyTorch implementation.") |
| flash_fft_available = False |
|
|
| try: |
| from flash_attn.modules.mlp import GatedMlp as MLP |
| triton_mlp = True |
| except ImportError as e: |
| print(f"Unable to import Triton-based MLP: {e}. Falling back to vanilla SwiGLU MLP instead.") |
| from modules import MLP |
| triton_mlp = False |
|
|
| try: |
| from flash_attn.ops.triton.layer_norm import RMSNorm |
| except ImportError as e: |
| print(f"Unable to import Triton-based RMSNorm: {e}. Falling back to PyTorch implementation.") |
| from torch.nn import RMSNorm |
|
|
| try: |
| from flash_attn.losses.cross_entropy import CrossEntropyLoss |
| except ImportError as e: |
| print(f"Unable to import Triton-based cross entropy loss: {e}. Falling back to PyTorch implementation.") |
| from torch.nn import CrossEntropyLoss |
|
|
| class Block(nn.Module): |
| def __init__(self, config, phi, n, flash_fft) -> None: |
| super(Block, self).__init__() |
| |
| self.rn_1 = RMSNorm(config.n_embd) |
| self.stu = STU(config, phi, n, flash_fft) |
| self.rn_2 = RMSNorm(config.n_embd) |
| self.attn = Attention(config) |
| self.rn_3 = RMSNorm(config.n_embd) |
| self.mlp = MLP( |
| config.n_embd, |
| config.n_embd * config.mlp_scale, |
| activation=F.silu, |
| bias1=config.bias, |
| bias2=config.bias, |
| ) if triton_mlp else MLP(config) |
| self.rn_4 = RMSNorm(config.n_embd) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.stu(self.rn_1(x)) |
| x = x + self.mlp(self.rn_2(x)) |
| x = x + self.attn(self.rn_3(x)) |
| x = x + self.mlp(self.rn_4(x)) |
| return x |
|
|
| class FlashSTU(PreTrainedModel): |
| config_class = FlashSTUConfig |
|
|
| def __init__(self, config) -> None: |
| super(FlashSTU, self).__init__(config) |
| self.config = config |
| self.n_layers = config.n_layers |
| self.n_embd = config.n_embd |
| self.mlp_scale = config.mlp_scale |
| self.seq_len = config.seq_len |
| self.n = nearest_power_of_two(self.seq_len * 2 - 1, round_up=True) |
| self.vocab_size = config.vocab_size |
| self.K = config.num_eigh |
| self.use_hankel_L = config.use_hankel_L |
| self.phi = get_spectral_filters(self.seq_len, self.K, self.use_hankel_L) |
| self.use_approx = config.use_approx |
| self.flash_fft = ( |
| FlashFFTConv(self.n, dtype=torch.bfloat16) |
| if config.use_flash_fft and flash_fft_available |
| else None |
| ) |
| self.dropout = config.dropout |
| self.bias = config.bias |
| self.loss_fn = CrossEntropyLoss() |
|
|
| self.flash_stu = nn.ModuleDict( |
| dict( |
| tok_emb=nn.Embedding(self.vocab_size, self.n_embd), |
| dropout=nn.Dropout(self.dropout), |
| hidden=nn.ModuleList( |
| [ |
| Block(self.config, self.phi, self.n, self.flash_fft) |
| for _ in range(self.n_layers) |
| ] |
| ), |
| rn_f=RMSNorm(config.n_embd) |
| ) |
| ) |
| self.lm_head = nn.Linear(self.n_embd, self.vocab_size, bias=self.bias) |
|
|
| self.std = (self.n_embd) ** -0.5 |
| self.apply(self._init_weights) |
| print("Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.tensor: |
| tok_emb = self.flash_stu.tok_emb(x) |
| x = self.flash_stu.dropout(tok_emb) |
|
|
| for block in self.flash_stu.hidden: |
| x = block(x) |
| x = self.flash_stu.rn_f(x) |
|
|
| y_hat = self.lm_head(x) |
| return y_hat |
|
|
| def _get_num_params(self): |
| n_params = sum(p.numel() for p in self.parameters()) |
| return n_params |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| if hasattr(module, "SCALE_INIT"): |
| self.std *= (2 * self.n_layers) ** -0.5 |
| torch.nn.init.normal_(module.weight, mean=0.0, std=self.std) |
| if module.bias is not None: |
| torch.nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| torch.nn.init.normal_(module.weight, mean=0.0, std=self.std) |
| elif isinstance(module, STU): |
| if self.use_approx: |
| torch.nn.init.xavier_normal_(module.M_inputs) |
| torch.nn.init.xavier_normal_(module.M_filters) |
| else: |
| torch.nn.init.xavier_normal_(module.M_phi_plus) |
| torch.nn.init.xavier_normal_(module.M_phi_minus) |
|
|
