forge-dl 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- forge_dl-0.1.0/LICENSE +21 -0
- forge_dl-0.1.0/PKG-INFO +207 -0
- forge_dl-0.1.0/README-pypi.md +159 -0
- forge_dl-0.1.0/README.md +588 -0
- forge_dl-0.1.0/forge/__init__.py +19 -0
- forge_dl-0.1.0/forge/autograd/__init__.py +20 -0
- forge_dl-0.1.0/forge/autograd/accelerate_backend.py +70 -0
- forge_dl-0.1.0/forge/autograd/engine.py +14 -0
- forge_dl-0.1.0/forge/autograd/fused_operations.py +120 -0
- forge_dl-0.1.0/forge/autograd/fusion.py +40 -0
- forge_dl-0.1.0/forge/autograd/grad_check.py +58 -0
- forge_dl-0.1.0/forge/autograd/mps_backend.py +52 -0
- forge_dl-0.1.0/forge/autograd/operations.py +1063 -0
- forge_dl-0.1.0/forge/compiler/__init__.py +23 -0
- forge_dl-0.1.0/forge/compiler/build.py +36 -0
- forge_dl-0.1.0/forge/compiler/codegen.py +99 -0
- forge_dl-0.1.0/forge/compiler/compiled_run.py +127 -0
- forge_dl-0.1.0/forge/compiler/fusion.py +62 -0
- forge_dl-0.1.0/forge/compiler/graph.py +108 -0
- forge_dl-0.1.0/forge/compiler/interpreter.py +109 -0
- forge_dl-0.1.0/forge/compiler/kernels.c +195 -0
- forge_dl-0.1.0/forge/dtype.py +26 -0
- forge_dl-0.1.0/forge/nn/__init__.py +23 -0
- forge_dl-0.1.0/forge/nn/layers.py +444 -0
- forge_dl-0.1.0/forge/nn/losses.py +122 -0
- forge_dl-0.1.0/forge/nn/module.py +92 -0
- forge_dl-0.1.0/forge/nn/parameter.py +13 -0
- forge_dl-0.1.0/forge/optim/__init__.py +8 -0
- forge_dl-0.1.0/forge/optim/optimizer.py +85 -0
- forge_dl-0.1.0/forge/serialization.py +60 -0
- forge_dl-0.1.0/forge/tensor.py +474 -0
- forge_dl-0.1.0/forge_dl.egg-info/PKG-INFO +207 -0
- forge_dl-0.1.0/forge_dl.egg-info/SOURCES.txt +36 -0
- forge_dl-0.1.0/forge_dl.egg-info/dependency_links.txt +1 -0
- forge_dl-0.1.0/forge_dl.egg-info/requires.txt +9 -0
- forge_dl-0.1.0/forge_dl.egg-info/top_level.txt +1 -0
- forge_dl-0.1.0/pyproject.toml +38 -0
- forge_dl-0.1.0/setup.cfg +4 -0
forge_dl-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Srihari Srinivasan
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
forge_dl-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,207 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: forge-dl
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A machine learning library and deep learning compiler written from scratch in Python
|
|
5
|
+
Author: Srihari Srinivasan
|
|
6
|
+
License: MIT License
|
|
7
|
+
|
|
8
|
+
Copyright (c) 2026 Srihari Srinivasan
|
|
9
|
+
|
|
10
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
11
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
12
|
+
in the Software without restriction, including without limitation the rights
|
|
13
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
14
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
15
|
+
furnished to do so, subject to the following conditions:
|
|
16
|
+
|
|
17
|
+
The above copyright notice and this permission notice shall be included in all
|
|
18
|
+
copies or substantial portions of the Software.
|
|
19
|
+
|
|
20
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
21
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
22
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
23
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
24
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
25
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
26
|
+
SOFTWARE.
|
|
27
|
+
|
|
28
|
+
Project-URL: Homepage, https://github.com/SrihariSr/Forge
|
|
29
|
+
Project-URL: Repository, https://github.com/SrihariSr/Forge
|
|
30
|
+
Keywords: machine-learning,autograd,compiler,from-scratch,neural-networks
|
|
31
|
+
Classifier: Development Status :: 3 - Alpha
|
|
32
|
+
Classifier: Intended Audience :: Developers
|
|
33
|
+
Classifier: Intended Audience :: Science/Research
|
|
34
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
35
|
+
Classifier: Programming Language :: Python :: 3
|
|
36
|
+
Classifier: Programming Language :: C
|
|
37
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
38
|
+
Requires-Python: >=3.10
|
|
39
|
+
Description-Content-Type: text/markdown
|
|
40
|
+
License-File: LICENSE
|
|
41
|
+
Provides-Extra: plots
|
|
42
|
+
Requires-Dist: matplotlib; extra == "plots"
|
|
43
|
+
Provides-Extra: bench
|
|
44
|
+
Requires-Dist: numpy; extra == "bench"
|
|
45
|
+
Provides-Extra: metal
|
|
46
|
+
Requires-Dist: pyobjc; extra == "metal"
|
|
47
|
+
Dynamic: license-file
|
|
48
|
+
|
|
49
|
+
# Forge
|
|
50
|
+
|
|
51
|
+
A machine learning library and deep learning compiler, written from scratch in Python.
|
|
52
|
+
|
|
53
|
+
No PyTorch, no TensorFlow, no NumPy in the core. Every operation, from the autograd engine to the attention mechanism, is implemented from first principles. On top of it sits a compiler that analyses computation graphs, fuses operations, and emits optimised C at run time.
|
|
54
|
+
|
|
55
|
+
Full documentation, derivations and benchmarks: **https://github.com/SrihariSr/Forge**
|
|
56
|
+
|
|
57
|
+
---
|
|
58
|
+
|
|
59
|
+
## Install
|
|
60
|
+
|
|
61
|
+
```bash
|
|
62
|
+
pip install forgeml
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
Python 3.10 or later. The core library has no dependencies.
|
|
66
|
+
|
|
67
|
+
Optional extras:
|
|
68
|
+
|
|
69
|
+
```bash
|
|
70
|
+
pip install "forgeml[plots]" # matplotlib, for the option pricing charts
|
|
71
|
+
pip install "forgeml[bench]" # numpy, for the compiler's BLAS comparison
|
|
72
|
+
pip install "forgeml[metal]" # pyobjc, for the Apple GPU backend
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
---
|
|
76
|
+
|
|
77
|
+
## Quick start
|
|
78
|
+
|
|
79
|
+
Tensors and automatic differentiation:
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
from forge import Tensor
|
|
83
|
+
|
|
84
|
+
a = Tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
85
|
+
b = Tensor([[5.0, 6.0], [7.0, 8.0]])
|
|
86
|
+
|
|
87
|
+
print(a + b) # element-wise addition
|
|
88
|
+
print(a @ b) # matrix multiplication
|
|
89
|
+
print(a.T) # transpose
|
|
90
|
+
|
|
91
|
+
x = Tensor([2.0, 3.0], requires_grad=True)
|
|
92
|
+
y = ((x * x) + x).sum()
|
|
93
|
+
y.backward()
|
|
94
|
+
print(x.grad) # dy/dx = 2x + 1
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
A transformer, built on that autograd engine:
|
|
98
|
+
|
|
99
|
+
```python
|
|
100
|
+
from forge.nn import GPT, CrossEntropyLoss
|
|
101
|
+
from forge.optim import Adam
|
|
102
|
+
|
|
103
|
+
model = GPT(vocab_size=65, embed_dim=128, num_heads=4,
|
|
104
|
+
ff_dim=256, num_layers=4, seq_len=32)
|
|
105
|
+
criterion = CrossEntropyLoss()
|
|
106
|
+
optimizer = Adam(model.parameters(), lr=0.002)
|
|
107
|
+
|
|
108
|
+
logits = model([[1, 2, 3, 4]])
|
|
109
|
+
loss = criterion(logits, [2, 3, 4, 5])
|
|
110
|
+
|
|
111
|
+
optimizer.zero_grad()
|
|
112
|
+
loss.backward()
|
|
113
|
+
optimizer.step()
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
---
|
|
117
|
+
|
|
118
|
+
## What is in it
|
|
119
|
+
|
|
120
|
+
**A differentiable tensor library.** Reverse-mode automatic differentiation on a define-by-run graph, a typed memory backend, and broadcasting that follows NumPy semantics without depending on NumPy. Every operation is checked against a numerical gradient computed by central finite differences.
|
|
121
|
+
|
|
122
|
+
**A neural network framework.** A `Module` system with automatic parameter registration, layers, optimisers and loss functions.
|
|
123
|
+
|
|
124
|
+
**A decoder-only transformer.** Multi-head causal self-attention, LayerNorm, GELU, learned positional encoding and pre-norm residual blocks, all built on the autograd engine above and gradient-checked individually.
|
|
125
|
+
|
|
126
|
+
**Hardware backends.** Matrix multiplication dispatches to the Apple Metal GPU, to Apple Accelerate BLAS on the CPU, or to a pure Python fallback, chosen by problem size.
|
|
127
|
+
|
|
128
|
+
**A deep learning compiler.** Builds a graph, plans which operations can share one pass over memory, generates C for each group, compiles it with gcc, and loads it back through ctypes while the program runs.
|
|
129
|
+
|
|
130
|
+
---
|
|
131
|
+
|
|
132
|
+
## The compiler
|
|
133
|
+
|
|
134
|
+
The C kernels are not built during installation, because they need a compiler and are architecture-specific. Build them once:
|
|
135
|
+
|
|
136
|
+
```bash
|
|
137
|
+
python -m forge.compiler.build
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
Importing `forge.compiler` before running that raises an error telling you to.
|
|
141
|
+
|
|
142
|
+
```python
|
|
143
|
+
from forge.compiler import placeholder, relu, compile_graph, run_compiled
|
|
144
|
+
|
|
145
|
+
a, b, c = placeholder("a"), placeholder("b"), placeholder("c")
|
|
146
|
+
graph = relu(a + b) * c
|
|
147
|
+
|
|
148
|
+
steps = compile_graph(graph) # analyse, fuse, generate C, compile
|
|
149
|
+
result = run_compiled(graph, steps, feeds)
|
|
150
|
+
```
|
|
151
|
+
|
|
152
|
+
For `relu(a + b) * c` the compiler writes this, and nothing else in the project wrote it:
|
|
153
|
+
|
|
154
|
+
```c
|
|
155
|
+
void fused_kernel(const float* in0, const float* in1, const float* in2,
|
|
156
|
+
float* out, int n){
|
|
157
|
+
for (int i = 0; i < n; i++) {
|
|
158
|
+
float t3 = (in0[i] + in1[i]);
|
|
159
|
+
float t4 = (t3 > 0.0f ? t3 : 0.0f);
|
|
160
|
+
out[i] = (t4 * in2[i]);
|
|
161
|
+
}
|
|
162
|
+
}
|
|
163
|
+
```
|
|
164
|
+
|
|
165
|
+
Three kernels become one. Three passes over memory become one. The intermediates live in registers and never reach main memory.
|
|
166
|
+
|
|
167
|
+
---
|
|
168
|
+
|
|
169
|
+
## Results
|
|
170
|
+
|
|
171
|
+
Measured on an Apple M4 Max. Every optimised path is verified to produce output identical to the unoptimised path before it is timed.
|
|
172
|
+
|
|
173
|
+
| Optimisation | Workload | Before | After | Speedup |
|
|
174
|
+
|---|---|---|---|---|
|
|
175
|
+
| Operator fusion | 16M elements | 8.0ms | 3.4ms | 2.35x |
|
|
176
|
+
| Blocked, threaded matmul | 768 x 768 | 348.0ms | 3.8ms | 91.5x |
|
|
177
|
+
| Both, 4-layer MLP | 256 x 256 | 43.0ms | 1.3ms | 33.6x |
|
|
178
|
+
|
|
179
|
+
How the matmul got there, with the same arithmetic throughout:
|
|
180
|
+
|
|
181
|
+
| Stage | Time | GFLOPS | This step | Cumulative |
|
|
182
|
+
|---|---|---|---|---|
|
|
183
|
+
| Naive triple loop | 350.6ms | 3 | | 1.0x |
|
|
184
|
+
| Loop reordering | 28.4ms | 32 | 12.3x | 12.3x |
|
|
185
|
+
| Cache blocking | 32.0ms | 28 | 0.89x | 11.0x |
|
|
186
|
+
| restrict | 32.0ms | 28 | 1.00x | 11.0x |
|
|
187
|
+
| Threading and tuning | 3.8ms | 238 | 8.4x | 92.0x |
|
|
188
|
+
|
|
189
|
+
Against OpenBLAS, called through NumPy on the same machine, at 768 x 768: 234 GFLOPS against 351, a gap of 1.5x.
|
|
190
|
+
|
|
191
|
+
A 550,977-parameter GPT trained on the tinyshakespeare corpus with data-parallel training across 8 CPU cores brought character-level cross-entropy from 3.66 to 1.44 over 33,113 steps.
|
|
192
|
+
|
|
193
|
+
---
|
|
194
|
+
|
|
195
|
+
## Limitations
|
|
196
|
+
|
|
197
|
+
The core library is pure Python, hence it is slower than a production framework.
|
|
198
|
+
|
|
199
|
+
The compiler handles the forward pass over six operations, with no autograd. Only matmul is threaded; the element-wise kernels are single-threaded, which is fine because they are memory-bound rather than compute-bound.
|
|
200
|
+
|
|
201
|
+
Metal and Accelerate backends require macOS and PyObjC. If not available, it defaults to pure Python.
|
|
202
|
+
|
|
203
|
+
---
|
|
204
|
+
|
|
205
|
+
Built by Srihari Srinivasan. MIT licensed.
|
|
206
|
+
|
|
207
|
+
The full write-up, including the mathematics behind each component and the bugs worth knowing about can be found at **https://github.com/SrihariSr/Forge**
|
|
@@ -0,0 +1,159 @@
|
|
|
1
|
+
# Forge
|
|
2
|
+
|
|
3
|
+
A machine learning library and deep learning compiler, written from scratch in Python.
|
|
4
|
+
|
|
5
|
+
No PyTorch, no TensorFlow, no NumPy in the core. Every operation, from the autograd engine to the attention mechanism, is implemented from first principles. On top of it sits a compiler that analyses computation graphs, fuses operations, and emits optimised C at run time.
|
|
6
|
+
|
|
7
|
+
Full documentation, derivations and benchmarks: **https://github.com/SrihariSr/Forge**
|
|
8
|
+
|
|
9
|
+
---
|
|
10
|
+
|
|
11
|
+
## Install
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
pip install forgeml
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
Python 3.10 or later. The core library has no dependencies.
|
|
18
|
+
|
|
19
|
+
Optional extras:
|
|
20
|
+
|
|
21
|
+
```bash
|
|
22
|
+
pip install "forgeml[plots]" # matplotlib, for the option pricing charts
|
|
23
|
+
pip install "forgeml[bench]" # numpy, for the compiler's BLAS comparison
|
|
24
|
+
pip install "forgeml[metal]" # pyobjc, for the Apple GPU backend
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## Quick start
|
|
30
|
+
|
|
31
|
+
Tensors and automatic differentiation:
|
|
32
|
+
|
|
33
|
+
```python
|
|
34
|
+
from forge import Tensor
|
|
35
|
+
|
|
36
|
+
a = Tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
37
|
+
b = Tensor([[5.0, 6.0], [7.0, 8.0]])
|
|
38
|
+
|
|
39
|
+
print(a + b) # element-wise addition
|
|
40
|
+
print(a @ b) # matrix multiplication
|
|
41
|
+
print(a.T) # transpose
|
|
42
|
+
|
|
43
|
+
x = Tensor([2.0, 3.0], requires_grad=True)
|
|
44
|
+
y = ((x * x) + x).sum()
|
|
45
|
+
y.backward()
|
|
46
|
+
print(x.grad) # dy/dx = 2x + 1
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
A transformer, built on that autograd engine:
|
|
50
|
+
|
|
51
|
+
```python
|
|
52
|
+
from forge.nn import GPT, CrossEntropyLoss
|
|
53
|
+
from forge.optim import Adam
|
|
54
|
+
|
|
55
|
+
model = GPT(vocab_size=65, embed_dim=128, num_heads=4,
|
|
56
|
+
ff_dim=256, num_layers=4, seq_len=32)
|
|
57
|
+
criterion = CrossEntropyLoss()
|
|
58
|
+
optimizer = Adam(model.parameters(), lr=0.002)
|
|
59
|
+
|
|
60
|
+
logits = model([[1, 2, 3, 4]])
|
|
61
|
+
loss = criterion(logits, [2, 3, 4, 5])
|
|
62
|
+
|
|
63
|
+
optimizer.zero_grad()
|
|
64
|
+
loss.backward()
|
|
65
|
+
optimizer.step()
|
|
66
|
+
```
|
|
67
|
+
|
|
68
|
+
---
|
|
69
|
+
|
|
70
|
+
## What is in it
|
|
71
|
+
|
|
72
|
+
**A differentiable tensor library.** Reverse-mode automatic differentiation on a define-by-run graph, a typed memory backend, and broadcasting that follows NumPy semantics without depending on NumPy. Every operation is checked against a numerical gradient computed by central finite differences.
|
|
73
|
+
|
|
74
|
+
**A neural network framework.** A `Module` system with automatic parameter registration, layers, optimisers and loss functions.
|
|
75
|
+
|
|
76
|
+
**A decoder-only transformer.** Multi-head causal self-attention, LayerNorm, GELU, learned positional encoding and pre-norm residual blocks, all built on the autograd engine above and gradient-checked individually.
|
|
77
|
+
|
|
78
|
+
**Hardware backends.** Matrix multiplication dispatches to the Apple Metal GPU, to Apple Accelerate BLAS on the CPU, or to a pure Python fallback, chosen by problem size.
|
|
79
|
+
|
|
80
|
+
**A deep learning compiler.** Builds a graph, plans which operations can share one pass over memory, generates C for each group, compiles it with gcc, and loads it back through ctypes while the program runs.
|
|
81
|
+
|
|
82
|
+
---
|
|
83
|
+
|
|
84
|
+
## The compiler
|
|
85
|
+
|
|
86
|
+
The C kernels are not built during installation, because they need a compiler and are architecture-specific. Build them once:
|
|
87
|
+
|
|
88
|
+
```bash
|
|
89
|
+
python -m forge.compiler.build
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
Importing `forge.compiler` before running that raises an error telling you to.
|
|
93
|
+
|
|
94
|
+
```python
|
|
95
|
+
from forge.compiler import placeholder, relu, compile_graph, run_compiled
|
|
96
|
+
|
|
97
|
+
a, b, c = placeholder("a"), placeholder("b"), placeholder("c")
|
|
98
|
+
graph = relu(a + b) * c
|
|
99
|
+
|
|
100
|
+
steps = compile_graph(graph) # analyse, fuse, generate C, compile
|
|
101
|
+
result = run_compiled(graph, steps, feeds)
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
For `relu(a + b) * c` the compiler writes this, and nothing else in the project wrote it:
|
|
105
|
+
|
|
106
|
+
```c
|
|
107
|
+
void fused_kernel(const float* in0, const float* in1, const float* in2,
|
|
108
|
+
float* out, int n){
|
|
109
|
+
for (int i = 0; i < n; i++) {
|
|
110
|
+
float t3 = (in0[i] + in1[i]);
|
|
111
|
+
float t4 = (t3 > 0.0f ? t3 : 0.0f);
|
|
112
|
+
out[i] = (t4 * in2[i]);
|
|
113
|
+
}
|
|
114
|
+
}
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
Three kernels become one. Three passes over memory become one. The intermediates live in registers and never reach main memory.
|
|
118
|
+
|
|
119
|
+
---
|
|
120
|
+
|
|
121
|
+
## Results
|
|
122
|
+
|
|
123
|
+
Measured on an Apple M4 Max. Every optimised path is verified to produce output identical to the unoptimised path before it is timed.
|
|
124
|
+
|
|
125
|
+
| Optimisation | Workload | Before | After | Speedup |
|
|
126
|
+
|---|---|---|---|---|
|
|
127
|
+
| Operator fusion | 16M elements | 8.0ms | 3.4ms | 2.35x |
|
|
128
|
+
| Blocked, threaded matmul | 768 x 768 | 348.0ms | 3.8ms | 91.5x |
|
|
129
|
+
| Both, 4-layer MLP | 256 x 256 | 43.0ms | 1.3ms | 33.6x |
|
|
130
|
+
|
|
131
|
+
How the matmul got there, with the same arithmetic throughout:
|
|
132
|
+
|
|
133
|
+
| Stage | Time | GFLOPS | This step | Cumulative |
|
|
134
|
+
|---|---|---|---|---|
|
|
135
|
+
| Naive triple loop | 350.6ms | 3 | | 1.0x |
|
|
136
|
+
| Loop reordering | 28.4ms | 32 | 12.3x | 12.3x |
|
|
137
|
+
| Cache blocking | 32.0ms | 28 | 0.89x | 11.0x |
|
|
138
|
+
| restrict | 32.0ms | 28 | 1.00x | 11.0x |
|
|
139
|
+
| Threading and tuning | 3.8ms | 238 | 8.4x | 92.0x |
|
|
140
|
+
|
|
141
|
+
Against OpenBLAS, called through NumPy on the same machine, at 768 x 768: 234 GFLOPS against 351, a gap of 1.5x.
|
|
142
|
+
|
|
143
|
+
A 550,977-parameter GPT trained on the tinyshakespeare corpus with data-parallel training across 8 CPU cores brought character-level cross-entropy from 3.66 to 1.44 over 33,113 steps.
|
|
144
|
+
|
|
145
|
+
---
|
|
146
|
+
|
|
147
|
+
## Limitations
|
|
148
|
+
|
|
149
|
+
The core library is pure Python, hence it is slower than a production framework.
|
|
150
|
+
|
|
151
|
+
The compiler handles the forward pass over six operations, with no autograd. Only matmul is threaded; the element-wise kernels are single-threaded, which is fine because they are memory-bound rather than compute-bound.
|
|
152
|
+
|
|
153
|
+
Metal and Accelerate backends require macOS and PyObjC. If not available, it defaults to pure Python.
|
|
154
|
+
|
|
155
|
+
---
|
|
156
|
+
|
|
157
|
+
Built by Srihari Srinivasan. MIT licensed.
|
|
158
|
+
|
|
159
|
+
The full write-up, including the mathematics behind each component and the bugs worth knowing about can be found at **https://github.com/SrihariSr/Forge**
|