Apex / ARCHITECTURE.md

Upload folder using huggingface_hub

1e4fb05 verified 2 days ago

14.1 kB

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.

Overview & High-Level Philosophy
Model Architecture Specification
Pre-Training Pipeline (Foundation Model)
Supervised Fine-Tuning (SFT) Alignment
Inference Engine & Multi-Turn Chat Memory
Dual-Layer Safety Guardrail Subsystem
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:

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.
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.
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.
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:

### 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:

Temperature Scaling (T=0.7): Softens the logit distribution: $\hat{z}_i = z_i / T$.
Repetition Penalty (alpha=1.2): Penalizes already generated tokens.
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.
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]

Layer 1 (Pre-Generation Filter): Evaluates input against high-risk categories via exact heuristics and unitary/toxic-bert.
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).

yadavkapil7155
/

Apex