# Architecture & System Design: Apex-64M Small Language Model (SLM)

Comprehensive technical documentation of the **63.8M Parameter Apex Small Language Model**, covering model architecture, pre-training token pipelines, Supervised Fine-Tuning (SFT) alignment, multi-turn memory inference, and guardrail subsystems.

---

## Table of Contents
1. [Overview & High-Level Philosophy](#1-overview--high-level-philosophy)
2. [Model Architecture Specification](#2-model-architecture-specification)
3. [Pre-Training Pipeline (Foundation Model)](#3-pre-training-pipeline-foundation-model)
4. [Supervised Fine-Tuning (SFT) Alignment](#4-supervised-fine-tuning-sft-alignment)
5. [Inference Engine & Multi-Turn Chat Memory](#5-inference-engine--multi-turn-chat-memory)
6. [Dual-Layer Safety Guardrail Subsystem](#6-dual-layer-safety-guardrail-subsystem)
7. [Parameter Accounting & Memory Footprint](#7-parameter-accounting--memory-footprint)

---

## 1. Overview & High-Level Philosophy

The objective of this project is to build an end-to-end generative AI assistant from scratch:
* **Zero pre-trained weight dependencies** during pre-training (trained from random weight initialization N(0, 0.02)).
* Optimized parameter scale (**63.82M parameters**) capable of training on consumer/free cloud hardware (Tesla T4 GPU) while running blazingly fast on standard laptop CPUs (<300 MB RAM).
* Complete 2-stage lifecycle: **Unsupervised Pre-Training** (C4 English corpus -> Base Model) followed by **Supervised Fine-Tuning** (Databricks Dolly 15k -> Chat Assistant).

```
+-----------------------------------------------------------------------------------+
|                                 FULL SYSTEM PIPELINE                              |
+-----------------------------------------------------------------------------------+
|                                                                                   |
|  [C4 English Dataset]                                                             |
|           |                                                                       |
|           v (tiktoken GPT-2 Tokenization & Continuous Packing)                    |
|  [Packed Memmap Cache (train.bin)]                                                |
|           |                                                                       |
|           v (Stage 1: Pre-Training ~60,000 steps, Cosine LR, FP16 AMP)            |
|  [Apex-64M Base Foundation Model: best_model.pt (~63.8M Params)]                  |
|           |                                                                       |
|           v (Stage 2: SFT Alignment on Dolly 15k, Prompt Masking, 1-Token Shift)  |
|  [Apex-64M Chat Assistant: sft_model.pt]                                          |
|           |                                                                       |
|           v (Inference Layer: Top-P, Repetition Penalty, Chat Memory)             |
|  [Safety Layer: Toxic-BERT + Keyword Filtering]                                   |
|           |                                                                       |
|           v                                                                       |
|  [End-User Multi-Turn Chat / Gradio Web UI]                                       |
+-----------------------------------------------------------------------------------+
```

---

## 2. Model Architecture Specification

The model is a **Decoder-Only Pre-LayerNorm Transformer** built in pure PyTorch with fused attention kernels.

```
                           +------------------------+
                           |   Input Token IDs (x)  |
                           +------------------------+
                                       |
                   +-------------------+-------------------+
                   |                                       |
                   v                                       v
         +--------------------+                  +--------------------+
         | Token Embed (wte)  |                  |  Pos Embed (wpe)   |
         |  [50257, 512]      |                  |    [512, 512]      |
         +--------------------+                  +--------------------+
                   |                                       |
                   +-------------------+-------------------+
                                       |
                                       v
                             +--------------------+
                             |  Dropout (p=0.1)   |
                             +--------------------+
                                       |
                                       v
                 +============================================+
                 |       Transformer Block (x 12 Layers)      |
                 |                                            |
                 |   +------------------------------------+   |
                 |   |          LayerNorm (ln_1)          |   |
                 |   +------------------------------------+   |
                 |                     |                      |
                 |                     v                      |
                 |   +------------------------------------+   |
                 |   | Causal Multi-Head Self-Attention   |   |
                 |   |  (8 Heads, d_k=64, FlashAttention) |   |
                 |   +------------------------------------+   |
                 |                     |                      |
                 |                     v                      |
                 |             [ + Residual Add ]             |
                 |                     |                      |
                 |                     v                      |
                 |   +------------------------------------+   |
                 |   |          LayerNorm (ln_2)          |   |
                 |   +------------------------------------+   |
                 |                     |                      |
                 |                     v                      |
                 |   +------------------------------------+   |
                 |   |  MLP (Linear 512 -> 2048 -> 512)   |   |
                 |   |  Activation: New-GELU Approximation|   |
                 |   +------------------------------------+   |
                 |                     |                      |
                 |                     v                      |
                 |             [ + Residual Add ]             |
                 +============================================+
                                       |
                                       v
                             +--------------------+
                             | Final LayerNorm    |
                             |      (ln_f)        |
                             +--------------------+
                                       |
                                       v
                             +--------------------+
                             |   LM Output Head   |
                             | (Tied Weights wte) |
                             +--------------------+
                                       |
                                       v
                             +--------------------+
                             | Logits [B, T, V]   |
                             +--------------------+
```

### Key Architectural Choices:
1. **Pre-Layer Normalization (Pre-LN):** LayerNorm is applied *before* the Multi-Head Attention and MLP sub-layers (rather than Post-LN as in original GPT). This provides stable gradient flow and eliminates warmup instability.
2. **PyTorch Fused Flash Attention (`F.scaled_dot_product_attention`):** Uses FlashAttention CUDA kernels when running on GPU, cutting memory complexity from O(N^2) to O(N) and boosting speed by 3x.
3. **Weight Tying (Weight Sharing):** The output projection matrix `lm_head.weight` is tied directly to `wte.weight` (50,257 x 512). This saves **25.73 Million parameters** (~40% of the total model parameter budget) while improving semantic convergence.
4. **New-GELU Activation:** Uses the accurate Gaussian Error Linear Unit approximation inside the feed-forward network.

---

## 3. Pre-Training Pipeline (Foundation Model)

### Tokenization & Data Packing:
* **Tokenizer:** OpenAI `tiktoken` BPE tokenizer (`gpt2` / `r50k_base`) with vocabulary size 50,257.
* **Continuous Token Packing:** Documents from C4 English are streamed, tokenized, and concatenated with the `<|endoftext|>` (50,256) separator into continuous `uint16` memory-mapped arrays (`train.bin`, `val.bin`).
* **Zero Padding Waste:** In pre-training, sequences are sliced into exact 512-token chunks with zero padding tokens.

### Target Alignment:
In `dataset.py`:
$$\mathbf{x} = \text{chunk}[0 : T-1], \quad \mathbf{y} = \text{chunk}[1 : T]$$
Position $t$ predicts token $t+1$.

### Pre-Training Hyperparameters:
* **Target Token Budget:** ~1.28 Billion tokens
* **Effective Batch Size:** 32 (8 micro-batch size x 4 gradient accumulation steps)
* **Optimizer:** AdamW (`beta1 = 0.9`, `beta2 = 0.95`, `weight_decay = 0.01`)
* **Precision:** FP16 Automatic Mixed Precision (`torch.autocast`) with `GradScaler`
* **Learning Rate Schedule:** Cosine decay from `6e-4` to `6e-5` across 60,000 steps.
* **Final Pre-Training Loss:** **`~4.02`** (Validation Perplexity ≈ 55.7).

---

## 4. Supervised Fine-Tuning (SFT) Alignment

To transform the raw completion model into an instructional chat assistant, the model is fine-tuned on **Databricks Dolly 15k** (15,011 human-written instruction pairs).

### Prompt Formatting Template:
```text
### Instruction:
{instruction}

### Context:
{optional context}

### Response:
{assistant response}<|endoftext|>
```

### SFT Loss Masking (Prompt Masking):
Standard causal language modeling calculates loss on the entire sequence. However, calculating loss on the user prompt teaches the model to memorize user questions rather than learning how to respond.

We apply **selective loss masking** with `ignore_index = -1`:
$$\mathbf{y}_{\text{target}}[t] = \begin{cases} -1 & \text{if } t \text{ is in Prompt} \\ \text{token ID} & \text{if } t \text{ is in Response} \\ -1 & \text{if } t \text{ is in Padding} \end{cases}$$

$$\mathcal{L} = -\frac{1}{N_{\text{resp}}} \sum_{t \in \text{Response}} \log P(y_t \mid x_{\le t})$$

### The 1-Token Autoregressive Shift:
To maintain autoregressive causal next-token alignment:
$$\mathbf{x}_{\text{input}} = \text{tokens}[0 : N-1], \quad \mathbf{y}_{\text{target}} = \text{labels}[1 : N]$$
The last prompt token (`### Response:\n`) at index $t$ directly targets the first token of the assistant's answer at index $t+1$.

### SFT Training Parameters:
* **Epochs:** 3 Epochs (1,326 optimizer steps)
* **Learning Rate:** `2e-5` with Cosine Annealing to `2e-6` (Conservative rate prevents catastrophic forgetting)
* **Dataset Filter:** Slices exceeding 512 tokens skipped to avoid mid-sentence cutoff (14,122 clean samples preserved).
* **Final SFT Loss:** **`~3.17`**

---

## 5. Inference Engine & Multi-Turn Chat Memory

### Generation Sampling Strategy:
To balance creativity and coherence while avoiding repetitive loops:
1. **Temperature Scaling (T=0.7):** Softens the logit distribution: $\hat{z}_i = z_i / T$.
2. **Repetition Penalty (alpha=1.2):** Penalizes already generated tokens.
3. **Top-P (Nucleus) Filtering (P=0.90):** Retains the smallest set of top tokens whose cumulative probability $\ge 0.90$, truncating the low-probability tail.
4. **Early Termination Safeguards:**
   * Stops immediately on `<|endoftext|>` (50,256).
   * Stops if model generates follow-up prompt headers (`### Instruction:`).

### Multi-Turn Memory Management:
The inference engine implements a **Sliding Window Multi-Turn Buffer**:
* Maintains a structured dialogue history array `conversation_history`.
* Appends previous turns into a unified context prompt.
* If prompt tokens exceed 380, it dynamically trims older turns from the left, ensuring 130+ tokens remain available for response generation within the 512 context limit.

---

## 6. Dual-Layer Safety Guardrail Subsystem

The inference pipeline incorporates an asynchronous dual-layer guardrail:

```
[User Input]
     |
     v
[Layer 1: Input Keyword Heuristic + Toxic-BERT Classifier]
     |
     +---> If Unsafe ---> [Return Safe Rejection Message]
     |
     v (If Safe)
[Apex-64M Neural Generation]
     |
     v
[Layer 2: Output Toxicity Verification]
     |
     +---> If Unsafe ---> [Return Neutral Fallback Message]
     |
     v (If Safe)
[Delivered Assistant Response]
```

1. **Layer 1 (Pre-Generation Filter):** Evaluates input against high-risk categories via exact heuristics and `unitary/toxic-bert`.
2. **Layer 2 (Post-Generation Moderation):** Scans model output prior to display to ensure generated content meets safety thresholds.

---

## 7. Parameter Accounting & Memory Footprint

### Layer-by-Layer Parameter Breakdown:

| Component | Dimensions / Formula | Parameters | % of Model |
| :--- | :--- | :--- | :--- |
| **Token Embeddings (`wte`)** | 50,257 x 512 | 25,731,584 | 40.32% |
| **Positional Embeddings (`wpe`)** | 512 x 512 | 262,144 | 0.41% |
| **12 Transformer Blocks:** | | | |
| - *Attention QKV Projections* | 3 x (512 x 512 + 512) | 787,968 x 12 = 9,455,616 | 14.82% |
| - *Attention Output Projection* | 512 x 512 + 512 | 262,656 x 12 = 3,151,872 | 4.94% |
| - *LayerNorm 1 (`ln_1`)* | 2 x 512 | 1,024 x 12 = 12,288 | 0.02% |
| - *MLP FC1 (Up-Projection)* | 512 x 2048 + 2048 | 1,050,624 x 12 = 12,607,488 | 19.75% |
| - *MLP FC2 (Down-Projection)* | 2048 x 512 + 512 | 1,049,088 x 12 = 12,589,056 | 19.72% |
| - *LayerNorm 2 (`ln_2`)* | 2 x 512 | 1,024 x 12 = 12,288 | 0.02% |
| **Final LayerNorm (`ln_f`)** | 2 x 512 | 1,024 | <0.01% |
| **LM Output Head (`lm_head`)** | Tied with `wte` | *0 (Reused)* | 0.00% |
| **Total Parameters** | | **63,823,360** | **100.0%** |

### Memory Footprint:
* **FP32 Storage:** ~255.3 MB
* **FP16 Storage:** ~127.6 MB
* **INT8 Quantized:** ~63.8 MB
* **Runtime RAM Requirement:** <300 MB (Runs on any mobile device or CPU).