from dataclasses import dataclass import math from typing import Optional import torch import torch.nn as nn from torch.nn import functional as F from .base import ( Block, NormFunc, configure_optimizers, estimate_mfu, get_init_fn, init_weights_simple, ) TIE_WEIGHTS = False SIMPLE_INIT = True Z_LOSS = True @dataclass class GPTConfig: n_layer: int = 24 n_head: int = 16 # query heads n_kv_head: Optional[int] = None d_head: int = 64 block_size: int = 4288 bias: bool = False dropout: float = 0.0 text_vocab_size: int = 60_032 text_codebook_size: int = 60_001 text_pad_token: int = 1 text_infer_token: int = 60_001 text_end_token: int = 60_002 semantic_vocab_size: int = 4032 semantic_codebook_size: int = 4000 semantic_n_codebooks: int = 1 semantic_pad_token: int = 4000 semantic_infer_token: int = 4001 semantic_rate_hz: int = 25 semantic_shift_factor: int = 50 coarse_vocab_size: int = 2112 coarse_codebook_size: int = 2048 coarse_n_codebooks: int = 12 coarse_pad_token: int = 2048 coarse_infer_token: int = 2049 coarse_rate_hz: int = 25 coarse_shift_factor: int = 5 t_text: int = 150 t_audio: int = 3008 t_memmap: int = 3008 use_rotary_pos_emb: bool = False attention_type: str = "torch" # "torch", "tao", "xformers" attention_sliding_window_size: int = -1 def __post_init__(self): # default to multi head attention if self.n_kv_head is None: self.n_kv_head = self.n_head @property def n_embd(self): """The width of the residual stream""" return self.n_head * self.d_head @property def use_learned_pos_emb(self): return not self.use_rotary_pos_emb class GPT(nn.Module): def __init__(self, config: GPTConfig): super().__init__() self.config = config model_dict = dict( wte_text=nn.Embedding(config.text_vocab_size, config.n_embd), ln_text=NormFunc(config.n_embd), wte_semantic=nn.ModuleList( [ nn.Embedding(config.semantic_vocab_size, config.n_embd) for _ in range(config.semantic_n_codebooks) ] ), ln_semantic=NormFunc(config.n_embd), wte_coarse=nn.ModuleList( [ nn.Embedding(config.coarse_vocab_size, config.n_embd) for _ in range(config.coarse_n_codebooks) ] ), ln_coarse=NormFunc(config.n_embd), drop=nn.Dropout(config.dropout), h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f=NormFunc(config.n_embd), ) if self.config.use_learned_pos_emb: model_dict["wpe"] = nn.Embedding(config.block_size, config.n_embd) self.transformer = nn.ModuleDict(model_dict) self.lm_heads = nn.ModuleList( [nn.Linear(config.n_embd, config.text_vocab_size, bias=False) for _ in range(1)] ) if TIE_WEIGHTS: for n in range(config.semantic_n_codebooks): self.transformer.wte_semantic[n].weight = self.lm_heads[n].weight for n in range(config.coarse_n_codebooks): n2 = config.semantic_n_codebooks + n self.transformer.wte_coarse[n].weight = self.lm_heads[n2].weight # init all weights if SIMPLE_INIT: self.apply(self._init_weights_simple) for pn, p in self.named_parameters(): if pn.endswith("c_proj.weight"): torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)) else: self._init_weights() print(f"number of parameters: {self.get_num_params()/1e6:.0f}M") def forward(self, x, y=None, audio_offset=0, return_logits=False, last_only=True): device = x.device b, ns, t = x.size() assert ns == 1 + self.config.semantic_n_codebooks + self.config.coarse_n_codebooks if y is not None: assert t == self.config.block_size _, _, t2 = y.size() assert t2 == self.config.t_text # embed text x_emb = self.transformer.wte_text(x[:, 0, :]) x_emb = self.transformer.ln_text(x_emb) # embed semantic for n in range(self.config.semantic_n_codebooks): x_emb += self.transformer.ln_semantic(self.transformer.wte_semantic[n](x[:, 1 + n, :])) # embed coarse for n in range(self.config.coarse_n_codebooks): n2 = 1 + n + self.config.semantic_n_codebooks x_emb += self.transformer.ln_coarse(self.transformer.wte_coarse[n](x[:, n2, :])) # x_emb (b, t, n_embd) if self.config.use_learned_pos_emb: pos = torch.arange(t, dtype=torch.long, device=device).unsqueeze(0) # shape (1, t) pos_emb = self.transformer.wpe(pos) # (1, t, n_embd) x_emb += pos_emb x = self.transformer.drop(x_emb) for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) x = x[:, audio_offset : audio_offset + self.config.t_text, :] if return_logits: if last_only: x = x[:, -1, :] text_logits_list = [] for n in range(1): text_logits_list.append(self.lm_heads[n](x)) text_logits = torch.stack(text_logits_list).swapaxes(0, 1) return text_logits loss_dict = {} if Z_LOSS: loss_dict["z_loss"] = 0 # projection logits = self.lm_heads[0](x) # flatten flat_logits = logits.reshape(-1, logits.size(-1)) flat_y = y[:, 0, :].reshape(-1) mask = flat_y != 1 loss_dict["text_0"] = F.cross_entropy(flat_logits[mask], flat_y[mask]) # loss_dict["text_0"] = F.cross_entropy( # logits.reshape(-1, logits.size(-1)), y[:, 0, :].reshape(-1) # ) if Z_LOSS: loss_dict["z_loss"] += (torch.logsumexp(logits, dim=-1) ** 2).mean() return loss_dict, logits def get_num_params(self, non_embedding=True): n_params = sum(p.numel() for p in self.parameters()) if non_embedding: for m in self.transformer.wte_semantic: n_params -= m.weight.numel() for m in self.transformer.wte_coarse: n_params -= m.weight.numel() if self.config.use_learned_pos_emb: n_params -= self.transformer.wpe.weight.numel() return n_params def _init_weights_simple(self, module): init_weights_simple(self, module) def _init_weights(self): # embeddings get_init_fn(self.config.n_embd, init_depth=None)(self.transformer.wte_text.weight) for module in self.transformer.wte_semantic: get_init_fn(self.config.n_embd, init_depth=None)(module.weight) for module in self.transformer.wte_coarse: get_init_fn(self.config.n_embd, init_depth=None)(module.weight) if self.config.use_learned_pos_emb: get_init_fn(self.config.n_embd, init_depth=None)(self.transformer.wpe.weight) # heads for module in self.lm_heads: get_init_fn(self.config.n_embd, init_depth=None)(module.weight) if module.bias is not None: torch.nn.init.zeros_(module.bias) # attention blocks for layer_idx, block in enumerate(self.transformer.h): # mlp module = block.mlp.c_fc get_init_fn(self.config.n_embd, init_depth=layer_idx + 1)(module.weight) if module.bias is not None: torch.nn.init.zeros_(module.bias) module = block.mlp.c_proj get_init_fn(block.mlp.embd_inner, init_depth=layer_idx + 1)(module.weight) if module.bias is not None: torch.nn.init.zeros_(module.bias) # attention for module in [block.attn.c_attn, block.attn.c_proj]: get_init_fn(self.config.n_embd, init_depth=layer_idx + 1)(module.weight) if module.bias is not None: torch.nn.init.zeros_(module.bias) def configure_optimizers(self, weight_decay, learning_rate, betas, device_type): return configure_optimizers(self, weight_decay, learning_rate, betas, device_type) def estimate_mfu(self, fwdbwd_per_iter, dt): return estimate_mfu(self, fwdbwd_per_iter, dt)