| --- |
| license: mit |
| datasets: |
| - Skylion007/openwebtext |
| language: |
| - en |
| metrics: |
| - perplexity |
| --- |
| ## Model Architecture & Training |
|
|
| Both models are **Decoder-only Transformers** (GPT-style) implemented in pure PyTorch. |
|
|
| | Hyperparameters | Main Model (Target) | Draft Model (Small) | Draft Model (Medium) | |
| | :--- | :--- | :--- | :--- | |
| | **Parameters** | ~150M | ~30M | ~70M | |
| | **Layers** | 12 | 2 | 6 | |
| | **Heads** | 12 | 4 | 8 | |
| | **Embedding Dim** | 768 | 256 | 512 | |
| | **Context Length** | 1024 | 1024 | 1024 | |
| | **Vocab Size** | 50304 | 50304 | 50304 | |
| | **Droupout** | 0.1 | 0.1 | 0.1 | |
| | **Dataset** | OpenWebText (Sample) | OpenWebText (Sample) | OpenWebText (Sample) | |
|
|
| ### Key Components |
|
|
| 1. **Token + Position Embeddings** |
| - Learned token embedding table: `(vocab_size, n_embd)` |
| - Rotary Position Embedding: parameter-free sinusoidal rotations on attention queries and keys |
|
|
| 2. **Decoder Blocks** |
| Each block consists of: |
| - **Masked Self-Attention** (MHA/MQA/GQA) |
| - **Key-Value (KV) Caching** for *O(1)* generation latency |
| - **Rotary Positional Embedding** to understand relative positions |
| - **Feed-Forward Network** (Linear → GELU → Linear → Dropout) |
| - Residual connection |
| - Pre-norm **RMSNorm** |
|
|
| 3. **Final RMSNorm** + **Language Modeling Head** |
| - Linear projection from `n_embd` → `vocab_size` to produce logits |
|
|
| ### Training Strategy |
|
|
| #### Main Model |
| The main model was trained in 4 distinct phases to hande hardware constraints and optimize convergence |
|
|
| | Phases | Focus | Max Iterations | Warmup Steps | Eval Iterations | Eval Interval | Accumulation Steps | Learning Rate | Weight Decay | Batch Size | |
| | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | |
| | **Phase 1** | **Initial Warmup** | 10_000 | 200 | 20 | 500 | 32 | 3e-4 | 0.1 | 8 | |
| | **Phase 2** | **Main Pre-Training** | 30_000 | 0 | 20 | 1_000 | 32 | 3e-4 | 0.15 | 8 | |
| | **Phase 3** | **Large-Batch Scaling** | 50_000 | 1_000 | 10 | 2_000 | 16 | 6e-5 | 0.1 | 16 | |
| | **Phase 4** | **Convergence** | 40_000 | 500 | 10 | 2_000 | 16 | 2e-5 | 0.08 | 16 | |
|
|
| #### Draft Model (Small) |
| The draft model (small) was trained only once |
|
|
| ```python |
| max_iters = 40_000 |
| warmup_steps = 1_000 |
| eval_iter = 20 |
| eval_interval = 1_000 |
| accumulation_steps = 16 |
| base_lr = 3e-4 |
| weight_decay=0.1 |
| batch_size = 16 |
| ``` |
|
|
| #### Draft Model (Medium) |
| The draft model (medium) was trained only once |
|
|
| ```python |
| max_iters = 40_000 |
| warmup_steps = 2_000 |
| eval_iters = 20 |
| eval_interval = 2_000 |
| accumulation_steps = 16 |
| base_lr = 3e-4 |
| weight_decay = 0.1 |
| ``` |
