@mineproness/tensorloop 2.1.55 → 2.1.57
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 +14 -20
- package/package.json +5 -2
- package/readme.md +907 -72
- package/src/Embedding/index.d.ts +5 -5
- package/src/Embedding/index.js +118 -33
- package/src/FFN/index.d.ts +5 -5
- package/src/FFN/index.js +34 -39
- package/src/Linear/index.d.ts +7 -8
- package/src/Linear/index.js +173 -49
- package/src/OperationManager.js +100 -0
- package/src/Relu/index.d.ts +5 -0
- package/src/Relu/index.js +103 -0
- package/src/SelfAttention/index.d.ts +7 -13
- package/src/SelfAttention/index.js +108 -55
- package/src/Softmax/index.d.ts +4 -3
- package/src/Softmax/index.js +192 -12
- package/src/Tokenizer/index.d.ts +1 -0
- package/src/Tokenizer/index.js +92 -7
- package/src/TransformerBlock/index.d.ts +6 -7
- package/src/TransformerBlock/index.js +24 -23
- package/src/getVec.js +28 -0
- package/src/shaders/CalculateAcc.wgsl +50 -0
- package/src/shaders/CalculateLoss.wgsl +35 -0
- package/src/shaders/ComputeDInput.wgsl +32 -0
- package/src/shaders/EB.wgsl +31 -0
- package/src/shaders/EF.wgsl +25 -0
- package/src/shaders/LB.wgsl +57 -0
- package/src/shaders/LF.wgsl +44 -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 +65 -0
- package/src/shaders/SoftmaxBackward.wgsl +31 -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/GELU/index.d.ts +0 -6
- package/src/GELU/index.js +0 -33
- package/src/GetConfigs.js +0 -35
- package/src/Model/index.d.ts +0 -10
- package/src/Model/index.js +0 -88
- package/src/PositionEmbedding/index.d.ts +0 -11
- package/src/PositionEmbedding/index.js +0 -49
- package/src/ReLU/index.d.ts +0 -5
- package/src/ReLU/index.js +0 -17
- package/src/gpiInit.js +0 -764
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
import { OperationManager } from "../OperationManager.js";
|
|
2
|
+
|
|
3
|
+
export class Relu {
|
|
4
|
+
embeddingSize;
|
|
5
|
+
paramBuffer;
|
|
6
|
+
oldParam;
|
|
7
|
+
forwardPipline;
|
|
8
|
+
gpu;
|
|
9
|
+
backwardPipline;
|
|
10
|
+
input;
|
|
11
|
+
constructor(device, embeddingSize) {
|
|
12
|
+
this.embeddingSize = embeddingSize
|
|
13
|
+
this.gpu = new OperationManager(device);
|
|
14
|
+
this.device = device;
|
|
15
|
+
this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/RF.wgsl")
|
|
16
|
+
this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/RB.wgsl")
|
|
17
|
+
}
|
|
18
|
+
async forward(input) {
|
|
19
|
+
const encoder = this.device.createCommandEncoder()
|
|
20
|
+
const tokenCount = input.size / (this.embeddingSize * 4)
|
|
21
|
+
if (!this.oldParam) {
|
|
22
|
+
this.oldParam = new Uint32Array([0, 0, 0]);
|
|
23
|
+
this.paramBuffer = await this.gpu.createUniFormBuffer(4 * 3);
|
|
24
|
+
}
|
|
25
|
+
const output = await this.gpu.createBuffer(input.size)
|
|
26
|
+
if (this.oldParam[1] !== tokenCount) {
|
|
27
|
+
this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
|
|
28
|
+
await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
|
|
29
|
+
}
|
|
30
|
+
const bind = await this.device.createBindGroup({
|
|
31
|
+
layout: await this.forwardPipline.getBindGroupLayout(0),
|
|
32
|
+
entries: [
|
|
33
|
+
{
|
|
34
|
+
binding: 0,
|
|
35
|
+
resource: {
|
|
36
|
+
buffer: input
|
|
37
|
+
}
|
|
38
|
+
},
|
|
39
|
+
{
|
|
40
|
+
binding: 1,
|
|
41
|
+
resource: {
|
|
42
|
+
buffer: this.paramBuffer
|
|
43
|
+
}
|
|
44
|
+
}, {
|
|
45
|
+
binding: 2,
|
|
46
|
+
resource: {
|
|
47
|
+
buffer: output
|
|
48
|
+
}
|
|
49
|
+
}
|
|
50
|
+
]
|
|
51
|
+
})
|
|
52
|
+
await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
|
|
53
|
+
this.device.queue.submit([encoder.finish()])
|
|
54
|
+
if(this.input) this.input.destroy()
|
|
55
|
+
this.input = input
|
|
56
|
+
return output
|
|
57
|
+
}
|
|
58
|
+
async backward(inputGradient) {
|
|
59
|
+
const encoder = this.device.createCommandEncoder()
|
|
60
|
+
const tokenCount = inputGradient.size / (this.embeddingSize * 4)
|
|
61
|
+
if (!this.oldParam) {
|
|
62
|
+
this.oldParam = new Uint32Array([0, 0, 0]);
|
|
63
|
+
this.paramBuffer = await this.gpu.createUniFormBuffer(4 * 3);
|
|
64
|
+
}
|
|
65
|
+
const output = await this.gpu.createBuffer(inputGradient.size)
|
|
66
|
+
if (this.oldParam[1] !== tokenCount) {
|
|
67
|
+
this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
|
|
68
|
+
await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
|
|
69
|
+
}
|
|
70
|
+
const bind = await this.device.createBindGroup({
|
|
71
|
+
layout: await this.backwardPipline.getBindGroupLayout(0),
|
|
72
|
+
entries: [
|
|
73
|
+
{
|
|
74
|
+
binding: 0,
|
|
75
|
+
resource: {
|
|
76
|
+
buffer: this.input
|
|
77
|
+
}
|
|
78
|
+
},
|
|
79
|
+
{
|
|
80
|
+
binding: 1,
|
|
81
|
+
resource: {
|
|
82
|
+
buffer: this.paramBuffer
|
|
83
|
+
}
|
|
84
|
+
}, {
|
|
85
|
+
binding: 2,
|
|
86
|
+
resource: {
|
|
87
|
+
buffer: output
|
|
88
|
+
}
|
|
89
|
+
}, {
|
|
90
|
+
binding: 3,
|
|
91
|
+
resource: {
|
|
92
|
+
buffer: inputGradient
|
|
93
|
+
}
|
|
94
|
+
}
|
|
95
|
+
]
|
|
96
|
+
})
|
|
97
|
+
await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
|
|
98
|
+
await this.device.queue.submit([encoder.finish()])
|
|
99
|
+
this.input.destroy()
|
|
100
|
+
inputGradient.destroy()
|
|
101
|
+
return output
|
|
102
|
+
}
|
|
103
|
+
}
|
|
@@ -1,15 +1,9 @@
|
|
|
1
1
|
export class SelfAttention{
|
|
2
|
-
|
|
3
|
-
|
|
4
|
-
k : Float32Array[]
|
|
5
|
-
input;
|
|
6
|
-
constructor(embeddingSize : number, configs: {
|
|
7
|
-
dirname: string
|
|
8
|
-
|
|
2
|
+
constructor(device : GPUAdapter,embeddingSize: number, configs : {
|
|
3
|
+
dirname: string
|
|
9
4
|
})
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
Save() : void
|
|
14
|
-
|
|
15
|
-
}
|
|
5
|
+
async forward(vec : GPUBuffer) : Promise<GPUBuffer>
|
|
6
|
+
async backward(dO : GPUBuffer, lr : number) : Promise<GPUBuffer>
|
|
7
|
+
async ClearInputCache() : Promise<void>
|
|
8
|
+
async Save() : Promise<void>
|
|
9
|
+
}
|
|
@@ -1,67 +1,120 @@
|
|
|
1
|
-
import {
|
|
2
|
-
import {
|
|
3
|
-
import {
|
|
1
|
+
import { Linear } from "../../index.js";
|
|
2
|
+
import { join } from 'path'
|
|
3
|
+
import { OperationManager } from "../OperationManager.js";
|
|
4
4
|
|
|
5
5
|
export class SelfAttention {
|
|
6
6
|
query;
|
|
7
|
-
key;
|
|
8
7
|
value;
|
|
9
|
-
|
|
10
|
-
k;
|
|
8
|
+
key;
|
|
11
9
|
q;
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
|
|
10
|
+
k;
|
|
11
|
+
v;
|
|
12
|
+
embeddingSize;
|
|
13
|
+
params;
|
|
14
|
+
attentionscoresShader;
|
|
15
|
+
softmaxForward;
|
|
17
16
|
gpu;
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
this.
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
this.q =
|
|
29
|
-
this.
|
|
30
|
-
this.
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
this.
|
|
34
|
-
|
|
35
|
-
|
|
17
|
+
input;
|
|
18
|
+
constructor(device, embeddingSize, configs) {
|
|
19
|
+
this.embeddingSize = embeddingSize;
|
|
20
|
+
const dd = [1, 2, 3].map((e) => {
|
|
21
|
+
return {
|
|
22
|
+
save: {
|
|
23
|
+
filename: [join(configs.dirname, `SELF-W-${e}.bin`), join(configs.dirname, `SELF-B-${e}.bin`)]
|
|
24
|
+
}
|
|
25
|
+
}
|
|
26
|
+
})
|
|
27
|
+
this.q = new Linear(device, embeddingSize, embeddingSize, dd[0])
|
|
28
|
+
this.k = new Linear(device, embeddingSize, embeddingSize, dd[1])
|
|
29
|
+
this.v = new Linear(device, embeddingSize, embeddingSize, dd[2])
|
|
30
|
+
this.gpu = new OperationManager(device)
|
|
31
|
+
this.attentionscoresShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/attentionSocres.wgsl")
|
|
32
|
+
this.softmaxForward = this.gpu.getShaderPipline("./src/shaders/SF.wgsl")
|
|
33
|
+
this.softmaxBackward = this.gpu.getShaderPipline("./src/shaders/SoftmaxBackward.wgsl")
|
|
34
|
+
this.outputForward = this.gpu.getShaderPipline("./src/shaders/attentionShader/CompueFinalForwardOutput.wgsl")
|
|
35
|
+
this.dVShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeDV.wgsl")
|
|
36
|
+
this.dAshader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeDA.wgsl")
|
|
37
|
+
this.dQshader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputedQ.wgsl")
|
|
38
|
+
this.dKShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeK.wgsl")
|
|
39
|
+
this.sumShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/sumGradient.wgsl")
|
|
36
40
|
}
|
|
41
|
+
async forward(vec) {
|
|
42
|
+
if (this.input) await this.input.destroy()
|
|
43
|
+
this.input = vec;
|
|
44
|
+
const tokenCount = vec.size / (this.embeddingSize * 4)
|
|
45
|
+
if (!this.asParams) this.asParams = await this.gpu.createUniFormBuffer(4 * 3);
|
|
46
|
+
if (!this.softmaxParams) this.softmaxParams = await this.gpu.createUniFormBuffer(4 * 3);
|
|
47
|
+
const encoder = await this.gpu.CreateEncoder()
|
|
48
|
+
const scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
|
|
49
|
+
await this.gpu.WriteBuffer(this.asParams, new Uint32Array([tokenCount, this.embeddingSize, 0]))
|
|
50
|
+
await this.gpu.WriteBuffer(this.softmaxParams, new Uint32Array([tokenCount, tokenCount, 0]))
|
|
51
|
+
if (this.query) this.query.destroy()
|
|
52
|
+
if (this.key) this.key.destroy()
|
|
53
|
+
if (this.value) this.value.destroy()
|
|
54
|
+
this.query = await this.q.forward(vec)
|
|
55
|
+
this.key = await this.k.forward(vec)
|
|
56
|
+
this.value = await this.v.forward(vec)
|
|
57
|
+
const bindAttention = await this.gpu.createBindGroup(this.attentionscoresShader, 0, [this.query, this.key, scroes, this.asParams])
|
|
58
|
+
await this.gpu.RunPipline(encoder, this.attentionscoresShader, bindAttention, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
|
|
59
|
+
if (this.weights) this.weights.destroy()
|
|
60
|
+
this.weights = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
|
|
61
|
+
const bindSF = await this.gpu.createBindGroup(this.softmaxForward, 0, [scroes, this.weights, this.softmaxParams]);
|
|
62
|
+
await this.gpu.RunPipline(encoder, this.softmaxForward, bindSF, tokenCount)
|
|
63
|
+
const finalOutput = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4)
|
|
64
|
+
const outputBind = await this.gpu.createBindGroup(this.outputForward, 0, [this.weights, this.value, finalOutput, this.asParams]);
|
|
65
|
+
await this.gpu.RunPipline(encoder, this.outputForward, outputBind, Math.ceil((tokenCount * this.embeddingSize) / 256))
|
|
37
66
|
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
const
|
|
44
|
-
const
|
|
45
|
-
const
|
|
46
|
-
const
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
67
|
+
await this.gpu.submitQueue(encoder)
|
|
68
|
+
scroes.destroy()
|
|
69
|
+
return finalOutput
|
|
70
|
+
}
|
|
71
|
+
async backward(dO, lr) {
|
|
72
|
+
const tokenCount = dO.size / (this.embeddingSize * 4);
|
|
73
|
+
const encoder = await this.gpu.CreateEncoder();
|
|
74
|
+
const dV = await this.gpu.createBuffer(dO.size);
|
|
75
|
+
const DvBind = await this.gpu.createBindGroup(this.dVShader, 0, [this.weights, dO, dV, this.asParams]);
|
|
76
|
+
await this.gpu.RunPipline(encoder, this.dVShader, DvBind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16));
|
|
77
|
+
const dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
|
|
78
|
+
const DaBind = await this.gpu.createBindGroup(this.dAshader, 0, [dO, this.value, dA, this.asParams]);
|
|
79
|
+
await this.gpu.RunPipline(encoder, this.dAshader, DaBind, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
|
|
80
|
+
const dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
|
|
81
|
+
const dSBind = await this.gpu.createBindGroup(this.softmaxBackward, 0, [this.weights, dA, dS, this.asParams])
|
|
82
|
+
await this.gpu.RunPipline(encoder, this.softmaxBackward, dSBind, Math.ceil((tokenCount * tokenCount) / 256))
|
|
83
|
+
const dQ = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
|
|
84
|
+
const dQbind = await this.gpu.createBindGroup(this.dQshader, 0, [dS, this.key, dQ, this.asParams]);
|
|
85
|
+
await this.gpu.RunPipline(encoder, this.dQshader, dQbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
|
|
86
|
+
const dK = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4)
|
|
87
|
+
const dKbind = await this.gpu.createBindGroup(this.dKShader, 0, [dS, this.query, dK, this.asParams]);
|
|
88
|
+
await this.gpu.RunPipline(encoder, this.dKShader, dKbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
|
|
89
|
+
await this.gpu.submitQueue(encoder);
|
|
90
|
+
const dq = await this.q.backward(dQ, lr)
|
|
91
|
+
const dk = await this.k.backward(dK, lr)
|
|
92
|
+
const dv = await this.v.backward(dV, lr)
|
|
93
|
+
const output = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4);
|
|
94
|
+
const encode = this.gpu.CreateEncoder()
|
|
95
|
+
const outputBind = await this.gpu.createBindGroup(this.sumShader, 0 , [dq, dk, dv, output, this.softmaxParams]);
|
|
96
|
+
await this.gpu.RunPipline(encode, this.sumShader, outputBind, Math.ceil((tokenCount * this.embeddingSize / 256)));
|
|
97
|
+
await this.gpu.submitQueue(encode)
|
|
98
|
+
dS.destroy()
|
|
99
|
+
dA.destroy()
|
|
100
|
+
dq.destroy()
|
|
101
|
+
dk.destroy()
|
|
102
|
+
dv.destroy()
|
|
103
|
+
dK.destroy()
|
|
104
|
+
dQ.destroy()
|
|
105
|
+
dV.destroy()
|
|
106
|
+
dO.destroy()
|
|
107
|
+
return output
|
|
108
|
+
|
|
56
109
|
}
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
110
|
+
async ClearInputCache(){
|
|
111
|
+
await this.q.ClearInputCache()
|
|
112
|
+
await this.k.ClearInputCache()
|
|
113
|
+
await this.v.ClearInputCache()
|
|
61
114
|
}
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
115
|
+
async Save(){
|
|
116
|
+
await this.q.Save()
|
|
117
|
+
await this.q.Save()
|
|
118
|
+
await this.q.Save()
|
|
66
119
|
}
|
|
67
120
|
}
|
package/src/Softmax/index.d.ts
CHANGED
|
@@ -1,5 +1,6 @@
|
|
|
1
1
|
export class Softmax{
|
|
2
|
-
|
|
3
|
-
|
|
4
|
-
|
|
2
|
+
constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number);
|
|
3
|
+
async forward(logit: GPUBuffer, isCpuReadBack) : Promise<Float32Array | GPUBuffer>;
|
|
4
|
+
async backward(probs : GPUBuffer, target : number[]) : Promise<GPUBuffer>
|
|
5
|
+
async crossEntropy(probs: GPUBuffer, target: number[]) : Promise<{loss: Float32Array, acc: Float32Array}>
|
|
5
6
|
}
|
package/src/Softmax/index.js
CHANGED
|
@@ -1,18 +1,198 @@
|
|
|
1
|
-
import {
|
|
1
|
+
import { log } from 'console';
|
|
2
|
+
import { OperationManager } from '../OperationManager.js'
|
|
2
3
|
|
|
3
|
-
|
|
4
|
-
|
|
4
|
+
export class Softmax {
|
|
5
|
+
device;
|
|
6
|
+
vocabSize;
|
|
7
|
+
forwardPipline;
|
|
8
|
+
embeddingSize;
|
|
5
9
|
gpu;
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
|
|
10
|
+
paramsForward;
|
|
11
|
+
oldParam;
|
|
12
|
+
bp
|
|
13
|
+
constructor(device, embeddingSize, vocabSize) {
|
|
14
|
+
this.vocabSize = vocabSize;
|
|
15
|
+
this.device = device;
|
|
16
|
+
this.embeddingSize = embeddingSize;
|
|
17
|
+
this.gpu = new OperationManager(device);
|
|
18
|
+
this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/SF.wgsl")
|
|
19
|
+
this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/SB.wgsl")
|
|
20
|
+
this.lossPipline = this.gpu.getShaderPipline("./src/shaders/CalculateLoss.wgsl")
|
|
21
|
+
this.accPipline = this.gpu.getShaderPipline("./src/shaders/CalculateAcc.wgsl")
|
|
22
|
+
this.battention = this.gpu.getShaderPipline("./src/shaders/SoftmaxBackward.wgsl")
|
|
23
|
+
this.reduceFn = this.gpu.getShaderPipline("./src/shaders/ReductionFunction.wgsl")
|
|
24
|
+
}
|
|
25
|
+
async forward(logits, cpuReadBack, encoderCPU) {
|
|
26
|
+
const encoder = encoderCPU ? encoderCPU : this.device.createCommandEncoder()
|
|
27
|
+
const tokenCount = logits.size / (this.vocabSize * 4)
|
|
28
|
+
if (!this.paramsForward) {
|
|
29
|
+
this.paramsForward = await this.gpu.createUniFormBuffer(4 * 3)
|
|
9
30
|
}
|
|
10
|
-
|
|
31
|
+
if (!this.oldParam) {
|
|
32
|
+
this.oldParam = new Uint32Array([this.vocabSize, 0, 0])
|
|
33
|
+
}
|
|
34
|
+
|
|
35
|
+
if (this.oldParam[1] !== tokenCount) { await this.gpu.WriteBuffer(this.paramsForward, new Uint32Array([this.vocabSize, tokenCount, 0])); this.oldParam = new Uint32Array([this.vocabSize, tokenCount, 0]); }
|
|
36
|
+
const output = await this.gpu.createBuffer(logits.size);
|
|
37
|
+
const bind = await this.device.createBindGroup({
|
|
38
|
+
layout: await this.forwardPipline.getBindGroupLayout(0),
|
|
39
|
+
entries: [
|
|
40
|
+
{
|
|
41
|
+
binding: 0,
|
|
42
|
+
resource: {
|
|
43
|
+
buffer: logits
|
|
44
|
+
}
|
|
45
|
+
},
|
|
46
|
+
{
|
|
47
|
+
binding: 1,
|
|
48
|
+
resource: {
|
|
49
|
+
buffer: output
|
|
50
|
+
}
|
|
51
|
+
}, {
|
|
52
|
+
binding: 2,
|
|
53
|
+
resource: {
|
|
54
|
+
buffer: this.paramsForward
|
|
55
|
+
}
|
|
56
|
+
}
|
|
57
|
+
]
|
|
58
|
+
})
|
|
59
|
+
await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((tokenCount * this.vocabSize) / 256))
|
|
60
|
+
if (!encoderCPU) {
|
|
61
|
+
await this.device.queue.submit([encoder.finish()])
|
|
62
|
+
logits.destroy()
|
|
63
|
+
}
|
|
64
|
+
if (cpuReadBack) {
|
|
65
|
+
const outputCPU = await this.gpu.readBuffer(output.size, output);
|
|
66
|
+
return outputCPU
|
|
67
|
+
}
|
|
68
|
+
return output
|
|
11
69
|
}
|
|
12
|
-
backward(
|
|
13
|
-
|
|
70
|
+
async backward(probs, target) {
|
|
71
|
+
const encoder = this.device.createCommandEncoder()
|
|
72
|
+
const tokenCount = probs.size / (this.vocabSize * 4)
|
|
73
|
+
if (!this.paramsForward) {
|
|
74
|
+
this.paramsForward = await this.gpu.createUniFormBuffer(4 * 3)
|
|
75
|
+
}
|
|
76
|
+
if (!this.oldParam) {
|
|
77
|
+
this.oldParam = new Uint32Array([this.vocabSize, 0, 0])
|
|
78
|
+
}
|
|
79
|
+
if (this.oldParam[1] !== tokenCount) {
|
|
80
|
+
await this.gpu.WriteBuffer(this.paramsForward, new Uint32Array([this.vocabSize, tokenCount, 0]));
|
|
81
|
+
this.oldParam = new Uint32Array([this.vocabSize, tokenCount, 0]);
|
|
82
|
+
}
|
|
83
|
+
const targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
|
|
84
|
+
await this.gpu.WriteBuffer(targetBuffer, new Float32Array(target))
|
|
85
|
+
const output = await this.gpu.createBuffer(probs.size)
|
|
86
|
+
const bind = await this.device.createBindGroup({
|
|
87
|
+
layout: await this.backwardPipline.getBindGroupLayout(0),
|
|
88
|
+
entries: [
|
|
89
|
+
{
|
|
90
|
+
binding: 0,
|
|
91
|
+
resource: {
|
|
92
|
+
buffer: probs
|
|
93
|
+
}
|
|
94
|
+
},
|
|
95
|
+
{
|
|
96
|
+
binding: 1,
|
|
97
|
+
resource: {
|
|
98
|
+
buffer: output
|
|
99
|
+
}
|
|
100
|
+
}, {
|
|
101
|
+
binding: 2,
|
|
102
|
+
resource: {
|
|
103
|
+
buffer: targetBuffer
|
|
104
|
+
}
|
|
105
|
+
}, {
|
|
106
|
+
binding: 3,
|
|
107
|
+
resource: {
|
|
108
|
+
buffer: this.paramsForward
|
|
109
|
+
}
|
|
110
|
+
}
|
|
111
|
+
]
|
|
112
|
+
})
|
|
113
|
+
await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.vocabSize) / 256))
|
|
114
|
+
this.device.queue.submit([encoder.finish()])
|
|
115
|
+
probs.destroy()
|
|
116
|
+
targetBuffer.destroy()
|
|
117
|
+
return output
|
|
14
118
|
}
|
|
15
|
-
crossEntropy(probs, target){
|
|
16
|
-
|
|
119
|
+
async crossEntropy(probs, target) {
|
|
120
|
+
const encoder = this.device.createCommandEncoder()
|
|
121
|
+
|
|
122
|
+
const tokenCount = probs.size / (this.vocabSize * 4)
|
|
123
|
+
const targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
|
|
124
|
+
const loss = await this.gpu.createBuffer(4 * tokenCount)
|
|
125
|
+
const acc = await this.gpu.createBuffer(4 * 2)
|
|
126
|
+
await this.gpu.WriteBuffer(targetBuffer, new Float32Array(target))
|
|
127
|
+
const bind = await this.device.createBindGroup({
|
|
128
|
+
layout: await this.lossPipline.getBindGroupLayout(0),
|
|
129
|
+
entries: [
|
|
130
|
+
{
|
|
131
|
+
binding: 0,
|
|
132
|
+
resource: {
|
|
133
|
+
buffer: probs
|
|
134
|
+
}
|
|
135
|
+
},
|
|
136
|
+
{
|
|
137
|
+
binding: 1,
|
|
138
|
+
resource: {
|
|
139
|
+
buffer: targetBuffer
|
|
140
|
+
}
|
|
141
|
+
}, {
|
|
142
|
+
binding: 2,
|
|
143
|
+
resource: {
|
|
144
|
+
buffer: loss
|
|
145
|
+
}
|
|
146
|
+
}, {
|
|
147
|
+
binding: 3,
|
|
148
|
+
resource: {
|
|
149
|
+
buffer: this.paramsForward
|
|
150
|
+
}
|
|
151
|
+
|
|
152
|
+
},
|
|
153
|
+
]
|
|
154
|
+
})
|
|
155
|
+
await this.gpu.RunPipline(encoder, this.lossPipline, bind, Math.ceil((this.vocabSize * tokenCount) / 256))
|
|
156
|
+
const outputLoss = await this.gpu.createBuffer(4)
|
|
157
|
+
const reduceBind = await this.gpu.createBindGroup(this.reduceFn, 0, [loss, outputLoss, this.paramsForward])
|
|
158
|
+
await this.gpu.RunPipline(encoder, this.reduceFn, reduceBind, 1);
|
|
159
|
+
const bind2 = await this.device.createBindGroup({
|
|
160
|
+
layout: await this.accPipline.getBindGroupLayout(0),
|
|
161
|
+
entries: [
|
|
162
|
+
{
|
|
163
|
+
binding: 0,
|
|
164
|
+
resource: {
|
|
165
|
+
buffer: this.paramsForward
|
|
166
|
+
}
|
|
167
|
+
},
|
|
168
|
+
{
|
|
169
|
+
binding: 1,
|
|
170
|
+
resource: {
|
|
171
|
+
buffer: probs
|
|
172
|
+
}
|
|
173
|
+
}, {
|
|
174
|
+
binding: 3,
|
|
175
|
+
resource: {
|
|
176
|
+
buffer: targetBuffer
|
|
177
|
+
}
|
|
178
|
+
},
|
|
179
|
+
{
|
|
180
|
+
binding: 2,
|
|
181
|
+
resource: {
|
|
182
|
+
buffer: acc
|
|
183
|
+
}
|
|
184
|
+
}
|
|
185
|
+
]
|
|
186
|
+
})
|
|
187
|
+
await this.gpu.RunPipline(encoder, this.accPipline, bind2, 1)
|
|
188
|
+
this.device.queue.submit([encoder.finish()])
|
|
189
|
+
targetBuffer.destroy()
|
|
190
|
+
loss.destroy()
|
|
191
|
+
const lossCpu = await this.gpu.readBuffer(outputLoss.size, outputLoss)
|
|
192
|
+
const accCPU = await this.gpu.readBuffer(acc.size, acc)
|
|
193
|
+
acc.destroy()
|
|
194
|
+
return { loss: lossCpu[0], acc: accCPU[0] }
|
|
17
195
|
}
|
|
18
|
-
|
|
196
|
+
|
|
197
|
+
}
|
|
198
|
+
|
package/src/Tokenizer/index.d.ts
CHANGED
package/src/Tokenizer/index.js
CHANGED
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
import { readFileSync } from 'fs'
|
|
1
|
+
import fs, { existsSync, readFileSync } from 'fs'
|
|
2
2
|
import { join, dirname, format } from 'path';
|
|
3
3
|
import { fileURLToPath } from 'url';
|
|
4
4
|
|
|
@@ -9,14 +9,20 @@ export class Tokenizer {
|
|
|
9
9
|
IdtoWord = {};
|
|
10
10
|
WordtoID = {};
|
|
11
11
|
EOSToken = 2
|
|
12
|
-
constructor() {
|
|
13
|
-
|
|
12
|
+
constructor(filename="") {
|
|
13
|
+
if (fs.existsSync(filename)) {
|
|
14
|
+
console.log("Warning: Tokenizer is loaded with a Custom Vocab File!")
|
|
15
|
+
console.log("")
|
|
16
|
+
this.vocabs = JSON.parse(readFileSync(filename, "utf-8"))
|
|
17
|
+
} else {
|
|
18
|
+
this.vocabs = JSON.parse(readFileSync(join(__dirname, "vocabs.json"), "utf-8"))
|
|
19
|
+
}
|
|
14
20
|
this.vocabSize = this.vocabs.length
|
|
15
21
|
this.IdtoWord = Object.fromEntries(this.vocabs.map((e, idx) => [idx, e]))
|
|
16
22
|
this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
|
|
17
23
|
}
|
|
18
24
|
encoder(text = "") {
|
|
19
|
-
const AddSpaces = text.toLocaleLowerCase().replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ")
|
|
25
|
+
const AddSpaces = text.toLocaleLowerCase().replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ").replaceAll("<" , " < ")
|
|
20
26
|
const splitedWords = AddSpaces.split(/\s+/).filter(e => e.length > 0)
|
|
21
27
|
const output = []
|
|
22
28
|
for (let a = 0; a < splitedWords.length; a++) {
|
|
@@ -34,9 +40,9 @@ export class Tokenizer {
|
|
|
34
40
|
return true
|
|
35
41
|
}
|
|
36
42
|
return false
|
|
37
|
-
|
|
43
|
+
|
|
38
44
|
}).reduce((prev, curr) => prev && curr, true)
|
|
39
|
-
|
|
45
|
+
if (wordsliced.length == 0) output.push("<unk>")
|
|
40
46
|
if (matched) {
|
|
41
47
|
wordsliced.map((e) => output.push(e))
|
|
42
48
|
output.push("$130")
|
|
@@ -70,6 +76,85 @@ export class Tokenizer {
|
|
|
70
76
|
}
|
|
71
77
|
}
|
|
72
78
|
}
|
|
73
|
-
return output.filter(e=> e.trim().length > 0).join(" ").replaceAll(/\s+([,.;'{}`!~\"?])/g, "$1")
|
|
79
|
+
return output.filter(e => e.trim().length > 0).join(" ").replaceAll(/\s+([,.;'{}`!~\"?])/g, "$1")
|
|
80
|
+
}
|
|
81
|
+
trainTokenizer(text = "ejrnienieijr", mergeLength = 5, vocabFileName) {
|
|
82
|
+
console.log("Training The Tokenizer....")
|
|
83
|
+
const arrText = text.toLocaleLowerCase().replaceAll("\r\n", "").replaceAll("\n", "").replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ").split(/\s+/).map(e => e.trim()).filter((e) => e.length > 0)
|
|
84
|
+
const reg = new RegExp(`.{0,${mergeLength}}`, "g")
|
|
85
|
+
const prevVocab = existsSync(vocabFileName) ? JSON.parse(readFileSync(vocabFileName, "utf-8")) : [];
|
|
86
|
+
const specal = ["<unk>",
|
|
87
|
+
"<PAD>",
|
|
88
|
+
"<EOS>",
|
|
89
|
+
"$12",
|
|
90
|
+
"$130",
|
|
91
|
+
"😀",
|
|
92
|
+
"😃",
|
|
93
|
+
"😄",
|
|
94
|
+
"😁",
|
|
95
|
+
"😆",
|
|
96
|
+
"😅",
|
|
97
|
+
"😂",
|
|
98
|
+
"🤣",
|
|
99
|
+
"😊",
|
|
100
|
+
"😎",
|
|
101
|
+
"❤️",
|
|
102
|
+
"👍",
|
|
103
|
+
"👎",
|
|
104
|
+
"🔥",
|
|
105
|
+
"💯",
|
|
106
|
+
"🎉",
|
|
107
|
+
"🚀",
|
|
108
|
+
"⭐",
|
|
109
|
+
"😭",
|
|
110
|
+
"😍",
|
|
111
|
+
"🥰",
|
|
112
|
+
"😡",
|
|
113
|
+
"🤔",
|
|
114
|
+
"😢",
|
|
115
|
+
"😮",
|
|
116
|
+
"👋",
|
|
117
|
+
"🙏",
|
|
118
|
+
"💀",
|
|
119
|
+
"🤖",
|
|
120
|
+
"👀",
|
|
121
|
+
"💕",
|
|
122
|
+
"💔",
|
|
123
|
+
"😂😂",
|
|
124
|
+
"😂😭",
|
|
125
|
+
"🔥🔥",
|
|
126
|
+
"❤️❤️",
|
|
127
|
+
"👍👍",
|
|
128
|
+
"<",
|
|
129
|
+
">",
|
|
130
|
+
".",
|
|
131
|
+
"{",
|
|
132
|
+
"}",
|
|
133
|
+
"?",
|
|
134
|
+
"!",
|
|
135
|
+
"@",
|
|
136
|
+
"$",
|
|
137
|
+
":",
|
|
138
|
+
"/",
|
|
139
|
+
"+",
|
|
140
|
+
"-",
|
|
141
|
+
"=",
|
|
142
|
+
"%",
|
|
143
|
+
"*",
|
|
144
|
+
"@"
|
|
145
|
+
]
|
|
146
|
+
const uniArr = [...new Set([...arrText, ...specal, ...prevVocab])];
|
|
147
|
+
const output = []
|
|
148
|
+
for(let a = 0 ; a < uniArr.length; a++){
|
|
149
|
+
let text = uniArr[a].match(reg).filter((e)=> e.length > 0);
|
|
150
|
+
output.push(...text)
|
|
151
|
+
}
|
|
152
|
+
let finalOutput = [...new Set([...output])]
|
|
153
|
+
fs.writeFileSync(vocabFileName, JSON.stringify(finalOutput, null, 2), "utf-8");
|
|
154
|
+
console.log("_________________________________")
|
|
155
|
+
console.log(`New VocabSize : ${finalOutput.length - prevVocab.length}`);
|
|
156
|
+
console.log("_________________________________")
|
|
157
|
+
console.log("Operation Completed!")
|
|
158
|
+
|
|
74
159
|
}
|
|
75
160
|
}
|