import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutput from .modules import Attention from .utils import nearest_power_of_two from .layers import AttentionLayer from .configuration_minitransformer import MiniTransformerConfig from .attn_masks import causal_mask from .attn_mods import generate_tanh_softcap from .rotary_emb import precompute_freqs_cis try: from liger_kernel.transformers.rms_norm import LigerRMSNorm as TritonNorm triton_norm = True except ImportError as e: print( f"Unable to import Triton-based RMSNorm: {e}. Falling back to PyTorch implementation." ) from torch.nn import RMSNorm triton_norm = False # Load the tokenizer from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "Hazan-Lab/Transformer_500M" tokenizer = AutoTokenizer.from_pretrained( model_name, trust_remote_code=True ) class MiniTransformer(PreTrainedModel): config_class = MiniTransformerConfig def __init__(self, config) -> None: super(MiniTransformer, self).__init__(config) self.num_layers = config.num_layers assert config.dim % config.num_heads == 0, f"dim ({self.dim}) must be divisible num_heads ({self.num_heads})" self.head_dim = config.dim // config.num_heads logit_softcap = generate_tanh_softcap(soft_cap=config.softcap) # From pytorch/pytorch#123411, we set persistent=True for torch.compile and PP compatibility self.register_buffer("freqs_cis", precompute_freqs_cis( head_dim=self.head_dim, max_seq_len=config.seq_len, theta=config.theta, ), persistent=True) self.tok_emb = nn.Embedding(config.vocab_size, config.dim) self.dropout = nn.Dropout(config.dropout) self.layers = nn.ModuleList() for _ in range(self.num_layers): layer = AttentionLayer(config, mask_mod=causal_mask, score_mod=logit_softcap) self.layers.append(layer) self.norm = nn.RMSNorm(config.dim) self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=config.bias) # self.tok_emb.weight = self.lm_head.weight self.std = (config.dim) ** -0.5 self.apply(self._init_weights) print("Model Parameter Count: %.2fM\n" % (self._get_num_params() / 1e6,)) def forward( self, input_ids: torch.Tensor, labels: torch.Tensor = None, **kwargs ) -> CausalLMOutput: # Compute embeddings tok_emb = self.tok_emb(input_ids) for layer in self.layers: tok_emb = layer(tok_emb, self.freqs_cis) # Normalize and project to vocabulary tok_emb = self.norm(tok_emb) logits = self.lm_head(tok_emb) loss = None if labels is not None: # Shift so that tokens predict the next token shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels[..., 1:].contiguous() loss_fct = nn.CrossEntropyLoss() loss = loss_fct( shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) ) return CausalLMOutput( loss=loss, logits=logits, ) def _get_num_params(self): n_params = sum(p.numel() for p in self.parameters()) if hasattr(self, "pos_emb") and self.pos_emb is not None: n_params -= self.pos_emb.weight.numel() if self.tok_emb.weight is self.lm_head.weight: n_params -= self.tok_emb.weight.numel() return n_params def _init_weights(self, module): if isinstance(module, nn.Linear): if hasattr(module, "SCALE_INIT"): self.std *= (2 * self.num_layers) ** -0.5 torch.nn.init.normal_(module.weight, mean=0.0, std=self.std) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=self.std) @staticmethod def top_k_top_p_filtering( logits: torch.Tensor, top_k: int = 50, top_p: float = 0.95, filter_value: float = float("-inf"), ): """ Filters a distribution of logits using top-k and/or nucleus (top-p) filtering. """ # top_k if top_k > 0: top_k = min(top_k, logits.size(-1)) # Remove all logits that are not in the top k indices_to_remove = logits < torch.topk(logits, top_k, dim=-1).values[:, -1, None] logits[indices_to_remove] = filter_value # top_p (nucleus) if 0 < top_p < 1.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # Remove tokens with cumulative probability above the threshold sorted_indices_to_remove = cumulative_probs > top_p # Shift the indices to the right to keep also the first token above the threshold sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone() sorted_indices_to_remove[:, 0] = False indices_to_remove = sorted_indices_to_remove.scatter( dim=1, index=sorted_indices, src=sorted_indices_to_remove ) logits[indices_to_remove] = filter_value return logits def generate( self, input_ids: torch.LongTensor, max_new_tokens: int = 50, temperature: float = 0.5, top_k: int = 50, top_p: float = 0.95, eos_token_id: int = None, pad_token_id: int = 0, **kwargs ): """ Naive token-by-token generation loop that uses top-k/top-p filtering and optional temperature. Args: input_ids (torch.LongTensor): shape (batch_size, sequence_length). max_new_tokens (int): max number of tokens to generate (beyond input_ids length). temperature (float): sampling temperature (>=0). top_k (int): Top-K sampling cutoff. top_p (float): Nucleus sampling cutoff. eos_token_id (int): If set, stop generation when this token is produced. pad_token_id (int): If set, can be used to pad sequences. (Not fully used here.) kwargs: Unused arguments (like num_beams) for compatibility. Returns: torch.LongTensor: shape (batch_size, sequence_length + generated_tokens). """ device = input_ids.device print("1=====================") print(tokenizer.decode(input_ids[0], skip_special_tokens=True)) print("1=====================") # We'll accumulate new tokens into generated_ids generated_ids = input_ids.clone() for _ in range(max_new_tokens): # Forward pass to get logits for the last token outputs = self.forward(generated_ids) logits = outputs.logits[:, -1, :] # shape: (batch_size, vocab_size) # Scale logits by temperature if temperature != 1.0: logits = logits / temperature # Filter logits using top-k and/or top-p logits = self.top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) # Convert to probabilities probabilities = F.softmax(logits, dim=-1) # Sample from the distribution next_token = torch.multinomial(probabilities, num_samples=1) # (batch_size, 1) # Append next token generated_ids = torch.cat([generated_ids, next_token], dim=1) # If eos_token_id is set and any sample produced it, we optionally could break early if eos_token_id is not None: # Check if all sequences in the batch ended # or if you want to do a more fine-grained approach if (next_token == eos_token_id).all(): break print("2=====================") print(tokenizer.decode(generated_ids[0], skip_special_tokens=True)) print("2=====================") return generated_ids