Create README.md

README.md

@@ -0,0 +1,368 @@
+---
+license: apache-2.0
+---
+
+# instinct.4B
+
+**[autonomousX/instinct-1.4B on Hugging Face](https://huggingface.co/autonomousX)**
+
+**[Training scripts and Logs on GITHUB](https://github.com/YADAV1825/AutonomousX-Instinct)**
+
+*instinct-1.4B is a fully reproducible, from-scratch trained 1.4B parameter language model trained on 70B tokens using TPU v4 infrastructure.*
+
+**instinct-1.4B** is a 1.4B parameter Large Language Model built entirely from scratch under the **AutonomousX** organization. 
+
+Compute for this project was supported by **[Google's TRC Program (TPU Research Cloud)](https://sites.research.google/trc/about/)**.
+
+---
+
+### 👨‍💻 Author Information
+**Rohit Yadav** B.Tech 3rd Year  
+Dr. B.R. Ambedkar National Institute of Technology (NIT) Jalandhar, India  
+**E-mail:** [yrohit1825@gmail.com](mailto:yrohit1825@gmail.com)  
+**LinkedIn:** [Rohit Yadav](https://www.linkedin.com/in/rohit-yadav-25535b256/)  
+**GitHub:** [YADAV1825](https://github.com/YADAV1825)
+
+**Research interests include:** Large Language Models, MultiModal Pipelines, Systems Programming, AI Infrastructure, Distributed Training.
+
+---
+# About AutonomousX
+
+AutonomousX focuses on open-source contributions aimed at building Large Language Models from scratch using custom training pipelines. Our work explores different training configurations including optimizers, datasets, and scalable TPU training using JAX and pmap. The goal is to provide transparent and reproducible implementations so that researchers, students, and developers can understand how modern LLMs are trained end-to-end.
+
+Due to the current scarcity of complete beginner-friendly guides for training LLMs on TPUs, especially using JAX, AutonomousX aims to bridge this gap by publishing full training pipelines, scripts, and documentation for the open-source community.
+
+Maintained by: Rohit Yadav | B.Tech NIT Jalandhar | yrohit1825@gmail.com | [Hugging_Face](https://huggingface.co/autonomousX)
+
+---
+
+### ⚠️ Disclaimer
+**This is a base model, not an SFT (Supervised Fine-Tuned) or RLHF (Reinforcement Learning from Human Feedback) model.** As a raw completion model, it may output undesired, biased, or nonsensical text. It is intended primarily for research and educational purposes.
+
+---
+
+### 📊 Model Overview
+
+| Attribute | Value |
+| :--- | :--- |
+| **Model Name** | instinct-1.4B |
+| **Organization** | AutonomousX |
+| **Parameters** | 1.4B |
+| **Vocabulary Size** | 50,304 |
+| **Dataset** | DOLMA |
+| **Tokenizer** | Pythia Tokenizer / BPE |
+| **Tokens Seen** | 70B |
+| **Training Hardware** | TPU v4-8 |
+| **Optimizer** | Lion |
+| **Architecture** | v-4 (128) pmap |
+| **Positional Embeddings**| RoPE |
+
+---
+
+### 🧠 Training Details
+
+instinct-1.4B was trained completely **from scratch** using **JAX/Flax on TPU v4-8 hardware**. 
+
+The training pipeline includes:
+* Dataset streaming from **DOLMA**.
+* **Pythia Tokenizer** with a 50,304 vocabulary size.
+* TPU optimized **JAX / Flax** training loop.
+* **Lion** optimizer for stable convergence.
+* Checkpointing and validation during training.
+* Rolling validation shard evaluation.
+
+#### 🔄 Training Workflow
+
+```mermaid
+graph TD
+    %% Dataset and Preparation
+    Data["Dataset: DOLMA\nRaw Text Data"]
+    Tokenizer["BPE Tokenizer\nVocabulary Construction"]
+    TokenizedData["Tokenized Data\nReady for Training"]
+
+    %% Model Architecture
+    Model["v-4 (128) pmap\nTransformer Decoder\n1.4B Parameters"]
+    RoPE["RoPE Positional Embeddings"]
+
+    %% Training Pipeline
+    Optimizer["Optimizer: Lion"]
+    ForwardPass["Forward Pass\nCompute Loss"]
+    BackwardPass["Backward Pass\nCompute Gradients"]
+    Update["Parameter Update"]
+
+    %% Logging and Checkpoints
+    Checkpoints["Model Checkpoints\nSaved up to 70B tokens"]
+    Logs["Training Logs\nLoss & Perplexity"]
+
+    %% Connections
+    Data --> Tokenizer
+    Tokenizer --> TokenizedData
+    TokenizedData --> ForwardPass
+    
+    Model --> ForwardPass
+    RoPE -.-> Model
+
+    ForwardPass --> BackwardPass
+    BackwardPass --> Optimizer
+    Optimizer --> Update
+    Update --> Model
+
+    Update --> Checkpoints
+    Update --> Logs
+```
+
+#### 📈 Training Curves
+The loss curves are saved in `training_log.txt` and `val_perplexity.txt`. Below is the visualization of the training progress:
+
+![image](https://cdn-uploads.huggingface.co/production/uploads/68bf07a31d80a360f1405b72/sWC0DlPjaqPoM6dsqNrA1.png)
+
+---
+
+### 🔄 Reproducibility
+
+The entire pipeline used to train the model is fully reproducible. This includes the dataset pipeline, tokenizer creation, model architecture, TPU training loop, and checkpointing system.
+
+**Full training pipeline repository:** [train.py](https://github.com/YADAV1825/instinct.4B/blob/main/train.py) (Example link format)
+
+---
+
+### 🚀 Run Inference (Colab TPU/GPU)
+
+The trained LLM inference script and model weights are available at: [autonomousX/instinct-1.4B on Hugging Face](https://huggingface.co/autonomousX).
+
+⚠️ **Disclaimer:** *This is a sample inference script originally written for the 0.5B (40B) variants. Please adjust parameters like `N_LAYERS`, `D_MODEL`, `N_HEADS`, `D_HEAD`, `D_FF`, etc., according to this specific model's architecture (1.4B) to run inference successfully.*
+
+A ready-to-run Google Colab TPU/GPU inference script is provided below. Simply open a notebook, set your runtime to TPU or GPU, and run it. *(Please be patient, it may take around 20 mins to run the model initialization).*
+
+<details>
+<summary>Click here to view the full Inference Code</summary>
+
+```python
+#please be patient It may take 20 mins to run the model
+# Install huggingface_hub if not installed
+!pip install -q huggingface_hub
+
+from huggingface_hub import snapshot_download
+
+repo_id = "autonomousX/instinct-1.4B"
+
+# Download entire repository
+local_path = snapshot_download(
+    repo_id=repo_id,
+    repo_type="model",
+    local_dir="TPU_1.4B",
+    local_dir_use_symlinks=False
+)
+
+print("Download complete!")
+print("Saved to:", local_path)
+
+# =========================
+# FAST 1.4B INFERENCE CELL
+# =========================
+
+import os
+import math
+import jax
+import jax.numpy as jnp
+from flax import linen as nn
+from flax.training import train_state, checkpoints
+import optax
+from transformers import AutoTokenizer
+
+# ---------------- CONFIG ----------------
+SEQ_LEN = 1024
+VOCAB_SIZE = 50304
+
+# NOTE: Adjust these parameters for your specific 1.4B architecture!
+N_LAYERS = 32
+D_MODEL = 1024
+N_HEADS = 16
+D_HEAD = 64
+D_FF = 4096
+ROTARY_PCT = 0.25
+
+CKPT_PATH = os.path.abspath("TPU_1.4B/checkpoint_0")
+
+# ---------------- RoPE ----------------
+def build_rope_cache(seq_len, head_dim, rotary_pct):
+    dim = int(head_dim * rotary_pct)
+    freqs = 1.0 / (10000 ** (jnp.arange(0, dim, 2) / dim))
+    pos = jnp.arange(seq_len)
+    angles = jnp.einsum("i,j->ij", pos, freqs)
+    return jnp.sin(angles), jnp.cos(angles)
+
+ROPE_SIN, ROPE_COS = build_rope_cache(SEQ_LEN, D_HEAD, ROTARY_PCT)
+
+def apply_rope(q, k):
+    dim = int(D_HEAD * ROTARY_PCT)
+    T = q.shape[1]
+
+    sin = ROPE_SIN[:T][None, :, None, :]
+    cos = ROPE_COS[:T][None, :, None, :]
+
+    q_rot, q_pass = q[..., :dim], q[..., dim:]
+    k_rot, k_pass = k[..., :dim], k[..., dim:]
+
+    q1, q2 = q_rot[..., ::2], q_rot[..., 1::2]
+    k1, k2 = k_rot[..., ::2], k_rot[..., 1::2]
+
+    q_rot = jnp.concatenate(
+        [q1 * cos - q2 * sin,
+         q1 * sin + q2 * cos],
+        axis=-1
+    )
+
+    k_rot = jnp.concatenate(
+        [k1 * cos - k2 * sin,
+         k1 * sin + k2 * cos],
+        axis=-1
+    )
+
+    return (
+        jnp.concatenate([q_rot, q_pass], axis=-1),
+        jnp.concatenate([k_rot, k_pass], axis=-1),
+    )
+
+# ---------------- MODEL ----------------
+class RMSNorm(nn.Module):
+    dim: int
+    eps: float = 1e-6
+    @nn.compact
+    def __call__(self, x):
+        scale = self.param("scale", nn.initializers.ones, (self.dim,))
+        norm = jnp.sqrt(jnp.mean(x**2, axis=-1, keepdims=True) + self.eps)
+        return x * (scale / norm)
+
+class Attention(nn.Module):
+    @nn.compact
+    def __call__(self, x, mask):
+        B, T, C = x.shape
+        qkv = nn.Dense(3 * C, use_bias=False, dtype=jnp.bfloat16)(x)
+        qkv = qkv.reshape(B, T, 3, N_HEADS, D_HEAD)
+
+        q = qkv[:, :, 0]
+        k = qkv[:, :, 1]
+        v = qkv[:, :, 2]
+
+        q, k = apply_rope(q, k)
+
+        att = jnp.einsum("bthd,bshd->bhts", q, k)
+        att = att / math.sqrt(D_HEAD)
+
+        mask = mask.astype(jnp.float32)
+        mask = (1.0 - mask) * -1e10
+        att = att + mask
+
+        att = nn.softmax(att.astype(jnp.float32), axis=-1)
+        att = att.astype(jnp.bfloat16)
+
+        out = jnp.einsum("bhts,bshd->bthd", att, v)
+        out = out.reshape(B, T, C)
+
+        return nn.Dense(C, use_bias=False, dtype=jnp.bfloat16)(out)
+
+class Block(nn.Module):
+    @nn.compact
+    def __call__(self, x, mask):
+        h = RMSNorm(D_MODEL)(x)
+        h = Attention()(h, mask)
+        x = x + h
+
+        h = RMSNorm(D_MODEL)(x)
+        h = nn.Dense(D_FF, dtype=jnp.bfloat16)(h)
+        h = nn.gelu(h)
+        h = nn.Dense(D_MODEL, dtype=jnp.bfloat16)(h)
+
+        return x + h
+
+class GPT(nn.Module):
+    @nn.compact
+    def __call__(self, input_ids):
+        batch, seq_len = input_ids.shape
+        mask = nn.attention.make_causal_mask(
+            jnp.ones((batch, seq_len), dtype=jnp.bool_)
+        )
+
+        x = nn.Embed(
+            VOCAB_SIZE,
+            D_MODEL,
+            embedding_init=nn.initializers.normal(0.02),
+            dtype=jnp.bfloat16,
+        )(input_ids)
+
+        RematBlock = nn.remat(Block)
+
+        for _ in range(N_LAYERS):
+            x = RematBlock()(x, mask)
+
+        x = RMSNorm(D_MODEL)(x)
+
+        return nn.Dense(
+            VOCAB_SIZE,
+            use_bias=False,
+            dtype=jnp.bfloat16
+        )(x)
+# ---------------- LOAD CHECKPOINT ----------------
+def create_state():
+    model = GPT()
+    rng = jax.random.PRNGKey(0)
+    params = model.init(rng, jnp.ones((1, SEQ_LEN), dtype=jnp.int32))
+    return train_state.TrainState.create(
+        apply_fn=model.apply,
+        params=params,
+        tx=optax.adamw(1e-4),
+    )
+
+state = create_state()
+state = checkpoints.restore_checkpoint(CKPT_PATH, state)
+params = state.params
+model = GPT()
+
+print("Checkpoint loaded.")
+
+@jax.jit
+def forward(params, input_ids):
+    return model.apply(params, input_ids)
+
+import jax.random as random
+
+def generate(params, input_ids, max_new_tokens=30, temperature=0.9, top_k=40):
+    rng = random.PRNGKey(0)
+
+    for _ in range(max_new_tokens):
+
+        logits = model.apply(params, input_ids)
+        logits = logits[:, -1, :]
+        logits = logits.astype(jnp.float32)
+
+        logits = logits / temperature
+
+        top_k_logits, top_k_indices = jax.lax.top_k(logits, top_k)
+        probs = jax.nn.softmax(top_k_logits, axis=-1)
+
+        rng, subkey = random.split(rng)
+        next_token_idx = random.categorical(subkey, jnp.log(probs))
+
+        next_token = jnp.take_along_axis(
+            top_k_indices,
+            next_token_idx[:, None],
+            axis=-1
+        )
+
+        input_ids = jnp.concatenate([input_ids, next_token], axis=1)
+
+    return input_ids
+# ---------------- RUN ----------------
+tokenizer = AutoTokenizer.from_pretrained("autonomousX/instinct-1.4B")
+
+prompt = "I am John,"
+tokens = tokenizer(prompt, return_tensors="np")
+input_ids = jnp.array(tokens["input_ids"], dtype=jnp.int32)
+
+output_ids = generate(params, input_ids, 200)
+
+print("\n=== GENERATED TEXT ===\n")
+print(tokenizer.decode(output_ids[0].tolist()))
+```
+</details>