CofeAI/FLM-Audio
132.9k
1import math2 3import torch4import torch.nn as nn5import torch.nn.functional as F6from transformers.configuration_utils import PretrainedConfig7 8 9class DepthGPTConfig(PretrainedConfig):10 def __init__(11 self,12 block_size: int = 8,13 vocab_size: int = 2049, # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency14 n_layer: int = 6,15 n_head: int = 16,16 n_embd: int = 1024,17 dropout: float = 0.0,18 bias: bool = False, # True: bias in Linears and LayerNorms, like GPT-2. False: a bit better and faster19 main_hidden_size = 1536,20 pad_token_id = 2048,21 use_cmlp = True,22 use_rmsnorm = False,23 use_swiglu = False24 ):25 """26 {27 "block_size": 8,28 "vocab_size": 2049,29 "n_layer": 6,30 "n_head": 16,31 "n_embd": 1024,32 "dropout": 0.0,33 "bias": false,34 "main_hidden_size": 1536,35 "pad_token_id": 2048,36 "use_cmlp": true37 }38 """39 # super().__init__(**kwargs)40 self.block_size = block_size41 self.vocab_size = vocab_size42 self.n_layer = n_layer43 self.n_head = n_head44 self.n_embd = n_embd45 self.dropout = dropout46 self.bias = bias47 self.main_hidden_size = main_hidden_size48 self.pad_token_id = pad_token_id49 self.use_cmlp = use_cmlp50 self.use_rmsnorm = use_rmsnorm51 self.use_swiglu = use_swiglu52 53################################################################################################54# GPT style55################################################################################################56 57class LayerNorm(nn.Module):58 """ LayerNorm but with an optional bias. PyTorch doesn't support simply bias=False """59 60 def __init__(self, ndim, bias):61 super().__init__()62 self.weight = nn.Parameter(torch.ones(ndim))63 self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None64 65 def forward(self, input):66 return F.layer_norm(input, self.weight.shape, self.weight, self.bias, 1e-5)67 68 69class RMSNorm(nn.Module):70 def __init__(self, dim: int, eps: float = 1e-6):71 super(RMSNorm, self).__init__()72 self.eps = eps73 self.weight = nn.Parameter(torch.ones(dim))74 75 def _norm(self, x):76 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)77 78 def forward(self, x):79 output = self._norm(x.float()).type_as(x)80 return output * self.weight81 82 83class CausalSelfAttention(nn.Module):84 85 def __init__(self, config):86 super().__init__()87 assert config.n_embd % config.n_head == 088 # key, query, value projections for all heads, but in a batch89 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)90 # output projection91 self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)92 # regularization93 self.attn_dropout = nn.Dropout(config.dropout)94 self.resid_dropout = nn.Dropout(config.dropout)95 self.n_head = config.n_head96 self.n_embd = config.n_embd97 self.dropout = config.dropout98 # flash attention make GPU go brrrrr but support is only in PyTorch >= 2.099 self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')100 if not self.flash:101 print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")102 # causal mask to ensure that attention is only applied to the left in the input sequence103 self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))104 .view(1, 1, config.block_size, config.block_size))105 106 def forward(self, x):107 B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)108 109 # calculate query, key, values for all heads in batch and move head forward to be the batch dim110 q, k, v = self.c_attn(x).split(self.n_embd, dim=2)111 k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)112 q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)113 v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)114 115 # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)116 if self.flash:117 # efficient attention using Flash Attention CUDA kernels118 y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)119 else:120 # manual implementation of attention121 att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))122 att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))123 att = F.softmax(att, dim=-1)124 att = self.attn_dropout(att)125 y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)126 y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side127 128 # output projection129 y = self.resid_dropout(self.c_proj(y))130 return y131 132 133class MLP(nn.Module):134 def __init__(self, config):135 super().__init__()136 self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)137 self.gelu = nn.GELU()138 self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)139 self.dropout = nn.Dropout(config.dropout)140 141 def forward(self, x):142 x = self.c_fc(x)143 x = self.gelu(x)144 x = self.c_proj(x)145 x = self.dropout(x)146 return x147 148 149class MLP_swiglu(nn.Module):150 def __init__(self, config):151 super().__init__()152 self.intermediate_size = int(8 * config.n_embd / 3)153 self.gate_proj = nn.Linear(config.n_embd, self.intermediate_size, bias=config.bias)154 self.up_proj = nn.Linear(config.n_embd, self.intermediate_size, bias=config.bias)155 self.down_proj = nn.Linear(self.intermediate_size, config.n_embd, bias=config.bias)156 self.act_fn = F.silu157 self.dropout = nn.Dropout(config.dropout)158 159 def forward(self, x):160 x = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))161 x = self.dropout(x)162 return x163 164class Block(nn.Module):165 166 def __init__(self, config):167 super().__init__()168 self.ln_1 = RMSNorm(config.n_embd) if config.use_rmsnorm else LayerNorm(config.n_embd, bias=config.bias)169 self.attn = CausalSelfAttention(config)170 self.ln_2 = RMSNorm(config.n_embd) if config.use_rmsnorm else LayerNorm(config.n_embd, bias=config.bias)171 mlp_cls = MLP_swiglu if config.use_swiglu else MLP172 self.mlp = mlp_cls(config)173 174 def forward(self, x):175 x = x + self.attn(self.ln_1(x))176 x = x + self.mlp(self.ln_2(x))177 return x178 179 180class BlockCMLP(nn.Module):181 182 def __init__(self, config):183 super().__init__()184 self.channel_size = config.block_size185 self.ln_1 = RMSNorm(config.n_embd) if config.use_rmsnorm else LayerNorm(config.n_embd, bias=config.bias)186 self.attn = CausalSelfAttention(config)187 self.ln_2 = RMSNorm(config.n_embd) if config.use_rmsnorm else LayerNorm(config.n_embd, bias=config.bias)188 mlp_cls = MLP_swiglu if config.use_swiglu else MLP189 self.mlps = nn.ModuleList([mlp_cls(config) for _ in range(self.channel_size)])190 191 assert self.channel_size == 8, f"DEBUG, self.channel_size={self.channel_size} != 8"192 193 def forward(self, x):194 _, channel_size, _ = x.shape195 # assert channel_size == self.channel_size196 x = x + self.attn(self.ln_1(x))197 198 xl = self.ln_2(x)199 x = x + torch.cat(200 [self.mlps[c](xl[:, c:c+1, :]) for c in range(self.channel_size)],201 dim=1202 )203 return x204 205 206class DepthGPT(nn.Module):207 208 def __init__(self, config):209 super().__init__()210 assert config.vocab_size is not None211 assert config.block_size is not None212 self.config = config213 self.num_channel = config.block_size214 215 self.linear_in = nn.Linear(config.main_hidden_size, config.n_embd * config.block_size, bias=False)216 217 block_cls = BlockCMLP if config.use_cmlp else Block218 self.transformer = nn.ModuleDict(dict(219 wtes = nn.ModuleList([nn.Embedding(config.vocab_size, config.n_embd) for _ in range(self.num_channel)]),220 wpe = nn.Embedding(self.num_channel, config.n_embd),221 drop = nn.Dropout(config.dropout),222 h = nn.ModuleList([block_cls(config) for _ in range(config.n_layer)]),223 ln_f = RMSNorm(config.n_embd) if config.use_rmsnorm else LayerNorm(config.n_embd, bias=config.bias),224 ))225 self.lm_heads = nn.ModuleList([nn.Linear(config.n_embd, config.vocab_size, bias=False) for _ in range(self.num_channel)])226 227 # with weight tying when using torch.compile() some warnings get generated:228 # "UserWarning: functional_call was passed multiple values for tied weights.229 # This behavior is deprecated and will be an error in future versions"230 # not 100% sure what this is, so far seems to be harmless. TODO investigate231 # self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying232 233 # init all weights234 self.apply(self._init_weights)235 # apply special scaled init to the residual projections, per GPT-2 paper236 for pn, p in self.named_parameters():237 if pn.endswith('c_proj.weight'):238 torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))239 240 # report number of parameters241 print("number of parameters: %.2fM" % (self.get_num_params()/1e6,))242 243 def get_num_params(self, non_embedding=False):244 """245 Return the number of parameters in the model.246 For non-embedding count (default), the position embeddings get subtracted.247 The token embeddings would too, except due to the parameter sharing these248 params are actually used as weights in the final layer, so we include them.249 """250 n_params = sum(p.numel() for p in self.parameters())251 if non_embedding:252 n_params -= self.transformer.wpe.weight.numel()253 return n_params254 255 def _init_weights(self, module):256 if isinstance(module, nn.Linear):257 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)258 if module.bias is not None:259 torch.nn.init.zeros_(module.bias)260 elif isinstance(module, nn.Embedding):261 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)262 263 def forward(self,264 main_hidden_states, # [seq, main_dim]265 audio_token_ids # [seq, 7]266 ):267 268 assert main_hidden_states.shape[0] == audio_token_ids.shape[0]269 in_audio_token_num = audio_token_ids.shape[-1]270 271 device = audio_token_ids.device272 273 audio_token_ids = F.pad(audio_token_ids, (1, 0), value=self.config.pad_token_id)274 275 x = torch.stack(276 [self.transformer.wtes[c](audio_token_ids[:, c]) for c in range(in_audio_token_num + 1)]277 ).transpose(0, 1) # [seq, in_audio_token_num]278 279 x += self.transformer.wpe(280 torch.arange(0, in_audio_token_num + 1, dtype=torch.long, device=device)281 ).unsqueeze(0) # position embeddings of shape (1, 8, depth_dim)282 283 main_hidden = self.linear_in(main_hidden_states).view(main_hidden_states.shape[0], self.config.block_size, -1)[:, :in_audio_token_num+1, :]284 x += main_hidden285 286 x = self.transformer.drop(x)287 for block in self.transformer.h:288 x = block(x)289 290 # [seq, 8, hidden]291 x = self.transformer.ln_f(x)292 293 # [seq, 8, hidden] (linear)-> [8, seq, vocab]294 x = torch.stack([self.lm_heads[c](x[:, c, :]) for c in range(x.shape[1])])295 296 # [8, seq, vocab] -> [seq, 8, vocab]297 x = x.transpose(0,1)298 299 return x300 def _initialize_weights(self, module):301 if isinstance(module, nn.Linear):302 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)303 if module.bias is not None:304 torch.nn.init.zeros_(module.bias)305 elif isinstance(module, nn.Embedding):306 torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)307 308 309if __name__ == "__main__":310 config = {311 "bias": False,312 "dropout": 0.0,313 "n_embd": 1024,314 "n_head": 16,315 "n_layer": 6,316 "use_cmlp": True,317 "use_rmsnorm": True,318 "use_swiglu": True,319 "main_hidden_size": 4096320 }321 model_config = DepthGPTConfig(**config)322 model = DepthGPT(config=model_config)323 324 main_hidden_states = torch.rand((1, 4096))325 decoded_audio_tokens = torch.empty((1, 0), dtype=torch.long, device=main_hidden_states.device)326 outputs = model(main_hidden_states, decoded_audio_tokens)327