@mineproness/tensorloop 2.1.65 → 2.1.68
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 +5 -1
- package/package.json +1 -1
- package/src/Embedding/index.js +2 -1
- package/src/FFN/index.js +2 -1
- package/src/Linear/index.js +2 -1
- package/src/Model/index.d.ts +15 -0
- package/src/Model/index.js +94 -0
- package/src/OperationManager.js +21 -2
- package/src/PositionEmbedding/index.d.ts +10 -0
- package/src/PositionEmbedding/index.js +70 -0
- package/src/SelfAttention/index.js +2 -1
- package/src/Softmax/index.js +1 -0
- package/src/Tokenizer/index.js +9 -5
- package/src/TransformerBlock/index.js +21 -5
- package/src/gpu.js +0 -0
- package/src/shaders/CalculateLoss.wgsl +1 -1
- package/src/shaders/LB.wgsl +4 -4
- package/src/shaders/PB.wgsl +25 -0
- package/src/shaders/PF.wgsl +21 -0
- package/src/shaders/SB.wgsl +0 -1
- package/src/shaders/add.wgsl +30 -0
package/index.js
CHANGED
|
@@ -6,6 +6,8 @@ import { FFN as ffn } from "./src/FFN/index.js";
|
|
|
6
6
|
import { Tokenizer as to } from './src/Tokenizer/index.js'
|
|
7
7
|
import {SelfAttention as ST} from './src/SelfAttention/index.js'
|
|
8
8
|
import {TransformerBlock as TT} from './src/TransformerBlock/index.js'
|
|
9
|
+
import { Model as m } from "./src/Model/index.js";
|
|
10
|
+
import {PositionEmbedding as PM } from './src/PositionEmbedding/index.js'
|
|
9
11
|
export const Embedding = E
|
|
10
12
|
export const Linear = L
|
|
11
13
|
export const Softmax = S;
|
|
@@ -13,4 +15,6 @@ export const RELU = r
|
|
|
13
15
|
export const FFN = ffn
|
|
14
16
|
export const Tokenizer = to;
|
|
15
17
|
export const SelfAttention = ST
|
|
16
|
-
export const TransformerBlock = TT;
|
|
18
|
+
export const TransformerBlock = TT;
|
|
19
|
+
export const Model = m;
|
|
20
|
+
export const PositionEmbedding = PM;
|
package/package.json
CHANGED
package/src/Embedding/index.js
CHANGED
|
@@ -134,6 +134,7 @@ export class Embedding {
|
|
|
134
134
|
await this.gpu.WriteBuffer(this.tokenBuffer, new Uint32Array(token));
|
|
135
135
|
await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * token.length) / 256))
|
|
136
136
|
this.device.queue.submit([encoder.finish()])
|
|
137
|
-
|
|
137
|
+
if(this.prevGra) this.prevGra.destroy()
|
|
138
|
+
this.prevGra = inputGradient;
|
|
138
139
|
}
|
|
139
140
|
}
|
package/src/FFN/index.js
CHANGED
|
@@ -29,10 +29,11 @@ export class FFN {
|
|
|
29
29
|
return final
|
|
30
30
|
}
|
|
31
31
|
async backward(inputGradient, lr) {
|
|
32
|
+
if(this.prevOUTGR) await this.prevOUTGR.destroy()
|
|
33
|
+
this.prevOUTGR = inputGradient;
|
|
32
34
|
const logit2 = await this.linear2.backward(inputGradient, lr)
|
|
33
35
|
const hidden = await this.hiddenLayer.backward(logit2)
|
|
34
36
|
const final = await this.linear1.backward(hidden, lr)
|
|
35
|
-
inputGradient.destroy()
|
|
36
37
|
return final
|
|
37
38
|
}
|
|
38
39
|
async Save(){
|
package/src/Linear/index.js
CHANGED
|
@@ -113,6 +113,8 @@ export class Linear {
|
|
|
113
113
|
return outputBuffer
|
|
114
114
|
}
|
|
115
115
|
async backward(dOutput, learningRate) {
|
|
116
|
+
if(this.pdOutput) await this.pdOutput.destroy();
|
|
117
|
+
this.pdOutput = dOutput
|
|
116
118
|
const encoder = this.device.createCommandEncoder()
|
|
117
119
|
const tokenCount = dOutput.size / (this.vocabSize * 4)
|
|
118
120
|
if (!this.oldParamb) {
|
|
@@ -173,7 +175,6 @@ export class Linear {
|
|
|
173
175
|
await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 256));
|
|
174
176
|
await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.vocabSize * tokenCount) / 256));
|
|
175
177
|
await this.device.queue.submit([encoder.finish()])
|
|
176
|
-
await dOutput.destroy()
|
|
177
178
|
this.input.destroy()
|
|
178
179
|
return outputBuffer
|
|
179
180
|
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
export class Model {
|
|
2
|
+
async forward(input: Number[], buffer: boolean): Promise<Float32Array | GPUBuffer>;
|
|
3
|
+
async backward(output: GPUBuffer, target: number[], lr: number): Promise<{ acc: number, loss: number }>
|
|
4
|
+
TrainTokenizer(dataset: string): void;
|
|
5
|
+
constructor(device :GPUAdapter, embeddingSize: number, configs: {
|
|
6
|
+
tokenizerName: string
|
|
7
|
+
blockSize: number,
|
|
8
|
+
dirname: string,
|
|
9
|
+
contextSize: number,
|
|
10
|
+
gpuName: string
|
|
11
|
+
|
|
12
|
+
})
|
|
13
|
+
async Save() : Promise<void>
|
|
14
|
+
async CreateContext()
|
|
15
|
+
}
|
|
@@ -0,0 +1,94 @@
|
|
|
1
|
+
import { Embedding, Linear, Softmax, TransformerBlock, Tokenizer, PositionEmbedding } from '../../index.js'
|
|
2
|
+
import { join } from 'path'
|
|
3
|
+
import fs from 'fs'
|
|
4
|
+
import { OperationManager } from '../OperationManager.js'
|
|
5
|
+
export class Model {
|
|
6
|
+
embedding;
|
|
7
|
+
posEmbedding;
|
|
8
|
+
linear;
|
|
9
|
+
softmax;
|
|
10
|
+
blocks = [];
|
|
11
|
+
blockSize;
|
|
12
|
+
dimSize;
|
|
13
|
+
tokenizer;
|
|
14
|
+
device;
|
|
15
|
+
configs;
|
|
16
|
+
tokenizerName;
|
|
17
|
+
op;
|
|
18
|
+
tokens;
|
|
19
|
+
constructor(device, embeddingSize = 64, configs) {
|
|
20
|
+
this.device = device;
|
|
21
|
+
this.dimSize = embeddingSize;
|
|
22
|
+
this.blockSize = configs.blockSize;
|
|
23
|
+
this.configs = configs
|
|
24
|
+
this.tokenizer = new Tokenizer(configs.tokenizerName)
|
|
25
|
+
this.tokenizerName = configs.tokenizerName
|
|
26
|
+
this.op = new OperationManager(device)
|
|
27
|
+
}
|
|
28
|
+
async CreateContext() {
|
|
29
|
+
this.embedding = new Embedding(this.device, this.dimSize, this.tokenizer.vocabSize, {
|
|
30
|
+
save: {
|
|
31
|
+
filename: join(this.configs.dirname, "./vec.bin")
|
|
32
|
+
}
|
|
33
|
+
})
|
|
34
|
+
this.linear = new Linear(this.device, this.dimSize, this.tokenizer.vocabSize, {
|
|
35
|
+
save: {
|
|
36
|
+
filename: [join(this.configs.dirname, "weights.bin"), join(this.configs.dirname, "bias.bin")]
|
|
37
|
+
}
|
|
38
|
+
})
|
|
39
|
+
this.softmax = new Softmax(this.device, this.dimSize, this.tokenizer.vocabSize);
|
|
40
|
+
for (let a = 0; a < this.blockSize; a++) {
|
|
41
|
+
console.log(`Load the Transfomer Block ${a + 1}...`)
|
|
42
|
+
if (!fs.existsSync(join(this.configs.dirname, `block-${a}`))) fs.mkdirSync(join(this.configs.dirname, `block-${a}`));
|
|
43
|
+
this.blocks[a] = new TransformerBlock(this.device, this.dimSize, {
|
|
44
|
+
dirname: join(this.configs.dirname, `block-${a}`)
|
|
45
|
+
})
|
|
46
|
+
}
|
|
47
|
+
this.posEmbedding = new PositionEmbedding(this.device, this.dimSize, this.configs.contextSize, {
|
|
48
|
+
save:{
|
|
49
|
+
filename: join(this.configs.dirname, "posVec.bin")
|
|
50
|
+
}
|
|
51
|
+
})
|
|
52
|
+
}
|
|
53
|
+
TrainTokenizer(text = "", mergeSize = 5) {
|
|
54
|
+
this.tokenizer.trainTokenizer(text, mergeSize, this.tokenizerName)
|
|
55
|
+
}
|
|
56
|
+
async ClearInputCache() {
|
|
57
|
+
await this.linear.ClearInputCache()
|
|
58
|
+
for (let a = 0; a < this.blockSize; a++) {
|
|
59
|
+
await this.blocks[a].ClearInputCache()
|
|
60
|
+
}
|
|
61
|
+
}
|
|
62
|
+
async forward(input, readBack) {
|
|
63
|
+
if(this.configs.contextSize < input.length) throw new Error("The Input is larger then contextSize")
|
|
64
|
+
this.tokens = input;
|
|
65
|
+
await this.ClearInputCache()
|
|
66
|
+
let vec = await this.embedding.forward(input)
|
|
67
|
+
vec = await this.posEmbedding.forward(vec)
|
|
68
|
+
for (let a = 0; a < this.blockSize; a++) {
|
|
69
|
+
vec = await this.blocks[a].forward(vec)
|
|
70
|
+
}
|
|
71
|
+
let logits = await this.linear.forward(vec);
|
|
72
|
+
let probs = await this.softmax.forward(logits, readBack)
|
|
73
|
+
return probs
|
|
74
|
+
}
|
|
75
|
+
async backward(outputGradient, target, lr = 0.02) {
|
|
76
|
+
const loss = await this.softmax.crossEntropy(outputGradient, target)
|
|
77
|
+
const dLogits = await this.softmax.backward(outputGradient, target);
|
|
78
|
+
let DInput = await this.linear.backward(dLogits, lr);
|
|
79
|
+
for (let a = this.blockSize - 1; a !== 0; a--) {
|
|
80
|
+
DInput = await this.blocks[a].backward(DInput, lr);
|
|
81
|
+
}
|
|
82
|
+
this.embedding.backward(this.tokens, DInput, lr);
|
|
83
|
+
this.posEmbedding.backward(DInput, lr)
|
|
84
|
+
return loss
|
|
85
|
+
}
|
|
86
|
+
async Save() {
|
|
87
|
+
await this.embedding.Save()
|
|
88
|
+
await this.linear.Save()
|
|
89
|
+
for (let a = this.blockSize - 1; a !== 0; a--) {
|
|
90
|
+
await this.blocks[a].Save();
|
|
91
|
+
}
|
|
92
|
+
await this.posEmbedding.Save()
|
|
93
|
+
}
|
|
94
|
+
}
|
package/src/OperationManager.js
CHANGED
|
@@ -8,9 +8,10 @@ export class OperationManager {
|
|
|
8
8
|
device;
|
|
9
9
|
matrixMatSHhaders;
|
|
10
10
|
params;
|
|
11
|
-
|
|
11
|
+
addShader;
|
|
12
12
|
constructor(device) {
|
|
13
13
|
this.device = device;
|
|
14
|
+
this.addShader = this.getShaderPipline("./src/shaders/add.wgsl")
|
|
14
15
|
}
|
|
15
16
|
async createBuffer(size) {
|
|
16
17
|
return this.device.createBuffer({
|
|
@@ -71,18 +72,21 @@ export class OperationManager {
|
|
|
71
72
|
pass.end()
|
|
72
73
|
}
|
|
73
74
|
|
|
74
|
-
async MakeACopyBuffer(
|
|
75
|
+
async MakeACopyBuffer(encoder2, sourceBuffer) {
|
|
76
|
+
const encoder = encoder2 ? encoder2 : await this.CreateEncoder()
|
|
75
77
|
const buffer = this.device.createBuffer({
|
|
76
78
|
size: sourceBuffer.size,
|
|
77
79
|
usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
|
|
78
80
|
})
|
|
79
81
|
encoder.copyBufferToBuffer(sourceBuffer, 0, buffer, 0, sourceBuffer.size);
|
|
82
|
+
if(!encoder2) await this.submitQueue(encoder)
|
|
80
83
|
return buffer
|
|
81
84
|
}
|
|
82
85
|
CreateEncoder() {
|
|
83
86
|
return this.device.createCommandEncoder()
|
|
84
87
|
}
|
|
85
88
|
createBindGroup(shader, groupNum, ent) {
|
|
89
|
+
|
|
86
90
|
return this.device.createBindGroup({
|
|
87
91
|
layout: shader.getBindGroupLayout(groupNum),
|
|
88
92
|
entries: ent.map((e, x) => {
|
|
@@ -96,5 +100,20 @@ export class OperationManager {
|
|
|
96
100
|
async submitQueue(encoder) {
|
|
97
101
|
this.device.queue.submit([encoder.finish()])
|
|
98
102
|
}
|
|
103
|
+
async add(a, b, size) {
|
|
104
|
+
if (!this.params) {
|
|
105
|
+
this.params = await this.createUniFormBuffer(4 * 4)
|
|
106
|
+
|
|
107
|
+
}
|
|
108
|
+
const encoder = await this.CreateEncoder()
|
|
109
|
+
const output = await this.createBuffer(a.size);
|
|
110
|
+
const bind = await this.createBindGroup(this.addShader, 0, [a, b, output, this.params]);
|
|
111
|
+
await this.WriteBuffer(this.params, new Uint32Array([size, 0]))
|
|
112
|
+
await this.RunPipline(encoder, this.addShader, bind, Math.ceil(size / 64));
|
|
113
|
+
await this.submitQueue(encoder)
|
|
114
|
+
// a.destroy()
|
|
115
|
+
// b.destroy()
|
|
116
|
+
return output
|
|
117
|
+
}
|
|
99
118
|
|
|
100
119
|
}
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
export class PositionEmbedding{
|
|
2
|
+
constructor(device: GPUBuffer ,embedding : number, contextSize: number, config:{
|
|
3
|
+
save: {
|
|
4
|
+
filename: string
|
|
5
|
+
}
|
|
6
|
+
});
|
|
7
|
+
async forward(input: GPUBuffer) : Promise<GPUBuffer>
|
|
8
|
+
async Save() : Promise<GPUBuffer>
|
|
9
|
+
async backward(inputGradient: GPUBuffer, lr: number) : Promise<void>
|
|
10
|
+
}
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
import { OperationManager } from "../OperationManager.js"
|
|
2
|
+
import { getVector, SaveVector } from "../getVec.js"
|
|
3
|
+
function GenerateVectors(embeddingSize, contextSize) {
|
|
4
|
+
const vec = new Float32Array(embeddingSize * contextSize).map(() => (Math.random() * 2 - 1) * 0.02)
|
|
5
|
+
return vec
|
|
6
|
+
}
|
|
7
|
+
|
|
8
|
+
export class PositionEmbedding {
|
|
9
|
+
vectors;
|
|
10
|
+
vecBuffer;
|
|
11
|
+
gpu;
|
|
12
|
+
embeddingSize;
|
|
13
|
+
contextSize;
|
|
14
|
+
params;
|
|
15
|
+
forwardShader;
|
|
16
|
+
backwardShader;
|
|
17
|
+
configs;
|
|
18
|
+
constructor(device, embeddingSize, contextSize, configs) {
|
|
19
|
+
this.embeddingSize = embeddingSize;
|
|
20
|
+
const total = embeddingSize * contextSize;
|
|
21
|
+
this.configs = configs;
|
|
22
|
+
this.contextSize = contextSize;
|
|
23
|
+
if (configs && configs.save) {
|
|
24
|
+
const vec = getVector(configs.save.filename, total);
|
|
25
|
+
if (vec) {
|
|
26
|
+
this.vectors = vec;
|
|
27
|
+
|
|
28
|
+
} else {
|
|
29
|
+
this.vectors = GenerateVectors(embeddingSize, contextSize)
|
|
30
|
+
SaveVector(this.configs.save.filename, this.vectors)
|
|
31
|
+
}
|
|
32
|
+
|
|
33
|
+
} else {
|
|
34
|
+
this.vectors = GenerateVectors(embeddingSize, contextSize)
|
|
35
|
+
|
|
36
|
+
}
|
|
37
|
+
this.gpu = new OperationManager(device)
|
|
38
|
+
this.forwardShader = this.gpu.getShaderPipline("./src/shaders/PF.wgsl")
|
|
39
|
+
this.backwardShader = this.gpu.getShaderPipline("./src/shaders/PB.wgsl")
|
|
40
|
+
}
|
|
41
|
+
async forward(input) {
|
|
42
|
+
const encoder = this.gpu.CreateEncoder()
|
|
43
|
+
if (!this.params) this.params = await this.gpu.createUniFormBuffer(4 * 4)
|
|
44
|
+
if (!this.vecBuffer) {
|
|
45
|
+
this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
|
|
46
|
+
await this.gpu.WriteBuffer(this.vecBuffer, this.vectors)
|
|
47
|
+
}
|
|
48
|
+
const output = await this.gpu.createBuffer(input.size);
|
|
49
|
+
const size = Math.ceil((input.size / 4) / 64);
|
|
50
|
+
const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input, output])
|
|
51
|
+
if (this.prevVec) this.prevVec.destroy();
|
|
52
|
+
this.prevVec = input;
|
|
53
|
+
await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
|
|
54
|
+
await this.gpu.submitQueue(encoder);
|
|
55
|
+
return output
|
|
56
|
+
}
|
|
57
|
+
async backward(outputGradient, lr){
|
|
58
|
+
const encoder = await this.gpu.CreateEncoder()
|
|
59
|
+
await this.gpu.WriteBuffer(this.params, new Float32Array([lr,0]))
|
|
60
|
+
const bindGroup = await this.gpu.createBindGroup(this.backwardShader, 0, [this.vecBuffer, outputGradient, this.params])
|
|
61
|
+
if(this.prevGradient) await this.prevGradient.destroy()
|
|
62
|
+
this.prevGradient = outputGradient;
|
|
63
|
+
await this.gpu.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((outputGradient.size / 4) / 64))
|
|
64
|
+
await this.gpu.submitQueue(encoder);
|
|
65
|
+
}
|
|
66
|
+
async Save() {
|
|
67
|
+
const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize, this.vecBuffer)
|
|
68
|
+
SaveVector(this.configs.save.filename, vec)
|
|
69
|
+
}
|
|
70
|
+
}
|
package/src/Softmax/index.js
CHANGED
package/src/Tokenizer/index.js
CHANGED
|
@@ -9,7 +9,7 @@ export class Tokenizer {
|
|
|
9
9
|
IdtoWord = {};
|
|
10
10
|
WordtoID = {};
|
|
11
11
|
EOSToken = 2
|
|
12
|
-
constructor(filename="") {
|
|
12
|
+
constructor(filename = "") {
|
|
13
13
|
if (fs.existsSync(filename)) {
|
|
14
14
|
console.log("Warning: Tokenizer is loaded with a Custom Vocab File!")
|
|
15
15
|
console.log("")
|
|
@@ -22,7 +22,7 @@ export class Tokenizer {
|
|
|
22
22
|
this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
|
|
23
23
|
}
|
|
24
24
|
encoder(text = "") {
|
|
25
|
-
const AddSpaces = text.toLocaleLowerCase().replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ").replaceAll("<"
|
|
25
|
+
const AddSpaces = text.toLocaleLowerCase().replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ").replaceAll("<", " < ")
|
|
26
26
|
const splitedWords = AddSpaces.split(/\s+/).filter(e => e.length > 0)
|
|
27
27
|
const output = []
|
|
28
28
|
for (let a = 0; a < splitedWords.length; a++) {
|
|
@@ -145,8 +145,8 @@ export class Tokenizer {
|
|
|
145
145
|
]
|
|
146
146
|
const uniArr = [...new Set([...arrText, ...specal, ...prevVocab])];
|
|
147
147
|
const output = []
|
|
148
|
-
for(let a = 0
|
|
149
|
-
let text = uniArr[a].match(reg).filter((e)=> e.length > 0);
|
|
148
|
+
for (let a = 0; a < uniArr.length; a++) {
|
|
149
|
+
let text = uniArr[a].match(reg).filter((e) => e.length > 0);
|
|
150
150
|
output.push(...text)
|
|
151
151
|
}
|
|
152
152
|
let finalOutput = [...new Set([...output])]
|
|
@@ -155,6 +155,10 @@ export class Tokenizer {
|
|
|
155
155
|
console.log(`New VocabSize : ${finalOutput.length - prevVocab.length}`);
|
|
156
156
|
console.log("_________________________________")
|
|
157
157
|
console.log("Operation Completed!")
|
|
158
|
-
|
|
158
|
+
this.vocabs = finalOutput;
|
|
159
|
+
this.vocabSize = finalOutput.length;
|
|
160
|
+
this.IdtoWord = Object.fromEntries(this.vocabs.map((e, idx) => [idx, e]))
|
|
161
|
+
this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
|
|
162
|
+
|
|
159
163
|
}
|
|
160
164
|
}
|
|
@@ -1,24 +1,40 @@
|
|
|
1
1
|
import {SelfAttention} from '../SelfAttention/index.js'
|
|
2
2
|
import {FFN} from '../FFN/index.js'
|
|
3
3
|
import {existsSync, mkdirSync} from 'fs'
|
|
4
|
-
|
|
4
|
+
import {OperationManager} from '../OperationManager.js'
|
|
5
5
|
export class TransformerBlock{
|
|
6
6
|
attention;
|
|
7
7
|
ffn;
|
|
8
|
+
gpu;
|
|
9
|
+
size;
|
|
8
10
|
constructor(device, embeddingSize, configs){
|
|
11
|
+
this.size = embeddingSize;
|
|
9
12
|
if(!existsSync(configs.dirname)) mkdirSync(configs.dirname);
|
|
10
13
|
this.attention = new SelfAttention(device, embeddingSize, configs)
|
|
11
14
|
this.ffn = new FFN(device, embeddingSize, configs)
|
|
15
|
+
this.gpu = new OperationManager(device)
|
|
12
16
|
}
|
|
13
17
|
async forward(input){
|
|
18
|
+
let context = input.size / (this.size * 4)
|
|
14
19
|
const xAtt = await this.attention.forward(input);
|
|
15
|
-
const
|
|
16
|
-
|
|
20
|
+
const re1 = await this.gpu.add(input, xAtt, this.size * context)
|
|
21
|
+
const xFF = await this.ffn.forward(re1);
|
|
22
|
+
const re2 = await this.gpu.add(re1, xFF, this.size * context)
|
|
23
|
+
xAtt.destroy()
|
|
24
|
+
xFF.destroy()
|
|
25
|
+
return re2
|
|
17
26
|
}
|
|
18
27
|
async backward(outGradient, lr){
|
|
28
|
+
const context = outGradient.size / (this.size * 4)
|
|
19
29
|
const dAtt = await this.ffn.backward(outGradient, lr);
|
|
20
|
-
const
|
|
21
|
-
|
|
30
|
+
const re1 = await this.gpu.add(outGradient, dAtt, this.size * context)
|
|
31
|
+
const dFF = await this.attention.backward(re1, lr)
|
|
32
|
+
const re2 = await this.gpu.add(re1, dFF , context)
|
|
33
|
+
outGradient.destroy()
|
|
34
|
+
re1.destroy()
|
|
35
|
+
dFF.destroy()
|
|
36
|
+
dAtt.destroy()
|
|
37
|
+
return re2
|
|
22
38
|
}
|
|
23
39
|
async ClearInputCache(){
|
|
24
40
|
await this.ffn.ClearInputCache()
|
package/src/gpu.js
ADDED
|
File without changes
|
package/src/shaders/LB.wgsl
CHANGED
|
@@ -42,16 +42,16 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
|
|
|
42
42
|
let inputValue = input[token * u32(params.inputSize) + inputIndex];
|
|
43
43
|
let gradient = outGradient[token * u32(params.vocabSize) + outputIndex];
|
|
44
44
|
let dW = inputValue * gradient;
|
|
45
|
-
dWeights += f32(dW);
|
|
45
|
+
dWeights += f32(max(-3.0, min(3.0, dW)));
|
|
46
46
|
}
|
|
47
|
-
let vClap = weight[index] - params.lr *
|
|
47
|
+
let vClap = weight[index] - params.lr * dWeights;
|
|
48
48
|
weight[index] = vClap;
|
|
49
49
|
if (inputIndex == 0u) {
|
|
50
50
|
var dBias = 0.0;
|
|
51
51
|
for (var token = 0u; token < u32(params.tokenCount); token++) {
|
|
52
|
-
dBias += f32(outGradient[token * u32(params.vocabSize) + outputIndex]);
|
|
52
|
+
dBias += f32(max(-3.0, min(3.0, outGradient[token * u32(params.vocabSize) + outputIndex])));
|
|
53
53
|
}
|
|
54
|
-
bias[outputIndex] = bias[outputIndex] - params.lr *
|
|
54
|
+
bias[outputIndex] = bias[outputIndex] - params.lr * dBias;
|
|
55
55
|
|
|
56
56
|
}
|
|
57
57
|
}
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
struct Params{
|
|
2
|
+
lr : f32
|
|
3
|
+
}
|
|
4
|
+
|
|
5
|
+
@group(0) @binding(0)
|
|
6
|
+
var <storage, read_write> posVec: array<f32>;
|
|
7
|
+
|
|
8
|
+
@group(0) @binding(1)
|
|
9
|
+
|
|
10
|
+
var <storage, read> outputGradient : array<f32>;
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@group(0) @binding(2)
|
|
14
|
+
|
|
15
|
+
var <uniform> params: Params;
|
|
16
|
+
|
|
17
|
+
@compute @workgroup_size(64)
|
|
18
|
+
|
|
19
|
+
fn main(@builtin(global_invocation_id) id : vec3<u32>){
|
|
20
|
+
let x = id.x;
|
|
21
|
+
if(x >= arrayLength(&outputGradient)){
|
|
22
|
+
return;
|
|
23
|
+
}
|
|
24
|
+
posVec[x] -= params.lr * outputGradient[x];
|
|
25
|
+
}
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
@group(0) @binding(0)
|
|
2
|
+
var <storage, read> posVec: array<f32>;
|
|
3
|
+
|
|
4
|
+
@group(0) @binding(1)
|
|
5
|
+
|
|
6
|
+
var <storage, read> vecs : array<f32>;
|
|
7
|
+
|
|
8
|
+
@group(0) @binding(2)
|
|
9
|
+
|
|
10
|
+
var <storage, read_write> output: array<f32>;
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@compute @workgroup_size(64)
|
|
14
|
+
|
|
15
|
+
fn main(@builtin(global_invocation_id) id : vec3<u32>){
|
|
16
|
+
let x = id.x;
|
|
17
|
+
if(x >= arrayLength(&vecs)){
|
|
18
|
+
return;
|
|
19
|
+
}
|
|
20
|
+
output[x] = vecs[x] + posVec[x];
|
|
21
|
+
}
|
package/src/shaders/SB.wgsl
CHANGED
|
@@ -32,7 +32,6 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
|
|
|
32
32
|
return;
|
|
33
33
|
}
|
|
34
34
|
var targetIndex = targetToken[token];
|
|
35
|
-
let total: u32 = params.tokenCount * params.vocabSize;
|
|
36
35
|
output[index] = probs[index];
|
|
37
36
|
if(vocabIndex == u32(targetIndex)){
|
|
38
37
|
output[index] -= 1;
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
struct Params{
|
|
2
|
+
size: u32
|
|
3
|
+
}
|
|
4
|
+
|
|
5
|
+
@group(0) @binding(0)
|
|
6
|
+
|
|
7
|
+
var <storage, read> a : array<f32>;
|
|
8
|
+
|
|
9
|
+
@group(0) @binding(1)
|
|
10
|
+
|
|
11
|
+
var <storage, read> b : array<f32>;
|
|
12
|
+
|
|
13
|
+
@group(0) @binding(2)
|
|
14
|
+
|
|
15
|
+
var <storage, read_write> output : array<f32>;
|
|
16
|
+
@group(0) @binding(3)
|
|
17
|
+
|
|
18
|
+
var <uniform> params : Params;
|
|
19
|
+
|
|
20
|
+
@compute @workgroup_size(64)
|
|
21
|
+
|
|
22
|
+
fn main(@builtin(global_invocation_id) id : vec3<u32>){
|
|
23
|
+
let index = id.x;
|
|
24
|
+
if(index >= params.size){
|
|
25
|
+
return;
|
|
26
|
+
}
|
|
27
|
+
let mainOut = a[index] + b[index];
|
|
28
|
+
output[index] = mainOut;
|
|
29
|
+
}
|
|
30
|
+
|