import math from typing import Tuple import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, PretrainedConfig from transformers.modeling_outputs import CausalLMOutput from typing import Callable from transformers.generation.utils import GenerationMixin from functools import partial from random import randrange import math from fast_hadamard_transform import hadamard_transform import torch from torch import nn, cat import torch.nn.functional as F from torch.nn import Module, Sequential from torch.utils._pytree import tree_flatten, tree_unflatten from einops import rearrange, repeat, reduce, einsum from einops.layers.torch import Rearrange, Reduce from .configuration_my_model import GPTConfig """ ein notation: b - batch d - feature dimension s - residual streams t - residual streams + num branch inputs f - number of fractions (division of feature dimension space) v - number of views for branch input """ # helper functions def exists(v): return v is not None def divisible_by(num, den): return (num % den) == 0 def default(v, d): return v if exists(v) else d def identity(t): return t def add(x, y): return x + y def sinkhorn_log(logits, num_iters=10, tau=0.05): n = logits.shape[-1] Z = logits / tau log_marginal = torch.full( (n,), -math.log(n), device=logits.device, dtype=logits.dtype ) u = torch.zeros(n, device=Z.device, dtype=Z.dtype) v = torch.zeros(n, device=Z.device, dtype=Z.dtype) for _ in range(num_iters): u = log_marginal - torch.logsumexp(Z + v.unsqueeze(0), dim=1) v = log_marginal - torch.logsumexp(Z + u.unsqueeze(1), dim=0) return torch.exp(Z + u.unsqueeze(1) + v.unsqueeze(0)) * n def zeropower_via_newtonschulz(X, steps=5, eps=1e-7, coeffs=(3.0, -3.2, 1.2)): a, b, c = coeffs X = X / (X.norm() + eps) transpose = False if X.shape[0] > X.shape[1]: X = X.T transpose = True for _ in range(steps): A = X @ X.T B = b * A + c * A @ A X = a * X + B @ X if transpose: X = X.T return X def orthostochastic_project( logits, ns_steps=5, ns_eps=1e-7, ns_coeffs=(3.0, -3.2, 1.2) ): O = zeropower_via_newtonschulz(logits, steps=ns_steps, eps=ns_eps, coeffs=ns_coeffs) return O.square() # main functions def get_expand_reduce_stream_functions( num_streams, add_stream_embed=False, dim=None, disable=False ): if num_streams == 1 or disable: return (nn.Identity(), nn.Identity()) if add_stream_embed: assert exists(dim), ( "`dim` must be passed into get_init_and_expand_reduce_stream_functions for returning an expansion function with stream embeddings added" ) expand_fn = StreamEmbed(num_streams, dim, expand_to_streams=True) else: expand_fn = Reduce( pattern="b ... -> (b s) ...", reduction="repeat", s=num_streams ) reduce_fn = Reduce(pattern="(b s) ... -> b ...", reduction="sum", s=num_streams) return expand_fn, reduce_fn def get_init_and_expand_reduce_stream_functions( num_streams, num_fracs=1, dim=None, add_stream_embed=False, disable=None ): disable = default(disable, num_streams == 1 and num_fracs == 1) hyper_conn_klass = HyperConnections if not disable else Residual init_hyper_conn_fn = partial(hyper_conn_klass, num_streams, num_fracs=num_fracs) expand_reduce_fns = get_expand_reduce_stream_functions( num_streams, add_stream_embed=add_stream_embed, dim=dim, disable=disable ) if exists(dim): init_hyper_conn_fn = partial(init_hyper_conn_fn, dim=dim) return (init_hyper_conn_fn, *expand_reduce_fns) # norms class RMSNorm(Module): def __init__(self, dim): super().__init__() self.scale = dim**0.5 self.gamma = nn.Parameter(torch.zeros(dim)) def forward(self, x): return F.normalize(x, dim=-1) * self.scale * (self.gamma + 1) # main classes # residual base class class Residual(Module): def __init__( self, *args, branch: Module | None = None, residual_transform: Module | None = None, **kwargs, ): super().__init__() self.branch = branch self.residual_transform = default(residual_transform, nn.Identity()) def width_connection(self, residuals): return residuals, residuals, dict() def depth_connection( self, branch_output, residuals, ): return branch_output + self.residual_transform(residuals) def decorate_branch(self, branch: Callable): assert not exists(self.branch), "branch was already wrapped on init" def forward_and_add_residual(residual, *args, **kwargs): branch_input, add_residual = self.forward(residual) branch_output = branch(branch_input, *args, **kwargs) residual = add_residual(branch_output) return residual return forward_and_add_residual def forward(self, residuals, *branch_args, **branch_kwargs): branch_input, residuals, residual_kwargs = self.width_connection(residuals) def add_residual_fn(branch_out): (branch_out, *rest), tree_spec = tree_flatten(branch_out) branch_out = self.depth_connection(branch_out, residuals, **residual_kwargs) return tree_unflatten((branch_out, *rest), tree_spec) if not exists(self.branch): return branch_input, add_residual_fn branch_output = self.branch(branch_input, *branch_args, **branch_kwargs) return add_residual_fn(branch_output) # hyper connection residual streams class HyperConnections(Module): def __init__( self, num_residual_streams, *, dim, branch: Module | None = None, layer_index=None, tanh=True, channel_first=False, dropout=0.0, residual_transform: Module | None = None, # to support resnet blocks where dimension in not equal to dimension out - usually a residual conv add_branch_out_to_residual=True, # will disable depth connections (weighted residual sum with beta) if set False num_input_views=1, # allow for the branch module to receive multiple input views, dimension placed on the very left (before batch) depth_residual_fn=add, num_fracs=1, # https://arxiv.org/abs/2503.14125 mhc=False, sinkhorn_iters=10, sinkhorn_tau=0.05, mhc_h_res_proj="sinkhorn", ns_steps=5, ns_eps=1e-7, ns_coeffs=(3.0, -3.2, 1.2), ): """ Appendix J, Algorithm2 in - https://arxiv.org/abs/2409.19606 """ super().__init__() self.branch = branch self.act = nn.Tanh() if tanh else nn.Identity() # frac-connections paper - num_fracs > 1 will be the `m` in their paper https://arxiv.org/abs/2503.14125 assert num_fracs >= 1 self.num_fracs = num_fracs self.has_fracs = num_fracs > 1 self.split_fracs = Rearrange("b ... (f d) -> b ... f d", f=num_fracs) self.merge_fracs = Rearrange("b ... f d -> b ... (f d)") assert divisible_by(dim, num_fracs), ( f"feature dimension ({dim}) must be divisible by the `num_fracs` ({num_fracs})" ) dim //= num_fracs # effective dim handled in dimension is feature dimension divided by num fractions # they used layernorm in paper, but rmsnorm is fine given what we know now self.norm = RMSNorm(dim) assert num_residual_streams > 0, "`num_residual_streams` must be greater than 0" self.num_residual_streams = num_residual_streams init_residual_index = ( default(layer_index, randrange(num_residual_streams)) % num_residual_streams ) # just choose one random residual stream if layer index not given # handle the parameter dimensions, which may require (num_residuals x num_fractions) - generalizing hyper + frac connections num_residual_streams_fracs = num_residual_streams * num_fracs num_input_views_fracs = num_input_views * num_fracs # width num residual streams assert num_input_views >= 1 self.num_input_views = num_input_views # width connection init_alpha0 = torch.zeros((num_residual_streams_fracs, num_input_views_fracs)) init_alpha0[init_residual_index, :] = 1.0 self.static_alpha = nn.Parameter( cat((init_alpha0, torch.eye(num_residual_streams_fracs)), dim=1) ) self.dynamic_alpha_fn = nn.Parameter( torch.zeros(dim, num_residual_streams_fracs + num_input_views_fracs) ) self.dynamic_alpha_scale = nn.Parameter(torch.ones(()) * 1e-2) # depth connection related (beta) self.add_branch_out_to_residual = add_branch_out_to_residual if add_branch_out_to_residual: self.static_beta = nn.Parameter(torch.ones(num_residual_streams_fracs)) dynamic_beta_shape = ( (dim,) if num_fracs == 1 else (dim, num_fracs) ) # preserve backwards compat self.dynamic_beta_fn = nn.Parameter(torch.zeros(dynamic_beta_shape)) self.dynamic_beta_scale = nn.Parameter(torch.ones(()) * 1e-2) # dropouts self.dropout = nn.Dropout(dropout) # channel first option self.channel_first = channel_first # maybe residual transform self.residual_transform = default(residual_transform, nn.Identity()) # maybe custom depth connection residual function # this is to prepare for gating the addition of the branch outputs to the residual streams # needed for memory lanes a la RMT / LMM self.depth_residual_fn = depth_residual_fn self.mhc = mhc self.sinkhorn_iters = sinkhorn_iters self.sinkhorn_tau = sinkhorn_tau self.mhc_h_res_proj = mhc_h_res_proj self.ns_steps = ns_steps self.ns_eps = ns_eps self.ns_coeffs = ns_coeffs if mhc: assert num_fracs == 1, "mhc currently requires num_fracs = 1" assert num_input_views == 1, "mhc currently requires num_input_views = 1" assert mhc_h_res_proj in ( "sinkhorn", "orthostochastic", ), "mhc_h_res_proj must be 'sinkhorn' or 'orthostochastic'" H_res_init = torch.full((num_residual_streams, num_residual_streams), -8.0) H_res_init.fill_diagonal_(0.0) self.H_res_logits = nn.Parameter(H_res_init) H_pre_init = torch.full((num_residual_streams,), -8.0) H_pre_init[init_residual_index] = 0.0 self.H_pre_logits = nn.Parameter(H_pre_init) if add_branch_out_to_residual: self.H_post_logits = nn.Parameter(torch.zeros(num_residual_streams)) def width_connection(self, residuals): streams = self.num_residual_streams maybe_transformed_residuals = self.residual_transform(residuals) # width connection # handle channel first if self.channel_first: residuals = rearrange(residuals, "b d ... -> b ... d") # split out fractions residuals = self.split_fracs(residuals) # split out streams residuals = rearrange(residuals, "(b s) ... d -> b ... s d", s=streams) if self.mhc: residuals_mixed_source = maybe_transformed_residuals if self.channel_first: residuals_mixed_source = rearrange( residuals_mixed_source, "b d ... -> b ... d" ) residuals_mixed_source = self.split_fracs(residuals_mixed_source) residuals_mixed_source = rearrange( residuals_mixed_source, "(b s) ... d -> b ... s d", s=streams ) if self.mhc_h_res_proj == "orthostochastic": H_res = orthostochastic_project( self.H_res_logits, ns_steps=self.ns_steps, ns_eps=self.ns_eps, ns_coeffs=self.ns_coeffs, ) else: H_res = sinkhorn_log( self.H_res_logits, self.sinkhorn_iters, self.sinkhorn_tau ) H_pre = F.softmax(self.H_pre_logits, dim=-1) H_post = None if self.add_branch_out_to_residual: H_post = F.softmax(self.H_post_logits, dim=-1) residuals_mixed = einsum( H_res, residuals_mixed_source, "s t, ... s d -> ... t d" ) branch_input = einsum(H_pre, residuals, "s, ... s d -> ... d") if getattr(self, "collect_stats", False): with torch.no_grad(): stats = dict( h_res_min=H_res.min(), h_res_row_sum=H_res.sum(dim=-1).mean(), h_res_col_sum=H_res.sum(dim=-2).mean(), h_pre_min=H_pre.min(), ) if H_post is not None: stats["h_post_min"] = H_post.min() self.last_stats = {k: v.detach() for k, v in stats.items()} if self.channel_first: branch_input = rearrange(branch_input, "b ... d -> b d ...") branch_input = self.merge_fracs(branch_input) return ( branch_input, maybe_transformed_residuals, dict(beta=H_post, residuals_mixed=residuals_mixed), ) # norm normed = self.norm(residuals) # alpha for weighted sum of residuals going into branch wc_weight = self.act(normed @ self.dynamic_alpha_fn) dynamic_alpha = wc_weight * self.dynamic_alpha_scale static_alpha = rearrange(self.static_alpha, "(f s) d -> f s d", s=streams) alpha = dynamic_alpha + static_alpha alpha = self.split_fracs( alpha ) # (batch, seq, fracs1, streams, fracs2, input + residual streams) # beta for weights from branch output back to residual streams beta = None if self.add_branch_out_to_residual: dc_weight = self.act(normed @ self.dynamic_beta_fn) if not self.has_fracs: dc_weight = rearrange(dc_weight, "... -> ... 1") dynamic_beta = dc_weight * self.dynamic_beta_scale static_beta = rearrange(self.static_beta, "... (s f) -> ... s f", s=streams) beta = dynamic_beta + static_beta if getattr(self, "collect_stats", False): with torch.no_grad(): num_input_views_fracs = self.num_input_views * self.num_fracs alpha_branch = alpha[..., :num_input_views_fracs] alpha_residual = alpha[..., num_input_views_fracs:] alpha_branch_abs_mean = alpha_branch.abs().mean() alpha_residual_abs_mean = alpha_residual.abs().mean() stats = dict( alpha_branch_mean=alpha_branch.mean(), alpha_branch_abs_mean=alpha_branch_abs_mean, alpha_residual_mean=alpha_residual.mean(), alpha_residual_abs_mean=alpha_residual_abs_mean, alpha_branch_residual_ratio=alpha_branch_abs_mean / (alpha_residual_abs_mean + 1e-8), ) if beta is not None: stats.update( beta_mean=beta.mean(), beta_abs_mean=beta.abs().mean(), beta_min=beta.min(), beta_max=beta.max(), ) self.last_stats = {k: v.detach() for k, v in stats.items()} mix_h = einsum(alpha, residuals, "... f1 s f2 t, ... f1 s d -> ... f2 t d") if self.num_input_views == 1: branch_input, residuals = mix_h[..., 0, :], mix_h[..., 1:, :] else: branch_input, residuals = ( mix_h[..., : self.num_input_views, :], mix_h[..., self.num_input_views :, :], ) branch_input = rearrange(branch_input, "b ... v d -> v b ... d") if self.channel_first: branch_input = rearrange(branch_input, "b ... d -> b d ...") # maybe merge fractions back branch_input = self.merge_fracs(branch_input) return branch_input, maybe_transformed_residuals, dict(beta=beta) def depth_connection(self, branch_output, residuals, *, beta, residuals_mixed=None): assert self.add_branch_out_to_residual # maybe split fractions branch_output = self.split_fracs(branch_output) # 'depth' connection if self.channel_first: branch_output = rearrange(branch_output, "b d ... -> b ... d") if self.mhc: assert residuals_mixed is not None assert beta is not None branch_to_streams = einsum(branch_output, beta, "b ... d, s -> b ... s d") output = residuals_mixed + branch_to_streams output = rearrange(output, "b ... s d -> (b s) ... d") output = self.merge_fracs(output) if self.channel_first: output = rearrange(output, "b ... d -> b d ...") return self.dropout(output) output = einsum( branch_output, beta, "b ... f1 d, b ... f1 s f2 -> b ... f2 s d" ) output = rearrange(output, "b ... s d -> (b s) ... d") # merge merge back fractions output = self.merge_fracs(output) # channel first if self.channel_first: output = rearrange(output, "b ... d -> b d ...") residuals = self.depth_residual_fn(output, residuals) return self.dropout(residuals) def decorate_branch(self, branch: Callable): assert not exists(self.branch), "branch was already wrapped on init" def forward_and_add_residual(residual, *args, **kwargs): branch_input, add_residual = self.forward(residual) branch_output = branch(branch_input, *args, **kwargs) residual = add_residual(branch_output) return residual return forward_and_add_residual def forward(self, residuals, *branch_args, **branch_kwargs): branch_input, residuals, residual_kwargs = self.width_connection(residuals) def add_residual_fn(branch_out): if not self.add_branch_out_to_residual: return branch_out (branch_out, *rest), tree_spec = tree_flatten(branch_out) branch_out = self.depth_connection(branch_out, residuals, **residual_kwargs) return tree_unflatten((branch_out, *rest), tree_spec) if not exists(self.branch): return branch_input, add_residual_fn branch_output = self.branch(branch_input, *branch_args, **branch_kwargs) return add_residual_fn(branch_output) HyperConnections.get_expand_reduce_stream_functions = staticmethod( get_expand_reduce_stream_functions ) HyperConnections.get_init_and_expand_reduce_stream_functions = staticmethod( get_init_and_expand_reduce_stream_functions ) # stream embed class StreamEmbed(Module): def __init__(self, num_streams, dim, channel_first=False, expand_to_streams=False): super().__init__() self.channel_first = channel_first self.num_streams = num_streams self.expand_to_streams = expand_to_streams self.stream_embed = nn.Parameter(torch.zeros(num_streams, dim)) def forward(self, residuals): if self.expand_to_streams: residuals = repeat(residuals, "b ... -> (b s) ...", s=self.num_streams) if self.channel_first: residuals = rearrange( residuals, "(b s) d ... -> b ... s d", s=self.num_streams ) else: residuals = rearrange( residuals, "(b s) ... d -> b ... s d", s=self.num_streams ) residuals = residuals + self.stream_embed if self.channel_first: residuals = rearrange( residuals, "b ... s d -> (b s) d ...", s=self.num_streams ) else: residuals = rearrange( residuals, "b ... s d -> (b s) ... d", s=self.num_streams ) return residuals # attention pool - taken from Enformer https://www.nature.com/articles/s41592-021-01252-x , in turn taken from somewhere else class AttentionPoolReduceStream(Module): def __init__(self, num_streams, dim, channel_first=False): super().__init__() self.num_streams = num_streams self.channel_first = channel_first self.to_attn_logits = nn.Linear(dim, dim, bias=False) self.to_attn_logits.weight.data.copy_(torch.eye(dim)) def forward(self, residuals): if self.channel_first: residuals = rearrange( residuals, "(b s) d ... -> b ... s d", s=self.num_streams ) else: residuals = rearrange( residuals, "(b s) ... d -> b ... s d", s=self.num_streams ) attn_logits = self.to_attn_logits(residuals) attn = attn_logits.softmax(dim=-2) residuals = reduce(residuals * attn, "b ... s d -> b ... d", "sum") if self.channel_first: residuals = rearrange(residuals, "b ... d -> b d ...") return residuals class Rotary(torch.nn.Module): def __init__(self, dim, base=10000): super().__init__() inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) self.seq_len_cached = None self.cos_cached = None self.sin_cached = None def forward(self, x): seq_len = x.shape[1] if seq_len != self.seq_len_cached: self.seq_len_cached = seq_len t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq) freqs = torch.outer(t, self.inv_freq).to(x.device) self.cos_cached = freqs.cos() self.sin_cached = freqs.sin() return self.cos_cached[None, :, None, :], self.sin_cached[None, :, None, :] def apply_rotary_emb(x, cos, sin): assert x.ndim == 4 # multihead attention d = x.shape[3]//2 x1 = x[..., :d] x2 = x[..., d:] y1 = x1 * cos + x2 * sin y2 = x1 * (-sin) + x2 * cos return torch.cat([y1, y2], 3) # class CasualSelfAttention(nn.Module): # def __init__(self, config): # super().__init__() # self.c_attn = nn.Linear(config.n_embd, 3* config.n_embd) # # self.c_proj = nn.Linear(config.n_embd, config.n_embd) # # self.c_proj.NANOGPT_SCALE_INIT = 1 # Attaching a flag/attribute for initialization # self.n_head = config.n_head # self.n_embd = config.n_embd # self.hada_scale = 1.0/ math.sqrt(config.n_embd) # self.rotary = Rotary(config.n_embd//config.n_head) # self.out_scale = nn.Parameter( # torch.ones(config.n_embd) / math.sqrt(2 * config.n_layer) # ) # self.out_bias = nn.Parameter(torch.zeros(config.n_embd)) # # causal mask to ensure that attention is only applied to the left in the input sequence # self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) # .view(1, 1, config.block_size, config.block_size)) # # @torch.compile # def fused_hadamard_output(self, y, B, T, C): # """Fuse reshape, hadamard, scaling operations""" # y = y.reshape(-1, C) # y = hadamard_transform(y, scale = self.hada_scale) # Batched transform on last dim # y = y * self.out_scale # Vector scaling # y = y + self.out_bias # <- add this # return y.view(B, T, C) # def forward(self, x): # B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) # # calculate query, key, values for all heads in batch and move head forward to be the batch dim # q, k, v = self.c_attn(x).split(self.n_embd, dim=2) # k = k.view(B, T, self.n_head, C // self.n_head) # (B, nh, T, hs) # q = q.view(B, T, self.n_head, C // self.n_head) # (B, nh, T, hs) # v = v.view(B, T, self.n_head, C // self.n_head) # (B, nh, T, hs) # cos, sin = self.rotary(q) # q = apply_rotary_emb(q, cos, sin) # k = apply_rotary_emb(k, cos, sin) # y = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True) # y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side # # output projection # # y: (B, T, C) # y = self.fused_hadamard_output(y,B,T,C) # # y2d = y.view(-1, C) # (B * T, C) # # y2d = hadamard_transform(y2d) # apply per-vector # # y = y2d.view(B, T, C) # restore original shape # # y = y * self.out_scale # return y class CasualSelfAttention(nn.Module): def __init__(self, config): super().__init__() self.c_attn = nn.Linear(config.n_embd, 3* config.n_embd) self.c_proj = nn.Linear(config.n_embd, config.n_embd) self.c_proj.NANOGPT_SCALE_INIT = 1 # Attaching a flag/attribute for initialization self.n_head = config.n_head self.n_embd = config.n_embd self.rotary = Rotary(config.n_embd//config.n_head) # causal mask to ensure that attention is only applied to the left in the input sequence self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) .view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) # calculate query, key, values for all heads in batch and move head forward to be the batch dim q, k, v = self.c_attn(x).split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, C // self.n_head) # (B, nh, T, hs) q = q.view(B, T, self.n_head, C // self.n_head) # (B, nh, T, hs) v = v.view(B, T, self.n_head, C // self.n_head) # (B, nh, T, hs) cos, sin = self.rotary(q) q = apply_rotary_emb(q, cos, sin) k = apply_rotary_emb(k, cos, sin) y = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side # output projection y = self.c_proj(y) return y class MLP(nn.Module): def __init__(self, config): super().__init__() # Set inner dim ~ 8/3 * d so that parameter count matches 4d-GELU inner_dim = int(4 * config.n_embd * 2 / 3) # inner_dim = 2048 # Round up to multiple of 256 for efficiency # inner_dim = ((inner_dim + 255) // 256) * 256 self.inner_dim = inner_dim # value and gate self.c_fc = nn.Linear(config.n_embd, 2 * inner_dim) self.c_proj = nn.Linear(inner_dim, config.n_embd) self.c_proj.NANOGPT_SCALE_INIT = 1 def forward(self, x): x_in = self.c_fc(x) x_gate, x_up = x_in.chunk(2, dim=-1) x = F.silu(x_gate) * x_up x = self.c_proj(x) return x class AttnBranch(nn.Module): def __init__(self, norm, attn): super().__init__() self.norm = norm self.attn = attn def forward(self, x): return self.attn(self.norm(x)) class Block(nn.Module): def __init__(self, config, layer_idx, init_hc): super().__init__() self.ln_1 = nn.RMSNorm(config.n_embd,eps = 1e-6) self.attn = CasualSelfAttention(config) self.ln_2 = nn.RMSNorm(config.n_embd,eps = 1e-6) self.mlp = MLP(config) self.attn_branch = AttnBranch(self.ln_1, self.attn) hc_kwargs = dict( mhc=config.mhc, sinkhorn_iters=config.sinkhorn_iters, sinkhorn_tau=config.sinkhorn_tau, mhc_h_res_proj=config.mhc_h_res_proj, ns_steps=config.ns_steps, ns_eps=config.ns_eps, ns_coeffs=config.ns_coeffs, ) self.hc_attn = init_hc( dim=config.n_embd, branch=self.attn_branch, layer_index=layer_idx * 2, **hc_kwargs, ) self.hc_mlp = init_hc( dim=config.n_embd, branch=nn.Sequential(self.ln_2, self.mlp), layer_index=layer_idx * 2 + 1, **hc_kwargs, ) def forward(self, x): x = self.hc_attn(x) x = self.hc_mlp(x) return x class GPTConfig(PretrainedConfig): model_type = "custom_gpt" def __init__( self, block_size=1024, vocab_size=50304, n_layer=12, n_head=12, n_embd=768, dropout=0.0, bias=True, hc_num_streams=1, hc_num_fracs=1, hc_disable=False, mhc=False, sinkhorn_iters=10, sinkhorn_tau=0.05, mhc_h_res_proj="sinkhorn", ns_steps=5, ns_eps=1e-7, ns_coeffs=(3.0, -3.2, 1.2), **kwargs, ): super().__init__(**kwargs) self.block_size = block_size self.vocab_size = vocab_size self.n_layer = n_layer self.n_head = n_head self.n_embd = n_embd self.dropout = dropout self.bias = bias self.hc_num_streams = hc_num_streams self.hc_num_fracs = hc_num_fracs self.hc_disable = hc_disable self.mhc = mhc self.sinkhorn_iters = sinkhorn_iters self.sinkhorn_tau = sinkhorn_tau self.mhc_h_res_proj = mhc_h_res_proj self.ns_steps = ns_steps self.ns_eps = ns_eps self.ns_coeffs = ns_coeffs # 🔑 HF compatibility aliases self.num_hidden_layers = n_layer self.num_attention_heads = n_head self.hidden_size = n_embd self.max_position_embeddings = block_size class GPT(PreTrainedModel, GenerationMixin): config_class = GPTConfig # config_class = MyGPTConfig def __init__(self, config): super().__init__(config) init_hc, expand_stream, reduce_stream = ( get_init_and_expand_reduce_stream_functions( config.hc_num_streams, num_fracs=config.hc_num_fracs, disable=config.hc_disable, ) ) self.expand_stream = expand_stream self.reduce_stream = reduce_stream self.transformer = nn.ModuleDict( dict( wte=nn.Embedding(config.vocab_size, config.n_embd), h=nn.ModuleList( [Block(config, i, init_hc) for i in range(config.n_layer)] ), ln_f = nn.RMSNorm(config.n_embd,eps = 1e-6) ) ) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight self.post_init() def prepare_inputs_for_generation( self, input_ids, past_key_values=None, **kwargs, ): # We do NOT use KV cache yet, so always feed full sequence return { "input_ids": input_ids, "past_key_values": None, } def forward( self, input_ids=None, attention_mask=None, # 👈 ADD THIS labels=None, past_key_values=None, use_cache=None, **kwargs, ): b, t = input_ids.size() assert t <= self.config.block_size pos = torch.arange(0, t, device=input_ids.device).unsqueeze(0) x = self.transformer.wte(input_ids) x = self.expand_stream(x) for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) x = self.reduce_stream(x) logits = self.lm_head(x) loss = None if labels is not None: loss = F.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ) return CausalLMOutput( loss=loss, logits=logits, )