@mineproness/tensorloop 2.1.54 → 2.1.56

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.
Files changed (51) hide show
  1. package/index.js +14 -20
  2. package/package.json +5 -2
  3. package/readme.md +907 -72
  4. package/src/Embedding/index.d.ts +5 -5
  5. package/src/Embedding/index.js +118 -33
  6. package/src/FFN/index.d.ts +5 -5
  7. package/src/FFN/index.js +34 -39
  8. package/src/Linear/index.d.ts +7 -8
  9. package/src/Linear/index.js +173 -49
  10. package/src/OperationManager.js +96 -0
  11. package/src/Relu/index.d.ts +5 -0
  12. package/src/Relu/index.js +103 -0
  13. package/src/SelfAttention/index.d.ts +7 -13
  14. package/src/SelfAttention/index.js +108 -55
  15. package/src/Softmax/index.d.ts +4 -3
  16. package/src/Softmax/index.js +192 -12
  17. package/src/Tokenizer/index.d.ts +1 -0
  18. package/src/Tokenizer/index.js +92 -7
  19. package/src/TransformerBlock/index.d.ts +6 -7
  20. package/src/TransformerBlock/index.js +24 -23
  21. package/src/getVec.js +28 -0
  22. package/src/shaders/CalculateAcc.wgsl +50 -0
  23. package/src/shaders/CalculateLoss.wgsl +35 -0
  24. package/src/shaders/ComputeDInput.wgsl +32 -0
  25. package/src/shaders/EB.wgsl +31 -0
  26. package/src/shaders/EF.wgsl +25 -0
  27. package/src/shaders/LB.wgsl +57 -0
  28. package/src/shaders/LF.wgsl +44 -0
  29. package/src/shaders/RB.wgsl +37 -0
  30. package/src/shaders/RF.wgsl +33 -0
  31. package/src/shaders/ReductionFunction.wgsl +28 -0
  32. package/src/shaders/SB.wgsl +36 -0
  33. package/src/shaders/SF.wgsl +65 -0
  34. package/src/shaders/SoftmaxBackward.wgsl +31 -0
  35. package/src/shaders/attentionShader/CompueFinalForwardOutput.wgsl +29 -0
  36. package/src/shaders/attentionShader/ComputeDA.wgsl +30 -0
  37. package/src/shaders/attentionShader/ComputeDV.wgsl +37 -0
  38. package/src/shaders/attentionShader/ComputeK.wgsl +30 -0
  39. package/src/shaders/attentionShader/ComputedQ.wgsl +31 -0
  40. package/src/shaders/attentionShader/attentionSocres.wgsl +40 -0
  41. package/src/shaders/attentionShader/sumGradient.wgsl +24 -0
  42. package/src/GELU/index.d.ts +0 -6
  43. package/src/GELU/index.js +0 -33
  44. package/src/GetConfigs.js +0 -35
  45. package/src/Model/index.d.ts +0 -7
  46. package/src/Model/index.js +0 -93
  47. package/src/PositionEmbedding/index.d.ts +0 -11
  48. package/src/PositionEmbedding/index.js +0 -49
  49. package/src/ReLU/index.d.ts +0 -5
  50. package/src/ReLU/index.js +0 -17
  51. package/src/gpiInit.js +0 -764
@@ -1,11 +1,11 @@
1
1
  export class Embedding{
2
- vectors: Float32Array[]
3
- constructor(embeddingSize: number, vocabSize: number, contextSize: number, configs: {
2
+ constructor(device: GPUBuffer ,embedding : number, vocabSize: number, config:{
3
+ cpuReadBack: boolean,
4
4
  save: {
5
5
  filename: string
6
6
  }
7
7
  })
8
- forward(token: number | number[]) : Float32Array[] | Float32Array[][]
9
- backward(inputGradient: number[] | Float32Array[], learingRate: number) : void
10
- Save() : void
8
+ async forward(tokens: number[]) : Promise<GPUBuffer>
9
+ async Save() : void
10
+ async backward(tokens: number[], inputGradient: GPUBuffer, lr: number) : Promise<void>
11
11
  }
@@ -1,53 +1,138 @@
1
- import fs, { existsSync } from 'fs'
2
- import { SaveVectors, LoadVectors } from '../GetConfigs.js'
3
- import { GPUEmbedding } from '../gpiInit.js'
4
- // import { GPUEmbedding } from '../GPU.js'
5
- // import { config } from 'process'
1
+ // import { config } from "process";
2
+ import { SaveVector, getVector } from '../getVec.js'
3
+ import { OperationManager } from "../OperationManager.js";
6
4
 
7
5
  function GenerateVectors(embeddingSize, vocabSize) {
8
- const vector = []
9
- for (let token = 0; token < vocabSize; token++) {
10
- vector[token] = new Float32Array(embeddingSize)
11
- for (let a = 0; a < embeddingSize; a++) {
12
- vector[token][a] = (Math.random() * 2 - 1) * 0.2
13
- }
14
- }
15
- return vector
6
+ const vec = new Float32Array(embeddingSize * vocabSize).map(() => (Math.random() * 2 - 1) * 0.02)
7
+ return vec
16
8
  }
17
9
 
18
10
  export class Embedding {
19
- vectors = []
20
- configs;
21
- embeddingSize = 0
11
+ embeddingSize;
12
+ vocabSize;
22
13
  gpu;
23
- constructor(embeddingSize, vocabSize, contextSize, configs) {
14
+ firstTrue;
15
+ vectors;
16
+ configs;
17
+ vectorsBuffer;
18
+ forwardPipline;
19
+ constructor(device, embeddingSize, vocabSize, configs) {
24
20
  this.configs = configs
25
- this.embeddingSize = embeddingSize
21
+ this.embeddingSize = embeddingSize;
22
+ this.vocabSize = vocabSize;
23
+ this.device = device
24
+ let total = embeddingSize * vocabSize;
26
25
  if (configs && configs.save) {
27
- if (existsSync(configs.save.filename)) {
28
- this.vectors = LoadVectors(configs.save.filename, embeddingSize)
26
+ const vec = getVector(configs.save.filename, total);
27
+ if (vec) {
28
+ this.vectors = vec;
29
+
29
30
  } else {
30
31
  this.vectors = GenerateVectors(embeddingSize, vocabSize)
31
- SaveVectors(this.vectors, this.configs.save.filename)
32
32
  }
33
+
33
34
  } else {
34
35
  this.vectors = GenerateVectors(embeddingSize, vocabSize)
35
36
 
36
- }
37
- this.gpu = new GPUEmbedding(this.vectors, vocabSize, embeddingSize, contextSize)
37
+ }
38
+ this.gpu = new OperationManager(device)
39
+ this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/EF.wgsl", {
40
+ embeddingSize: embeddingSize
41
+ })
42
+ this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/EB.wgsl", {
43
+ embeddingSize: embeddingSize
44
+ })
45
+ SaveVector(this.configs.save.filename, this.vectors)
38
46
  }
39
- forward(token) {
40
- return this.gpu.forward(token)
47
+ async forward(input) {
48
+ const encoder = this.device.createCommandEncoder()
49
+ if (!this.vectorsBuffer) {
50
+ this.vectorsBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize);
51
+ await this.gpu.WriteBuffer(this.vectorsBuffer, this.vectors)
52
+ }
53
+ const tokenBuffer = await this.gpu.createBuffer(Uint32Array.BYTES_PER_ELEMENT * input.length);
54
+ await this.gpu.WriteBuffer(tokenBuffer, new Uint32Array(input))
55
+ const outputBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length);
56
+ const bind = await this.device.createBindGroup({
57
+ layout: await this.forwardPipline.getBindGroupLayout(0),
58
+ entries: [
59
+ {
60
+ binding: 0,
61
+ resource: {
62
+ buffer: this.vectorsBuffer
63
+ }
64
+ },
65
+ {
66
+ binding: 1,
67
+ resource: {
68
+ buffer: outputBuffer
69
+ }
70
+ }, {
71
+ binding: 2,
72
+ resource: {
73
+ buffer: tokenBuffer
74
+ }
75
+ }
76
+ ]
77
+ })
78
+ await this.gpu.RunPipline(encoder,this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
79
+ this.device.queue.submit([encoder.finish()])
80
+ tokenBuffer.destroy()
81
+ if (this.configs.cpuReadBack) {
82
+ const output = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length, outputBuffer);
83
+ outputBuffer.destroy()
84
+ return output
85
+ }
86
+ return outputBuffer
87
+
41
88
  }
42
- backward(inputGradient, learingRate){
43
- this.gpu.backward(inputGradient, learingRate)
89
+ async Save() {
90
+ const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize, this.vectorsBuffer)
91
+ SaveVector(this.configs.save.filename, vec)
44
92
  }
45
- Save() {
46
- const vectors = []
47
- for(let i = 0; i < this.gpu.vectors.length; i++){
48
- const out = this.gpu.vectors[i].toArray()
49
- vectors.push(...out)
93
+ async backward(token, inputGradient, lr) {
94
+ const encoder = this.device.createCommandEncoder()
95
+ if (!this.oldParam) {
96
+ this.oldParam = new Float32Array([0, 0]);
97
+ this.param = await this.gpu.createUniFormBuffer(this.oldParam.byteLength);
98
+ this.tokenBuffer = await this.gpu.createBuffer(4 * token.length)
50
99
  }
51
- SaveVectors(vectors.map((e)=> new Float32Array(e)), this.configs.save.filename)
100
+ if (this.oldParam[1] !== lr) {
101
+ this.oldParam = new Float32Array([lr, 0]);
102
+ this.gpu.WriteBuffer(this.param, this.oldParam)
103
+ }
104
+ const bind = await this.device.createBindGroup({
105
+ layout: await this.backwardPipline.getBindGroupLayout(0),
106
+ entries: [
107
+ {
108
+ binding: 0,
109
+ resource: {
110
+ buffer: this.vectorsBuffer
111
+ }
112
+ },
113
+ {
114
+ binding: 1,
115
+ resource: {
116
+ buffer: inputGradient
117
+ }
118
+ }, {
119
+ binding: 2,
120
+ resource: {
121
+ buffer: this.tokenBuffer
122
+ }
123
+
124
+ }, {
125
+ binding: 3,
126
+ resource: {
127
+ buffer: this.param
128
+ }
129
+
130
+ }
131
+ ]
132
+ })
133
+ await this.gpu.WriteBuffer(this.tokenBuffer, new Float32Array(token));
134
+ await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * token.length) / 256))
135
+ this.device.queue.submit([encoder.finish()])
136
+ inputGradient.destroy()
52
137
  }
53
138
  }
@@ -1,9 +1,9 @@
1
1
  export class FFN{
2
- constructor(embeddingSize: number, hiddenLayer: "gelu" | "relu" , configs: {
2
+ constructor(device: GPUAdapter, embeddingSize: number, configs: {
3
3
  dirname: string
4
4
  })
5
- forward(input: Float32Array[]) : Float32Array[]
6
- backward(dInput: Float32Array[], LearningRate: number) : Float32Array[]
7
- Save() : void;
8
- ClearInputCache() : void
5
+ async forward(vec: GPUBuffer) : Promise<GPUBuffer>
6
+ async backward(inputGradient: GPUBuffer, lr: number) : Promise<GPUBuffer>
7
+ async Save() : Promise<void>
8
+ async ClearInputCache() : Promise<void>
9
9
  }
package/src/FFN/index.js CHANGED
@@ -1,51 +1,46 @@
1
- import { Linear } from '../Linear/index.js'
2
- import { GELU } from '../GELU/index.js'
3
- import { ReLU } from '../ReLU/index.js'
1
+ import { RELU, Linear } from "../../index.js";
4
2
  import { join } from 'path'
5
- import { GPUHiddenLayer } from '../gpiInit.js';
3
+ import { OperationManager } from "../OperationManager.js";
4
+
5
+
6
6
  export class FFN {
7
7
  linear1;
8
8
  linear2;
9
- hidden;
10
- hiddenName;
11
- constructor(embeddingSize, hiddenLayer = "gelu", configs) {
12
- this.hiddenName = hiddenLayer
13
- this.hidden = hiddenLayer.toLocaleLowerCase() == "gelu" ? new GELU() : new GPUHiddenLayer()
14
- if (configs) {
15
- this.linear1 = new Linear(embeddingSize, 4 * embeddingSize, {
16
- save: {
17
- filename: [join(configs.dirname, "ffnWeights.bin"), join(configs.dirname, "ffnBias.bin")]
18
- }
19
- })
20
- this.linear2 = new Linear(4 * embeddingSize, embeddingSize, {
21
- save: {
22
- filename: [join(configs.dirname, "ffnWeights1.bin"), join(configs.dirname, "ffnBias1.bin")]
23
- }
24
- })
25
- } else {
26
- this.linear1 = new Linear(embeddingSize, 4 * embeddingSize)
27
- this.linear2 = new Linear(4 * embeddingSize, embeddingSize)
28
- }
9
+ hiddenLayer;
10
+ Op
11
+ constructor(device, embeddingSize, configs) {
12
+ this.Op = new OperationManager(device)
13
+ this.linear1 = new Linear(device, embeddingSize, embeddingSize * 4, {
14
+ save: {
15
+ filename: [join(configs.dirname, "FFN_weights1.bin"), join(configs.dirname, "FFN_bias1.bin")]
16
+ }
17
+ });
18
+ this.hiddenLayer = new RELU(device, embeddingSize * 4);
19
+ this.linear2 = new Linear(device, embeddingSize * 4, embeddingSize, {
20
+ save: {
21
+ filename: [join(configs.dirname, "FFN_weights2.bin"), join(configs.dirname, "FFN_bias2.bin")]
22
+ }
23
+ })
29
24
  }
30
- forward(input) {
31
- const out1 = this.linear1.forward(input)
32
- let hidden = this.hidden.forward(out1)
33
- const result = this.linear2.forward(hidden)
34
- return result
25
+ async forward(vec) {
26
+ const logit1 = await this.linear1.forward(vec)
27
+ const hidden = await this.hiddenLayer.forward(logit1)
28
+ const final = await this.linear2.forward(hidden)
29
+ return final
35
30
  }
36
- backward(dInput, LearningRate) {
37
- const out2 = this.linear2.backward(dInput, LearningRate);
38
- const dHidden = this.hidden.backward(out2)
39
- const result = this.linear1.backward(dHidden, LearningRate)
40
- return result
31
+ async backward(inputGradient, lr) {
32
+ const logit2 = await this.linear2.backward(inputGradient, lr)
33
+ const hidden = await this.hiddenLayer.backward(logit2)
34
+ const final = await this.linear1.backward(hidden, lr)
35
+ inputGradient.destroy()
36
+ return final
41
37
  }
42
- Save() {
43
- this.linear1.Save()
44
- this.linear2.Save()
38
+ async Save(){
39
+ await this.linear1.Save()
40
+ await this.linear2.Save()
45
41
  }
46
- ClearInputCache() {
42
+ async ClearInputCache(){
47
43
  this.linear1.ClearInputCache()
48
44
  this.linear2.ClearInputCache()
49
- this.hidden.ClearInputCache()
50
45
  }
51
46
  }
@@ -1,13 +1,12 @@
1
- export class Linear{
2
- Weights : Float32Array[]
3
- Bias: Float32Array
4
- constructor(embeddingSize : number, vocabSize : number, configs: {
1
+ export class Linear {
2
+ constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number, configs: {
3
+ cpuReadBack: boolean,
5
4
  save: {
6
5
  filename: string[]
7
6
  }
8
7
  })
9
- forward(input : Float32Array[]) : Float32Array[]
10
- backward(outGradient : number[] , learingRate: number) : Float32Array[]
11
- Save() : void
12
- ClearInputCache() : void
8
+ async forward(vec: GPUBuffer): Promise<GPUBuffer | Float32Array>
9
+ async backward(dOutput: GPUBuffer, lr: number): Promise<GPUBuffer>
10
+ async Save(): Promise<void>
11
+ async ClearInputCache() : Promise<void>
13
12
  }
@@ -1,67 +1,191 @@
1
- import { LoadBias, LoadVectors, SaveBias, SaveVectors } from '../GetConfigs.js'
2
- import { existsSync } from 'fs'
3
- import { GPULinear } from '../gpiInit.js'
4
- function GenerateWeightsBias(embeddingSize, vocabSize) {
5
- const vectors = []
6
- const Bias = new Float32Array(vocabSize)
7
- for (let a = 0; a < vocabSize; a++) {
8
- vectors[a] = new Float32Array(embeddingSize)
9
- for (let t = 0; t < embeddingSize; t++) {
10
- vectors[a][t] = (Math.random() * 2 - 1) * 0.02
11
- }
12
- Bias[a] = 0
13
- }
14
- return { vectors, Bias }
1
+
2
+ import { create, globals } from "webgpu";
3
+ import fs from 'fs'
4
+ import { buffer } from "stream/consumers";
5
+ import { SaveVector, getVector } from '../getVec.js'
6
+ import { OperationManager } from "../OperationManager.js";
7
+ Object.assign(globalThis, globals)
8
+ const navigator = { gpu: create([]) }
9
+
10
+
11
+ function GenerateVectors(embeddingSize, vocabSize) {
12
+ const vec = new Float32Array(embeddingSize * vocabSize).map(() => (Math.random() * 2 - 1) * 0.02)
13
+ return vec
14
+ }
15
+
16
+ function GenerateBias(vocabSize) {
17
+ const vec = new Float32Array(vocabSize)
18
+ return vec
15
19
  }
16
20
 
17
21
  export class Linear {
18
- Weights = []
19
- Bias = []
22
+ device;
23
+ embeddingSize;
24
+ vocabSize;
25
+ configs;
26
+ bias;
27
+ weights;
28
+ forwardPipline;
20
29
  configs;
21
- input = [];
30
+ input;
31
+ wBuffer;
32
+ bBuffer;
22
33
  gpu;
23
- constructor(embeddingSize, vocabSize, configs) {
34
+ paramForward;
35
+ oldParam;
36
+ constructor(device, embeddingSize, vocabSize, configs) {
37
+ this.configs = configs
38
+ this.embeddingSize = embeddingSize;
39
+ this.vocabSize = vocabSize;
40
+ this.configs = configs;
41
+ this.device = device
42
+ let total = embeddingSize * vocabSize;
24
43
  if (configs && configs.save) {
25
- this.configs = configs
26
- if (existsSync(configs.save.filename[0]) && existsSync(configs.save.filename[1])) {
27
- const Vectors = LoadVectors(configs.save.filename[0], embeddingSize)
28
- const Bias = LoadBias(configs.save.filename[1])
29
- this.Weights = Vectors;
30
- this.Bias = Bias
44
+ const w = getVector(configs.save.filename[0], total);
45
+ const b = getVector(configs.save.filename[1], total, true)
46
+ if (!w && !b) {
47
+ this.bias = GenerateBias(vocabSize)
48
+ this.weights = GenerateVectors(embeddingSize, vocabSize)
49
+ SaveVector(configs.save.filename[0], this.weights)
50
+ SaveVector(configs.save.filename[1], this.bias)
31
51
  } else {
32
- const { vectors, Bias } = GenerateWeightsBias(embeddingSize, vocabSize)
33
- this.Weights = vectors;
34
- this.Bias = Bias;
35
- SaveVectors(this.Weights, configs.save.filename[0])
36
- SaveBias(this.Bias, configs.save.filename[1])
52
+ this.bias = b;
53
+ this.weights = w;
37
54
  }
38
55
  } else {
39
- const { vectors, Bias } = GenerateWeightsBias(embeddingSize, vocabSize)
40
- this.Weights = vectors;
41
- this.Bias = Bias;
56
+ this.bias = GenerateBias(vocabSize)
57
+ this.weights = GenerateVectors(embeddingSize, vocabSize)
42
58
  }
43
-
59
+ this.gpu = new OperationManager(device);
60
+ this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/LF.wgsl")
61
+ this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/LB.wgsl")
62
+ this.dInputBackward = this.gpu.getShaderPipline("./src/shaders/ComputeDInput.wgsl")
44
63
  }
45
- Save() {
46
- const {bias, weights} = this.gpu.Save()
47
- SaveVectors(weights, this.configs.save.filename[0])
48
- SaveBias(bias, this.configs.save.filename[1])
64
+ async forward(vec) {
65
+ const encoder = this.device.createCommandEncoder()
66
+ const tokenCount = vec.size / (this.embeddingSize * 4)
67
+ if (!this.wBuffer) {
68
+
69
+ this.wBuffer = await this.gpu.createBuffer(this.embeddingSize * this.vocabSize * 4)
70
+ await this.gpu.WriteBuffer(this.wBuffer, this.weights)
71
+ this.bBuffer = await this.gpu.createBuffer(this.vocabSize * 4)
72
+ await this.gpu.WriteBuffer(this.bBuffer, this.bias)
73
+ this.paramForward = await this.gpu.createUniFormBuffer(4 * 4)
74
+ this.paramBackward = await this.gpu.createUniFormBuffer(4 * 5)
75
+ }
76
+ if (!this.oldParam) {
77
+ this.oldParam = new Uint32Array([this.embeddingSize, this.vocabSize, 0, 0])
78
+ }
79
+ if (this.oldParam[2] !== tokenCount) {
80
+ await this.gpu.WriteBuffer(this.paramForward, new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount, 0]))
81
+ this.oldParam = new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount, 0]);
82
+ }
83
+ const outputBuffer = await this.gpu.createBuffer(4 * this.vocabSize * tokenCount)
84
+ const bind = await this.device.createBindGroup({
85
+ layout: await this.forwardPipline.getBindGroupLayout(0),
86
+ entries: [
87
+ {
88
+ binding: 0,
89
+ resource: this.wBuffer
90
+ },
91
+ {
92
+ binding: 1,
93
+ resource: outputBuffer
94
+ },
95
+ {
96
+ binding: 2,
97
+ resource: vec
98
+ },
99
+ {
100
+ binding: 3,
101
+ resource: this.bBuffer
102
+ },
103
+ {
104
+ binding: 4,
105
+ resource: this.paramForward
106
+ }
107
+ ]
108
+ })
109
+ await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.vocabSize * tokenCount) / 256))
110
+ if (this.input) this.input.destroy()
111
+ this.input = await this.gpu.MakeACopyBuffer(encoder, vec);
112
+ this.device.queue.submit([encoder.finish()])
113
+ return outputBuffer
49
114
  }
50
- forward(input) {
51
- if (!this.gpu) {
52
- this.gpu = new GPULinear(this.Weights, this.Bias, input.length)
115
+ async backward(dOutput, learningRate) {
116
+ const encoder = this.device.createCommandEncoder()
117
+ const tokenCount = dOutput.size / (this.vocabSize * 4)
118
+ if (!this.oldParamb) {
119
+ this.oldParamb = new Float32Array([this.vocabSize, 0, 0, this.embeddingSize, 0])
53
120
  }
54
- for (let a = 0; a < input.length; a++) {
55
- this.input[a] = input;
121
+ if (this.oldParamb[1] !== tokenCount || this.oldParamb[2] !== learningRate) {
122
+ await this.gpu.WriteBuffer(this.paramBackward, new Float32Array([this.vocabSize, tokenCount, learningRate, this.embeddingSize, 0]));
123
+ this.oldParamb = new Float32Array([this.vocabSize, tokenCount, learningRate, this.embeddingSize, 0]);
56
124
  }
57
- return this.gpu.forward(input)
125
+ const outputBuffer = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4);
126
+ const bind = await this.device.createBindGroup({
127
+ layout: await this.backwardPipline.getBindGroupLayout(0),
128
+ entries: [
129
+ {
130
+ binding: 0,
131
+ resource: this.bBuffer
132
+ },
133
+ {
134
+ binding: 1,
135
+ resource: this.wBuffer
136
+ },
137
+ {
138
+ binding: 2,
139
+ resource: dOutput
140
+ },
141
+ {
142
+ binding: 3,
143
+ resource: this.input
144
+ },
145
+ {
146
+ binding: 4,
147
+ resource: this.paramBackward
148
+ }
149
+ ]
150
+ })
151
+ const bind2 = await this.device.createBindGroup({
152
+ layout: await this.dInputBackward.getBindGroupLayout(0),
153
+ entries: [
154
+ {
155
+ binding: 0,
156
+ resource: this.wBuffer
157
+ },
158
+ {
159
+ binding: 1,
160
+ resource: dOutput
161
+ },
162
+ {
163
+ binding: 3,
164
+ resource: this.paramBackward
165
+ },
166
+ {
167
+ binding: 4,
168
+ resource: outputBuffer
169
+ }
170
+ ]
171
+ })
172
+ // console.log((await this.gpu.readBuffer(this.input.size, this.input))[0])
173
+ await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 256));
174
+ await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.vocabSize * tokenCount) / 256));
175
+ await this.device.queue.submit([encoder.finish()])
176
+ await dOutput.destroy()
177
+ this.input.destroy()
178
+ return outputBuffer
179
+
58
180
  }
59
- backward(out, lr){
60
- return this.gpu.backward(out, this.input, lr)
181
+ async Save() {
182
+ const b = await this.gpu.readBuffer(this.bBuffer.size, this.bBuffer);
183
+ const w = await this.gpu.readBuffer(this.wBuffer.size, this.wBuffer);
184
+ SaveVector(this.configs.save.filename[0], w)
185
+ SaveVector(this.configs.save.filename[1], b)
61
186
  }
62
- ClearInputCache() {
63
- this.input = []
187
+ async ClearInputCache() {
188
+ if (this.input) await this.input.destroy();
64
189
  }
65
- }
66
-
67
190
 
191
+ }
@@ -0,0 +1,96 @@
1
+ import fs from 'fs'
2
+
3
+ export class OperationManager {
4
+ device;
5
+ matrixMatSHhaders;
6
+ params;
7
+ matrixParams;
8
+ constructor(device) {
9
+ this.device = device;
10
+ }
11
+ async createBuffer(size) {
12
+ return this.device.createBuffer({
13
+ size,
14
+ usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
15
+ })
16
+ }
17
+ async createUniFormBuffer(size) {
18
+ return this.device.createBuffer({
19
+ size,
20
+ usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.UNIFORM
21
+ })
22
+ }
23
+ async readBuffer(size, buffer) {
24
+ const readBuffer = this.device.createBuffer({
25
+ size: size,
26
+ usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
27
+ })
28
+ const encoder = this.device.createCommandEncoder()
29
+ encoder.copyBufferToBuffer(buffer, 0, readBuffer, 0, size);
30
+ this.device.queue.submit([await encoder.finish()])
31
+ await readBuffer.mapAsync(GPUMapMode.READ);
32
+ const mapped = readBuffer.getMappedRange();
33
+ const result = new Float32Array(
34
+ mapped.slice(0)
35
+ );
36
+ readBuffer.unmap()
37
+ readBuffer.destroy()
38
+ return result
39
+ }
40
+ getShaderPipline(filename, value = { "embedding": 123 }, entryPoint = "main") {
41
+ let readShaderCode = fs.readFileSync(filename, "utf-8");
42
+ if (value) {
43
+ for (const a in value) {
44
+ readShaderCode = readShaderCode.replaceAll(`%${a}%`, value[a])
45
+ }
46
+ };
47
+ const compiledKernel = this.device.createShaderModule({
48
+ code: readShaderCode
49
+ });
50
+ const compiledPipLine = this.device.createComputePipeline({
51
+ layout: "auto",
52
+ compute: {
53
+ module: compiledKernel,
54
+ entryPoint: entryPoint
55
+ }
56
+ });
57
+ return compiledPipLine;
58
+ }
59
+ async WriteBuffer(buffer, data) {
60
+ this.device.queue.writeBuffer(buffer, 0, data)
61
+ }
62
+ async RunPipline(encoder, pipline, binding, dispatchWorkgroups, y = 1, z = 1) {
63
+ const pass = encoder.beginComputePass();
64
+ pass.setPipeline(pipline)
65
+ pass.setBindGroup(0, binding)
66
+ pass.dispatchWorkgroups(dispatchWorkgroups, y , z)
67
+ pass.end()
68
+ }
69
+
70
+ async MakeACopyBuffer(encoder, sourceBuffer) {
71
+ const buffer = this.device.createBuffer({
72
+ size: sourceBuffer.size,
73
+ usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
74
+ })
75
+ encoder.copyBufferToBuffer(sourceBuffer, 0, buffer, 0, sourceBuffer.size);
76
+ return buffer
77
+ }
78
+ CreateEncoder(){
79
+ return this.device.createCommandEncoder()
80
+ }
81
+ createBindGroup(shader,groupNum,ent){
82
+ return this.device.createBindGroup({
83
+ layout: shader.getBindGroupLayout(groupNum),
84
+ entries: ent.map((e, x)=>{
85
+ return {
86
+ binding : x,
87
+ resource: e
88
+ }
89
+ })
90
+ })
91
+ }
92
+ async submitQueue(encoder){
93
+ this.device.queue.submit([encoder.finish()])
94
+ }
95
+
96
+ }
@@ -0,0 +1,5 @@
1
+ export class Relu{
2
+ constructor(device : GPUAdapter, embeddingSize : number);
3
+ async forward(input: GPUBuffer) : Promise<GPUBuffer>
4
+ async backward(outGradinet: GPUBuffer) : Promise<GPUBuffer>
5
+ }