fixes

config.json

@@ -16,14 +16,14 @@
   "global_bsz": 524288,
   "bsz": 2,
   "warmup_steps": 1907,
-  "eval_peruse_alibiiod": 50,
+  "eval_period": 50,
   "save_period": 500,
   "max_lr": 3.0e-4,
   "min_lr": 3.0e-5,
   "max_norm": 1.0,
   "dilation": 1,
-  "fsdp": true,
-  "ddp": false,
+  "fsdp": false,
+  "ddp": true,
   "mixed_precision": true,
   "torch_dtype": "bfloat16",
   "cpu_offload": false,
@@ -48,4 +48,4 @@
   "theta": 10000.0,
   "use_alibi": false,
   "torch_compile": false
-}
+}

configuration_minitransformer.py

@@ -1,6 +1,5 @@
 import torch
 from transformers import PretrainedConfig, AutoConfig
-
 class MiniTransformerConfig(PretrainedConfig):
     model_type = "minitransformer"
 
@@ -11,6 +10,7 @@ class MiniTransformerConfig(PretrainedConfig):
         num_heads: int = 8,
         num_layers: int = 12,
         seq_len: int = 8192,
+        weight_tying: bool = True,
         window_size: int = 8192,
         vocab_size: int = 200064,
         mlp_scale: int = 12,
@@ -18,7 +18,7 @@ class MiniTransformerConfig(PretrainedConfig):
         dropout: float = 0.0,
         softcap: float = 50.0,
         theta: float = 10_000.0,
-        use_alibi: bool = False,
+        use_alibi: bool = False, # Default to RoPE
         torch_dtype: torch.dtype = torch.bfloat16,
         device: torch.device = None,
         **kwargs,
@@ -29,6 +29,7 @@ class MiniTransformerConfig(PretrainedConfig):
         self.num_heads = num_heads
         self.num_layers = num_layers
         self.seq_len = seq_len
+        self.weight_tying = weight_tying
         self.window_size = window_size
         self.vocab_size = vocab_size
         self.hidden_size = dim
@@ -40,5 +41,4 @@ class MiniTransformerConfig(PretrainedConfig):
         self.theta = theta
         self.use_alibi = use_alibi
         self.torch_dtype = torch_dtype
-        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')  # Store as string
-
+        self.device = device

modeling_minitransformer.py

@@ -1,46 +1,199 @@
+import json
+import math
+
 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 transformers import PreTrainedModel, PretrainedConfig
 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
+    from flash_attn import flash_attn_func
 except ImportError as e:
     print(
-        f"Unable to import Triton-based RMSNorm: {e}. Falling back to PyTorch implementation."
+        f"Unable to import Triton-based flash attention: {e}. No alternative currently available."
+    )
+
+
+def precompute_freqs_cis(head_dim: int, max_seq_len: int, theta: float = 10000.0):    
+    # For half the dimensions, build the scale factor:
+    freq_seq = torch.arange(0, head_dim, 2).float() / head_dim
+    freqs = 1.0 / (theta ** freq_seq)
+
+    # Outer product with positions
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    angles = torch.outer(t, freqs)
+    
+    # Build a complex exponential e^{i * theta}
+    freqs_cis = torch.polar(
+        torch.ones_like(angles),
+        angles
+    )
+    return freqs_cis
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    """
+    x is [B, n_heads, seq_len, head_dim_as_complex],
+    so we want to broadcast freqs_cis from [max_seq_len, half_dim]
+    to [1, 1, seq_len, half_dim].
+    """
+    seq_len = x.shape[2]
+    freqs_cis = freqs_cis[:seq_len]  # slice down to current seq_len
+    return freqs_cis.view(1, 1, seq_len, -1)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Convert real -> complex by grouping last dim in pairs
+    # shape => [B, n_heads, seq_len, head_dim//2, 2] => complex => [B, n_heads, seq_len, head_dim//2]
+    xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+
+    # Broadcast the frequencies to match [B, n_heads, seq_len, head_dim//2]
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_complex)
+
+    # Multiply => apply rotation
+    xq_complex = xq_complex * freqs_cis
+    xk_complex = xk_complex * freqs_cis
+
+    # Convert back to real => shape [B, n_heads, seq_len, head_dim]
+    xq_out = torch.view_as_real(xq_complex).reshape(*xq.shape)
+    xk_out = torch.view_as_real(xk_complex).reshape(*xk.shape)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+def nearest_power_of_two(x: int, round_up: bool = False) -> int:
+    return (
+        1 << math.floor(math.log2(x)) if not round_up else 1 << math.ceil(math.log2(x))
     )
-    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 Attention(nn.Module):
+    def __init__(self, config):
+        super(Attention, self).__init__()
+        self.dim, self.num_heads = config.dim, config.num_heads
+        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
+
+        self.c_attn = nn.Linear(self.dim, 3*self.dim, bias=config.bias)
+        self.c_proj = nn.Linear(config.dim, config.dim, bias=config.bias)
+        self.c_proj.SCALE_INIT = 1
+
+        self.alibi_slopes = self._get_alibi_slopes(self.num_heads) if config.use_alibi else None
+        self.window_size = config.window_size
+        self.softcap = config.softcap
+
+        self.dropout = config.dropout
+        self.resid_dropout = nn.Dropout(self.dropout)
+
+    def _generate_slopes(self, n: int):
+            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            return [start * (start**i) for i in range(n)]
+
+    def _get_alibi_slopes(self, num_heads: int, interpolation_factor: float = 0.25):
+        # If n_heads is a power of 2, generate slopes directly
+        if math.log2(num_heads).is_integer():
+            slopes = self._generate_slopes(num_heads)
+        else:
+            # Get slopes for the nearest power of two
+            n = nearest_power_of_two(num_heads, round_up=False)
+            slopes_power_of_two = self._generate_slopes(n)
+
+            # Generate extra slopes
+            extra_slopes = self._generate_slopes(2 * n)
+            extra_slopes_trunc = extra_slopes[0::2][: num_heads - n]
+            slopes = slopes_power_of_two + extra_slopes_trunc
+        slopes = torch.tensor(slopes, device=torch.device("cuda"))
+        slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
+        return slopes
+
+    def forward(
+        self,
+        x: torch.Tensor = None,
+        q: torch.Tensor = None,
+        k: torch.Tensor = None,
+        v: torch.Tensor = None,
+        freqs_cis: torch.Tensor = None,
+    ) -> torch.Tensor:
+        if x is not None:
+            q = k = v = x
+        if any(t is None for t in [q, k, v]):
+            raise ValueError("Must provide either x for self-attention or q/k/v for cross-attention.")
+
+        bsz, q_len, dim = q.shape
+        _, k_len, _ = k.shape
+        _, v_len, _ = v.shape
+
+        qkv = self.c_attn(x)
+        q, k, v = torch.chunk(qkv, 3, dim=2)
+
+        q = q.view(bsz, q_len, self.num_heads, self.head_dim)
+        k = k.view(bsz, k_len, self.num_heads, self.head_dim)
+        v = v.view(bsz, v_len, self.num_heads, self.head_dim)
+
+        if self.alibi_slopes is None: # Use either ALiBi or RoPE
+            q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis)
+
+        y = flash_attn_func(  # https://arxiv.org/pdf/2307.08691
+            q=q, k=k, v=v,
+            dropout_p=self.dropout if self.training else 0.0,
+            causal=True,
+            window_size=(self.window_size, 0), # Set to config.seq_len if full attention
+            alibi_slopes=self.alibi_slopes, # https://arxiv.org/pdf/2108.12409
+            softcap=self.softcap,  # https://arxiv.org/pdf/2408.00118
+        )
+
+        y = y.contiguous().view(bsz, q_len, -1)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, config) -> None:
+        super(AttentionLayer, self).__init__()
+        self.attn_norm = nn.RMSNorm(config.dim)
+        self.attn = Attention(config=config)
+        self.mlp_norm = nn.RMSNorm(config.dim)
+        self.mlp = MLP(config)
+
+    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor=None) -> torch.Tensor:
+        x = x + self.attn(x=self.attn_norm(x), freqs_cis=freqs_cis)
+        x = x + self.mlp(self.mlp_norm(x))
+        return x
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        # https://arxiv.org/pdf/2002.05202
+        super().__init__()
+        self.hidden_size = config.dim
+        self.intermediate_size = config.dim * config.mlp_scale
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.bias)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.bias)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        gate = self.gate_proj(x)
+        gate = F.gelu(gate, approximate="tanh")
+        up = self.up_proj(x)
+        fuse = gate * up
+        outputs = self.down_proj(fuse)
+        outputs = self.dropout(outputs)
+        return outputs
 
 class MiniTransformer(PreTrainedModel):
+
     config_class = MiniTransformerConfig
 
     def __init__(self, config) -> None:
-        super(MiniTransformer, self).__init__(config)
+        super(Transformer, 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(
@@ -54,55 +207,36 @@ class MiniTransformer(PreTrainedModel):
 
         self.layers = nn.ModuleList()
         for _ in range(self.num_layers):
-            layer = AttentionLayer(config, mask_mod=causal_mask, score_mod=logit_softcap)
+            layer = AttentionLayer(config=config)
             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
+        if config.weight_tying:
+            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)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tok_emb = self.tok_emb(x)
+        x = self.dropout(tok_emb)
 
         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,
-        )
+            x = layer(x, self.freqs_cis)
+
+        y_hat = self.lm_head(self.norm(x))
+
+        return y_hat
 
     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):
@@ -114,105 +248,10 @@ class MiniTransformer(PreTrainedModel):
                 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
+        elif isinstance(module, Attention):
+            torch.nn.init.xavier_normal_(module.c_attn.weight)
+            torch.nn.init.xavier_normal_(module.c_proj.weight)
+            if module.c_attn.bias is not None:
+                torch.nn.init.zeros_(module.c_attn.bias)
+            if module.c_proj.bias is not None:
+                torch.nn.init.zeros_(module.c_proj.bias)