@imaduddin/fastbrain 0.0.0-stage β 1.0.1
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.
- package/index.js +20 -0
- package/package.json +29 -6
- package/readme.md +959 -0
- package/src/Embedding/index.d.ts +21 -0
- package/src/Embedding/index.js +140 -0
- package/src/FFN/index.d.ts +23 -0
- package/src/FFN/index.js +45 -0
- package/src/Linear/index.d.ts +31 -0
- package/src/Linear/index.js +197 -0
- package/src/OperationManager.js +169 -0
- package/src/PositionEmbedding/index.d.ts +30 -0
- package/src/PositionEmbedding/index.js +65 -0
- package/src/RMSNorm/index.d.ts +30 -0
- package/src/RMSNorm/index.js +95 -0
- package/src/Relu/index.d.ts +25 -0
- package/src/Relu/index.js +108 -0
- package/src/SelfAttention/index.d.ts +28 -0
- package/src/SelfAttention/index.js +138 -0
- package/src/Softmax/index.d.ts +26 -0
- package/src/Softmax/index.js +224 -0
- package/src/Tokenizer/index.d.ts +12 -0
- package/src/Tokenizer/index.js +164 -0
- package/src/Tokenizer/vocabs.json +19618 -0
- package/src/TransformerBlock/index.d.ts +28 -0
- package/src/TransformerBlock/index.js +78 -0
- package/src/getVec.js +30 -0
- package/src/gpu.js +0 -0
- package/src/residualConnection/index.js +32 -0
- package/src/shaders/CalculateAcc.wgsl +50 -0
- package/src/shaders/CalculateLoss.wgsl +30 -0
- package/src/shaders/ComputeDInput.wgsl +32 -0
- package/src/shaders/EB.wgsl +33 -0
- package/src/shaders/EF.wgsl +25 -0
- package/src/shaders/LB.wgsl +75 -0
- package/src/shaders/LF.wgsl +60 -0
- package/src/shaders/Norm/rmsBackward.wgsl +65 -0
- package/src/shaders/Norm/rmsForward.wgsl +30 -0
- package/src/shaders/Norm/sumSquare.wgsl +38 -0
- package/src/shaders/Norm/updateGamma.wgsl +52 -0
- package/src/shaders/PB.wgsl +25 -0
- package/src/shaders/PF.wgsl +19 -0
- package/src/shaders/RB.wgsl +37 -0
- package/src/shaders/RF.wgsl +33 -0
- package/src/shaders/ReductionFunction.wgsl +28 -0
- package/src/shaders/SB.wgsl +36 -0
- package/src/shaders/SF.wgsl +75 -0
- package/src/shaders/SoftmaxBackward.wgsl +31 -0
- package/src/shaders/add.wgsl +30 -0
- package/src/shaders/attentionShader/CompueFinalForwardOutput.wgsl +29 -0
- package/src/shaders/attentionShader/ComputeDA.wgsl +30 -0
- package/src/shaders/attentionShader/ComputeDV.wgsl +37 -0
- package/src/shaders/attentionShader/ComputeK.wgsl +30 -0
- package/src/shaders/attentionShader/ComputedQ.wgsl +31 -0
- package/src/shaders/attentionShader/attentionSocres.wgsl +40 -0
- package/src/shaders/attentionShader/sumGradient.wgsl +24 -0
- package/src/shaders/calculateTotal.wgsl +26 -0
- package/README.md +0 -3
package/readme.md
ADDED
|
@@ -0,0 +1,959 @@
|
|
|
1
|
+
# fastBrain
|
|
2
|
+
|
|
3
|
+
**GPU-accelerated neural network and Transformer library for JavaScript using WebGPU.**
|
|
4
|
+
|
|
5
|
+
fastBrain is a JavaScript library for building and training neural networks and GPT-style Transformer models directly with **WebGPU**.
|
|
6
|
+
|
|
7
|
+
It is designed around GPU computation while keeping the API relatively simple, making it possible to experiment with embeddings, linear layers, self-attention, Transformer blocks, softmax, and training without relying on Python-based ML frameworks.
|
|
8
|
+
|
|
9
|
+
> π§ **Status:** Experimental / actively developing
|
|
10
|
+
|
|
11
|
+
---
|
|
12
|
+
|
|
13
|
+
## β¨ Features
|
|
14
|
+
|
|
15
|
+
* β‘ WebGPU-accelerated computation
|
|
16
|
+
* π§ Transformer architecture components
|
|
17
|
+
* π€ Token embeddings
|
|
18
|
+
* π’ Linear layers
|
|
19
|
+
* π― Softmax and cross-entropy
|
|
20
|
+
* π Forward and backward propagation
|
|
21
|
+
* π§© Transformer blocks
|
|
22
|
+
* πΎ Model weight saving
|
|
23
|
+
* π¦ JavaScript / Node.js
|
|
24
|
+
* π₯οΈ GPU computation through WebGPU
|
|
25
|
+
* π οΈ Designed for building LLM experiments from scratch
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## π§ Transformer Architecture
|
|
30
|
+
|
|
31
|
+
fastBrain can be used to construct a GPT-style language model using individual components.
|
|
32
|
+
|
|
33
|
+
A typical model can look like:
|
|
34
|
+
|
|
35
|
+
```text
|
|
36
|
+
Token IDs
|
|
37
|
+
β
|
|
38
|
+
βΌ
|
|
39
|
+
ββββββββββββ
|
|
40
|
+
β Tokenizerβ
|
|
41
|
+
ββββββ¬ββββββ
|
|
42
|
+
β
|
|
43
|
+
βΌ
|
|
44
|
+
ββββββββββββ
|
|
45
|
+
β Embeddingβ
|
|
46
|
+
ββββββ¬ββββββ
|
|
47
|
+
β
|
|
48
|
+
βΌ
|
|
49
|
+
βββββββββββββββββββ
|
|
50
|
+
β TransformerBlockβ
|
|
51
|
+
β Self-Attention β
|
|
52
|
+
β + FFN β
|
|
53
|
+
ββββββββββ¬βββββββββ
|
|
54
|
+
β
|
|
55
|
+
βΌ
|
|
56
|
+
βββββββββββββββββββ
|
|
57
|
+
β TransformerBlockβ
|
|
58
|
+
ββββββββββ¬βββββββββ
|
|
59
|
+
β
|
|
60
|
+
βΌ
|
|
61
|
+
βββββββββββββββββββ
|
|
62
|
+
β TransformerBlockβ
|
|
63
|
+
ββββββββββ¬βββββββββ
|
|
64
|
+
β
|
|
65
|
+
βΌ
|
|
66
|
+
ββββββββββββ
|
|
67
|
+
β Linear β
|
|
68
|
+
ββββββ¬ββββββ
|
|
69
|
+
β
|
|
70
|
+
βΌ
|
|
71
|
+
ββββββββββββ
|
|
72
|
+
β Softmax β
|
|
73
|
+
ββββββ¬ββββββ
|
|
74
|
+
β
|
|
75
|
+
βΌ
|
|
76
|
+
Token Probabilities
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
The components can be connected manually, allowing you to experiment with the architecture instead of using a single black-box model.
|
|
80
|
+
|
|
81
|
+
---
|
|
82
|
+
|
|
83
|
+
## π¦ Installation
|
|
84
|
+
|
|
85
|
+
Install fastBrain with npm:
|
|
86
|
+
|
|
87
|
+
```bash
|
|
88
|
+
npm install @imaduddin/fastbrain
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
fastBrain requires a runtime with WebGPU support.
|
|
92
|
+
|
|
93
|
+
For Node.js WebGPU environments, you can use a WebGPU implementation such as:
|
|
94
|
+
|
|
95
|
+
```bash
|
|
96
|
+
npm install webgpu
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
---
|
|
100
|
+
|
|
101
|
+
## π Basic Example
|
|
102
|
+
|
|
103
|
+
The following example creates a small language model using:
|
|
104
|
+
|
|
105
|
+
* Tokenizer
|
|
106
|
+
* 64-dimensional embeddings
|
|
107
|
+
* 3 Transformer blocks
|
|
108
|
+
* Linear output layer
|
|
109
|
+
* Softmax
|
|
110
|
+
* Cross-entropy loss
|
|
111
|
+
* Backpropagation
|
|
112
|
+
* GPU computation
|
|
113
|
+
|
|
114
|
+
```js
|
|
115
|
+
import {
|
|
116
|
+
Embedding,
|
|
117
|
+
Linear,
|
|
118
|
+
Tokenizer,
|
|
119
|
+
TransformerBlock,
|
|
120
|
+
Softmax
|
|
121
|
+
} from "@imaduddin/fastBrain"
|
|
122
|
+
|
|
123
|
+
import { create, globals } from "webgpu"
|
|
124
|
+
import fs from "fs"
|
|
125
|
+
|
|
126
|
+
Object.assign(globalThis, globals)
|
|
127
|
+
|
|
128
|
+
const navigator = {
|
|
129
|
+
gpu: create(["adapter=NVIDIA GeForce 920M"])
|
|
130
|
+
}
|
|
131
|
+
|
|
132
|
+
const contextSize = 24
|
|
133
|
+
const lr = 0.02
|
|
134
|
+
|
|
135
|
+
async function train() {
|
|
136
|
+
|
|
137
|
+
// Load tokenizer
|
|
138
|
+
const tokenizer = new Tokenizer("./vocab.json")
|
|
139
|
+
|
|
140
|
+
// Initialize WebGPU
|
|
141
|
+
const adapter = await navigator.gpu.requestAdapter({
|
|
142
|
+
powerPreference: "low-power"
|
|
143
|
+
})
|
|
144
|
+
|
|
145
|
+
const device = await adapter.requestDevice()
|
|
146
|
+
|
|
147
|
+
// Token embedding
|
|
148
|
+
const embedding = new Embedding(
|
|
149
|
+
device,
|
|
150
|
+
64,
|
|
151
|
+
tokenizer.vocabSize,
|
|
152
|
+
{
|
|
153
|
+
cpuReadBack: false,
|
|
154
|
+
|
|
155
|
+
save: {
|
|
156
|
+
filename: "model/vec.bin"
|
|
157
|
+
}
|
|
158
|
+
}
|
|
159
|
+
)
|
|
160
|
+
|
|
161
|
+
// Output projection
|
|
162
|
+
const linear = new Linear(
|
|
163
|
+
device,
|
|
164
|
+
64,
|
|
165
|
+
tokenizer.vocabSize,
|
|
166
|
+
{
|
|
167
|
+
save: {
|
|
168
|
+
filename: [
|
|
169
|
+
"./model/bias.bin",
|
|
170
|
+
"./model/weights.bin"
|
|
171
|
+
]
|
|
172
|
+
}
|
|
173
|
+
}
|
|
174
|
+
)
|
|
175
|
+
|
|
176
|
+
// Transformer blocks
|
|
177
|
+
const block = new TransformerBlock(
|
|
178
|
+
device,
|
|
179
|
+
64,
|
|
180
|
+
{
|
|
181
|
+
dirname: "./model/block1"
|
|
182
|
+
}
|
|
183
|
+
)
|
|
184
|
+
|
|
185
|
+
const block1 = new TransformerBlock(
|
|
186
|
+
device,
|
|
187
|
+
64,
|
|
188
|
+
{
|
|
189
|
+
dirname: "./model/block2"
|
|
190
|
+
}
|
|
191
|
+
)
|
|
192
|
+
|
|
193
|
+
const block2 = new TransformerBlock(
|
|
194
|
+
device,
|
|
195
|
+
64,
|
|
196
|
+
{
|
|
197
|
+
dirname: "./model/block3"
|
|
198
|
+
}
|
|
199
|
+
)
|
|
200
|
+
|
|
201
|
+
// Output softmax
|
|
202
|
+
const softmax = new Softmax(
|
|
203
|
+
device,
|
|
204
|
+
64,
|
|
205
|
+
tokenizer.vocabSize
|
|
206
|
+
)
|
|
207
|
+
|
|
208
|
+
// Load dataset
|
|
209
|
+
const tokens = fs
|
|
210
|
+
.readFileSync("./dataset.txt", "utf-8")
|
|
211
|
+
.split("\r\n")
|
|
212
|
+
.filter(e => e.length > 0)
|
|
213
|
+
.map(e => e.trim() + "<EOS>")
|
|
214
|
+
.map(e => tokenizer.encoder(e).tokenIDs)
|
|
215
|
+
.flat()
|
|
216
|
+
|
|
217
|
+
// Create next-token prediction pairs
|
|
218
|
+
const trainingPairs = []
|
|
219
|
+
|
|
220
|
+
for (let i = 0; i < tokens.length; i++) {
|
|
221
|
+
|
|
222
|
+
trainingPairs.push({
|
|
223
|
+
inputIDX: tokens.slice(
|
|
224
|
+
i,
|
|
225
|
+
i + contextSize
|
|
226
|
+
),
|
|
227
|
+
|
|
228
|
+
targetIDX: tokens.slice(
|
|
229
|
+
i + 1,
|
|
230
|
+
i + contextSize + 1
|
|
231
|
+
)
|
|
232
|
+
})
|
|
233
|
+
|
|
234
|
+
}
|
|
235
|
+
|
|
236
|
+
const start = 0
|
|
237
|
+
const end = 8000
|
|
238
|
+
|
|
239
|
+
// Training
|
|
240
|
+
for (let epoch = 0; epoch < 100; epoch++) {
|
|
241
|
+
|
|
242
|
+
let totalLoss = 0
|
|
243
|
+
let totalAcc = 0
|
|
244
|
+
|
|
245
|
+
for (let i = start; i < end; i++) {
|
|
246
|
+
|
|
247
|
+
// -------------------------
|
|
248
|
+
// Forward pass
|
|
249
|
+
// -------------------------
|
|
250
|
+
|
|
251
|
+
let vec = await embedding.forward(
|
|
252
|
+
trainingPairs[i].inputIDX
|
|
253
|
+
)
|
|
254
|
+
|
|
255
|
+
vec = await block.forward(vec)
|
|
256
|
+
|
|
257
|
+
vec = await block1.forward(vec)
|
|
258
|
+
|
|
259
|
+
vec = await block2.forward(vec)
|
|
260
|
+
|
|
261
|
+
const logits = await linear.forward(vec)
|
|
262
|
+
|
|
263
|
+
const probs = await softmax.forward(logits)
|
|
264
|
+
|
|
265
|
+
// -------------------------
|
|
266
|
+
// Loss
|
|
267
|
+
// -------------------------
|
|
268
|
+
|
|
269
|
+
const { loss, acc } =
|
|
270
|
+
await softmax.crossEntropy(
|
|
271
|
+
probs,
|
|
272
|
+
trainingPairs[i].targetIDX
|
|
273
|
+
)
|
|
274
|
+
|
|
275
|
+
totalLoss += loss
|
|
276
|
+
totalAcc += acc
|
|
277
|
+
|
|
278
|
+
// -------------------------
|
|
279
|
+
// Backward pass
|
|
280
|
+
// -------------------------
|
|
281
|
+
|
|
282
|
+
const dLogits =
|
|
283
|
+
await softmax.backward(
|
|
284
|
+
probs,
|
|
285
|
+
trainingPairs[i].targetIDX
|
|
286
|
+
)
|
|
287
|
+
|
|
288
|
+
let dInput =
|
|
289
|
+
await linear.backward(
|
|
290
|
+
dLogits,
|
|
291
|
+
lr
|
|
292
|
+
)
|
|
293
|
+
|
|
294
|
+
dInput =
|
|
295
|
+
await block2.backward(
|
|
296
|
+
dInput,
|
|
297
|
+
lr
|
|
298
|
+
)
|
|
299
|
+
|
|
300
|
+
dInput =
|
|
301
|
+
await block1.backward(
|
|
302
|
+
dInput,
|
|
303
|
+
lr
|
|
304
|
+
)
|
|
305
|
+
|
|
306
|
+
dInput =
|
|
307
|
+
await block.backward(
|
|
308
|
+
dInput,
|
|
309
|
+
lr
|
|
310
|
+
)
|
|
311
|
+
|
|
312
|
+
embedding.backward(
|
|
313
|
+
trainingPairs[i].inputIDX,
|
|
314
|
+
dInput,
|
|
315
|
+
lr
|
|
316
|
+
)
|
|
317
|
+
|
|
318
|
+
// Clear cached tensors
|
|
319
|
+
await block.ClearInputCache()
|
|
320
|
+
await block1.ClearInputCache()
|
|
321
|
+
await block2.ClearInputCache()
|
|
322
|
+
await linear.ClearInputCache()
|
|
323
|
+
}
|
|
324
|
+
|
|
325
|
+
// -------------------------
|
|
326
|
+
// Save model
|
|
327
|
+
// -------------------------
|
|
328
|
+
|
|
329
|
+
await embedding.Save()
|
|
330
|
+
await linear.Save()
|
|
331
|
+
await block.Save()
|
|
332
|
+
await block1.Save()
|
|
333
|
+
await block2.Save()
|
|
334
|
+
|
|
335
|
+
console.log(
|
|
336
|
+
`Epoch ${epoch + 1} | ` +
|
|
337
|
+
`Loss: ${totalLoss / (end - start)} | ` +
|
|
338
|
+
`Accuracy: ${totalAcc / (end - start)}`
|
|
339
|
+
)
|
|
340
|
+
}
|
|
341
|
+
}
|
|
342
|
+
|
|
343
|
+
train()
|
|
344
|
+
```
|
|
345
|
+
|
|
346
|
+
---
|
|
347
|
+
|
|
348
|
+
# π€ Tokenization
|
|
349
|
+
|
|
350
|
+
fastBrain can work with a custom tokenizer.
|
|
351
|
+
|
|
352
|
+
For example:
|
|
353
|
+
|
|
354
|
+
```js
|
|
355
|
+
const tokenizer = new Tokenizer("./vocab.json")
|
|
356
|
+
```
|
|
357
|
+
|
|
358
|
+
Encode text into token IDs:
|
|
359
|
+
|
|
360
|
+
```js
|
|
361
|
+
const result = tokenizer.encoder(
|
|
362
|
+
"Hello world"
|
|
363
|
+
)
|
|
364
|
+
|
|
365
|
+
console.log(result.tokenIDs)
|
|
366
|
+
```
|
|
367
|
+
|
|
368
|
+
The resulting token IDs can then be passed directly into the embedding layer:
|
|
369
|
+
|
|
370
|
+
```js
|
|
371
|
+
const vectors = await embedding.forward(
|
|
372
|
+
result.tokenIDs
|
|
373
|
+
)
|
|
374
|
+
```
|
|
375
|
+
|
|
376
|
+
This makes it possible to use your own tokenizer and vocabulary instead of requiring a predefined tokenizer.
|
|
377
|
+
|
|
378
|
+
---
|
|
379
|
+
|
|
380
|
+
# π§© Embedding
|
|
381
|
+
|
|
382
|
+
The `Embedding` layer converts token IDs into dense vectors.
|
|
383
|
+
|
|
384
|
+
Example:
|
|
385
|
+
|
|
386
|
+
```js
|
|
387
|
+
const embedding = new Embedding(
|
|
388
|
+
device,
|
|
389
|
+
64,
|
|
390
|
+
tokenizer.vocabSize
|
|
391
|
+
)
|
|
392
|
+
```
|
|
393
|
+
|
|
394
|
+
Here:
|
|
395
|
+
|
|
396
|
+
```text
|
|
397
|
+
64
|
|
398
|
+
```
|
|
399
|
+
|
|
400
|
+
is the embedding dimension.
|
|
401
|
+
|
|
402
|
+
For a vocabulary of `V` tokens, the embedding matrix contains:
|
|
403
|
+
|
|
404
|
+
```text
|
|
405
|
+
V Γ 64
|
|
406
|
+
```
|
|
407
|
+
|
|
408
|
+
parameters.
|
|
409
|
+
|
|
410
|
+
Forward pass:
|
|
411
|
+
|
|
412
|
+
```js
|
|
413
|
+
const vectors = await embedding.forward(tokens)
|
|
414
|
+
```
|
|
415
|
+
|
|
416
|
+
---
|
|
417
|
+
|
|
418
|
+
# π§± TransformerBlock
|
|
419
|
+
|
|
420
|
+
Transformer blocks combine the main components required for Transformer-based models.
|
|
421
|
+
|
|
422
|
+
```js
|
|
423
|
+
const block = new TransformerBlock(
|
|
424
|
+
device,
|
|
425
|
+
64,
|
|
426
|
+
{
|
|
427
|
+
dirname: "./model/block1"
|
|
428
|
+
}
|
|
429
|
+
)
|
|
430
|
+
```
|
|
431
|
+
|
|
432
|
+
Multiple blocks can be stacked:
|
|
433
|
+
|
|
434
|
+
```js
|
|
435
|
+
vec = await block.forward(vec)
|
|
436
|
+
|
|
437
|
+
vec = await block1.forward(vec)
|
|
438
|
+
|
|
439
|
+
vec = await block2.forward(vec)
|
|
440
|
+
```
|
|
441
|
+
|
|
442
|
+
And during training:
|
|
443
|
+
|
|
444
|
+
```js
|
|
445
|
+
dInput = await block2.backward(
|
|
446
|
+
dInput,
|
|
447
|
+
lr
|
|
448
|
+
)
|
|
449
|
+
|
|
450
|
+
dInput = await block1.backward(
|
|
451
|
+
dInput,
|
|
452
|
+
lr
|
|
453
|
+
)
|
|
454
|
+
|
|
455
|
+
dInput = await block.backward(
|
|
456
|
+
dInput,
|
|
457
|
+
lr
|
|
458
|
+
)
|
|
459
|
+
```
|
|
460
|
+
|
|
461
|
+
This allows fastBrain to build deeper Transformer networks by stacking blocks.
|
|
462
|
+
|
|
463
|
+
---
|
|
464
|
+
|
|
465
|
+
# π― Linear Layer
|
|
466
|
+
|
|
467
|
+
The `Linear` layer projects the Transformer representation into vocabulary space.
|
|
468
|
+
|
|
469
|
+
```js
|
|
470
|
+
const linear = new Linear(
|
|
471
|
+
device,
|
|
472
|
+
64,
|
|
473
|
+
tokenizer.vocabSize
|
|
474
|
+
)
|
|
475
|
+
```
|
|
476
|
+
|
|
477
|
+
For example:
|
|
478
|
+
|
|
479
|
+
```text
|
|
480
|
+
Embedding size = 64
|
|
481
|
+
Vocabulary = 3,000
|
|
482
|
+
|
|
483
|
+
64 β 3,000
|
|
484
|
+
```
|
|
485
|
+
|
|
486
|
+
The output contains one logit for every vocabulary token.
|
|
487
|
+
|
|
488
|
+
---
|
|
489
|
+
|
|
490
|
+
# π Softmax
|
|
491
|
+
|
|
492
|
+
The `Softmax` layer converts logits into probabilities.
|
|
493
|
+
|
|
494
|
+
```js
|
|
495
|
+
const probs = await softmax.forward(logits)
|
|
496
|
+
```
|
|
497
|
+
|
|
498
|
+
The predicted probability distribution can then be evaluated using cross-entropy:
|
|
499
|
+
|
|
500
|
+
```js
|
|
501
|
+
const { loss, acc } =
|
|
502
|
+
await softmax.crossEntropy(
|
|
503
|
+
probs,
|
|
504
|
+
targetTokens
|
|
505
|
+
)
|
|
506
|
+
```
|
|
507
|
+
|
|
508
|
+
---
|
|
509
|
+
|
|
510
|
+
# π Backpropagation
|
|
511
|
+
|
|
512
|
+
fastBrain supports manually controlled backpropagation.
|
|
513
|
+
|
|
514
|
+
The gradients flow backwards through the network:
|
|
515
|
+
|
|
516
|
+
```text
|
|
517
|
+
Softmax
|
|
518
|
+
β
|
|
519
|
+
Linear
|
|
520
|
+
β
|
|
521
|
+
TransformerBlock 3
|
|
522
|
+
β
|
|
523
|
+
TransformerBlock 2
|
|
524
|
+
β
|
|
525
|
+
TransformerBlock 1
|
|
526
|
+
β
|
|
527
|
+
Embedding
|
|
528
|
+
```
|
|
529
|
+
|
|
530
|
+
In code:
|
|
531
|
+
|
|
532
|
+
```js
|
|
533
|
+
const dLogits =
|
|
534
|
+
await softmax.backward(
|
|
535
|
+
probs,
|
|
536
|
+
targetIDX
|
|
537
|
+
)
|
|
538
|
+
|
|
539
|
+
let dInput =
|
|
540
|
+
await linear.backward(
|
|
541
|
+
dLogits,
|
|
542
|
+
lr
|
|
543
|
+
)
|
|
544
|
+
|
|
545
|
+
dInput =
|
|
546
|
+
await block2.backward(
|
|
547
|
+
dInput,
|
|
548
|
+
lr
|
|
549
|
+
)
|
|
550
|
+
|
|
551
|
+
dInput =
|
|
552
|
+
await block1.backward(
|
|
553
|
+
dInput,
|
|
554
|
+
lr
|
|
555
|
+
)
|
|
556
|
+
|
|
557
|
+
dInput =
|
|
558
|
+
await block.backward(
|
|
559
|
+
dInput,
|
|
560
|
+
lr
|
|
561
|
+
)
|
|
562
|
+
|
|
563
|
+
embedding.backward(
|
|
564
|
+
inputIDX,
|
|
565
|
+
dInput,
|
|
566
|
+
lr
|
|
567
|
+
)
|
|
568
|
+
```
|
|
569
|
+
|
|
570
|
+
This gives you direct control over the training process.
|
|
571
|
+
|
|
572
|
+
---
|
|
573
|
+
|
|
574
|
+
# β‘ WebGPU
|
|
575
|
+
|
|
576
|
+
fastBrain is designed around WebGPU.
|
|
577
|
+
|
|
578
|
+
Example:
|
|
579
|
+
|
|
580
|
+
```js
|
|
581
|
+
import { create, globals } from "webgpu"
|
|
582
|
+
|
|
583
|
+
Object.assign(globalThis, globals)
|
|
584
|
+
|
|
585
|
+
const navigator = {
|
|
586
|
+
gpu: create([
|
|
587
|
+
"adapter=NVIDIA GeForce 920M"
|
|
588
|
+
])
|
|
589
|
+
}
|
|
590
|
+
```
|
|
591
|
+
|
|
592
|
+
Then request the GPU device:
|
|
593
|
+
|
|
594
|
+
```js
|
|
595
|
+
const adapter =
|
|
596
|
+
await navigator.gpu.requestAdapter()
|
|
597
|
+
|
|
598
|
+
const device =
|
|
599
|
+
await adapter.requestDevice()
|
|
600
|
+
```
|
|
601
|
+
|
|
602
|
+
The device is passed to fastBrain layers:
|
|
603
|
+
|
|
604
|
+
```js
|
|
605
|
+
const embedding =
|
|
606
|
+
new Embedding(device, 64, vocabSize)
|
|
607
|
+
|
|
608
|
+
const linear =
|
|
609
|
+
new Linear(device, 64, vocabSize)
|
|
610
|
+
|
|
611
|
+
const block =
|
|
612
|
+
new TransformerBlock(device, 64)
|
|
613
|
+
```
|
|
614
|
+
|
|
615
|
+
This lets the layers perform their GPU computations through WebGPU.
|
|
616
|
+
|
|
617
|
+
---
|
|
618
|
+
|
|
619
|
+
|
|
620
|
+
|
|
621
|
+
# π€ Training the Tokenizer
|
|
622
|
+
|
|
623
|
+
fastBrain includes a tokenizer that can learn a vocabulary directly from your dataset.
|
|
624
|
+
|
|
625
|
+
```js
|
|
626
|
+
import fs from "fs"
|
|
627
|
+
import { Tokenizer } from "./index.js"
|
|
628
|
+
|
|
629
|
+
const tokenizer = new Tokenizer("./vocab.json")
|
|
630
|
+
|
|
631
|
+
tokenizer.trainTokenizer(
|
|
632
|
+
fs.readFileSync("./dataset.txt", "utf-8"),
|
|
633
|
+
6,
|
|
634
|
+
"./vocab.json"
|
|
635
|
+
)
|
|
636
|
+
```
|
|
637
|
+
|
|
638
|
+
## `trainTokenizer()`
|
|
639
|
+
|
|
640
|
+
```js
|
|
641
|
+
tokenizer.trainTokenizer(
|
|
642
|
+
text,
|
|
643
|
+
maxMergeLength,
|
|
644
|
+
outputPath
|
|
645
|
+
)
|
|
646
|
+
```
|
|
647
|
+
|
|
648
|
+
| Parameter | Description |
|
|
649
|
+
| ---------------- | ---------------------------------------------------- |
|
|
650
|
+
| `text` | Training text used to build the tokenizer vocabulary |
|
|
651
|
+
| `maxMergeLength` | Maximum length used when creating token merges |
|
|
652
|
+
| `outputPath` | Location where the generated vocabulary is saved |
|
|
653
|
+
|
|
654
|
+
### Example
|
|
655
|
+
|
|
656
|
+
```js
|
|
657
|
+
tokenizer.trainTokenizer(
|
|
658
|
+
fs.readFileSync("./dataset.txt", "utf-8"),
|
|
659
|
+
6,
|
|
660
|
+
"./vocab.json"
|
|
661
|
+
)
|
|
662
|
+
```
|
|
663
|
+
|
|
664
|
+
Here:
|
|
665
|
+
|
|
666
|
+
```text
|
|
667
|
+
6
|
|
668
|
+
```
|
|
669
|
+
|
|
670
|
+
is the **maximum merge length**.
|
|
671
|
+
|
|
672
|
+
The tokenizer learns merges from the dataset and saves the resulting vocabulary to:
|
|
673
|
+
|
|
674
|
+
```text
|
|
675
|
+
vocab.json
|
|
676
|
+
```
|
|
677
|
+
|
|
678
|
+
### Tokenizer Workflow
|
|
679
|
+
|
|
680
|
+
```text
|
|
681
|
+
dataset.txt
|
|
682
|
+
β
|
|
683
|
+
βΌ
|
|
684
|
+
βββββββββββββββββββ
|
|
685
|
+
β trainTokenizer()β
|
|
686
|
+
β β
|
|
687
|
+
β maxMergeLength=6β
|
|
688
|
+
ββββββββββ¬βββββββββ
|
|
689
|
+
β
|
|
690
|
+
βΌ
|
|
691
|
+
vocab.json
|
|
692
|
+
β
|
|
693
|
+
βΌ
|
|
694
|
+
ββββββββββββββ
|
|
695
|
+
β Tokenizer β
|
|
696
|
+
βββββββ¬βββββββ
|
|
697
|
+
β
|
|
698
|
+
βΌ
|
|
699
|
+
tokenizer.encoder()
|
|
700
|
+
β
|
|
701
|
+
βΌ
|
|
702
|
+
Token IDs
|
|
703
|
+
```
|
|
704
|
+
|
|
705
|
+
Once the vocabulary has been created, load it with:
|
|
706
|
+
|
|
707
|
+
```js
|
|
708
|
+
const tokenizer = new Tokenizer("./vocab.json")
|
|
709
|
+
```
|
|
710
|
+
|
|
711
|
+
Then encode text:
|
|
712
|
+
|
|
713
|
+
```js
|
|
714
|
+
const result = tokenizer.encoder(
|
|
715
|
+
"Hello world"
|
|
716
|
+
)
|
|
717
|
+
|
|
718
|
+
console.log(result.tokenIDs)
|
|
719
|
+
```
|
|
720
|
+
|
|
721
|
+
The resulting token IDs can be passed directly into fastBrain's embedding layer:
|
|
722
|
+
|
|
723
|
+
```js
|
|
724
|
+
const vectors = await embedding.forward(
|
|
725
|
+
result.tokenIDs
|
|
726
|
+
)
|
|
727
|
+
```
|
|
728
|
+
|
|
729
|
+
This allows you to create a tokenizer vocabulary from your own dataset and then use that vocabulary for Transformer training.
|
|
730
|
+
|
|
731
|
+
|
|
732
|
+
---
|
|
733
|
+
|
|
734
|
+
|
|
735
|
+
|
|
736
|
+
|
|
737
|
+
# πΎ Saving Models
|
|
738
|
+
|
|
739
|
+
fastBrain layers can save their parameters to files.
|
|
740
|
+
|
|
741
|
+
Embedding:
|
|
742
|
+
|
|
743
|
+
```js
|
|
744
|
+
const embedding = new Embedding(
|
|
745
|
+
device,
|
|
746
|
+
64,
|
|
747
|
+
vocabSize,
|
|
748
|
+
{
|
|
749
|
+
save: {
|
|
750
|
+
filename: "model/vec.bin"
|
|
751
|
+
}
|
|
752
|
+
}
|
|
753
|
+
)
|
|
754
|
+
```
|
|
755
|
+
|
|
756
|
+
Linear:
|
|
757
|
+
|
|
758
|
+
```js
|
|
759
|
+
const linear = new Linear(
|
|
760
|
+
device,
|
|
761
|
+
64,
|
|
762
|
+
vocabSize,
|
|
763
|
+
{
|
|
764
|
+
save: {
|
|
765
|
+
filename: [
|
|
766
|
+
"./model/bias.bin",
|
|
767
|
+
"./model/weights.bin"
|
|
768
|
+
]
|
|
769
|
+
}
|
|
770
|
+
}
|
|
771
|
+
)
|
|
772
|
+
```
|
|
773
|
+
|
|
774
|
+
Transformer blocks can use their own directories:
|
|
775
|
+
|
|
776
|
+
```js
|
|
777
|
+
const block = new TransformerBlock(
|
|
778
|
+
device,
|
|
779
|
+
64,
|
|
780
|
+
{
|
|
781
|
+
dirname: "./model/block1"
|
|
782
|
+
}
|
|
783
|
+
)
|
|
784
|
+
```
|
|
785
|
+
|
|
786
|
+
A model can therefore be organized like:
|
|
787
|
+
|
|
788
|
+
```text
|
|
789
|
+
model/
|
|
790
|
+
βββ vec.bin
|
|
791
|
+
βββ weights.bin
|
|
792
|
+
βββ bias.bin
|
|
793
|
+
β
|
|
794
|
+
βββ block1/
|
|
795
|
+
β βββ ...
|
|
796
|
+
β
|
|
797
|
+
βββ block2/
|
|
798
|
+
β βββ ...
|
|
799
|
+
β
|
|
800
|
+
βββ block3/
|
|
801
|
+
βββ ...
|
|
802
|
+
```
|
|
803
|
+
|
|
804
|
+
---
|
|
805
|
+
|
|
806
|
+
# π§ͺ Example Model Configuration
|
|
807
|
+
|
|
808
|
+
The example above uses:
|
|
809
|
+
|
|
810
|
+
| Component | Configuration |
|
|
811
|
+
| ------------------- | --------------------------: |
|
|
812
|
+
| Embedding dimension | 64 |
|
|
813
|
+
| Context length | 24 |
|
|
814
|
+
| Transformer blocks | 3 |
|
|
815
|
+
| Vocabulary | `tokenizer.vocabSize` |
|
|
816
|
+
| Learning rate | 0.02 |
|
|
817
|
+
| Training pairs | 8,000 |
|
|
818
|
+
| Epochs | 100 |
|
|
819
|
+
| Optimizer | SGD-style parameter updates |
|
|
820
|
+
|
|
821
|
+
The vocabulary size is determined dynamically from the tokenizer:
|
|
822
|
+
|
|
823
|
+
```js
|
|
824
|
+
tokenizer.vocabSize
|
|
825
|
+
```
|
|
826
|
+
|
|
827
|
+
So the same architecture can be used with different vocabulary sizes.
|
|
828
|
+
|
|
829
|
+
---
|
|
830
|
+
|
|
831
|
+
# π Example Project
|
|
832
|
+
|
|
833
|
+
A simple project can look like:
|
|
834
|
+
|
|
835
|
+
```text
|
|
836
|
+
my-model/
|
|
837
|
+
β
|
|
838
|
+
βββ dataset.txt
|
|
839
|
+
βββ vocab.json
|
|
840
|
+
βββ train.js
|
|
841
|
+
β
|
|
842
|
+
βββ model/
|
|
843
|
+
βββ vec.bin
|
|
844
|
+
βββ weights.bin
|
|
845
|
+
βββ bias.bin
|
|
846
|
+
β
|
|
847
|
+
βββ block1/
|
|
848
|
+
βββ block2/
|
|
849
|
+
βββ block3/
|
|
850
|
+
```
|
|
851
|
+
|
|
852
|
+
Run training with:
|
|
853
|
+
|
|
854
|
+
```bash
|
|
855
|
+
node train.js
|
|
856
|
+
```
|
|
857
|
+
|
|
858
|
+
---
|
|
859
|
+
|
|
860
|
+
# π οΈ Design Philosophy
|
|
861
|
+
|
|
862
|
+
fastBrain is intended to make the internals of neural networks accessible from JavaScript.
|
|
863
|
+
|
|
864
|
+
Instead of hiding everything behind a single high-level training function, fastBrain exposes the individual operations:
|
|
865
|
+
|
|
866
|
+
```js
|
|
867
|
+
embedding.forward()
|
|
868
|
+
transformer.forward()
|
|
869
|
+
linear.forward()
|
|
870
|
+
softmax.forward()
|
|
871
|
+
|
|
872
|
+
softmax.backward()
|
|
873
|
+
linear.backward()
|
|
874
|
+
transformer.backward()
|
|
875
|
+
embedding.backward()
|
|
876
|
+
```
|
|
877
|
+
|
|
878
|
+
This makes it useful for:
|
|
879
|
+
|
|
880
|
+
* Learning how neural networks work
|
|
881
|
+
* Experimenting with Transformer architectures
|
|
882
|
+
* Building small language models
|
|
883
|
+
* WebGPU experiments
|
|
884
|
+
* GPU programming experiments
|
|
885
|
+
* Creating custom neural-network layers
|
|
886
|
+
* Understanding forward and backward propagation
|
|
887
|
+
|
|
888
|
+
---
|
|
889
|
+
|
|
890
|
+
# π§ Building an LLM From Scratch
|
|
891
|
+
|
|
892
|
+
fastBrain can be used as a foundation for experimenting with GPT-style language models.
|
|
893
|
+
|
|
894
|
+
A simplified training pipeline is:
|
|
895
|
+
|
|
896
|
+
```text
|
|
897
|
+
Dataset
|
|
898
|
+
β
|
|
899
|
+
βΌ
|
|
900
|
+
Tokenizer
|
|
901
|
+
β
|
|
902
|
+
βΌ
|
|
903
|
+
Token IDs
|
|
904
|
+
β
|
|
905
|
+
βΌ
|
|
906
|
+
Embedding
|
|
907
|
+
β
|
|
908
|
+
βΌ
|
|
909
|
+
Transformer Blocks
|
|
910
|
+
β
|
|
911
|
+
βΌ
|
|
912
|
+
Linear Projection
|
|
913
|
+
β
|
|
914
|
+
βΌ
|
|
915
|
+
Softmax
|
|
916
|
+
β
|
|
917
|
+
βΌ
|
|
918
|
+
Cross Entropy
|
|
919
|
+
β
|
|
920
|
+
βΌ
|
|
921
|
+
Backpropagation
|
|
922
|
+
β
|
|
923
|
+
ββββββββββββββββΊ Update weights
|
|
924
|
+
```
|
|
925
|
+
|
|
926
|
+
The architecture can then be expanded with additional Transformer blocks, larger embeddings, larger vocabularies, and other normalization or optimization techniques.
|
|
927
|
+
|
|
928
|
+
---
|
|
929
|
+
|
|
930
|
+
# π¬ Experimental
|
|
931
|
+
|
|
932
|
+
fastBrain is still an experimental project.
|
|
933
|
+
|
|
934
|
+
The API and internal implementation may change as new features are added.
|
|
935
|
+
|
|
936
|
+
The project is primarily intended for experimentation, learning, and building custom GPU-accelerated neural-network systems.
|
|
937
|
+
|
|
938
|
+
---
|
|
939
|
+
|
|
940
|
+
# π License
|
|
941
|
+
|
|
942
|
+
|
|
943
|
+
|
|
944
|
+
MIT License
|
|
945
|
+
|
|
946
|
+
|
|
947
|
+
---
|
|
948
|
+
|
|
949
|
+
## π¨βπ» Author
|
|
950
|
+
|
|
951
|
+
Created by **mineproness**.
|
|
952
|
+
|
|
953
|
+
fastBrain is part of the broader **@imaduddin** ecosystem of JavaScript AI/ML experiments.
|
|
954
|
+
|
|
955
|
+
---
|
|
956
|
+
|
|
957
|
+
# β If you find fastBrain interesting
|
|
958
|
+
|
|
959
|
+
Give the project a star β and experiment with building your own neural networks and Transformer models directly in JavaScript.
|