bobbydeveloper2014 commited on
Commit
be5cbce
·
verified ·
1 Parent(s): 4d74557

Upload README.md with huggingface_hub

Browse files
Files changed (1) hide show
  1. README.md +3 -201
README.md CHANGED
@@ -1,205 +1,7 @@
1
- ---
2
- license: apache-2.0
3
- language:
4
- - en
5
- pipeline_tag: text-generation
6
- tags:
7
- - hybrid-ssm
8
- - mha
9
- - moe
10
- - state-space-models
11
- - recurrent-gated-expert
12
- - arg
13
- - hse
14
- - b300
15
- - runpod
16
- - instruct
17
- - chat
18
- base_model: Phase-Technologies/forge-500m-base
19
- ---
20
-
21
- # FORGE-500M-Chat: First-Principles Optimal Recurrent-Gated Expert Architecture
22
-
23
- <div align="center">
24
-
25
- [![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
26
- [![Parameters](https://img.shields.io/badge/Parameters-500M%20Total%20%7C%20320M%20Active-green.svg)]()
27
- [![Hardware](https://img.shields.io/badge/Trained%20On-4%C3%97%20NVIDIA%20B300%20SXM6-76B900.svg?logo=nvidia)]()
28
- [![Architecture](https://img.shields.io/badge/Architecture-Hybrid%20SSM%20%2B%20MHA%20%2B%20HSE--MoE-orange.svg)]()
29
- [![Training Tokens](https://img.shields.io/badge/Tokens-1.66B%20Pretrain%20%2B%20Aligned-purple.svg)]()
30
-
31
- </div>
32
-
33
- ---
34
-
35
- ## Executive Summary
36
-
37
- **FORGE-500M-Chat** is a state-of-the-art hybrid sub-billion parameter conversational language model designed from first principles by **Soham Pal**. It solves the compute accessibility gap in AI research by replacing quadratic self-attention with linear-time recurrent state-space recurrence, paired with hierarchical sparse expert mixture-of-experts (HSE-MoE).
38
-
39
- The model was pretrained on **1.66 Billion tokens** from scratch across 4× NVIDIA B300 SXM6 MIG compute slices, aligned via **Supervised Fine-Tuning (SFT)** on multi-turn dialogue with per-token loss masking, and polished using **Direct Preference Optimization (DPO)**.
40
-
41
- The full console training trajectory from Step 1 to Step 821 is archived in this repository under [`training_journey.log`](./training_journey.log).
42
 
43
  ---
44
 
45
- ## Model Architecture Specifications
46
-
47
- FORGE employs an asymmetric, multi-layer hybrid topology designed for maximum parameter utility and linear sequence scalability:
48
-
49
- | Specification | Value | Technical Description |
50
- | :--- | :--- | :--- |
51
- | **Total Parameters** | **500,019,616** | 500.0M total weights (FP16/BF16 storage: ~1.0 GB) |
52
- | **Active Parameters** | **~320,000,000** | Active per-token forward pass via Sparse HSE-MoE |
53
- | **Embedding Size** | 211,444,736 | 206,464 vocabulary × 1,024 hidden dimension |
54
- | **Hidden Dimension ($d_{\text{model}}$)** | 1,024 | Intermediate representation width |
55
- | **Total Transformer Layers** | 18 | Interleaved hybrid recurrence + global attention |
56
- | **SSM Recurrence Layers** | 15 (83.3%) | **Adaptive Recurrent Gating (ARG)** linear state space |
57
- | **Attention Layers** | 3 (16.7%) | **Global Multi-Head Attention (MHA)** with RoPE + YaRN |
58
- | **Attention Heads** | 8 | 8 Query heads, 8 Key/Value heads |
59
- | **MoE FFN Layers** | 9 | **Hierarchical Sparse Expert (HSE)** (32 experts, top-2 active) |
60
- | **Dense FFN Layers** | 9 | Dense SwiGLU FFN with $8/3 \times d_{\text{model}}$ expansion |
61
- | **Context Window** | 4,096 | Native context (scalable to 16,384 with YaRN) |
62
- | **Tokenizer** | CRAYON Turbo | 206,464 vocabulary (BPE with subword byte-fallback) |
63
-
64
- ```
65
- Input Tokens (B, T)
66
- │
67
- ▼
68
- [CRAYON Embedding (206k Vocab, 1024 d_model)]
69
- │
70
- ┌────┴───────────────────────────────────────┐
71
- │ Layer 0..17: Interleaved Blocks │
72
- │ ├─ 15× ARG Layers: Linear SSM Recurrence │
73
- │ ├─ 3× Global MHA: Long-Range Attention │
74
- │ ├─ 9× SwiGLU Dense FFN │
75
- │ └─ 9× Hierarchical Sparse Expert (MoE) │
76
- └────┬───────────────────────────────────────┘
77
- ▼
78
- [Final RMSNorm + Chunked LM Projection Head]
79
- │
80
- ▼
81
- Target Logits / Probabilities
82
- ```
83
-
84
- ---
85
-
86
- ## The Full Training Journey & Historical Narrative
87
-
88
- The pretraining and post-training of FORGE-500M represents an empirical milestone in lean, high-efficiency AI development:
89
-
90
- ### 1. The Pretraining Stage
91
- - **Cluster**: 4× NVIDIA B300 SXM6 MIG slices (`1g.34gb`, 18 SMs each, CUDA 13.2, Driver 595.91).
92
- - **Global Batch**: 1,048,576 tokens per step (32 micro-batches × 32,768 tokens per batch).
93
- - **Physical Throughput**: 7,337 unique tokens per second (~26.4 Million tokens/hour).
94
- - **Step 1**: Initial Cross-Entropy Loss: **12.512**, Initial Perplexity: **> 150,000**.
95
- - **Step 821**: Final Loss: **4.6473**, Final Perplexity: **104.30**.
96
- - **Cumulative Pretrain Tokens**: **1,660,944,384 tokens** (1.66 Billion tokens).
97
- - **Challenges Overcome**:
98
- - Survived multiple cloud host network disruptions and PCIe MIG interface resets via resilient checkpoint auto-recovery.
99
- - Successfully resolved PyTorch 2.5 `nn.Module` dynamic attribute restrictions with custom parameter wrapping (`convert_to_attr_dict`).
100
- - Implemented chunked linear projection (`chunk_size=1024`) on the 206,464-token vocabulary to eliminate VRAM allocation spikes.
101
-
102
- ### 2. Supervised Fine-Tuning (SFT)
103
- - **Dataset**: Curated mixture of ShareGPT dialogues, UltraChat multi-turn instructions, CodeFeedback Python programming pairs, MetaMath mathematical reasoning, and OpenOrca synthetic data.
104
- - **Protocol**: Chat template formatting with assistant turn loss masking (no loss on system or user prompts).
105
- - **Optimization**: PyTorch DistributedDataParallel (DDP) across all 4 GPUs with cosine learning rate decay ($2\times 10^{-5} \to 2\times 10^{-6}$) and BF16 autocast.
106
- - **Result**: Loss dropped rapidly to **~4.0**, establishing clean turn-taking and instruction responsiveness.
107
-
108
- ### 3. Direct Preference Optimization (DPO)
109
- - **Dataset**: Multi-domain preference pairs from UltraFeedback, Anthropic HH-RLHF, and HelpSteer2.
110
- - **Protocol**: Trainable policy model aligned against frozen BF16 SFT reference model with KL divergence regularization ($\beta = 0.1$).
111
- - **Result**: Achieved positive reward margin on chosen vs. rejected completions, eliminating repetitive degradation and improving refusal of hazardous requests.
112
-
113
- ---
114
-
115
- ## Complete Training Console Logs Preserved
116
-
117
- In accordance with our commitment to total scientific reproducibility, the complete console log file (`training_journey.log`) containing every training step from Step 1 to Step 821, memory footprints, step durations, and loss curves is provided directly in the repository root:
118
-
119
- - **Filename**: [`training_journey.log`](./training_journey.log)
120
- - **Size**: ~5.9 MB
121
- - **Coverage**: Steps 1 through 821, SFT, and DPO alignment logs.
122
-
123
- ---
124
-
125
- ## How to Run Inference
126
-
127
- ### 1. Google Colab (1-Click Quickstart)
128
-
129
- In a Google Colab notebook (T4 / A100 GPU), paste and run this single cell:
130
-
131
- ```bash
132
- # 1. Clone repository & enter directory
133
- !git clone https://huggingface.co/Phase-Technologies/FORGE-500M-Chat
134
- %cd FORGE-500M-Chat
135
-
136
- # 2. Install dependencies
137
- !pip install -q torch einops triton transformers huggingface_hub
138
-
139
- # 3. Generate text immediately
140
- !python run_inference.py --prompt "<|bos|><|usr|>Explain how quantum computers differ from classical computers.<|asst|>" --max_new_tokens 128
141
- ```
142
-
143
- ---
144
-
145
- ### 2. Pure Python API (Colab / Jupyter / Local)
146
-
147
- ```python
148
- import sys
149
- import torch
150
- from pathlib import Path
151
- from huggingface_hub import snapshot_download
152
-
153
- # 1. Download model repository
154
- repo_dir = Path(snapshot_download(repo_id="Phase-Technologies/FORGE-500M-Chat"))
155
- sys.path.insert(0, str(repo_dir))
156
-
157
- from tokenizer.crayon_wrapper import ForgeTokenizer
158
- from config import ForgeModelConfig
159
- from model.forge_model import build_forge_500m
160
-
161
- tokenizer = ForgeTokenizer.from_pretrained(str(repo_dir / "tokenizer"))
162
- config = ForgeModelConfig.from_json(str(repo_dir / "model_config.json"))
163
- model = build_forge_500m(config)
164
-
165
- # 3. Load BF16 weights
166
- weights = torch.load(repo_dir / "pytorch_model.bin", map_location="cpu")
167
- model.load_state_dict({k: v.to(torch.bfloat16) for k, v in weights.items()})
168
- model = model.to("cuda" if torch.cuda.is_available() else "cpu").to(torch.bfloat16)
169
- model.eval()
170
-
171
- # 4. Formatted Chat Prompt
172
- prompt = "<|bos|><|usr|> Explain how quantum computers differ from classical computers in simple terms.<|asst|>"
173
-
174
- input_ids = torch.tensor([tokenizer.encode(prompt, add_bos=False, add_eos=False)], device=model.device)
175
-
176
- # 5. Autoregressive Generation
177
- with torch.inference_mode():
178
- output_ids = model.generate(
179
- input_ids=input_ids,
180
- max_new_tokens=128,
181
- temperature=0.7,
182
- top_p=0.9,
183
- eos_token_id=tokenizer.special_tokens["<|eos|>"],
184
- )
185
-
186
- print(tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True))
187
- ```
188
-
189
- ### 3. Special Conversation Tokens
190
-
191
- | Token | String | Function |
192
- | :--- | :--- | :--- |
193
- | **BOS** | `<|bos|>` | Beginning of sequence |
194
- | **EOS** | `<|eos|>` | End of sequence / turn termination |
195
- | **SYS** | `<|sys|>` | System instruction block prefix |
196
- | **USR** | `<|usr|>` | User prompt prefix |
197
- | **ASST** | `<|asst|>` | Assistant completion prefix |
198
- | **PAD** | `<|pad|>` | Sequence padding token |
199
-
200
- ---
201
-
202
- ## a note from the creator.....
203
-
204
- honestly i really wanted to train this beauty for full 15b tokens..... that was the real plan from day one to see how far this custom architecture could push..... but running multiple top tier gpus burns money so fast and i just ran completely out of budget..... literally had to stop pretraining early at 821 steps..... then quickly did sft and dpo right on the pod before shutting it down..... huge shoutout and love to runpod.io for providing the insane compute and infrastructure that made training this beast possible..... every single cent went into this little 500m model and seeing it come alive from scratch with custom cuda kernels and our own tokenizer was pure magic..... maybe one day with better funding we will train the full 3b or 15b token run..... until then i hope you guys enjoy playing with this model as much as i loved building it..... thank you so much to everyone who supported this crazy journey.....
205
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
 
2
  ---
3
 
4
+ ## Run with nimformers-framework
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
 
6
+ This model can be run with **nimformers-framework**: https://github.com/bobbyshop-vui/nimformers-framework
7
+ The FIR graph for it is published in [forge-500m-chat-fir](https://huggingface.co/bobbydeveloper2014/forge-500m-chat-fir).