""" All the layers included here aim to be already optimized by torch itself, without requiring triton/cuda kernels (explicitly), in fact they try to only use matmul and Linear operations which are arguably already optimized implicitly in torch. """ import time import torch from Vathos._basics import * from typing import Tuple, Optional, Union, List import re from Vathos.complexity import combine_big_o_product, combine_big_o_sum import math import torch.nn.functional as F from tqdm import tqdm from Vathos._spatials import * class Block1d(Layer): def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, norm=nn.LayerNorm): super().__init__() self.spatial_mixer = spatial_mixer self.channel_mixer = channel_mixer self.norm1 = norm(spatial_mixer.d_model) self.norm2 = norm(spatial_mixer.d_model) def forward(self, x: torch.Tensor, ve=None): if ve is not None: x = x + self.spatial_mixer(self.norm1(x), ve) else: x = x + self.spatial_mixer(self.norm1(x)) x = x + self.channel_mixer(self.norm2(x)) return x def generate(self, x: torch.Tensor, ve=None): h = self.norm1(x) if self.spatial_mixer.has_custom_generate(): spatial_out = _call_with_ve(self.spatial_mixer.generate, h, ve) else: spatial_out = _call_with_ve(self.spatial_mixer, h, ve) x = x + spatial_out if self.channel_mixer.has_custom_generate(): x = x + self.channel_mixer.generate(self.norm2(x)) else: x = x + self.channel_mixer(self.norm2(x)) return x def _call_with_ve(fn, h, ve): """Call fn(h, ve=ve) if it accepts ve, otherwise fn(h). Used by Block.generate so that spatial mixers without ve support stay drop-in.""" if ve is None: return fn(h) try: return fn(h, ve=ve) except TypeError: return fn(h) class CBlock1d(Layer): def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, norm=nn.LayerNorm): super().__init__() self.spatial_mixer = spatial_mixer self.channel_mixer = channel_mixer self.norm1 = norm(spatial_mixer.d_model) self.norm2 = norm(spatial_mixer.d_model) self.l = nn.Parameter(torch.tensor([1.0, 1.0, 1.0, 1.0])) def forward(self, x: torch.Tensor, ve=None): if ve is not None: x = x * self.l[0] + self.spatial_mixer(self.norm1(x), ve) * self.l[1] else: x = x * self.l[0] + self.spatial_mixer(self.norm1(x)) * self.l[1] x = x * self.l[2] + self.channel_mixer(self.norm2(x)) * self.l[3] return x def generate(self, x: torch.Tensor, ve=None): h = self.norm1(x) if self.spatial_mixer.has_custom_generate(): spatial_out = _call_with_ve(self.spatial_mixer.generate, h, ve) else: spatial_out = _call_with_ve(self.spatial_mixer, h, ve) x = x * self.l[0] + spatial_out * self.l[1] if self.channel_mixer.has_custom_generate(): x = x * self.l[2] + self.channel_mixer.generate(self.norm2(x)) * self.l[3] else: x = x * self.l[2] + self.channel_mixer(self.norm2(x)) * self.l[3] return x # TODO: Smear class SmearBlock1d(Layer): def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, norm=nn.LayerNorm): super().__init__() self.spatial_mixer = spatial_mixer self.channel_mixer = channel_mixer self.norm1 = norm(spatial_mixer.d_model) self.norm2 = norm(spatial_mixer.d_model) def forward(self, x: torch.Tensor): x = x + self.spatial_mixer(self.norm1(x)) x = x + self.channel_mixer(self.norm2(x)) return x def generate(self, x: torch.Tensor): # Use generate if available, otherwise forward if self.spatial_mixer.has_custom_generate(): x = x + self.spatial_mixer.generate(self.norm1(x)) else: x = x + self.spatial_mixer(self.norm1(x)) if self.channel_mixer.has_custom_generate(): x = x + self.channel_mixer.generate(self.norm2(x)) else: x = x + self.channel_mixer(self.norm2(x)) return x class JBlock1d(Layer): def __init__(self, d_model, channel_mixer: nn.Module, spatial_mixer: nn.Module, norm=nn.LayerNorm): super().__init__() self.channel_mixer = channel_mixer self.spatial_mixer = spatial_mixer self.norm = norm(d_model) def forward(self, x: torch.Tensor, ve=None): h = self.norm(x) if ve is not None: return x + self.spatial_mixer(h, ve) + self.channel_mixer(h) return x + self.spatial_mixer(h) + self.channel_mixer(h) def generate(self, x: torch.Tensor, ve=None): h = self.norm(x) sm_fn = self.spatial_mixer.generate if (hasattr(self.spatial_mixer, 'has_custom_generate') and self.spatial_mixer.has_custom_generate()) else self.spatial_mixer cm_fn = self.channel_mixer.generate if (hasattr(self.channel_mixer, 'has_custom_generate') and self.channel_mixer.has_custom_generate()) else self.channel_mixer spatial_out = _call_with_ve(sm_fn, h, ve) channel_out = cm_fn(h) return x + spatial_out + channel_out class ExpandingBlock1d(Layer): def __init__(self, d_model, channel_mixer: Layer, spatial_mixer: Layer, expand=1): super().__init__() self.spatial_mixer = spatial_mixer self.channel_mixer = channel_mixer self.expander = nn.Linear(d_model, d_model * expand, bias=False) self.contractor = nn.Linear(d_model * expand, d_model, bias=False) self.norm1 = nn.LayerNorm(spatial_mixer.d_model) self.norm2 = nn.LayerNorm(spatial_mixer.d_model) def forward(self, x: torch.Tensor): x = x + self.contractor(self.spatial_mixer(self.norm1(self.expander(x)))) x = x + self.channel_mixer(self.norm2(x)) return x class Renamer: def __init__(self, constructor, renames: dict): super().__init__() self.constructor = constructor self.renames = renames def __call__(self, *args, **kwargs): renamed_kwargs = {} for key, value in kwargs.items(): if key in self.renames: renamed_kwargs[self.renames[key]] = value else: renamed_kwargs[key] = value return self.constructor(*args, renamed_kwargs) class BlockStack(Layer): def __init__(self, blocks: Tuple[Block1d | Layer]): super().__init__() self.blocks = blocks self.stack = nn.ModuleList(blocks) def forward(self, x: torch): for block in self.stack: x = block(x) return x class DepthwiseCausalConv1d(Layer): __name__ = "CausalConv1d" __complexity__ = "O(L d k)" def __init__(self, d_model, k=3): super().__init__() self.d_model = d_model self.k = k self.pad = k - 1 self.K = nn.Parameter(torch.randn(d_model, k) / (k ** 0.5)) def forward(self, x): b, L, d = x.shape x = x.transpose(1, 2) x_pad = F.pad(x, (self.pad, 0)) out = F.conv1d(x_pad, self.K.unsqueeze(1), groups=d) return out.transpose(1, 2) class CausalConv1d(Layer): __name__ = "CausalConv1d" __complexity__ = "O(L d k)" def __init__(self, d_model, k=3, groups=None, outproj=False): super().__init__() self.d_model = d_model self.k = k self.pad = k - 1 self.groups = groups if groups is not None else d_model self.outproj = nn.Linear(groups, d_model, bias=False) if outproj else Identity() if self.groups != d_model and not outproj: flag( "Output channel numbers will be the number of groups, use outproj=True if you want to enable a linear projection to go back to d_model") self.K = nn.Parameter(torch.randn(d_model, k) / (k ** 0.5)) def forward(self, x): x = x.transpose(1, 2) x_pad = F.pad(x, (self.pad, 0)) out = F.conv1d(x_pad, self.K.unsqueeze(1), groups=self.groups) return out.transpose(1, 2) class LSTM(Layer): __name__ = "LSTM" __complexity__ = "O(L d^2)" def __init__(self, d_model, d_hidden, bidirectional=False, dropout=0.1, n_layer=1): super().__init__() self.bidirectional = bidirectional self.d_model = d_model self.d_hidden = d_hidden self.projout = nn.Linear(d_hidden, d_model) if d_model != d_hidden else Identity() self.LSTM = nn.LSTM(d_model, d_hidden, bidirectional=bidirectional, batch_first=True, num_layers=n_layer) self.dropout = nn.Dropout(dropout) def forward(self, x): out, _ = self.LSTM(x) out = self.dropout(out) return self.projout(out) class GRU(Layer): __name__ = "GRU" __complexity__ = "O(L d^2)" def __init__(self, d_model, d_hidden, bidirectional=False, dropout=0.1, n_layer=1): super().__init__() self.bidirectional = bidirectional self.d_model = d_model self.d_hidden = d_hidden self.projout = nn.Linear(d_hidden, d_model) if d_model != d_hidden else Identity() self.GRU = nn.GRU(d_model, d_hidden, bidirectional=bidirectional, batch_first=True, num_layers=n_layer) self.dropout = nn.Dropout(dropout) def forward(self, x): out, _ = self.GRU(x) out = self.dropout(out) return self.projout(out) class LinearMixer(Layer): __name__ = "Linear Mixer" __complexity__ = "O(L^2 d)" def __init__(self, d_model, max_len=1000, causal=True): super().__init__() self.causal = causal self.d_model = d_model self.W = nn.Parameter(torch.randn(max_len, max_len) * 0.01) def forward(self, x): x = x / math.sqrt(self.d_model) if self.causal: return torch.tril(self.W[:x.shape[1], :x.shape[1]]) @ x else: return self.W[:x.shape[1], :x.shape[1]] @ x class MLPMixer(Layer): __name__ = "Linear Mixer" __complexity__ = "O(L^2 d)" def __init__(self, d_model, max_len=1000, L_expand=1, causal=False, activation=nn.GELU): super().__init__() if causal: if L_expand < 1: raise ValueError("For the MLPMixer to be causal, L_expand must be greater than 1") if L_expand > 1: raise NotImplementedError( "Causal MLP Mixer with L_expand is currently not implemented") # TODO: Causal Lex self.causal = causal self.d_model = d_model self.W1 = nn.Parameter(torch.randn(max_len * L_expand, max_len) / max_len) self.W2 = nn.Parameter(torch.randn(max_len, max_len * L_expand) / max_len) self.activation = activation() def forward(self, x): x = x / math.sqrt(self.d_model) return self.W2 @ self.activation(self.W1 @ x) class ShortConvGatedMixer(Layer): def __init__(self, d_model, mixer, mixer_params, k=4, activation=nn.Sigmoid, k1=None, k2=None): super().__init__() if k1 is None: k1 = k2 = k self.mixer = mixer(mixer_params) self.conv_1 = DepthwiseCausalConv1d(d_model, k=k1) self.conv_2 = DepthwiseCausalConv1d(d_model, k=k2) self.activation = activation() def forward(self, x): g1 = self.conv_1(x) g2 = self.activation(self.conv_2(x)) return self.mixer(g1) * g2 + (1 - g2) * x class SmearGate(Layer): """Smear gate (Modded-NanoGPT 2025-09 record): gate input-dependent + skip token-1. Per ogni posizione t > 0: x[t] += smear_lambda * sigmoid(W_gate @ x[t, :gate_input_dim]) * x[t-1] Gate sparse (default 12 dim su d_model), input-dependent. Token 0 invariato. smear_lambda init=0 ⇒ no-op iniziale. Receptive field = 1 (lookback=1 in fast gen). """ __name__ = "SmearGate" def __init__(self, d_model: int, gate_input_dim: int = 12): super().__init__() self.gate_input_dim = gate_input_dim self.gate = nn.Linear(gate_input_dim, d_model, bias=False) self.smear_lambda = nn.Parameter(torch.zeros(1)) def _init_weights(self): # Identity-at-init: gate orthogonal; smear_lambda già 0 da __init__ # ⇒ shifted = 0 ⇒ no-op iniziale indipendentemente dal gate. nn.init.orthogonal_(self.gate.weight) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: [B, L, D]. Caso L==1: nessuno "x[t-1]" disponibile → no-op. if x.size(1) < 2: return x gate = torch.sigmoid(self.gate(x[:, 1:, :self.gate_input_dim])) # [B, L-1, D] prev = x[:, :-1, :] # [B, L-1, D] shifted = self.smear_lambda * gate * prev # [B, L-1, D] return torch.cat([x[:, :1, :], x[:, 1:, :] + shifted], dim=1) class Embedder(Layer): __name__ = "SymbolicEmbedder" __complexity__ = "O(L d)" def __init__(self, vocab_size, d_model: int): super().__init__() self.vocab_size = vocab_size self.d_model = d_model self.frozen = False self.embedding = nn.Embedding(vocab_size, d_model) def freeze(self): if self.frozen: flag("Trying to freeze an already frozen Embedder") else: self.embedding.weight.requires_grad = False self.frozen = True def unfreeze(self): if self.frozen: self.embedding.weight.requires_grad = True self.frozen = False else: flag("Trying to unfreeze an already unfrozen Embedder") def forward(self, x): return self.embedding(x) def finetune(self): self.freeze() class EasyEmbedder(Layer): __name__ = "SymbolicEmbedder" __complexity__ = "O(L d)" def __init__(self, vocab_size, d_model: int, dropout=0.13): super().__init__() self.vocab_size = vocab_size self.d_model = d_model self.embedding = nn.Embedding(vocab_size, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.dropout(self.embedding(x)) class HybridAttentionBlock1d(Layer): def __init__(self, d_model, sec_mixer, sec_params, attn_params, n_attn=1, n_sec=3, channel_mixer=MLP, channel_params=None): super().__init__() self.n_sec = n_sec self.n_attn = n_attn self.d_model = d_model self.attn_params = attn_params self.sec_params = sec_params self.sec_mixer = sec_mixer self.channel_mixer = channel_mixer self.channel_params = channel_params if channel_params is None and isinstance(channel_mixer, MLP): self.channel_params = {'expand': 2} else: self.channel_params = channel_params self.attn_blocks = nn.ModuleList([ Block1d( channel_mixer=self.channel_mixer(d_model=d_model, **self.channel_params), spatial_mixer=MultiheadAttentionMixer(d_model=d_model, **self.attn_params) ) for _ in range(self.n_attn) ]) self.sec_blocks = nn.ModuleList([ Block1d( channel_mixer=self.channel_mixer(d_model=d_model, **self.channel_params), spatial_mixer=self.sec_mixer(d_model=d_model, **self.sec_params) ) for _ in range(self.n_sec) ]) def forward(self, x): for sec in self.sec_blocks: x = sec(x) for attn in self.attn_blocks: x = attn(x) return x ######################################################################################################################## # VISION ######################################################################################################################## class PatchEmbedder(Layer): __name__ = "PatchEmbedder" def __init__( self, vocab_size=None, d_model: int = 768, img_size: Union[int, Tuple[int, int]] = 224, patch_size: Union[int, Tuple[int, int]] = 16, in_chans: int = 3, flatten: bool = True, norm_layer: Optional[Layer] = None, cls: bool = False, ): super().__init__() if isinstance(patch_size, int): patch_size = (patch_size, patch_size) if isinstance(img_size, int): img_size = (img_size, img_size) self.img_size = img_size self.patch_size = patch_size self.in_chans = in_chans self.embed_dim = d_model self.flatten = flatten self.frozen = False self.norm = norm_layer if norm_layer is not None else None self.use_cls = cls # ← NEW self.proj = nn.Conv2d(in_chans, d_model, kernel_size=patch_size, stride=patch_size) if self.use_cls: assert flatten, "CLS token requires flatten=True" self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model)) # ← NEW self.grid_size = None if img_size is not None: self.grid_size = (math.ceil(img_size[0] / patch_size[0]), math.ceil(img_size[1] / patch_size[1])) self._num_patches = None if self.grid_size is not None: self._num_patches = self.grid_size[0] * self.grid_size[1] def freeze(self): if self.frozen: flag("Trying to freeze an already frozen PatchEmbedder") return for p in self.proj.parameters(): p.requires_grad = False if self.norm is not None: for p in self.norm.parameters(): p.requires_grad = False if self.use_cls: # ← FREEZE CLS TOKEN TOO self.cls_token.requires_grad = False self.frozen = True def unfreeze(self): if not self.frozen: flag("Trying to unfreeze an already unfrozen PatchEmbedder") return for p in self.proj.parameters(): p.requires_grad = True if self.norm is not None: for p in self.norm.parameters(): p.requires_grad = True if self.use_cls: self.cls_token.requires_grad = True self.frozen = False def num_patches(self, H: Optional[int] = None, W: Optional[int] = None) -> int: if H is None or W is None: if self.img_size is not None: H, W = self.img_size else: raise ValueError("Must provide H and W or set img_size during init.") ph, pw = self.patch_size Hp = math.ceil(H / ph) Wp = math.ceil(W / pw) return Hp * Wp def _pad_to_patch_multiple(self, x: torch.Tensor) -> torch.Tensor: _, _, H, W = x.shape ph, pw = self.patch_size pad_h = (ph - H % ph) % ph pad_w = (pw - W % pw) % pw if pad_h == 0 and pad_w == 0: return x return nn.functional.pad(x, (0, pad_w, 0, pad_h), mode='constant', value=0) def forward(self, x: torch.Tensor) -> torch.Tensor: assert x.dim() == 4, "Input must be 4D (B, C, H, W)" assert x.dtype in (torch.float32, torch.float16, torch.bfloat16), "Input dtype must be float" B, C, H, W = x.shape assert C == self.in_chans, f"Expected {self.in_chans} channels, got {C}" x = self._pad_to_patch_multiple(x) x = self.proj(x) # (B, embed_dim, H_p, W_p) H_p, W_p = x.shape[2], x.shape[3] self._num_patches = H_p * W_p if self.flatten: x = x.flatten(2).transpose(1, 2) # (B, N, embed_dim) if self.use_cls: cls_tok = self.cls_token.expand(B, -1, -1) # (B, 1, D) x = torch.cat((cls_tok, x), dim=1) # (B, 1+N, D) if self.norm is not None: x = self.norm(x) else: if self.norm is not None: x = self.norm(x) return x class MeanClassificationHead(Layer): __name__ = "MeanClassificationHead" def __init__(self, d_model, vocab_size): super().__init__() self.proj = nn.Linear(d_model, vocab_size) def forward(self, x): x = x.mean(dim=1) return self.proj(x) class ClsHead(Layer): __name__ = "Cls Head" def __init__(self, d_model, vocab_size): super().__init__() self.proj = nn.Linear(d_model, vocab_size) def forward(self, x): x = x[:, -1:, :] return self.proj(x)[:, 0, :] class MultiHeadUnembedder(Layer): __name__ = "UnbiasedLinear" __complexity__ = "O(L d^2)" def __init__(self, d_model, vocab_size, k=4): super(MultiHeadUnembedder, self).__init__() self.linear = nn.Linear(d_model, vocab_size * k, bias=False) def forward(self, x): return self.linear(x) class VathosModel(Layer): __name__ = "VathosModel" def __init__(self): super().__init__() self._losses = [] self._losses_dict = {} self._losses_per_epoch = [] self._losses_per_epoch_dict = {} self._losses_this_epoch = [] self._metrics = dict() self._metrics_this_epoch = dict() self._metrics_per_epoch = dict() self.autosave = True self.autosave_overwrite = True self.best_loss = float('inf') self.checkpoints = 0 self.steps = 0 self.steps_per_epoch = 0 self.epochs = 0 self.finetuning = False self.name = 'VathosModel_defaultName' def flag_not_training(self): if not self.training: flag("You are not training") def register_loss(self, loss: float) -> bool: """Registra loss + NaN/Inf detection. Ritorna True se loss finita, False se NaN/Inf (con warning). Mantiene `self._nan_streak` (consecutive NaN counter). """ is_bad = (loss != loss) or (loss == float('inf')) or (loss == float('-inf')) if is_bad: self._nan_streak = getattr(self, "_nan_streak", 0) + 1 import warnings warnings.warn( f"VathosModel.register_loss: loss non-finita ({loss}) a step " f"{self.steps}. nan_streak={self._nan_streak}.", RuntimeWarning, stacklevel=2, ) else: self._nan_streak = 0 self._losses_dict[self.steps] = loss self._losses.append(loss) self._losses_this_epoch.append(loss) self.steps += 1 return not is_bad def get_last_loss(self): return self._losses[-1] @property def nan_streak(self) -> int: return getattr(self, "_nan_streak", 0) def get_mean_loss(self, epoch=True): if epoch: return np.mean(self._losses_this_epoch) else: return np.mean(self._losses) def register_metrics(self, metrics: dict): for metric in metrics: if metric in self._metrics: self._metrics[metric].append(metrics[metric]) self._metrics_this_epoch[metric].append(metrics[metric]) else: flag(f"Registering a new metric {metric}") self._metrics[metric] = [metrics[metric]] self._metrics_this_epoch[metric] = [metrics[metric]] def register_epoch(self): self.epochs += 1 self._losses_per_epoch.append(np.mean(self._losses_this_epoch)) self._losses_per_epoch_dict[self.steps] = np.mean(self._losses_this_epoch) self._losses_this_epoch = [] for metric in self._metrics_this_epoch: if metric in self._metrics_per_epoch: self._metrics_per_epoch[metric].append(np.mean(self._metrics_this_epoch[metric])) else: self._metrics_per_epoch[metric] = [np.mean(self._metrics_this_epoch[metric])] self._metrics_this_epoch[metric] = [] if self._losses_per_epoch[-1] < self.best_loss: self.best_loss = self._losses_per_epoch[-1] if self.autosave: self.checkpoints += 1 if self.autosave_overwrite: self.save_checkpoint(f'{self.name}-checkpoint.pt') else: self.save_checkpoint(f'{self.name}-checkpoint-{self.checkpoints}.pt') def save_state_dict(self, path): torch.save(self.state_dict(), path) def plot_losses(self): print(self.steps_per_epoch) plt.plot( list(self._losses_dict.keys()), list(self._losses_dict.values()), label="Losses", linewidth=1) plt.plot( list(self._losses_per_epoch_dict.keys()), list(self._losses_per_epoch_dict.values()), label="Losses Per Epoch", linewidth=2) plt.xlabel("steps") plt.ylabel("loss") plt.title("Model Losses per Steps") plt.show() def plot_metrics(self, figsize=(12, 8)): n_metrics = len(self._metrics_per_epoch) n_plots = 1 + n_metrics n_cols = 2 n_rows = (n_plots + 1) // 2 fig, axes = plt.subplots(n_rows, n_cols, figsize=figsize) if n_plots == 1: axes = np.array([axes]) else: axes = axes.flatten() ax = axes[0] ax.plot( list(self._losses_dict.keys()), list(self._losses_dict.values()), label="Losses", linewidth=1, alpha=0.6) ax.plot( list(self._losses_per_epoch_dict.keys()), list(self._losses_per_epoch_dict.values()), label="Losses Per Epoch", linewidth=2) ax.set_xlabel("Steps") ax.set_ylabel("Loss") ax.set_title("Training Loss") ax.legend() ax.grid(True, alpha=0.3) for idx, (metric_name, metric_values) in enumerate(self._metrics_per_epoch.items(), start=1): ax = axes[idx] if metric_name in self._metrics and len(self._metrics[metric_name]) > 0: step_indices = list(range(len(self._metrics[metric_name]))) ax.plot(step_indices, self._metrics[metric_name], label=f"{metric_name}", linewidth=1, alpha=0.6) epoch_indices = list(range(len(metric_values))) ax.plot(epoch_indices, metric_values, label=f"{metric_name} Per Epoch", linewidth=2, marker='o') ax.set_xlabel("Steps/Epochs") ax.set_ylabel(metric_name) ax.set_title(f"Metric: {metric_name}") ax.legend() ax.grid(True, alpha=0.3) # Hide any unused subplots for idx in range(n_plots, len(axes)): axes[idx].set_visible(False) plt.tight_layout() plt.show() def save_checkpoint(self, path): """Save complete model checkpoint including training state""" checkpoint = { 'model_state_dict': self.state_dict(), 'losses': self._losses, 'losses_dict': self._losses_dict, 'losses_per_epoch': self._losses_per_epoch, 'losses_per_epoch_dict': self._losses_per_epoch_dict, 'losses_this_epoch': self._losses_this_epoch, 'metrics': self._metrics, 'metrics_this_epoch': self._metrics_this_epoch, 'metrics_per_epoch': self._metrics_per_epoch, 'best_loss': self.best_loss, 'checkpoints': self.checkpoints, 'steps': self.steps, 'steps_per_epoch': self.steps_per_epoch, 'epochs': self.epochs, 'autosave': self.autosave, 'autosave_overwrite': self.autosave_overwrite, } torch.save(checkpoint, path) def load_checkpoint(self, path): """Load complete model checkpoint including training state""" checkpoint = torch.load(path, weights_only=False) self.load_state_dict(checkpoint['model_state_dict']) # Restore training state self._losses = checkpoint['losses'] self._losses_dict = checkpoint['losses_dict'] self._losses_per_epoch = checkpoint['losses_per_epoch'] self._losses_per_epoch_dict = checkpoint['losses_per_epoch_dict'] self._losses_this_epoch = checkpoint['losses_this_epoch'] self._metrics = checkpoint['metrics'] self._metrics_this_epoch = checkpoint['metrics_this_epoch'] self._metrics_per_epoch = checkpoint['metrics_per_epoch'] self.best_loss = checkpoint['best_loss'] self.checkpoints = checkpoint['checkpoints'] self.steps = checkpoint['steps'] self.steps_per_epoch = checkpoint['steps_per_epoch'] self.epochs = checkpoint['epochs'] self.autosave = checkpoint['autosave'] self.autosave_overwrite = checkpoint['autosave_overwrite'] def finetune(self): self.finetuning = True def train(self, *args, **kwargs): self.finetuning = False return super().train(*args, **kwargs) def eval(self): self.finetuning = False return super().eval() ######################################################################################################################## # Assemblers ######################################################################################################################## class SequenceModel(VathosModel): __name__ = "SequenceModel" def __init__(self, vocab_size: int, d_model: int, n_layers: int, max_len=1024, pos_encoder: bool | None | Layer | nn.Module = None, embedder=Embedder, embedder_args: dict = None, unembedder=UnbiasedLinear, unembedder_args=None, channel_mixer=MLP, spatial_mixer: Layer | nn.Module = MultiheadAttentionMixer, channel_args: dict = None, spatial_args: dict = None, name='', pad='none', baseblock=Block1d, baseblock_args=None, dropout=0.1, weight_tying=False, norm=nn.LayerNorm, d_modifiers: List | None = None, unet_skips=False, unet_weights=False ): super().__init__() self.pad = pad flag( "If you need to use any RoPE or alternative positional encodings which operate directly in the spatial mixer, be sure to call activate it in the spatial_args (e.g. rope=True)", 2) if channel_args is None and channel_mixer is MLP: channel_args = {"expand": 2, "activation": nn.GELU, "depth": 2} if spatial_args is None: spatial_args = {} if channel_args is None: channel_args = {} if embedder_args is None: embedder_args = {} if unembedder_args is None: unembedder_args = {'input_features': d_model, 'output_features': vocab_size} if baseblock_args is None: baseblock_args = {} if d_modifiers is None: d_modifiers = [1 for _ in range(d_model)] else: print( f"{NUM} Initialized SequenceModel with d_model structure: {[int(d_model * d) for d in d_modifiers]}{RES}") self.pipe = {} self.name = name self.baseblock = baseblock self.spatial_mixer = spatial_mixer self.channel_mixer = channel_mixer self.baseblock_args = baseblock_args self.spatial_args = spatial_args self.channel_args = channel_args self.vocab_size = vocab_size self.max_len = max_len self.d_model = d_model self.n_layers = n_layers self.unet_skips = unet_skips if self.unet_skips: self.skip_weights = nn.Parameter(torch.zeros(self.n_layers // 2), requires_grad=unet_weights) if d_modifiers is not None: flag("Unet Skips detected, assure d_modifier are symmetrical to make it work!") self.embedder = embedder(vocab_size=vocab_size, d_model=d_model, **embedder_args) self.pos_encoder = pos_encoder(d_model, max_len=max_len) if pos_encoder not in (True, False, None) else \ (SinusoidalPositionalEncoding(d_model, max_len=max_len) if pos_encoder is True else nn.Identity()) self.blocks = nn.ModuleList([ self.baseblock( d_model=int(d_model * d_modifiers[i]), channel_mixer=channel_mixer(d_model=int(d_model * d_modifiers[i]), **channel_args), spatial_mixer=spatial_mixer(d_model=int(d_model * d_modifiers[i]), **spatial_args), norm=norm, **baseblock_args ) for i in range(n_layers) ]) self._piped_blocks = None self.norm = nn.LayerNorm(d_model) self.unembedder = unembedder(**unembedder_args) if weight_tying and hasattr(self.embedder, 'embedding') and hasattr(self.unembedder, 'linear'): self.unembedder.linear.weight = self.embedder.embedding.weight elif weight_tying: raise TypeError( "Automatic weight tying is only possible if the the Embedder has a 'embedding' attribute and Unembedder has a linear attribute" "You should manually do weight tying if you aim to use specific layer:" "\n e.g. model.unembedder.yourmodule.weight = model.embedder.youembeddings.weight is an auto weight tying example") self.embedder_complexity = embedder.__complexity__ if hasattr(embedder, "__complexity__") else "O(L d)" self.unembedder_complexity = unembedder.__complexity__ if hasattr(unembedder, "__complexity__") else "O(L d)" self.spatial_complextiy = spatial_mixer.__complexity__ if hasattr(spatial_mixer, "__complexity__") else "O(L d)" self.channel_complexity = channel_mixer.__complexity__ if hasattr(channel_mixer, "__complexity__") else "O(L d)" self.runned = False self.embed_scale = math.sqrt(d_model) self._init_weights() def _init_weights(self): """ Initialize weights following GPT-style conventions: - Small std for embeddings - Xavier/Kaiming for linear layers - Scaled initialization for residual projections """ std_embed = 0.02 try: nn.init.normal_(self.embedder.embedding.weight, mean=0.0, std=std_embed) except: pass try: nn.init.normal_(self.unembedder.weight, mean=0.0, std=std_embed) except: pass for module in self.modules(): if isinstance(module, nn.Linear): std_init = 0.02 nn.init.normal_(module.weight, mean=0.0, std=std_init) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) for block_idx, block in enumerate(self.blocks): for name, module in block.named_modules(): if isinstance(module, nn.Linear): depth_scale = (2.0 * self.n_layers) ** -0.5 if 'out' in name.lower() or 'proj' in name.lower() or '.l2' in name or '.g2' in name: with torch.no_grad(): module.weight.data *= depth_scale def forward(self, x: torch.LongTensor, unembed=True): # B, L = x.size(0), x.size(1) # x = self.embedder(x) * self.embed_scale x = self.pos_encoder(x) skips = [] half_layers = self.n_layers // 2 for i, block in enumerate(self.blocks): if self.unet_skips: if i < half_layers: skips.append(x) elif i >= self.n_layers - half_layers: skip_idx = self.n_layers - 1 - i x = x + self.skip_weights[skip_idx] * skips[skip_idx] x = block(x) if unembed: x = self.norm(self.unembedder(x)) return x def insert_block(self, idx, module): self.blocks.insert(idx, module) def append(self, module): self.blocks.append(module) @torch.no_grad() def _clear_all_caches(self): """Clear KV caches in all attention layers""" for block in self.blocks: for module in block.modules(): if hasattr(module, 'clear_cache'): module.clear_cache() def forward(self, x: torch.LongTensor, unembed=True): B, L = x.size(0), x.size(1) x = self.embedder(x) * self.embed_scale x = self.pos_encoder(x) if self.pad == 'sqrt': n = int((x.shape[1] ** 0.5) + 0.999999) x = F.pad(x, (0, 0, 0, n ** 2 - L), mode="constant", value=0) for block in self.blocks: x = block(x) if unembed: x = self.unembedder(x) return x[:, :L, :] def _sample_token(self, logits, temperature=1.0, top_p=1.0, top_k=None): logits = logits / (temperature + 1e-8) if top_k is not None and top_k > 0: top_k = min(top_k, logits.size(-1)) v, _ = torch.topk(logits, top_k) logits[logits < v[:, [-1]]] = float('-inf') if top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] = float('-inf') probs = F.softmax(logits, dim=-1) return torch.multinomial(probs, num_samples=1) def generate(self, *args, **kwargs): if not 'custom_generate' in kwargs: kwargs['custom_generate'] = False if kwargs['custom_generate']: del kwargs['custom_generate'] return self.custom_generate(*args, **kwargs) else: del kwargs['custom_generate'] return self.simple_generate(*args, **kwargs) @torch.no_grad() def simple_generate(self, prompt: torch.Tensor, max_len=100, temperature=1.0, top_p=1.0, top_k=50, token_end=None, repetition_penalty=1.0): self.eval() if prompt.dim() == 1: prompt = prompt.unsqueeze(0) generated = prompt.clone() pbar = tqdm(range(max_len), desc="Simple Gen") for _ in pbar: logits = self.forward(generated, unembed=True) next_token_logits = logits[:, -1, :] # Apply repetition penalty if repetition_penalty != 1.0: next_token_logits = self._apply_repetition_penalty( next_token_logits, generated, repetition_penalty ) next_token = self._sample_token(next_token_logits, temperature, top_p, top_k) generated = torch.cat([generated, next_token], dim=1) if token_end is not None and (next_token == token_end).all(): break return generated @torch.no_grad() def custom_generate(self, prompt: torch.Tensor, max_len=100, temperature=1.0, top_p=1.0, top_k=50, token_end=None, repetition_penalty=1.0): self.eval() self._clear_all_caches() if prompt.dim() == 1: prompt = prompt.unsqueeze(0) x = self.embedder(prompt) * self.embed_scale if self.pos_encoder is not None and not isinstance(self.pos_encoder, nn.Identity): x = self.pos_encoder(x) for block in self.blocks: if block.has_custom_generate(): x = block.generate(x) else: x = block(x) generated = prompt.clone() logits = self.unembedder(x) next_token_logits = logits[:, -1, :] # Apply repetition penalty if repetition_penalty != 1.0: next_token_logits = self._apply_repetition_penalty( next_token_logits, generated, repetition_penalty ) next_token = self._sample_token(next_token_logits, temperature, top_p, top_k) generated = torch.cat([generated, next_token], dim=1) pbar = tqdm(range(max_len - 1), desc="Fast Gen") for _ in pbar: x_t = self.embedder(next_token) * self.embed_scale current_pos = generated.shape[1] - 1 if isinstance(self.pos_encoder, SinusoidalPositionalEncoding): pe_slice = self.pos_encoder.pe[current_pos: current_pos + 1].unsqueeze(0) x_t = x_t + pe_slice for block in self.blocks: if block.has_custom_generate(): x_t = block.generate(x_t) else: x_t = block(x_t) logits = self.unembedder(x_t) next_token_logits = logits[:, -1, :] # Apply repetition penalty if repetition_penalty != 1.0: next_token_logits = self._apply_repetition_penalty( next_token_logits, generated, repetition_penalty ) next_token = self._sample_token(next_token_logits, temperature, top_p, top_k) generated = torch.cat([generated, next_token], dim=1) if token_end is not None and (next_token == token_end).all(): break self._clear_all_caches() return generated def _apply_repetition_penalty(self, logits: torch.Tensor, generated: torch.Tensor, repetition_penalty: float) -> torch.Tensor: """ Apply repetition penalty to logits based on previously generated tokens. Args: logits: Shape (batch_size, vocab_size) generated: Shape (batch_size, seq_len) - previously generated tokens repetition_penalty: Penalty factor (> 1.0 discourages repetition) Returns: Modified logits with repetition penalty applied """ batch_size = logits.shape[0] for i in range(batch_size): # Get unique tokens in the generated sequence for this batch item unique_tokens = generated[i].unique() # Apply penalty: divide logits by penalty if positive, multiply if negative for token_id in unique_tokens: if logits[i, token_id] > 0: logits[i, token_id] /= repetition_penalty else: logits[i, token_id] *= repetition_penalty return logits def summary(self): complexities = [self.channel_complexity, self.spatial_complextiy, self.embedder_complexity, self.unembedder_complexity] print(f'{NUM}VATHOS{RES} {self.name} Summary:') print(f"{NUM}SequenceModel{RES}(d_model={NUM}{self.d_model}{RES}, n_layer={NUM}{self.n_layers}{RES})") print(f"\t - {NUM}VOCAB_SIZE:{RES}: {NUM}{self.vocab_size}{RES}") print(f"\t - {NUM}D_MODEL:{RES}: {NUM}{self.d_model}{RES}") print(f"\t - {NUM}N_LAYERS:{RES}: {NUM}{self.n_layers}{RES}") print("") print(f"\t - {NUM}Embedder{RES}: {getname(self.embedder)} - {NUM}{self.embedder_complexity}{RES}") print(f"\t - {NUM}Unembedder{RES}: {getname(self.unembedder)} - {NUM}{self.unembedder_complexity}{RES}") print( f"\t - {NUM}Spatial Mixer{RES}: {getname(self.spatial_mixer)}({self.spatial_args}) - {NUM}{self.spatial_complextiy}{RES}") print( f"\t - {NUM}Channel Mixer{RES}: {getname(self.channel_mixer)}({self.channel_args}) - {NUM}{self.channel_complexity}{RES}") print(f"Num Parameters: {NUM}{sum([p.numel() for p in self.parameters()]):_}{RES}") print(f"Num Trainable Parameters: {NUM}{sum([p.numel() for p in self.parameters() if p.requires_grad]):_}{RES}") print(f"Total Complexity: {NUM}{combine_big_o_sum(complexities)}{RES}") def finetune(self): flag( "Finetune simply checks for finetune() methods in spatial mixers, channel mixers, embedder and unembedder, " "if a finetune method is not available, the module/Layer will be left as it is") super().finetune() if hasattr(self.embedder, "finetune"): self.embedder.finetune() if hasattr(self.unembedder, "finetune"): self.unembedder.finetune() for block in self.blocks: if hasattr(block, "finetune"): block.channel_mixer.finetune() block.spatial_mixer.finetune() class ModdedFormer(VathosModel): def __init__(self, vocab_size: int, embed_dim: int, d_models: List[int], spatials: List[Builder], expand=3, M_dims: List[int] | None = None, weights_tying=True, baseblock=Block1d, norm=RMSNorm, ffn_act=ReLU2, unet_skips=False, max_len=2400, zeroskip=False, UDLP=VariableUDLP, skips: List[int | None] | None = None, value_embeddings: List[int | None] | None = None, ve_type: str = 'scalar', ve_gate_dim: int | None = None, learnable_pe: bool = False, input_projections: bool | int=False, smear_gate: None | nn.Module | Layer=False): super().__init__() self.vocab_size = vocab_size self.embed_dim = embed_dim assert d_models[0] == embed_dim assert max(d_models) == min(d_models), "Variables d_models is WIP and not working" self.d_models = d_models n_layers = self.n_layer = len(d_models) self.input_projections = input_projections self.smear_gate = smear_gate if M_dims is not None: self.M_dims = M_dims else: self.M_dims = [d * expand for d in d_models] self.embedder = Embedder(vocab_size, embed_dim) self.unembedder = UnbiasedLinear(d_models[-1], vocab_size) self.spatials = spatials self.norm = norm self.unet = unet_skips self.max_len = max_len self.zeroskip = zeroskip self.skips = skips if skips is not None else [None] * n_layers assert len(self.skips) == n_layers self.skip_lambdas = nn.ParameterDict() for source_idx, target_idx in enumerate(self.skips): if target_idx is not None: assert target_idx > source_idx assert target_idx < n_layers self.skip_lambdas[f"route_{source_idx}_to_{target_idx}"] = nn.Parameter(torch.zeros(1)) if weights_tying: assert d_models[-1] == embed_dim self.unembedder.linear.weight = self.embedder.embedding.weight blocks = [] for i in range(n_layers): in_dim = embed_dim if i == 0 else self.d_models[i - 1] out_dim = self.d_models[i] blocks.append( baseblock( self.d_models[i], UDLP(in_dim, d_output=out_dim, M=self.M_dims[i], activation=ffn_act), self.spatials[i](self.d_models[i]), norm=norm )) self.blocks = nn.ModuleList(blocks) if zeroskip: self.zeroskip_params = nn.ParameterList([nn.Parameter(torch.tensor([0.0])) for _ in range(n_layers)]) if input_projections > 0: if not zeroskip: raise ValueError("Zeroskip must be activated when using input_projections") self.inputs_projection_linears = nn.Linear(in_features=d_models[0], out_features=d_models[0]*input_projections, bias=False) self.inputs_projection_params = nn.ParameterList([nn.Parameter(torch.tensor([0.0])) for _ in range(n_layers) for i in range(input_projections)]) assert ve_type in ('scalar', 'gate'), f"ve_type must be 'scalar' or 'gate', got {ve_type}" self.ve_type = ve_type self.ve_gate_dim = ve_gate_dim if ve_gate_dim is not None else embed_dim self.value_embeddings_cfg = value_embeddings # Value embeddings: ModuleDict/ParameterDict come prima (per leggibilità + checkpoint legacy), # MA in più costruiamo strutture INDEXED-BY-LAYER per il forward (no str() runtime, no dict # lookup → compile-friendly). if value_embeddings is not None: assert len(value_embeddings) == n_layers, "value_embeddings length must equal n_layers" unique_groups = sorted(set(v for v in value_embeddings if v is not None)) self.ve_embeddings = nn.ModuleDict({ str(g): nn.Embedding(vocab_size, embed_dim) for g in unique_groups }) if ve_type == 'scalar': self.ve_scales = nn.ParameterDict({ str(i): nn.Parameter(torch.zeros(1)) for i, v in enumerate(value_embeddings) if v is not None }) else: # gate self.ve_gates = nn.ModuleDict({ str(i): nn.Linear(self.ve_gate_dim, embed_dim, bias=False) for i, v in enumerate(value_embeddings) if v is not None }) # Strutture parallele indexed-by-layer per il forward (riferiscono gli stessi param). self._ve_per_layer = nn.ModuleList([ self.ve_embeddings[str(v)] if v is not None else nn.Identity() for v in value_embeddings ]) self._has_ve_per_layer = tuple(v is not None for v in value_embeddings) self._has_any_ve = any(self._has_ve_per_layer) if ve_type == 'scalar': self._ve_scale_per_layer = nn.ParameterList([ self.ve_scales[str(i)] if value_embeddings[i] is not None else nn.Parameter(torch.zeros(1), requires_grad=False) for i in range(n_layers) ]) self._ve_gate_per_layer = nn.ModuleList([nn.Identity() for _ in range(n_layers)]) else: self._ve_scale_per_layer = nn.ParameterList([ nn.Parameter(torch.zeros(1), requires_grad=False) for _ in range(n_layers) ]) self._ve_gate_per_layer = nn.ModuleList([ self.ve_gates[str(i)] if value_embeddings[i] is not None else nn.Identity() for i in range(n_layers) ]) else: self.ve_embeddings = nn.ModuleDict() self._has_ve_per_layer = tuple([False] * n_layers) self._has_any_ve = False self._ve_per_layer = nn.ModuleList([nn.Identity() for _ in range(n_layers)]) self._ve_scale_per_layer = nn.ParameterList() self._ve_gate_per_layer = nn.ModuleList() # --------------------------------------------------------------------- # Skip routing: pre-computed plan + ParameterList (no string keys, no dict mutato) # --------------------------------------------------------------------- # Per ogni layer i: # - _skip_save[i]: bool, se salvare l'output come sorgente di skip # - _skip_consume_at[i]: tuple di (source_idx, gate_idx_in_list) da consumare qui self._skip_save = tuple(s is not None for s in self.skips) consume_lists = [[] for _ in range(n_layers)] gate_param_count = 0 # Mapping (source, target) -> gate index nella ParameterList self._skip_gate_map = {} for source_idx, target_idx in enumerate(self.skips): if target_idx is not None: self._skip_gate_map[(source_idx, target_idx)] = gate_param_count consume_lists[target_idx].append((source_idx, gate_param_count)) gate_param_count += 1 self._skip_consume_at = tuple(tuple(lst) for lst in consume_lists) self._has_any_skip = gate_param_count > 0 # ParameterList parallela: stessi tensori della ParameterDict legacy, indicizzati posizionalmente. if self._has_any_skip: self._skip_gates_list = nn.ParameterList([ self.skip_lambdas[f"route_{s}_to_{t}"] for (s, t), _ in sorted(self._skip_gate_map.items(), key=lambda kv: kv[1]) ]) else: self._skip_gates_list = nn.ParameterList() self.learnable_pe = learnable_pe if self.learnable_pe: self.pos_emb = nn.Parameter(torch.zeros(1, max_len, embed_dim)) self.final_norm = RMSNorm(d_models[-1]) self._init_weights() # Flag globale fast-path: nessuna feature opzionale → loop pulito senza branch. self._has_any_zeroskip = bool(zeroskip) self._has_input_projections = int(input_projections) > 0 self._has_smear_gate = bool(smear_gate) self._has_learnable_pe = bool(learnable_pe) self._fast_path = not any([ self._has_any_ve, self._has_any_skip, self._has_any_zeroskip, self._has_input_projections, self._has_smear_gate, self._has_learnable_pe, ]) def _apply_ve_weight(self, ve: torch.Tensor, x: torch.Tensor, layer_idx: int) -> torch.Tensor: """Legacy helper (mantenuto per backward compat). Il forward usa la versione indicizzata.""" if self.ve_type == 'scalar': return ve * self.ve_scales[str(layer_idx)] else: gate = 2.0 * torch.sigmoid(self.ve_gates[str(layer_idx)](x[..., :self.ve_gate_dim])) return ve * gate def _compute_ve_indexed(self, input_ids: torch.Tensor, x: torch.Tensor, i: int): """Versione fast: ModuleList[i] invece di ModuleDict[str(group_id)].""" ve = self._ve_per_layer[i](input_ids) if self.ve_type == 'scalar': return ve * self._ve_scale_per_layer[i] # gate type return ve * (2.0 * torch.sigmoid(self._ve_gate_per_layer[i](x[..., :self.ve_gate_dim]))) # ------------------------------------------------------------------------- # Forward — fast path se niente feature opzionali, slow path altrimenti # ------------------------------------------------------------------------- def forward(self, x): # FAST PATH: standard transformer pretrain, no skip / no VE / no zeroskip / no PE / no smear. # È il caso di gran lunga più comune (es. PiCO 2 / Guido-1 baseline). if self._fast_path: x = self.embedder(x) for block in self.blocks: x = block(x) return self.unembedder(self.final_norm(x)) # SLOW PATH: include tutte le feature opzionali. Anche qui niente string keys / dict mutati: # le strutture sono pre-computate al __init__ e indexed-by-layer. input_ids = x # alias per VE lookup x0 = self.embedder(x) if self._has_input_projections: xs = self.inputs_projection_linears(x0).chunk(self.input_projections, dim=-1) if self._has_learnable_pe: x0 = x0 + self.pos_emb[:, :x.size(1), :] if self._has_smear_gate: x0 = self.smear_gate(x0) x = x0 # Buffer di skip: lista Python di lunghezza fissa = n_layers. Allocata solo se serve. skip_buffer = [None] * self.n_layer if self._has_any_skip else None for i, block in enumerate(self.blocks): # VE — indexed lookup (no string) ve = None if self._has_ve_per_layer[i]: ve = self._compute_ve_indexed(input_ids, x, i) # Block forward + zeroskip variants if self._has_any_zeroskip and self._has_input_projections: proj_params = self.inputs_projection_params[ i * self.input_projections: (i + 1) * self.input_projections ] x = block(x, ve=ve) + x0 * self.zeroskip_params[i] + sum( [xj * p for xj, p in zip(xs, proj_params)] ) elif self._has_any_zeroskip: x = block(x, ve=ve) + x0 * self.zeroskip_params[i] else: x = block(x, ve=ve) # Skip routing — pre-computed plan, indexed gates (no dict, no string) if self._has_any_skip: consume = self._skip_consume_at[i] for source_idx, gate_idx in consume: x = x + self._skip_gates_list[gate_idx] * skip_buffer[source_idx] if self._skip_save[i]: skip_buffer[i] = x return self.unembedder(self.final_norm(x)) @torch.no_grad() def _clear_all_caches(self): for block in self.blocks: for module in block.modules(): if hasattr(module, 'clear_cache'): module.clear_cache() def _embed_input(self, ids: torch.Tensor, pos_start: int) -> torch.Tensor: x0 = self.embedder(ids) if self.learnable_pe: L = ids.size(1) x0 = x0 + self.pos_emb[:, pos_start:pos_start + L, :] if self.smear_gate: x0 = self.smear_gate(x0) return x0 def _blocks_generate(self, x: torch.Tensor, x0: torch.Tensor, xs, input_ids: torch.Tensor) -> torch.Tensor: """ Mirror of forward()'s block loop, but calls block.generate(...) so that attention mixers populate / consume their KV cache. Works for both prefill (L>1) and single-token increments (L=1). """ active_skips = {} for i, block in enumerate(self.blocks): ve = None if self.value_embeddings_cfg is not None and self.value_embeddings_cfg[i] is not None: group_id = self.value_embeddings_cfg[i] ve = self.ve_embeddings[str(group_id)](input_ids) ve = self._apply_ve_weight(ve, x, i) if self.zeroskip and self.input_projections > 0: proj_params = self.inputs_projection_params[ i * self.input_projections: (i + 1) * self.input_projections ] x = block.generate(x, ve=ve) + x0 * self.zeroskip_params[i] + sum( [xj * p for xj, p in zip(xs, proj_params)] ) elif self.zeroskip: x = block.generate(x, ve=ve) + x0 * self.zeroskip_params[i] else: x = block.generate(x, ve=ve) for source_idx, target_idx in enumerate(self.skips): if target_idx == i: gate = self.skip_lambdas[f"route_{source_idx}_to_{target_idx}"] x = x + gate * active_skips[source_idx] del active_skips[source_idx] if self.skips[i] is not None: active_skips[i] = x return x @torch.no_grad() def generate(self, prompt: torch.Tensor, max_len: int = 100, temperature: float = 1.0, top_k: int | None = None, top_p: float = 1.0, repetition_penalty: float = 1.0, token_end=None, simple: bool = False): """ KV-cached AR generation. Prefills the prompt in one shot (every spatial mixer's .generate accepts L>1 and populates its kv_cache), then walks token-by-token reusing the cached keys/values. simple=True falls back to simple_ar_generate (O(L^2)) — useful for regression-checking the cached path. """ if simple: return simple_ar_generate( self, prompt, max_len=max_len, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, token_end=token_end, ) self.eval() self._clear_all_caches() if self.value_embeddings_cfg is not None: flag("ModdedFormer.generate: value_embeddings only work with spatial mixers " "whose .generate accepts a 've=' kwarg (e.g. MultiheadAttentionMixer). " "On other mixers ve is silently dropped during generation.", 2) if prompt.dim() == 1: prompt = prompt.unsqueeze(0) device = next(self.parameters()).device prompt = prompt.to(device) L = prompt.size(1) x0 = self._embed_input(prompt, pos_start=0) xs = (self.inputs_projection_linears(x0).chunk(self.input_projections, dim=-1) if self.input_projections > 0 else None) x = self._blocks_generate(x0, x0, xs, prompt) logits = self.unembedder(self.final_norm(x))[:, -1, :] if repetition_penalty != 1.0: logits = apply_repetition_penalty(logits, prompt, repetition_penalty) next_token = sample_next_token(logits, temperature, top_k, top_p) generated = torch.cat([prompt, next_token], dim=1) pos = L pbar = tqdm(range(max_len - 1), desc="ModdedFormer fast gen") for _ in pbar: x0_t = self._embed_input(next_token, pos_start=pos) xs_t = (self.inputs_projection_linears(x0_t).chunk(self.input_projections, dim=-1) if self.input_projections > 0 else None) x_t = self._blocks_generate(x0_t, x0_t, xs_t, next_token) logits = self.unembedder(self.final_norm(x_t))[:, -1, :] if repetition_penalty != 1.0: logits = apply_repetition_penalty(logits, generated, repetition_penalty) next_token = sample_next_token(logits, temperature, top_k, top_p) generated = torch.cat([generated, next_token], dim=1) pos += 1 if token_end is not None and (next_token == token_end).all(): break self._clear_all_caches() return generated def _init_weights(self): std_embed = 0.02 try: nn.init.normal_(self.embedder.embedding.weight, mean=0.0, std=std_embed) except: pass try: nn.init.normal_(self.unembedder.weight, mean=0.0, std=std_embed) except: pass if self.learnable_pe: nn.init.normal_(self.pos_emb, mean=0.0, std=std_embed) for block_idx, block in enumerate(self.blocks): block.channel_mixer._init_weights() try: block.spatial_mixer._init_weights() except: pass def summary(self): print(f"Vathos {NUM}ModdedFormer{RES} Summary") print(f"Embedding dim: {self.embed_dim}") print(f"Learnable PE: {self.learnable_pe}") print(f"Skips: {self.skips}") for i in range(self.n_layer): if i == 0: in_dim = self.embed_dim else: in_dim = self.d_models[i - 1] out_dim = self.d_models[i] M = self.M_dims[i] print(f"Layer {i}: {in_dim} -> {out_dim}") print(f"Block: {i}: {self.blocks[i]}") print(f"\tFFN Dimension: {M}") print(f"\tFFN Activation: {self.blocks[i].channel_mixer.activation}") print(f"Num Parameters: {NUM}{sum([p.numel() for p in self.parameters()]):_}{RES}") print(f"Num Trainable Parameters: {NUM}{sum([p.numel() for p in self.parameters() if p.requires_grad]):_}{RES}") ####################################################################################################################### """def test_causality(module=MTransformer(8, 4, 2)): torch.manual_seed(42) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = module x = torch.randn(2, 20, model.d_model, device=device) print(f"\nInput shape: {x.shape}") with torch.no_grad(): output = model(x) model.eval() with torch.no_grad(): full_output = model(x) for t in range(1, 20): prefix_output = model(x[:, :t, :]) max_diff = (full_output[:, :t, :] - prefix_output).abs().max().item() if max_diff > 1e-5: print(f"Causality violated") break else: print("Causality verified") def test_causality_symbolic(module=SequenceModel(128, 16, 4, 2, pos_encoder=True)): torch.manual_seed(42) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = module x = torch.randint(0, model.vocab_size, (2, model.max_len), device=device) print(f"\nInput shape: {x.shape}") with torch.no_grad(): output = model(x) model.eval() with torch.no_grad(): full_output = model(x, unembed=False) for t in range(1, 20): prefix_output = model(x[:, :t], unembed=False) max_diff = (full_output[:, :t] - prefix_output).abs().max().item() if max_diff > 1e-5: print(f"Causality violated") break else: print("Causality verified") def test_symbolic_model(model): vocab_size = model.vocab_size length = model.max_len x = torch.randint(0, vocab_size, (2, length)) print(f"Input shape {x.shape}") x = model(x) print(f"Output shape {x.shape}") print(f"Bounds: min:{x.min()}, max:{x.max()}, mean:{x.mean()}, std:{x.std()}, sum_of_a_vector: {x[0, 0, :].sum()}")""" ######################################################################################################################## NAMES = { "MLP": MLP, "Attention": MultiheadAttentionMixer, "MHA": MultiheadAttentionMixer, "Embed": Embedder, "EMBED": Embedder, "E": Embedder, "PE": SinusoidalPositionalEncoding, "PatchEmbedder": PatchEmbedder, "ClassificationHead": MeanClassificationHead, "MCH": MeanClassificationHead, "ClsHead": ClsHead, } def expand_architecture(arch_string): pattern = r"(\(.+?\))x(\d+)" def replacer(match): content = match.group(1) count = int(match.group(2)) return " -> ".join([content] * count) return re.sub(pattern, replacer, arch_string) def assemble(code, d_model=512): code = code.strip() if "x" in code: code = expand_architecture(code) divided = code.split("->") for arch in divided: arch = arch.strip() if arch.startswith("("): layers = [] for subarch in arch[1:-1].split(","): subarch = subarch.strip() if subarch in NAMES: layer = NAMES[subarch](d_model=d_model) else: raise ValueError(f"Unknown Layer: {subarch}") layers.append(subarch) block = Block1d(*layers) else: if arch in NAMES: layer = NAMES[arch] elif hasattr(nn, arch): layer = getattr(nn, arch) else: raise ValueError(f"Unknown Layer: {arch}") def wrap(module): return tWrapper(module) def get_builder(layer, params): pass ######################################################################################################################## # Pre Built ######################################################################################################################## GQA = GroupedQueryAttention GQANOO = GroupedQueryAttentionNOO GQANOV = GroupedQueryAttentionNOV Attention = MultiheadAttentionMixer AttentionNOV = MultiheadAttentionMixerNOV NOVLa2 = MultiheadAttentionMixerNOVLa2 # ============================================================================= # HIGH-LEVEL MODELS # ============================================================================= # # Modelli pronti all'uso, ottimizzati per training/inference efficienti. A # differenza di SequenceModel / ModdedFormer (pensati per studio + ablation), # qui ogni Python branch è pre-risolto al __init__, niente ParameterDict / # ModuleDict string-keyed, niente dict mutati nel forward. Tutto torch.compile- # friendly. # # d_model è fissa su tutti i layer (sceltà per semplicità + compile efficiency). # ============================================================================= class BlockX0(Layer): """Pre-norm block + x0-skip nativo. Forward(x, x0) -> x + spatial(norm1(x)) + channel(norm2(x)) + x0_lambda * x0 `x0_lambda` init=0 ⇒ comportamento identico a Block1d al primo step. """ __name__ = "BlockX0" def __init__(self, d_model: int, channel_mixer: Layer, spatial_mixer: Layer, norm=RMSNorm): super().__init__() self.spatial_mixer = spatial_mixer self.channel_mixer = channel_mixer self.norm1 = norm(d_model) self.norm2 = norm(d_model) self.x0_lambda = nn.Parameter(torch.zeros(1)) def forward(self, x: torch.Tensor, x0: torch.Tensor) -> torch.Tensor: x = x + self.spatial_mixer(self.norm1(x)) x = x + self.channel_mixer(self.norm2(x)) return x + self.x0_lambda * x0 def generate(self, x: torch.Tensor, x0: torch.Tensor) -> torch.Tensor: h = self.norm1(x) sm = self.spatial_mixer x = x + (sm.generate(h) if sm.has_custom_generate() else sm(h)) h = self.norm2(x) cm = self.channel_mixer x = x + (cm.generate(h) if cm.has_custom_generate() else cm(h)) return x + self.x0_lambda * x0 class PiCOBlock(BlockX0): """BlockX0 + value-embedding nativo (off di default, zero costo se off). Quando `ve_enabled=True`, forward accetta un tensor `ve` esterno e lo passa al spatial_mixer pre-scalato per `ve_scale` (Parameter init=0 ⇒ no-op iniziale). """ __name__ = "PiCOBlock" def __init__(self, d_model: int, channel_mixer: Layer, spatial_mixer: Layer, norm=RMSNorm, ve_enabled: bool = False): super().__init__(d_model, channel_mixer, spatial_mixer, norm) self.ve_enabled = ve_enabled if ve_enabled: self.ve_scale = nn.Parameter(torch.zeros(1)) def forward(self, x: torch.Tensor, x0: torch.Tensor, ve: torch.Tensor = None) -> torch.Tensor: if self.ve_enabled and ve is not None: x = x + self.spatial_mixer(self.norm1(x), ve=ve * self.ve_scale) else: x = x + self.spatial_mixer(self.norm1(x)) x = x + self.channel_mixer(self.norm2(x)) return x + self.x0_lambda * x0 def generate(self, x: torch.Tensor, x0: torch.Tensor, ve: torch.Tensor = None) -> torch.Tensor: h = self.norm1(x) sm = self.spatial_mixer if self.ve_enabled and ve is not None: scaled = ve * self.ve_scale x = x + (sm.generate(h, ve=scaled) if sm.has_custom_generate() else sm(h, ve=scaled)) else: x = x + (sm.generate(h) if sm.has_custom_generate() else sm(h)) h = self.norm2(x) cm = self.channel_mixer x = x + (cm.generate(h) if cm.has_custom_generate() else cm(h)) return x + self.x0_lambda * x0 class PiCOFormer(VathosModel): """High-level transformer. Ultra-efficient, modular, x0-skip nativo. Caratteristiche: - `d_model` fissa su tutti i layer. - Spatial mixers heterogeneous (lista di `Builder`, uno per layer). - x0-skip nativo via `BlockX0` / `PiCOBlock`. - Value embeddings opzionali (zero cost se off, indexed-by-layer). - Smear gate opzionale (Modded-NanoGPT style) su x0. - Final logit softcap. - Tied embeddings di default. Args: vocab_size: vocab size. d_model: residual stream dim (fissa). n_layers: numero blocchi. spatials: `Builder` (condiviso) o `List[Builder]` di lunghezza n_layers. channel: `Builder` channel mixer (default VariableUDLP M=4*d_model, ReLU²). norm: classe norm (default RMSNorm). ve_groups: `List[int|None]` di lunghezza n_layers. `None` = no VE su quel layer; int = id gruppo (gruppi condividono lo stesso `nn.Embedding`). `None` globalmente = niente VE (default, niente memory cost). smear_gate: `nn.Module` o `None`. Applicato a x0 dopo l'embedding. logit_softcap: softcap finale. 0 disabilita. tied_embeddings: tying input/output embedding (default True). """ __name__ = "PiCOFormer" def __init__(self, vocab_size: int, d_model: int, n_layers: int, spatials, channel=None, norm=RMSNorm, ve_groups=None, smear_gate: nn.Module = None, smear_gate_lookback: int = 0, logit_softcap: float = 30.0, tied_embeddings: bool = True): super().__init__() self.vocab_size = vocab_size self.d_model = d_model self.n_layers = n_layers self.softcap = logit_softcap # Spatials: lista di Builder o Builder singolo (condiviso) spatials_list = spatials if isinstance(spatials, list) else [spatials] * n_layers assert len(spatials_list) == n_layers, "len(spatials) must equal n_layers" # Channel default = VariableUDLP, M=4*d_model, ReLU² if channel is None: channel = Builder(VariableUDLP, d_output=d_model, M=4 * d_model, activation=ReLU2) # VE: pre-compute embedding condiviso per group + indice per layer self._has_any_ve = ve_groups is not None and any(g is not None for g in ve_groups) if self._has_any_ve: assert len(ve_groups) == n_layers, "len(ve_groups) must equal n_layers" unique_groups = sorted({g for g in ve_groups if g is not None}) self.ve_embeddings = nn.ModuleList([nn.Embedding(vocab_size, d_model) for _ in unique_groups]) group_to_idx = {g: i for i, g in enumerate(unique_groups)} self._ve_idx_per_layer = tuple( group_to_idx[g] if g is not None else -1 for g in ve_groups ) else: self.ve_embeddings = nn.ModuleList() self._ve_idx_per_layer = (-1,) * n_layers self._ve_enabled_per_layer = tuple(idx >= 0 for idx in self._ve_idx_per_layer) # Embeddings + final norm + unembedder self.embedder = Embedder(vocab_size, d_model) self.unembedder = UnbiasedLinear(d_model, vocab_size) if tied_embeddings: self.unembedder.linear.weight = self.embedder.embedding.weight self.final_norm = norm(d_model) # Blocks self.blocks = nn.ModuleList([ PiCOBlock( d_model, channel_mixer=channel(d_model), spatial_mixer=spatials_list[i](d_model), norm=norm, ve_enabled=self._ve_enabled_per_layer[i], ) for i in range(n_layers) ]) # Smear gate opzionale (nn.Module [B,L,D] -> [B,L,D]). # smear_gate_lookback: 0 = stateless (per-token pointwise) — il fast path # generate() lo applica al singolo token. Se >0, il fast path mantiene una # rolling window di `lookback` token raw embeddings e ricomputa smear su # `window + new_token` (necessario per conv causali kernel>1, RNN, ecc.). self.smear_gate = smear_gate self._smear_lookback = int(smear_gate_lookback) self._x_window = None self._smear_gate_warned = False self._init_weights() def _init_weights(self): # Embeddings: normal std=0.02 (standard). LM head tied di default ⇒ stesso peso. std_embed = 0.02 nn.init.normal_(self.embedder.embedding.weight, mean=0.0, std=std_embed) if self.unembedder.linear.weight is not self.embedder.embedding.weight: nn.init.normal_(self.unembedder.linear.weight, mean=0.0, std=std_embed) for emb in self.ve_embeddings: nn.init.normal_(emb.weight, mean=0.0, std=std_embed) # Sub-modules: ciascuno definisce la propria policy di identity-init. if self.smear_gate is not None and hasattr(self.smear_gate, "_init_weights"): self.smear_gate._init_weights() for block in self.blocks: if hasattr(block.channel_mixer, "_init_weights"): block.channel_mixer._init_weights() if hasattr(block.spatial_mixer, "_init_weights"): block.spatial_mixer._init_weights() # x0_lambda (BlockX0) e ve_scale (PiCOBlock) sono già zero-init da nn.Parameter(torch.zeros(1)). def _compute_ves(self, input_ids: torch.Tensor): """Pre-computa la lista di tensori VE per ogni layer (None se non abilitato).""" if not self._has_any_ve: return (None,) * self.n_layers # Computa unique embeddings UNA volta, poi distribuisce per layer unique_ves = [emb(input_ids) for emb in self.ve_embeddings] return tuple(unique_ves[idx] if idx >= 0 else None for idx in self._ve_idx_per_layer) def forward(self, x: torch.Tensor) -> torch.Tensor: input_ids = x x0 = self.embedder(x) if self.smear_gate is not None: x0 = self.smear_gate(x0) ves = self._compute_ves(input_ids) x = x0 for i, block in enumerate(self.blocks): x = block(x, x0, ve=ves[i]) logits = self.unembedder(self.final_norm(x)) if self.softcap > 0: logits = self.softcap * torch.tanh(logits / self.softcap) return logits @torch.no_grad() def _clear_caches(self): for block in self.blocks: for mod in block.modules(): if hasattr(mod, "clear_cache"): mod.clear_cache() self._x_window = None @torch.no_grad() def generate(self, prompt: torch.Tensor, max_len: int = 100, temperature: float = 1.0, top_k: int | None = None, top_p: float = 1.0, repetition_penalty: float = 1.0, token_end=None, simple: bool = False) -> torch.Tensor: """KV-cached AR generation. Prefill in un colpo, poi token-by-token.""" if simple: return simple_ar_generate( self, prompt, max_len=max_len, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, token_end=token_end, ) self.eval() self._clear_caches() # Warning una tantum se smear_gate è non-trivial ma lookback=0 if (self.smear_gate is not None and not isinstance(self.smear_gate, nn.Identity) and self._smear_lookback == 0 and not self._smear_gate_warned): import warnings warnings.warn( "PiCOFormer.generate(simple=False): smear_gate non-Identity con " "smear_gate_lookback=0. Se il modulo ha state temporale (conv " "kernel>1, RNN, ecc.), il fast path divergerà da simple=True. " "Setta smear_gate_lookback= o usa simple=True.", RuntimeWarning, stacklevel=2, ) self._smear_gate_warned = True if prompt.dim() == 1: prompt = prompt.unsqueeze(0) device = next(self.parameters()).device prompt = prompt.to(device) # --- Prefill -------------------------------------------------------- x0_raw = self.embedder(prompt) if self.smear_gate is not None: x0 = self.smear_gate(x0_raw) if self._smear_lookback > 0: # Salva ultimi `lookback` raw embeddings come window (esclude i nuovi) self._x_window = x0_raw[:, -self._smear_lookback:, :].clone() else: x0 = x0_raw ves = self._compute_ves(prompt) x = x0 for i, block in enumerate(self.blocks): x = block.generate(x, x0, ve=ves[i]) logits = self.unembedder(self.final_norm(x))[:, -1, :] if self.softcap > 0: logits = self.softcap * torch.tanh(logits / self.softcap) if repetition_penalty != 1.0: logits = apply_repetition_penalty(logits, prompt, repetition_penalty) next_token = sample_next_token(logits, temperature, top_k, top_p) generated = torch.cat([prompt, next_token], dim=1) # --- Step-by-step --------------------------------------------------- pbar = tqdm(range(max_len - 1), desc="PiCOFormer fast gen") for _ in pbar: x0_t_raw = self.embedder(next_token) if self.smear_gate is not None: if self._smear_lookback > 0: # combina window + new, applica smear, prendi ultimo token combined = torch.cat([self._x_window, x0_t_raw], dim=1) smeared = self.smear_gate(combined) x0_t = smeared[:, -1:, :] # aggiorna window: tieni gli ultimi `lookback` raw embeddings self._x_window = combined[:, -self._smear_lookback:, :] else: x0_t = self.smear_gate(x0_t_raw) else: x0_t = x0_t_raw ves_t = self._compute_ves(next_token) x_t = x0_t for i, block in enumerate(self.blocks): x_t = block.generate(x_t, x0_t, ve=ves_t[i]) logits = self.unembedder(self.final_norm(x_t))[:, -1, :] if self.softcap > 0: logits = self.softcap * torch.tanh(logits / self.softcap) if repetition_penalty != 1.0: logits = apply_repetition_penalty(logits, generated, repetition_penalty) next_token = sample_next_token(logits, temperature, top_k, top_p) generated = torch.cat([generated, next_token], dim=1) if token_end is not None and (next_token == token_end).all(): break self._clear_caches() return generated if __name__ == "__main__": d_models = [64 for i in range(4)] m_dims = [128, 64, 32, 16] attn = Builder(GQA, n_heads=4, n_kv_heads=2) model = ModdedFormer( vocab_size=100, embed_dim=64, d_models=d_models, spatials=[attn for _ in range(len(d_models))], M_dims=m_dims, input_projections=3, zeroskip=True ) model.summary() out = model( torch.randint(0, 99, (2, 128)) ) print(out.shape) model.profile()