Elizesel/kronos-btc-bot
0
1import math2 3from einops import rearrange, reduce4import torch5import torch.nn as nn6from torch.autograd import Function7import torch.nn.functional as F8 9 10class DifferentiableEntropyFunction(Function):11 @staticmethod12 def forward(ctx, zq, basis, K, eps):13 zb = (zq + 1) / 214 zi = ((zb * basis).sum(-1)).to(torch.int64)15 cnt = torch.scatter_reduce(torch.zeros(2 ** K, device=zq.device, dtype=zq.dtype),16 0,17 zi.flatten(),18 torch.ones_like(zi.flatten()).to(zq.dtype),19 'sum')20 prob = (cnt + eps) / (cnt + eps).sum()21 H = -(prob * torch.log(prob)).sum()22 ctx.save_for_backward(zq, zi, prob)23 ctx.K = K24 return H25 26 @staticmethod27 def backward(ctx, grad_output):28 zq, zi, prob = ctx.saved_tensors29 grad_array = -grad_output * (torch.log(prob) + 1) / zi.numel() / ctx.K30 reord_grad = grad_array[zi.flatten()].reshape(zi.shape)31 grad_input = reord_grad.unsqueeze(-1) * zq32 return grad_input, None, None, None, None33 34 35def codebook_entropy(zq, basis, K, eps=1e-4):36 return DifferentiableEntropyFunction.apply(zq, basis, K, eps)37 38 39class BinarySphericalQuantizer(nn.Module):40 def __init__(self, embed_dim, beta, gamma0, gamma, zeta,41 input_format='bchw',42 soft_entropy=True, group_size=9,43 persample_entropy_compute='analytical',44 cb_entropy_compute='group',45 l2_norm=True,46 inv_temperature=1):47 """48 Paper link: https://arxiv.org/pdf/2406.07548.pdf49 Here we use the official implementation of the BinarySphericalQuantizer.50 """51 super().__init__()52 self.embed_dim = embed_dim53 self.beta = beta # loss weight for commit loss54 self.gamma0 = gamma0 # loss weight for entropy penalty55 self.gamma = gamma # loss weight for entropy penalty56 self.zeta = zeta # loss weight for entire entropy penalty57 self.input_format = input_format58 assert self.embed_dim % group_size == 0, "embed_dim must be divisible by group_size"59 self.num_groups = self.embed_dim // group_size60 self.group_size = group_size61 assert persample_entropy_compute in ['group', 'analytical'], "persample_entropy_compute must be either 'group' or 'analytical'"62 assert cb_entropy_compute in ['group', 'nce'], "cb_entropy_compute must be either 'group' or 'nce'"63 self.persample_entropy_compute = persample_entropy_compute64 self.cb_entropy_compute = cb_entropy_compute65 self.l2_norm = l2_norm66 self.inv_temperature = inv_temperature67 68 self.register_buffer('basis', 2 ** torch.arange(embed_dim - 1, -1, -1))69 self.register_buffer('group_basis', 2 ** torch.arange(group_size - 1, -1, -1))70 71 self.num_dimensions = 2 ** embed_dim72 self.bits_per_index = embed_dim73 74 # we only need to keep the codebook portion up to the group size75 # because we approximate the H loss with this subcode76 group_codes = torch.arange(2 ** self.group_size)77 group_codebook = self.indexes_to_codes(group_codes).float()[:, -group_size:]78 self.register_buffer('group_codebook', group_codebook, persistent=False)79 80 self.soft_entropy = soft_entropy # soft_entropy: Sec 3.2 of https://arxiv.org/pdf/1911.05894.pdf81 82 def quantize(self, z):83 assert z.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {z.shape[-1]}"84 85 zhat = torch.where(z > 0,86 torch.tensor(1, dtype=z.dtype, device=z.device),87 torch.tensor(-1, dtype=z.dtype, device=z.device))88 return z + (zhat - z).detach()89 90 def forward(self, z, collect_metrics=True):91 # if self.input_format == 'bchw':92 # z = rearrange(z, 'b c h w -> b h w c')93 zq = self.quantize(z)94 95 q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.96 97 zq = zq * q_scale98 99 if not collect_metrics:100 return zq, zq.new_zeros(()), {}101 102 indices = self.codes_to_indexes(zq.detach())103 group_indices = self.codes_to_group_indexes(zq.detach())104 if not self.training:105 used_codes = torch.unique(indices, return_counts=False)106 else:107 used_codes = None108 109 if self.soft_entropy:110 persample_entropy, cb_entropy, avg_prob = self.soft_entropy_loss(z)111 entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy112 else:113 zb_by_sample = ((zq + 1) / 2).reshape(z.shape[0], -1, z.shape[-1]).to(torch.float32)114 persample_entropy = self.get_hard_per_sample_entropy(zb_by_sample)115 cb_entropy = codebook_entropy(zq, self.basis, self.embed_dim)116 entropy_penalty = self.gamma0 * persample_entropy - self.gamma * cb_entropy117 118 # commit loss119 commit_loss = self.beta * torch.mean(((zq.detach() - z) ** 2).sum(dim=-1))120 121 # if self.input_format == 'bchw':122 # zq = rearrange(zq, 'b h w c -> b c h w')123 124 return (125 zq,126 commit_loss + self.zeta * entropy_penalty / self.inv_temperature,127 {"H": cb_entropy, "used_codes": used_codes, "indices": indices, "group_indices": group_indices,128 "avg_prob": avg_prob}129 )130 131 def soft_entropy_loss(self, z):132 # if we divide the code in subgroups of size group_size, the codebook will be of size 2 ** group_size133 # the sub-code is the last group_size bits of the full code134 group_code_book = self.group_codebook / (self.embed_dim ** 0.5 if self.l2_norm else 1)135 divided_z = rearrange(z, '... (g c) -> ... g c', c=self.group_size)136 137 # we calculate the distance between the divided_z and the codebook for each subgroup138 distance = - 2 * torch.einsum('... g c, d c ->... g d', divided_z, group_code_book)139 prob = (-distance * self.inv_temperature).softmax(dim=-1)140 if self.persample_entropy_compute == 'analytical':141 if self.l2_norm:142 p = torch.sigmoid(-4 * z / (self.embed_dim ** 0.5) * self.inv_temperature)143 else:144 p = torch.sigmoid(-4 * z * self.inv_temperature)145 prob = torch.stack([p, 1 - p], dim=-1)146 per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()147 else:148 per_sample_entropy = self.get_entropy(prob, dim=-1, normalize=False).sum(dim=-1).mean()149 150 # macro average of the probability of each subgroup151 avg_prob = reduce(prob, '... g d ->g d', 'mean')152 codebook_entropy = self.get_entropy(avg_prob, dim=-1, normalize=False)153 154 # the approximation of the entropy is the sum of the entropy of each subgroup155 return per_sample_entropy, codebook_entropy.sum(), avg_prob156 157 def get_hard_per_sample_entropy(self, zb_by_sample):158 probs_per_dim = zb_by_sample.sum(1) / zb_by_sample.shape[1]159 persample_entropy = - probs_per_dim * torch.log(probs_per_dim + 1e-8) - (1 - probs_per_dim) * torch.log(1 - probs_per_dim + 1e-8)160 persample_entropy = persample_entropy.sum(-1)161 return persample_entropy.mean()162 163 def codes_to_indexes(self, zhat):164 """Converts a `code` to an index in the codebook.165 Args:166 zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}167 """168 assert zhat.shape[-1] == self.embed_dim, f"Expected {self.embed_dim} dimensions, got {zhat.shape[-1]}"169 return ((zhat + 1) / 2 * self.basis).sum(axis=-1).to(torch.int64)170 171 def codes_to_group_indexes(self, zhat):172 """Converts a `code` to a list of indexes (in groups) in the codebook.173 Args:174 zhat: A tensor of shape (B, ..., C) containing the codes. must be in {-1, 1}175 """176 zhat_in_group = rearrange(zhat, 'b ... (g c) -> b ... g c', c=self.group_size)177 return ((zhat_in_group + 1) / 2 * self.group_basis).sum(axis=-1).to(torch.int64)178 179 def indexes_to_codes(self, indices):180 """Inverse of `indexes_to_codes`."""181 indices = indices.unsqueeze(-1)182 codes_non_centered = torch.remainder(183 torch.floor_divide(indices, self.basis), 2184 )185 return codes_non_centered * 2 - 1186 187 def group_indexes_to_codes(self, group_indices):188 """Inverse of `group_indexes_to_codes`."""189 group_indices = group_indices.unsqueeze(-1)190 codes_non_centered = torch.remainder(191 torch.floor_divide(group_indices, self.group_basis), 2192 )193 codes_non_centered = rearrange(codes_non_centered, 'b ... g c -> b ... (g c)')194 return codes_non_centered * 2 - 1195 196 def get_entropy(self, count, dim=-1, eps=1e-4, normalize=True):197 if normalize:198 probs = (count + eps) / (count + eps).sum(dim=dim, keepdim=True)199 else:200 probs = count201 H = -(probs * torch.log(probs + 1e-8)).sum(dim=dim)202 return H203 204 def get_group_codebook_entry(self, group_indices):205 z_q = self.group_indexes_to_codes(group_indices)206 q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.207 z_q = z_q * q_scale208 if self.input_format == 'bchw':209 h, w = int(z_q.shape[1] ** 0.5)210 assert h * w == z_q.shape[1], 'Invalid sequence length'211 z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)212 return z_q213 214 def get_codebook_entry(self, indices):215 z_q = self.indexes_to_codes(indices)216 q_scale = 1. / (self.embed_dim ** 0.5) if self.l2_norm else 1.217 z_q = z_q * q_scale218 if self.input_format == 'bchw':219 h, w = int(z_q.shape[1] ** 0.5)220 assert h * w == z_q.shape[1], 'Invalid sequence length'221 z_q = rearrange(z_q, 'b (h w) c -> b c h w', h=h)222 return z_q223 224 225class BSQuantizer(nn.Module):226 227 def __init__(self, s1_bits, s2_bits, beta, gamma0, gamma, zeta, group_size):228 super().__init__()229 self.codebook_dim = s1_bits + s2_bits230 self.s1_bits = s1_bits231 self.s2_bits = s2_bits232 self.bsq = BinarySphericalQuantizer(self.codebook_dim, beta, gamma0, gamma, zeta, group_size=group_size)233 234 def bits_to_indices(self, bits):235 bits = (bits >= 0).to(torch.long)236 indices = 2 ** torch.arange(237 0,238 bits.shape[-1],239 1,240 dtype=torch.long,241 device=bits.device,242 )243 return (bits * indices).sum(-1)244 245 def forward(self, z, half=False, collect_metrics=True):246 z = F.normalize(z, dim=-1)247 quantized, bsq_loss, metrics = self.bsq(z, collect_metrics=collect_metrics)248 if half:249 q_pre = quantized[:, :, :self.s1_bits]250 q_post = quantized[:, :, self.s1_bits:]251 z_indices = [self.bits_to_indices(q_pre), self.bits_to_indices(q_post)]252 else:253 z_indices = self.bits_to_indices(quantized)254 return bsq_loss, quantized, z_indices255 256 257class RMSNorm(torch.nn.Module):258 def __init__(self, dim: int, eps: float = 1e-5):259 super().__init__()260 self.eps = eps261 self.weight = nn.Parameter(torch.ones(dim))262 263 def _norm(self, x):264 return x * torch.rsqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps)265 266 def forward(self, x):267 output = self._norm(x.float()).type_as(x)268 return output * self.weight269 270 271class FeedForward(nn.Module):272 def __init__(self, d_model, ff_dim, ffn_dropout_p=0.0):273 super().__init__()274 275 self.w1 = nn.Linear(d_model, ff_dim, bias=False)276 self.w3 = nn.Linear(d_model, ff_dim, bias=False)277 self.w2 = nn.Linear(ff_dim, d_model, bias=False)278 self.ffn_dropout = nn.Dropout(ffn_dropout_p)279 280 def forward(self, x):281 return self.ffn_dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))282 283 284class RotaryPositionalEmbedding(nn.Module):285 def __init__(self, dim):286 super().__init__()287 inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))288 self.register_buffer("inv_freq", inv_freq)289 self.seq_len_cached = None290 self.cos_cached = None291 self.sin_cached = None292 293 def _update_cos_sin_cache(self, x, seq_len):294 if seq_len != self.seq_len_cached:295 self.seq_len_cached = seq_len296 t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)297 freqs = torch.einsum('i,j->ij', t, self.inv_freq)298 emb = torch.cat((freqs, freqs), dim=-1).to(x.device)299 self.cos_cached = emb.cos()[None, None, :, :]300 self.sin_cached = emb.sin()[None, None, :, :]301 return self.cos_cached, self.sin_cached302 303 def forward(self, q, k):304 cos, sin = self._update_cos_sin_cache(q, q.shape[-2])305 return (306 (q * cos) + (self._rotate_half(q) * sin),307 (k * cos) + (self._rotate_half(k) * sin),308 )309 310 def _rotate_half(self, x):311 x1, x2 = x.chunk(2, dim=-1)312 return torch.cat((-x2, x1), dim=-1)313 314 315class MultiHeadAttentionWithRoPE(nn.Module):316 def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout_p=0.0):317 super().__init__()318 self.d_model = d_model319 self.n_heads = n_heads320 self.head_dim = d_model // n_heads321 322 self.q_proj = nn.Linear(d_model, d_model)323 self.k_proj = nn.Linear(d_model, d_model)324 self.v_proj = nn.Linear(d_model, d_model)325 self.out_proj = nn.Linear(d_model, d_model)326 self.rotary = RotaryPositionalEmbedding(self.head_dim)327 self.attn_dropout_p = attn_dropout_p328 self.resid_dropout = nn.Dropout(resid_dropout_p)329 330 def forward(self, x, key_padding_mask=None):331 batch_size, seq_len, _ = x.shape332 333 q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)334 k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)335 v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)336 337 q, k = self.rotary(q, k)338 339 if key_padding_mask is not None:340 attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2) # [batch, 1, 1, seq_len]341 attn_mask = attn_mask.expand(-1, self.n_heads, seq_len, -1) # [batch, n_heads, q_len, k_len]342 else:343 attn_mask = None344 345 attn_output = F.scaled_dot_product_attention(346 q, k, v,347 attn_mask=attn_mask,348 dropout_p=self.attn_dropout_p if self.training else 0.0,349 is_causal=True350 )351 352 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)353 return self.resid_dropout(self.out_proj(attn_output))354 355 356class MultiHeadCrossAttentionWithRoPE(nn.Module):357 def __init__(self, d_model, n_heads, attn_dropout_p=0.0, resid_dropout=0.0):358 super().__init__()359 self.d_model = d_model360 self.n_heads = n_heads361 self.head_dim = d_model // n_heads362 363 self.q_proj = nn.Linear(d_model, d_model)364 self.k_proj = nn.Linear(d_model, d_model)365 self.v_proj = nn.Linear(d_model, d_model)366 self.out_proj = nn.Linear(d_model, d_model)367 self.rotary = RotaryPositionalEmbedding(self.head_dim)368 self.attn_dropout_p = attn_dropout_p369 self.resid_dropout = nn.Dropout(resid_dropout)370 371 def forward(self, query, key, value, key_padding_mask=None):372 batch_size, q_len, _ = query.shape373 _, seq_len, _ = key.shape374 375 q = self.q_proj(query).view(batch_size, q_len, self.n_heads, self.head_dim).transpose(1, 2)376 k = self.k_proj(key).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)377 v = self.v_proj(value).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)378 379 q, k = self.rotary(q, k)380 381 if key_padding_mask is not None:382 attn_mask = key_padding_mask.unsqueeze(1).unsqueeze(2)383 attn_mask = attn_mask.expand(-1, self.n_heads, q_len, -1)384 else:385 attn_mask = None386 387 is_causal_flag = self.training388 389 attn_output = F.scaled_dot_product_attention(390 q, k, v,391 attn_mask=attn_mask,392 dropout_p=self.attn_dropout_p if self.training else 0.0,393 is_causal=is_causal_flag394 )395 396 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.d_model)397 return self.resid_dropout(self.out_proj(attn_output))398 399 400class HierarchicalEmbedding(nn.Module):401 def __init__(self, s1_bits, s2_bits, d_model=256):402 super().__init__()403 self.s1_bits = s1_bits404 self.s2_bits = s2_bits405 406 vocab_s1 = 2 ** s1_bits407 vocab_s2 = 2 ** s2_bits408 409 self.emb_s1 = nn.Embedding(vocab_s1, d_model)410 self.emb_s2 = nn.Embedding(vocab_s2, d_model)411 self.d_model = d_model412 self.fusion_proj = nn.Linear(d_model * 2, d_model)413 414 nn.init.normal_(self.emb_s1.weight, mean=0, std=d_model ** -0.5)415 nn.init.normal_(self.emb_s2.weight, mean=0, std=d_model ** -0.5)416 417 def split_token(self, token_ids: torch.Tensor, s2_bits: int):418 """Inputs:419 token_ids (torch.Tensor): Composite token IDs of shape [batch_size, seq_len] or [N], each in range [0, 2^(s1_bits + s2_bits) - 1].420 s2_bits (int): Number of low bits used for the fine token (s2).421 """422 assert isinstance(s2_bits, int) and s2_bits > 0, "s2_bits must be a positive integer"423 424 t = token_ids.long()425 mask = (1 << s2_bits) - 1426 s2_ids = t & mask # extract low bits427 s1_ids = t >> s2_bits # extract high bits428 return s1_ids, s2_ids429 430 def forward(self, token_ids):431 """Inputs:432 token_ids:433 - tuple or list: (s1_ids, s2_ids), each of shape [batch_size, seq_len], or434 - torch.Tensor: composite token IDs of shape [batch_size, seq_len], which will be split into (s1_ids, s2_ids) internally.435 Output: [batch_size, seq_len, d_model]436 """437 if isinstance(token_ids, tuple) or isinstance(token_ids, list):438 s1_ids, s2_ids = token_ids439 else:440 s1_ids, s2_ids = self.split_token(token_ids, self.s2_bits)441 s1_emb = self.emb_s1(s1_ids) * math.sqrt(self.d_model)442 s2_emb = self.emb_s2(s2_ids) * math.sqrt(self.d_model)443 return self.fusion_proj(torch.cat([s1_emb, s2_emb], dim=-1))444 445 446class DependencyAwareLayer(nn.Module):447 def __init__(self, d_model, n_heads=4, attn_dropout_p=0.0, resid_dropout=0.0):448 super().__init__()449 self.cross_attn = MultiHeadCrossAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout)450 self.norm = RMSNorm(d_model)451 452 def forward(self, hidden_states, sibling_embed, key_padding_mask=None):453 """hidden_states: [batch, seq_len, d_model]454 sibling_embed: Embedding from another subtoken455 """456 attn_out = self.cross_attn(457 query=sibling_embed,458 key=hidden_states,459 value=hidden_states,460 key_padding_mask=key_padding_mask461 )462 return self.norm(hidden_states + attn_out)463 464 465class TransformerBlock(nn.Module):466 def __init__(self, d_model, n_heads, ff_dim=1024, ffn_dropout_p=0.0, attn_dropout_p=0.0, resid_dropout_p=0.0):467 super().__init__()468 self.norm1 = RMSNorm(d_model)469 self.self_attn = MultiHeadAttentionWithRoPE(d_model, n_heads, attn_dropout_p, resid_dropout_p)470 self.norm2 = RMSNorm(d_model)471 self.ffn = FeedForward(d_model, ff_dim, ffn_dropout_p)472 473 def forward(self, x, key_padding_mask=None):474 residual = x475 x = self.norm1(x)476 attn_out = self.self_attn(x, key_padding_mask=key_padding_mask)477 x = residual + attn_out478 479 residual = x480 x = self.norm2(x)481 ffn_out = self.ffn(x)482 x = residual + ffn_out483 return x484 485 486class DualHead(nn.Module):487 def __init__(self, s1_bits, s2_bits, d_model):488 super().__init__()489 self.vocab_s1 = 2 ** s1_bits490 self.vocab_s2 = 2 ** s2_bits491 self.proj_s1 = nn.Linear(d_model, self.vocab_s1)492 self.proj_s2 = nn.Linear(d_model, self.vocab_s2)493 494 def compute_loss(self, s1_logits, s2_logits, s1_targets, s2_targets, padding_mask=None):495 if padding_mask is not None:496 valid_mask = (padding_mask == 0)497 s1_logits = s1_logits[valid_mask]498 s2_logits = s2_logits[valid_mask]499 s1_targets = s1_targets[valid_mask]500 s2_targets = s2_targets[valid_mask]501 ce_s1 = F.cross_entropy(s1_logits, s1_targets)502 ce_s2 = F.cross_entropy(s2_logits, s2_targets)503 else:504 ce_s1 = F.cross_entropy(s1_logits.reshape(-1, self.vocab_s1), s1_targets.reshape(-1))505 ce_s2 = F.cross_entropy(s2_logits.reshape(-1, self.vocab_s2), s2_targets.reshape(-1))506 ce_loss = (ce_s1 + ce_s2) / 2507 return ce_loss, ce_s1, ce_s2508 509 def forward(self, x):510 return self.proj_s1(x)511 512 def cond_forward(self, x2):513 return self.proj_s2(x2)514 515 516class FixedEmbedding(nn.Module):517 def __init__(self, c_in, d_model):518 super(FixedEmbedding, self).__init__()519 520 w = torch.zeros(c_in, d_model).float()521 w.require_grad = False522 523 position = torch.arange(0, c_in).float().unsqueeze(1)524 div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()525 526 w[:, 0::2] = torch.sin(position * div_term)527 w[:, 1::2] = torch.cos(position * div_term)528 529 self.emb = nn.Embedding(c_in, d_model)530 self.emb.weight = nn.Parameter(w, requires_grad=False)531 532 def forward(self, x):533 return self.emb(x).detach()534 535 536class TemporalEmbedding(nn.Module):537 def __init__(self, d_model, learn_pe):538 super(TemporalEmbedding, self).__init__()539 540 minute_size = 60541 hour_size = 24542 weekday_size = 7543 day_size = 32544 month_size = 13545 546 Embed = FixedEmbedding if not learn_pe else nn.Embedding547 self.minute_embed = Embed(minute_size, d_model)548 self.hour_embed = Embed(hour_size, d_model)549 self.weekday_embed = Embed(weekday_size, d_model)550 self.day_embed = Embed(day_size, d_model)551 self.month_embed = Embed(month_size, d_model)552 553 def forward(self, x):554 x = x.long()555 556 minute_x = self.minute_embed(x[:, :, 0])557 hour_x = self.hour_embed(x[:, :, 1])558 weekday_x = self.weekday_embed(x[:, :, 2])559 day_x = self.day_embed(x[:, :, 3])560 month_x = self.month_embed(x[:, :, 4])561 562 return hour_x + weekday_x + day_x + month_x + minute_x563 564 565 566 567 568 569 570 571