Upload model_tiny.py with huggingface_hub

model_tiny.py

@@ -0,0 +1,220 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def get_alibi_slopes(num_heads: int) -> torch.Tensor:
+    n = 2 ** math.ceil(math.log2(num_heads))
+    m = torch.tensor([2 ** (-(i + 1)) for i in range(n)])
+    m = m[:num_heads]
+    return m
+
+
+def build_alibi_bias(num_heads: int, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    slopes = get_alibi_slopes(num_heads).to(device=device, dtype=dtype)
+    pos = torch.arange(seq_len, device=device, dtype=dtype)
+    mask = pos.view(1, seq_len) - pos.view(seq_len, 1)
+    mask = mask.abs().neg().unsqueeze(0).unsqueeze(0)
+    bias = mask * slopes.view(-1, 1, 1)
+    return bias
+
+
+class ALiBiAttention(nn.Module):
+    def __init__(self, hidden: int, num_heads: int, num_kv_heads: int):
+        super().__init__()
+        self.hidden = hidden
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = hidden // num_heads
+        self.num_groups = num_heads // num_kv_heads
+
+        self.q_proj = nn.Linear(hidden, hidden, bias=False)
+        self.k_proj = nn.Linear(hidden, num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden, num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(hidden, hidden, bias=False)
+        self._alibi_bias = None
+        self._alibi_seq_len = 0
+
+    def _get_alibi(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        if self._alibi_bias is None or seq_len > self._alibi_seq_len:
+            bias = build_alibi_bias(self.num_heads, seq_len, device, dtype)
+            causal = torch.triu(torch.full((seq_len, seq_len), float('-inf'), device=device, dtype=dtype), diagonal=1)
+            self._alibi_bias = bias + causal
+            self._alibi_seq_len = seq_len
+        return self._alibi_bias[:, :, :seq_len, :seq_len]
+
+    def forward(self, x):
+        B, T, _ = x.shape
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        if self.num_groups > 1:
+            k = k.repeat_interleave(self.num_groups, dim=1)
+            v = v.repeat_interleave(self.num_groups, dim=1)
+
+        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        scores = scores + self._get_alibi(T, x.device, scores.dtype)
+
+        attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(x.dtype)
+        out = (attn @ v).transpose(1, 2).contiguous().view(B, T, self.hidden)
+        return self.o_proj(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, hidden: int, intermediate: int):
+        super().__init__()
+        self.gate = nn.Linear(hidden, intermediate, bias=False)
+        self.up = nn.Linear(hidden, intermediate, bias=False)
+        self.down = nn.Linear(intermediate, hidden, bias=False)
+
+    def forward(self, x):
+        return self.down(F.gelu(self.gate(x)) * self.up(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, hidden: int, intermediate: int, num_heads: int, num_kv_heads: int):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(hidden)
+        self.attn = ALiBiAttention(hidden, num_heads, num_kv_heads)
+        self.ln2 = nn.LayerNorm(hidden)
+        self.mlp = MLP(hidden, intermediate)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class TinyModel(nn.Module):
+    def __init__(self, vocab_size=4096, hidden=128, intermediate=512,
+                 num_layers=2, num_heads=8, num_kv_heads=4, max_seq_len=2048,
+                 tie_weights=True):
+        super().__init__()
+        self.hidden = hidden
+        self.max_seq_len = max_seq_len
+        self.token_embed = nn.Embedding(vocab_size, hidden)
+        self.blocks = nn.ModuleList([
+            TransformerBlock(hidden, intermediate, num_heads, num_kv_heads)
+            for _ in range(num_layers)
+        ])
+        self.ln_f = nn.LayerNorm(hidden)
+        self.lm_head = nn.Linear(hidden, vocab_size, bias=False)
+        if tie_weights:
+            self.lm_head.weight = self.token_embed.weight
+
+    def forward(self, input_ids, labels=None):
+        x = self.token_embed(input_ids)
+        for block in self.blocks:
+            x = block(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        return logits, loss
+
+    def generate(self, input_ids, max_new_tokens=128, temperature=0.7, top_p=0.9):
+        self.eval()
+        for _ in range(max_new_tokens):
+            seq_len = input_ids.size(1)
+            if seq_len > self.max_seq_len:
+                input_ids = input_ids[:, -self.max_seq_len:]
+            with torch.no_grad():
+                logits, _ = self.forward(input_ids)
+            logits = logits[:, -1, :] / temperature
+
+            if top_p < 1.0:
+                sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True)
+                cum_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+                cutoff = cum_probs > top_p
+                cutoff[..., 1:] = cutoff[..., :-1].clone()
+                cutoff[..., 0] = False
+                logits[~cutoff] = float('-inf')
+
+            probs = F.softmax(logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_token], dim=1)
+
+            if next_token.item() == 2:
+                break
+        return input_ids
+
+
+class LoRALinear(nn.Module):
+    def __init__(self, original: nn.Linear, r: int = 8, alpha: float = 16, dropout: float = 0.0):
+        super().__init__()
+        self.linear = original
+        self.r = r
+        self.alpha = alpha
+        self.scaling = alpha / r
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        in_dim, out_dim = original.in_features, original.out_features
+        self.lora_A = nn.Parameter(torch.zeros(r, in_dim))
+        self.lora_B = nn.Parameter(torch.zeros(out_dim, r))
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        self.linear.weight.requires_grad = False
+        if self.linear.bias is not None:
+            self.linear.bias.requires_grad = False
+
+    def forward(self, x):
+        return self.linear(x) + self.dropout(x) @ self.lora_A.T @ self.lora_B.T * self.scaling
+
+
+def apply_lora(model, target_modules=None, r=8, alpha=16, dropout=0.0):
+    if target_modules is None:
+        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate", "up", "down"]
+    lora_params = 0
+    for name, module in model.named_modules():
+        if not isinstance(module, nn.Linear):
+            continue
+        key = name.split(".")[-1]
+        if key not in target_modules:
+            continue
+        parent = model
+        parts = name.split(".")
+        for p in parts[:-1]:
+            parent = getattr(parent, p)
+        lora = LoRALinear(module, r=r, alpha=alpha, dropout=dropout)
+        setattr(parent, parts[-1], lora)
+        lora_params += 2 * r * module.in_features + module.out_features * r
+    n = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"LoRA applied: {lora_params:,} trainable params (total trainable: {n:,})")
+    return model
+
+
+def count_params(model):
+    return sum(p.numel() for p in model.parameters())
+
+
+def create_model():
+    model = TinyModel(
+        vocab_size=4096, hidden=128, intermediate=512,
+        num_layers=2, num_heads=8, num_kv_heads=4,
+        max_seq_len=2048, tie_weights=True,
+    )
+    n = count_params(model)
+    print(f"TinyModel: {n:,} params ({n/1e6:.2f}M)")
+    return model
+
+
+if __name__ == "__main__":
+    m = create_model()
+    x = torch.randint(0, 100, (1, 16))
+    logits, loss = m(x, labels=x)
+    print(f"Forward OK: logits {logits.shape}, loss {loss.item():.4f}")
+    gen = m.generate(x, max_new_tokens=10)
+    print(f"Generate OK: {gen.shape}")
+    m2 = create_model()
+    apply_lora(m2)
+    x2 = torch.randint(0, 100, (1, 16))
+    logits2, _ = m2(x2)
+    print(f"LoRA forward OK: {logits2.shape}")