.

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

MANIFEST.in

@@ -0,0 +1,4 @@
+include *.json
+include README.md
+include pytorch_model.bin
+recursive-include gemma3_hf *.py

README.md

@@ -0,0 +1,97 @@
+---
+license: mit
+language: en
+tags:
+- gemma3
+- slm
+- tinystories
+model_type: gemma3
+---
+
+# Gemma3 270M: Small Language Model Implementation from Scratch
+
+This repository contains a **PyTorch implementation** of the **Google's Gemma3 model** with **270 million parameters**, trained from scratch on the **TinyStories** dataset.
+
+[Github Repo Link](https://github.com/ShubhamWaghmare11/Gemma3_270M_Scratch)
+
+## About
+
+This model is based on Google DeepMind's Gemma3 architecture and was built from the ground up to explore training dynamics, architecture design, and generation quality of small LLMs. It includes advanced components such as:
+
+- Sliding Window Attention (512-token window)
+
+- Rotary Positional Embeddings (RoPE)
+
+- RMSNorm for stable training
+
+- Grouped Key-Value Attention (1 KV group)
+
+
+## Model Architecture
+
+- Parameters: 270M total (170M embedding + 100M transformer)
+- Layers: 18 transformer blocks
+- Attention Heads: 4 Query Heads, 1 KV Group
+- Hidden Dimension: 2048
+- Embedding Dimension: 640
+- Head Dimension: 256
+- Vocabulary Size: 50,257 (GPT-2 tokenizer)
+- Context Length: 32,768 tokens (trained with 128 block size)
+- Sliding Window: 512 tokens
+
+
+
+## Training Details
+
+- Dataset: [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) by Roneneldan
+
+- Steps: 150,000 steps (not epochs)
+
+- Batch Size: 32
+
+- Loss Function: Cross-Entropy
+
+- Optimizer: AdamW
+
+- LR Scheduler: Linear Warmup + Cosine Decay
+
+- Hardware: Single NVIDIA A100 GPU
+
+
+
+
+
+
+## Requirements
+
+```bash
+pip install git+https://huggingface.co/Shubhamw11/Gemma-270M-TinyStories
+```
+
+## How to use
+
+You can now load and use the model from the `gemma3_tinystories` library:
+
+```python
+from gemma3_tinystories import HFGemma3Model, Gemma3Config
+import tiktoken
+import torch
+
+config = Gemma3Config.from_pretrained("Shubhamw11/Gemma-270M-TinyStories")
+model = HFGemma3Model.from_pretrained("Shubhamw11/Gemma-270M-TinyStories", config=config).model
+tokenizer = tiktoken.get_encoding("gpt2")
+```
+
+# Generate text
+```python
+
+#define the device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+input_text = "Once upon a time, there was a little"
+context = torch.tensor(tokenizer.encode(input_text), dtype=torch.long).unsqueeze(0).to(device)
+model.to(device)
+response = model.generate(context, max_new_tokens=200, temperature=1.1, top_k=5)
+
+print(tokenizer.decode(response.squeeze().tolist()))
+```

config.json

@@ -0,0 +1,44 @@
+{
+  "context_length": 32768,
+  "dtype": "bfloat16",
+  "emb_dim": 640,
+  "finetuning_task": null,
+  "head_dim": 256,
+  "hidden_dim": 2048,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "n_heads": 4,
+  "n_kv_groups": 1,
+  "n_layers": 18,
+  "num_labels": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "output_past": true,
+  "pruned_heads": {},
+  "qk_norm": true,
+  "query_pre_attn_scalar": 256,
+  "rope_base": 1000000.0,
+  "rope_local_base": 10000.0,
+  "sliding_window": 512,
+  "torchscript": false,
+  "use_bfloat16": false,
+  "vocab_size": 50257
+}

gemma3_tinystories/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_gemma3 import HFGemma3DPONegative, Gemma3Config
+
+__all__ = ["HFGemma3DPONegative", "Gemma3Config"]

gemma3_tinystories/modeling_gemma3.py

@@ -0,0 +1,229 @@
+from transformers import PreTrainedModel, PretrainedConfig
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ------------------- Config -------------------
+class Gemma3Config(PretrainedConfig):
+    model_type = "gemma3"
+
+    def __init__(
+        self,
+        vocab_size=50257,
+        emb_dim=512,
+        n_layers=12,
+        n_heads=8,
+        n_kv_groups=2,
+        head_dim=None,
+        context_length=1024,
+        sliding_window=64,
+        rope_base=10000,
+        rope_local_base=10000,
+        layer_types=None,
+        qk_norm=False,
+        query_pre_attn_scalar=None,
+        dtype="float32",
+        hidden_dim=2048,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.emb_dim = emb_dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_groups = n_kv_groups
+        self.head_dim = head_dim
+        self.context_length = context_length
+        self.sliding_window = sliding_window
+        self.rope_base = rope_base
+        self.rope_local_base = rope_local_base
+        self.layer_types = layer_types or ["sliding_attention"] * n_layers
+        self.qk_norm = qk_norm
+        self.query_pre_attn_scalar = query_pre_attn_scalar
+        self.dtype = dtype
+        self.hidden_dim = hidden_dim
+
+# ------------------- Utilities -------------------
+def compute_rope_params(head_dim, theta_base=10_000, context_length=4096, dtype=torch.float32):
+    assert head_dim % 2 == 0, "Embedding dimension must be even"
+    inv_freq = 1.0 / (theta_base ** (torch.arange(0, head_dim, 2, dtype=dtype) / head_dim))
+    positions = torch.arange(context_length, dtype=dtype)
+    angles = positions[:, None] * inv_freq[None, :]
+    angles = torch.cat([angles, angles], dim=1)
+    return torch.cos(angles), torch.sin(angles)
+
+def apply_rope(x, cos, sin):
+    batch_size, num_heads, seq_len, head_dim = x.shape
+    x1, x2 = x[..., :head_dim//2], x[..., head_dim//2:]
+    cos = cos[:seq_len, :].unsqueeze(0).unsqueeze(0)
+    sin = sin[:seq_len, :].unsqueeze(0).unsqueeze(0)
+    rotated = torch.cat((-x2, x1), dim=-1)
+    return ((x * cos) + (rotated * sin)).to(dtype=x.dtype)
+
+class RMSNorm(nn.Module):
+    def __init__(self, emb_dim, eps=1e-6, bias=False):
+        super().__init__()
+        self.eps = eps
+        self.scale = nn.Parameter(torch.zeros(emb_dim))
+        self.shift = nn.Parameter(torch.zeros(emb_dim)) if bias else None
+
+    def forward(self, x):
+        input_dtype = x.dtype
+        x_f = x.float()
+        var = x_f.pow(2).mean(dim=-1, keepdim=True)
+        x_norm = x_f * torch.rsqrt(var + self.eps)
+        out = x_norm * (1.0 + self.scale.float())
+        if self.shift is not None:
+            out = out + self.shift.float()
+        return out.to(input_dtype)
+
+# ------------------- Core Layers -------------------
+class GroupedQueryAttention(nn.Module):
+    def __init__(self, d_in, num_heads, num_kv_groups, head_dim=None, qk_norm=False,
+                 query_pre_attn_scalar=None, dtype=torch.float32):
+        super().__init__()
+        if isinstance(dtype, str):
+            dtype = getattr(torch, dtype)
+        self.num_heads = num_heads
+        self.num_kv_groups = num_kv_groups
+        self.group_size = num_heads // num_kv_groups
+        if head_dim is None:
+            head_dim = d_in // num_heads
+        self.head_dim = head_dim
+        self.d_out = num_heads * head_dim
+
+        self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)
+        self.W_key = nn.Linear(d_in, num_kv_groups*head_dim, bias=False, dtype=dtype)
+        self.W_value = nn.Linear(d_in, num_kv_groups*head_dim, bias=False, dtype=dtype)
+        self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)
+
+        self.q_norm = RMSNorm(head_dim) if qk_norm else None
+        self.k_norm = RMSNorm(head_dim) if qk_norm else None
+        self.scaling = (query_pre_attn_scalar**-0.5 if query_pre_attn_scalar else head_dim**-0.5)
+
+    def forward(self, x, mask, cos, sin):
+        b, num_tokens, _ = x.shape
+        q = self.W_query(x).view(b, num_tokens, self.num_heads, self.head_dim).transpose(1,2)
+        k = self.W_key(x).view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1,2)
+        v = self.W_value(x).view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1,2)
+        if self.q_norm: q = self.q_norm(q)
+        if self.k_norm: k = self.k_norm(k)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        k, v = k.repeat_interleave(self.group_size, dim=1), v.repeat_interleave(self.group_size, dim=1)
+        q = q * self.scaling
+        attn_scores = q @ k.transpose(2,3)
+        attn_scores = attn_scores.masked_fill(mask, float('-inf'))
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        context = (attn_weights @ v).transpose(1,2).reshape(b, num_tokens, self.d_out)
+        return self.out_proj(context)
+
+class FeedForward(nn.Module):
+    def __init__(self, cfg):
+        dtype = cfg["dtype"]
+        if isinstance(dtype, str): dtype = getattr(torch, dtype)
+        super().__init__()
+        self.fc1 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False, dtype=dtype)
+        self.fc2 = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False, dtype=dtype)
+        self.fc3 = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False, dtype=dtype)
+    def forward(self, x):
+        return self.fc3(F.gelu(self.fc1(x), approximate="tanh") * self.fc2(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, cfg, attn_type):
+        super().__init__()
+        self.attn_type = attn_type
+        self.att = GroupedQueryAttention(cfg["emb_dim"], cfg["n_heads"], cfg["n_kv_groups"],
+                                        head_dim=cfg["head_dim"], qk_norm=cfg["qk_norm"],
+                                        query_pre_attn_scalar=cfg["query_pre_attn_scalar"],
+                                        dtype=cfg["dtype"])
+        self.ff = FeedForward(cfg)
+        self.input_layernorm = RMSNorm(cfg["emb_dim"])
+        self.post_attention_layernorm = RMSNorm(cfg["emb_dim"])
+        self.pre_feedforward_layernorm = RMSNorm(cfg["emb_dim"])
+        self.post_feedforward_layernorm = RMSNorm(cfg["emb_dim"])
+
+    def forward(self, x, mask_global, mask_local, cos_global, sin_global, cos_local, sin_local):
+        shortcut = x
+        x = self.input_layernorm(x)
+        attn_mask = mask_local if self.attn_type=="sliding_attention" else mask_global
+        cos, sin = (cos_local, sin_local) if self.attn_type=="sliding_attention" else (cos_global, sin_global)
+        x_attn = self.att(x, attn_mask, cos, sin)
+        x_attn = self.post_attention_layernorm(x_attn)
+        x = shortcut + x_attn
+        shortcut = x
+        x_ffn = self.pre_feedforward_layernorm(x)
+        x_ffn = self.ff(x_ffn)
+        x = shortcut + self.post_feedforward_layernorm(x_ffn)
+        return x
+
+# ------------------- Gemma3 Model -------------------
+class Gemma3Model(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        assert cfg["layer_types"] is not None and len(cfg["layer_types"])==cfg["n_layers"]
+        dtype = cfg["dtype"]
+        if isinstance(dtype, str): dtype = getattr(torch, dtype)
+        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"], dtype=dtype)
+        self.blocks = nn.ModuleList([TransformerBlock(cfg, t) for t in cfg["layer_types"]])
+        self.final_norm = RMSNorm(cfg["emb_dim"])
+        self.out_head = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False, dtype=dtype)
+        self.cfg = cfg
+        cos_l, sin_l = compute_rope_params(cfg["head_dim"], cfg["rope_local_base"], cfg["context_length"])
+        cos_g, sin_g = compute_rope_params(cfg["head_dim"], cfg["rope_base"], cfg["context_length"])
+        self.register_buffer("cos_local", cos_l, persistent=False)
+        self.register_buffer("sin_local", sin_l, persistent=False)
+        self.register_buffer("cos_global", cos_g, persistent=False)
+        self.register_buffer("sin_global", sin_g, persistent=False)
+
+    def _create_masks(self, seq_len, device):
+        ones = torch.ones((seq_len, seq_len), dtype=torch.bool, device=device)
+        mask_global = torch.triu(ones, 1)
+        far_past = torch.triu(ones, diagonal=self.cfg["sliding_window"]).T
+        mask_local = mask_global | far_past
+        return mask_global, mask_local
+
+    def forward(self, input_ids, targets=None):
+        b, seq_len = input_ids.shape
+        x = self.tok_emb(input_ids) * (self.cfg["emb_dim"] ** 0.5)
+        mask_global, mask_local = self._create_masks(seq_len, x.device)
+        for block in self.blocks:
+            x = block(x, mask_global, mask_local, self.cos_global, self.sin_global, self.cos_local, self.sin_local)
+        x = self.final_norm(x)
+        logits = self.out_head(x.to(getattr(torch, self.cfg["dtype"])))
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        for _ in range(max_new_tokens):
+            ctx_len = self.cfg["context_length"]
+            idx_cond = idx if idx.size(1) <= ctx_len else idx[:, -ctx_len:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float('-inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx
+
+# ------------------- Hugging Face Wrapper -------------------
+class HFGemma3DPONegative(PreTrainedModel):
+    config_class = Gemma3Config
+    def __init__(self, config):
+        super().__init__(config)
+        cfg_dict = config.to_dict()
+        if isinstance(cfg_dict.get("dtype"), str):
+            cfg_dict["dtype"] = cfg_dict["dtype"]
+        self.model = Gemma3Model(cfg_dict)
+
+    def forward(self, *args, **kwargs):
+        return self.model(*args, **kwargs)
+    
+    @torch.no_grad()
+    def generate(self, *args, **kwargs):
+        return self.model.generate(*args, **kwargs)
+

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ecfb01576a4b1228338d9d6f8e7c60a70df2824b30fe306b189f38153ac21639
+size 658711843

setup.py

@@ -0,0 +1,23 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="gemma3_tinystories_dpo",  
+    version="0.1.1",               
+    author="Shubham Waghmare",
+    description="DPO-aligned Gemma3-270M for negative sentiment narratives",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    url="https://huggingface.co/Shubhamw11/gemma-3-270m-dpo-negative", 
+    packages=find_packages(),
+    install_requires=[
+        "torch",
+        "transformers",
+        "tiktoken",
+    ],
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires='>=3.8',
+)