@imaduddin/fastbrain 1.0.2 → 1.0.4

package/index.js

@@ -8,6 +8,7 @@ import {SelfAttention as ST} from './src/SelfAttention/index.js'
 import {TransformerBlock as TT} from './src/TransformerBlock/index.js'
 import {PositionEmbedding as PM } from './src/PositionEmbedding/index.js'
 import { RMSNorm as RM } from "./src/RMSNorm/index.js";
+import {model as M} from './src/model/index.js'
 export const Embedding = E
 export const Linear = L
 export const Softmax = S;
@@ -17,4 +18,5 @@ export const Tokenizer = to;
 export const SelfAttention = ST
 export const TransformerBlock = TT;
 export const PositionEmbedding = PM;
-export const RMSNorm = RM
+export const RMSNorm = RM
+export const model = M

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@imaduddin/fastbrain",
-  "version": "1.0.2",
+  "version": "1.0.4",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called fastBrain.",
   "keywords": [
     "ai",

package/src/Embedding/index.js

@@ -45,8 +45,8 @@ export class Embedding {
         })
         SaveVector(this.configs.save.filename, this.vectors)
     }
-    async forward(input, encoderMain) {
-        const encoder = encoderMain ? encoderMain :  this.device.createCommandEncoder()
+    async forward(input, encoder2) {
+        const encoder = encoder2 ? encoder2 : this.device.createCommandEncoder()
         if (!this.vectorsBuffer) {
             this.vectorsBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize);
             await this.gpu.WriteBuffer(this.vectorsBuffer, this.vectors)
@@ -86,7 +86,7 @@ export class Embedding {
             ]
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
-        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
+        if(!encoder2) await this.device.queue.submit([encoder.finish()])
         return this.outputBuffer
 
     }
@@ -94,8 +94,8 @@ export class Embedding {
         const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize, this.vectorsBuffer)
         SaveVector(this.configs.save.filename, vec)
     }
-    async backward(inputGradient, lr, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+    async backward(inputGradient, lr, en2) {
+        const encoder = en2 ? en2 : this.device.createCommandEncoder()
         if (!this.oldParam) {
             this.oldParam = new Float32Array([0, 0]);
             this.param = await this.gpu.createUniFormBuffer(this.oldParam.byteLength);
@@ -134,7 +134,7 @@ export class Embedding {
             ]
         })
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * (this.tokenBuffer.size / 4)) / 256))
-        if(!encoderMain) this.device.queue.submit([encoder.finish()])
+        if(!en2) this.device.queue.submit([encoder.finish()])
     }
 
 }

package/src/FFN/index.js

@@ -28,10 +28,10 @@ export class FFN {
         const final = await this.linear2.forward(hidden, encoderMain)
         return final
     }
-    async backward(inputGradient, lr, encoderMain) {
-        const logit2 = await this.linear2.backward(inputGradient, lr, encoderMain)
-        const hidden = await this.hiddenLayer.backward(logit2, encoderMain)
-        const final = await this.linear1.backward(hidden, lr, encoderMain)
+    async backward(inputGradient, lr, encoder) {
+        const logit2 = await this.linear2.backward(inputGradient, lr, encoder)
+        const hidden = await this.hiddenLayer.backward(logit2,encoder)
+        const final = await this.linear1.backward(hidden, lr, encoder)
         return final
     }
     async Save() {

package/src/Linear/index.js

@@ -110,7 +110,7 @@ export class Linear {
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil(this.vocabSize / 16), Math.ceil(tokenCount / 4))
 
-        if (!encoderMain) this.device.queue.submit([encoder.finish()])
+        if(!encoderMain) this.device.queue.submit([encoder.finish()])
         this.input = vec;
         return this.outputBuffer
     }
@@ -179,10 +179,9 @@ export class Linear {
                 }
             ]
         })
-        // console.log((await this.gpu.readBuffer(this.input.size, this.input))[0])
-        await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 256));
-        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil(this.vocabSize));
-        if (!encoderMain) await this.device.queue.submit([encoder.finish()])
+        await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 64));
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, this.vocabSize);
+        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
         return this.gradientBuffer
 
     }

package/src/PositionEmbedding/index.js

@@ -38,8 +38,8 @@ export class PositionEmbedding {
         this.forwardShader = this.gpu.getShaderPipline("./src/shaders/PF.wgsl")
         this.backwardShader = this.gpu.getShaderPipline("./src/shaders/PB.wgsl")
     }
-    async forward(input, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.gpu.CreateEncoder()
+    async forward(input, en2) {
+        const encoder =  en2 ? en2 : this.gpu.CreateEncoder()
         if (!this.params) this.params = await this.gpu.createUniFormBuffer(4 * 4)
         if (!this.vecBuffer) {
             this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
@@ -48,15 +48,15 @@ export class PositionEmbedding {
         const size = Math.ceil((input.size / 4) / 64);
         const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input.mainBuffer])
         await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
-        if(!encoderMain) await this.gpu.submitQueue(encoder);
+        if(!en2) await this.gpu.submitQueue(encoder);
         return input
     }
-    async backward(outputGradient, lr, encoderMain){
-        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder()
+    async backward(outputGradient, lr, en2){
+        const encoder = en2 ? en2 : await this.gpu.CreateEncoder()
         await this.gpu.WriteBuffer(this.params, new Float32Array([lr,0]))
         const bindGroup = await this.gpu.createBindGroup(this.backwardShader, 0, [this.vecBuffer, outputGradient.mainBuffer, this.params])
         await this.gpu.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((outputGradient.size / 4) / 64))
-        if(!encoderMain) await this.gpu.submitQueue(encoder);
+        if(!en2) await this.gpu.submitQueue(encoder);
     }
     async Save() {
         const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize, this.vecBuffer)

package/src/RMSNorm/index.js

@@ -45,7 +45,7 @@ export class RMSNorm {
       this.gamma = new Float32Array(embeddingSize).fill(1)
     }
   }
-  async forward(vec, encoderMain) {
+  async forward(vec, en) {
     if (!this.gammabuffer) {
       this.gammabuffer = await this.op.createBuffer(this.embeddingSize * 4);
       await this.op.WriteBuffer(this.gammabuffer, this.gamma);
@@ -57,34 +57,30 @@ export class RMSNorm {
       this.rmsReduction.destroy()
       this.rmsReduction = await this.op.createBuffer(4 * context)
     }
-    if (this.output.size !== vec.size) {
-      this.output.destroy()
-      this.output = await this.op.ManageBuffer(vec.size)
-    }
-    const encoder = encoderMain ? encoderMain : await this.op.CreateEncoder()
+    if (this.output) this.output.destroy();
+    this.output = await this.op.ManageBuffer(vec.size)
+    const encoder = en ? en : await this.op.CreateEncoder()
     const bind1 = await this.op.createBindGroup(this.reductionShader, 0, [vec.mainBuffer, this.rmsReduction])
     const bind2 = await this.op.createBindGroup(this.forwardShader, 0, [this.rmsReduction, vec.mainBuffer, this.output.mainBuffer, this.gammabuffer])
     await this.op.RunPipline(encoder, this.reductionShader, bind1, context)
     await this.op.RunPipline(encoder, this.forwardShader, bind2, Math.ceil((this.embeddingSize * context) / 64))
-    if (!encoderMain) await this.op.submitQueue(encoder);
+    if(!en) await this.op.submitQueue(encoder);
     this.input = vec;
     return this.output
 
   }
-  async backward(outGradient, lr, encoderMain) {
+  async backward(outGradient, lr, en) {
     const contextSize = outGradient.size / (this.embeddingSize * 4);
-    if (!this.gradientOutput) this.gradientOutput = await this.op.ManageBuffer(this.embeddingSize * 4 * contextSize)
-    if (this.gradientOutput.size !== outGradient.size) {
-      this.gradientOutput = await this.op.ManageBuffer(this.embeddingSize * 4 * contextSize)
-    }
-    if (!this.params) this.params = await this.op.createBuffer(4 * 2);
+    if (this.gradientOutput) this.gradientOutput.destroy()
+    this.gradientOutput = await this.op.ManageBuffer(this.embeddingSize * 4 * contextSize)
+    if (!this.params) this.params = await this.op.createUniFormBuffer(4 * 2);
     await this.op.WriteBuffer(this.params, new Float32Array([lr, 0]))
-    const encoder = encoderMain ? encoderMain : this.op.CreateEncoder()
+    const encoder = en ? en : this.op.CreateEncoder()
     const bind1 = await this.op.createBindGroup(this.backwardShader, 0, [this.input.mainBuffer, outGradient.mainBuffer, this.rmsReduction, this.gammabuffer, this.gradientOutput.mainBuffer])
     await this.op.RunPipline(encoder, this.backwardShader, bind1, contextSize);
     const bind2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input.mainBuffer, outGradient.mainBuffer, this.rmsReduction, this.gammabuffer, this.params])
     await this.op.RunPipline(encoder, this.updateGamma, bind2, this.embeddingSize);
-    if (!encoderMain) this.op.submitQueue(encoder);
+    if(!en) this.op.submitQueue(encoder);
     return this.gradientOutput
   }
   async Save() {

package/src/Relu/index.js

@@ -15,8 +15,8 @@ export class Relu {
         this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/RF.wgsl")
         this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/RB.wgsl")
     }
-    async forward(input, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+    async forward(input, en) {
+        const encoder =  en ? en : this.device.createCommandEncoder()
         const tokenCount = input.mainBuffer.size / (this.embeddingSize * 4)
         if (!this.oldParam) {
             this.oldParam = new Uint32Array([0, 0, 0]);
@@ -54,12 +54,12 @@ export class Relu {
             ]
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
-        if(!encoderMain) this.device.queue.submit([encoder.finish()])
+        if(!en) this.device.queue.submit([encoder.finish()])
         this.input = input
         return this.output
     }
-    async backward(inputGradient, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+    async backward(inputGradient,en) {
+        const encoder = en ? en :  this.device.createCommandEncoder()
         const tokenCount = inputGradient.size / (this.embeddingSize * 4)
         if (!this.oldParam) {
             this.oldParam = new Uint32Array([0, 0, 0]);
@@ -102,7 +102,7 @@ export class Relu {
             ]
         })
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
-        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
+        if(!en) await this.device.queue.submit([encoder.finish()])
         return this.gradOut
     }
 }

package/src/SelfAttention/index.js

@@ -38,12 +38,12 @@ export class SelfAttention {
         this.dKShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/ComputeK.wgsl")
         this.sumShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/sumGradient.wgsl")
     }
-    async forward(vec, encoderMain) {
+    async forward(vec, encoder2) {
         this.input = vec.mainBuffer;
         const tokenCount = vec.mainBuffer.size / (this.embeddingSize * 4)
         if (!this.asParams) this.asParams = await this.gpu.createUniFormBuffer(4 * 3);
         if (!this.softmaxParams) this.softmaxParams = await this.gpu.createUniFormBuffer(4 * 3);
-        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder()
+        const encoder = encoder2 ? encoder2 :  await this.gpu.CreateEncoder()
         if (!this.scroes) {
             this.scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
         }
@@ -53,9 +53,9 @@ export class SelfAttention {
         }
         await this.gpu.WriteBuffer(this.asParams, new Uint32Array([tokenCount, this.embeddingSize, 0]))
         await this.gpu.WriteBuffer(this.softmaxParams, new Uint32Array([tokenCount, tokenCount, 0]))
-        this.query = await this.q.forward(vec, encoderMain)
-        this.key = await this.k.forward(vec, encoderMain)
-        this.value = await this.v.forward(vec, encoderMain)
+        this.query = await this.q.forward(vec, encoder)
+        this.key = await this.k.forward(vec, encoder)
+        this.value = await this.v.forward(vec, encoder)
         const bindAttention = await this.gpu.createBindGroup(this.attentionscoresShader, 0, [this.query.mainBuffer, this.key.mainBuffer, this.scroes, this.asParams])
         await this.gpu.RunPipline(encoder, this.attentionscoresShader, bindAttention, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
         if (this.weights) this.weights.destroy()
@@ -71,12 +71,12 @@ export class SelfAttention {
         await this.finalOutput.count()
         const outputBind = await this.gpu.createBindGroup(this.outputForward, 0, [this.weights, this.value.mainBuffer, this.finalOutput.mainBuffer, this.asParams]);
         await this.gpu.RunPipline(encoder, this.outputForward, outputBind, Math.ceil((tokenCount * this.embeddingSize) / 256))
-        if(!encoderMain) await this.gpu.submitQueue(encoder)
+        if(!encoder2) await this.gpu.submitQueue(encoder)
         return this.finalOutput
     }
-    async backward(dO, lr, encoderMain) {
+    async backward(dO, lr,encoder2) {
         const tokenCount = dO.size / (this.embeddingSize * 4);
-        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder();
+        const encoder = encoder2 ? encoder2 :  await this.gpu.CreateEncoder();
         if (!this.dV) this.dV = await this.gpu.createBuffer(dO.size);
         if (!this.dA) this.dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
         if (!this.dS) this.dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
@@ -112,9 +112,9 @@ export class SelfAttention {
         await this.gpu.RunPipline(encoder, this.dQshader, dQbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
         const dKbind = await this.gpu.createBindGroup(this.dKShader, 0, [this.dS, this.query.mainBuffer, this.dK, this.asParams]);
         await this.gpu.RunPipline(encoder, this.dKShader, dKbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
-        const dq = await this.q.backward({ mainBuffer: this.dQ, size: this.dQ.size}, lr, encoderMain)
-        const dk = await this.k.backward({ mainBuffer: this.dK, size: this.dK.size}, lr, encoderMain)
-        const dv = await this.v.backward({ mainBuffer: this.dV, size: this.dV.size}, lr, encoderMain)
+        const dq = await this.q.backward({ mainBuffer: this.dQ, size: this.dQ.size}, lr, encoder)
+        const dk = await this.k.backward({ mainBuffer: this.dK, size: this.dK.size}, lr, encoder)
+        const dv = await this.v.backward({ mainBuffer: this.dV, size: this.dV.size}, lr, encoder)
         if(!this.gradOut) this.gradOut = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
         if(this.gradOut.size !== tokenCount * 4 * this.embeddingSize || this.gradOut.isDestroyed() || this.gradOut.size == 0){
             this.gradOut.destroy()
@@ -126,7 +126,7 @@ export class SelfAttention {
         }
         const outputBind = await this.gpu.createBindGroup(this.sumShader, 0, [dq.mainBuffer, dk.mainBuffer, dv.mainBuffer, this.gradOut.mainBuffer, this.asParams]);
         await this.gpu.RunPipline(encoder, this.sumShader, outputBind, Math.ceil((tokenCount * this.embeddingSize / 256)));
-        if(!encoderMain) await this.gpu.submitQueue(encoder)
+        if(!encoder2) await this.gpu.submitQueue(encoder)
         return this.gradOut
 
     }

package/src/Softmax/index.js

@@ -26,8 +26,8 @@ export class Softmax {
         this.reduceFn = this.gpu.getShaderPipline("./src/shaders/ReductionFunction.wgsl")
         this.totalCalculate = this.gpu.getShaderPipline('./src/shaders/calculateTotal.wgsl')
     }
-    async forward(logits, encoderCPU) {
-        const encoder = encoderCPU ? encoderCPU : this.device.createCommandEncoder()
+    async forward(logits, en) {
+        const encoder = en ? en : this.device.createCommandEncoder()
         const tokenCount = logits.size / (this.vocabSize * 4)
         if (!this.paramsForward) {
             this.paramsForward = await this.gpu.createUniFormBuffer(4 * 3)
@@ -65,13 +65,12 @@ export class Softmax {
             ]
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, tokenCount)
-        if (!encoderCPU) {
-            await this.device.queue.submit([encoder.finish()])
-        }
+        if(!en) await this.device.queue.submit([encoder.finish()])
+
         return this.output
     }
-    async backward(probs, target, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+    async backward(probs, target,encoder2) {
+        const encoder = encoder2 ? encoder2 : this.device.createCommandEncoder()
         const tokenCount = probs.size / (this.vocabSize * 4)
         if (!this.paramsForward) {
             this.paramsForward = await this.gpu.createUniFormBuffer(4 * 3)
@@ -88,6 +87,11 @@ export class Softmax {
             this.targetBuffer.destroy()
             this.targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
         }
+        if (!this.dLogits) this.dLogits = await this.gpu.ManageBuffer(probs.size);
+        if (this.dLogits.size !== probs.size) {
+            this.dLogits.destroy()
+            this.dLogits = await this.gpu.ManageBuffer(probs.size);
+        }
         await this.gpu.WriteBuffer(this.targetBuffer, new Float32Array(target))
         const bind = await this.device.createBindGroup({
             layout: await this.backwardPipline.getBindGroupLayout(0),
@@ -107,12 +111,18 @@ export class Softmax {
                     resource: {
                         buffer: this.paramsForward
                     }
+                },
+                {
+                    binding: 4,
+                    resource: {
+                        buffer: this.dLogits.mainBuffer
+                    }
                 }
             ]
         })
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.vocabSize) / 256))
-        if (!encoderMain) this.device.queue.submit([encoder.finish()])
-        return probs
+        if(!encoder2) await this.gpu.submitQueue(encoder)
+        return this.dLogits
     }
     async crossEntropy(probs, target, encoderMain) {
         const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder();

package/src/TransformerBlock/index.js

@@ -54,13 +54,13 @@ export class TransformerBlock {
         return residual2;
     }
 
-    async backward(outGradient, lr, encoderMain) {
-        const dAtt = await this.ffn.backward(outGradient, lr, encoderMain);
-        const dNorm2 = await this.rmsnorm2.backward(dAtt, lr)
-        const dResidual1 = await this.rback2.add(outGradient, dNorm2, encoderMain)
-        const dFF = await this.attention.backward(dResidual1, lr, encoderMain)
-        const dNorm = await this.rmsnorm.backward(dFF, lr)
-        const dResidual2 = await this.rback1.add(dNorm, dResidual1, encoderMain)
+    async backward(outGradient, lr, encoder) {
+        const dAtt = await this.ffn.backward(outGradient, lr, encoder);
+        const dNorm2 = await this.rmsnorm2.backward(dAtt, lr, encoder)
+        const dResidual1 = await this.rback2.add(outGradient, dNorm2, encoder)
+        const dFF = await this.attention.backward(dResidual1, lr, encoder)
+        const dNorm = await this.rmsnorm.backward(dFF, lr, encoder)
+        const dResidual2 = await this.rback1.add(dNorm, dResidual1 ,encoder)
         return dResidual2;
     }
 

package/src/model/index.d.ts

@@ -0,0 +1,43 @@
+import { Embedding } from '../Embedding/index.ts'
+import { Linear } from '../Linear/index.ts'
+import { PositionEmbedding } from '../PositionEmbedding/index.ts'
+import { TransformerBlock } from '../TransformerBlock/index.ts'
+import { Softmax } from '../Softmax/index.ts'
+import { Tokenizer } from '../Tokenizer/index.ts'
+
+
+
+export class model {
+    tokenizer: Tokenizer
+    embedding: Embedding
+    linear: Linear
+    PosEmbedding: PositionEmbedding
+    blocks: TransformerBlock[]
+    softmax: Softmax
+    embeddingSize: number
+    loopSize: number
+    contextSize: number
+    constructor(device: GPUAdapter, configs: {
+        embeddingSize: number,
+        contextSize: number,
+        blockSize: number,
+        dirname: string,
+        vocabFile: string
+    })
+    getLoss : Function
+    getAcc: Function
+    clearLoss : Function
+    async forward(input: number[], target: number[]): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async backward(x: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, target: number[], lr : number)
+    async Save() : Promise<void>
+}

package/src/model/index.js

@@ -0,0 +1,87 @@
+import { Embedding } from "../Embedding/index.js"
+import { Linear } from "../Linear/index.js"
+import { PositionEmbedding } from "../PositionEmbedding/index.js"
+import { Softmax } from "../Softmax/index.js"
+import { Tokenizer } from "../Tokenizer/index.js"
+import { join } from 'path'
+import { TransformerBlock } from "../TransformerBlock/index.js"
+import { OperationManager } from '../OperationManager.js'
+export class model {
+    tokenizer
+    embedding
+    linear
+    PosEmbedding
+    blocks = []
+    softmax
+    embeddingSize
+    contextSize
+    op;
+    clearLoss;
+    getLoss;
+    getAcc;
+    constructor(device, configs) {
+        this.contextSize = configs.contextSize
+        this.embeddingSize = configs.embeddingSize;
+        this.tokenizer = new Tokenizer(configs.vocabFile)
+        this.op = new OperationManager(device)
+        this.embedding = new Embedding(device, this.embeddingSize, this.tokenizer.vocabSize, {
+            save: {
+                filename: join(configs.dirname, "vec.bin")
+            }
+        })
+        this.softmax = new Softmax(device, this.embeddingSize, this.tokenizer.vocabSize);
+        this.PosEmbedding = new PositionEmbedding(device, this.embeddingSize, configs.contextSize, {
+            save: {
+                filename: join(configs.dirname, "posVec.bin")
+            }
+        })
+        this.linear = new Linear(device, this.embeddingSize, this.tokenizer.vocabSize, {
+            save: {
+                filename: [join(configs.dirname, "w.bin"), join(configs.dirname, "b.bin")]
+            }
+        })
+        for (let a = 0; a < configs.blockSize; a++) {
+            this.blocks[a] = new TransformerBlock(device, this.embeddingSize, {
+                dirname: join(configs.dirname, `block-${a + 1}`)
+            })
+        }
+    }
+    async forward(input, target) {
+        const encoder = await this.op.CreateEncoder()
+        if (input.length > this.contextSize) throw new Error("The input get large then ContextSize")
+        let vec = await this.embedding.forward(input, encoder)
+        vec = await this.PosEmbedding.forward(vec, encoder)
+        for (let i = 0; i < this.blocks.length; i++) {
+            vec = await this.blocks[i].forward(vec, encoder)
+        }
+
+        let logits = await this.linear.forward(vec, encoder)
+        let probs = await this.softmax.forward(logits, encoder)
+        await this.op.submitQueue(encoder)
+        const { acc, clear, loss } = await this.softmax.crossEntropy(probs, target)
+        this.getAcc = acc;
+        this.getLoss = loss;
+        this.clearLoss = clear;
+        return probs
+    }
+    async backward(x, target, lr) {
+        const encoder = await this.op.CreateEncoder()
+        let dLogits = await this.softmax.backward(x, target, encoder)
+        let dInput = await this.linear.backward(dLogits, lr, encoder)
+        for (let i = this.blocks.length - 1; i >= 0; i--) {
+            dInput = await this.blocks[i].backward(dInput, lr, encoder)
+        };
+        await this.embedding.backward(dInput, lr)
+        await this.PosEmbedding.backward(dInput, lr)
+        await this.op.submitQueue(encoder)
+
+    }
+    async Save() {
+        await this.PosEmbedding.Save()
+        await this.embedding.Save()
+        await this.linear.Save()
+        for (let i = this.blocks.length - 1; i >= 0; i--) {
+             await this.blocks[i].Save()
+        };
+    }
+}

package/src/residualConnection/index.js

@@ -10,11 +10,11 @@ export class residualConnection {
         this.gpu = new OperationManager(device)
         this.addShader = this.gpu.getShaderPipline("./src/shaders/add.wgsl")
     }
-    async add(a, b, encoderMain) {
+    async add(a, b, en) {
         if (!this.params) {
             this.params = await this.gpu.createUniFormBuffer(4 * 4)
         }
-        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder()
+        const encoder = en ? en : await this.gpu.CreateEncoder()
         if (!this.output) {
             this.output = await this.gpu.ManageBuffer(a.size);
         }
@@ -22,11 +22,10 @@ export class residualConnection {
             this.output.destroy()
             this.output = await this.gpu.ManageBuffer(a.size);
         }
-        await this.output.count()
         const bind = await this.gpu.createBindGroup(this.addShader, 0, [a.mainBuffer, b.mainBuffer, this.output.mainBuffer, this.params]);
         await this.gpu.WriteBuffer(this.params, new Uint32Array([(a.size  / 4), 0]))
         await this.gpu.RunPipline(encoder, this.addShader, bind, Math.ceil((a.size / 4) / 64));
-        if(!encoderMain) await this.gpu.submitQueue(encoder)
+        if(!en) await this.gpu.submitQueue(encoder)
         return this.output
     }
 }

package/src/shaders/ComputeDInput.wgsl

@@ -1,3 +1,4 @@
+
 struct Params {
     vocabSize: u32,
     tokenCount: u32,
@@ -6,27 +7,40 @@ struct Params {
 }
 
 @group(0) @binding(0)
-var<storage, read_write> weight: array<f32>;
+var<storage, read> weight: array<f32>;
+
 @group(0) @binding(1)
 var<storage, read> outGradient: array<f32>;
+
 @group(0) @binding(3)
 var<uniform> params: Params;
+
 @group(0) @binding(4)
 var<storage, read_write> dInput: array<f32>;
+
 @compute @workgroup_size(256)
 fn main(@builtin(global_invocation_id) id: vec3<u32>) {
     let index = id.x;
     let total = params.tokenCount * params.inputSize;
-    if (index >= u32(total)) {
+
+    if (index >= total) {
         return;
     }
-    let token =index / u32(params.inputSize);
-    let inputIndex = index % u32(params.inputSize);
-    var gradInput = 0.0;
-    for (var i = 0u; i < u32(params.vocabSize); i++) {
-        let w = weight[i * u32(params.inputSize) + inputIndex];
-        let grad = outGradient[token * u32(params.vocabSize) + i];
-        gradInput += grad * w;
+
+    let inputSize = params.inputSize;
+    let vocabSize = params.vocabSize;
+
+    let token = index / inputSize;
+    let inputIndex = index % inputSize;
+
+    let weightBase = inputIndex;
+    let gradBase = token * vocabSize;
+    var sum = 0.0;
+
+    for (var i = 0u; i < vocabSize; i++) {
+        let wIndex = i * inputSize + weightBase;
+        let gIndex = gradBase + i;
+        sum += outGradient[gIndex] * weight[wIndex];
     }
-    dInput[index] = gradInput;
-}
+    dInput[index] = sum;
+}

package/src/shaders/LB.wgsl

@@ -1,3 +1,4 @@
+
 struct Params {
     vocabSize: u32,
     tokenCount: u32,
@@ -20,48 +21,44 @@ var<storage, read> input: array<f32>;
 @group(0) @binding(4)
 var<uniform> params: Params;
 
+const WG_SIZE: u32 = 256u;
+
 var<workgroup> sharedBias: array<f32, 256>;
 
 @compute @workgroup_size(256)
-fn main(
-    @builtin(workgroup_id) wid: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>
-) {
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
     let outputIndex = wid.x;
     let tid = lid.x;
-
     let vocab = params.vocabSize;
     let tokens = params.tokenCount;
     let inputSize = params.inputSize;
-
-    for (
-        var inputIndex = tid;
-        inputIndex < inputSize;
-        inputIndex += 256u
-    ) {
+    if (outputIndex >= vocab) {
+        return;
+    }
+    let gradBase = outputIndex;
+    let weightBase = outputIndex * inputSize;
+    for (var j = tid; j < inputSize; j += WG_SIZE) {
         var grad = 0.0;
-
         for (var t = 0u; t < tokens; t++) {
-            grad += input[t * inputSize + inputIndex]
-                  * outGradient[t * vocab + outputIndex];
+            let inputIndex = t * inputSize + j;
+            let outputGradIndex = t * vocab + gradBase;
+            grad += input[inputIndex] * outGradient[outputGradIndex];
         }
 
         if (tokens > 0u) {
-            weight[outputIndex * inputSize + inputIndex] -=
-                params.lr * grad / f32(tokens);
+            let idx = weightBase + j;
+            weight[idx] -= params.lr * (grad / f32(tokens));
         }
     }
 
     var biasGrad = 0.0;
-
-    for (var t = tid; t < tokens; t += 256u) {
-        biasGrad += outGradient[t * vocab + outputIndex];
+    for (var t = tid; t < tokens; t += WG_SIZE) {
+        biasGrad += outGradient[t * vocab + gradBase];
     }
 
     sharedBias[tid] = biasGrad;
     workgroupBarrier();
-
-    for (var stride = 128u; stride > 0u; stride >>= 1u) {
+    for (var stride = WG_SIZE / 2u; stride > 0u; stride >>= 1u) {
         if (tid < stride) {
             sharedBias[tid] += sharedBias[tid + stride];
         }
@@ -69,7 +66,6 @@ fn main(
     }
 
     if (tid == 0u && tokens > 0u) {
-        bias[outputIndex] -=
-            params.lr * sharedBias[0] / f32(tokens);
+        bias[outputIndex] -= params.lr * (sharedBias[0] / f32(tokens));
     }
-}
+}

package/src/shaders/Norm/updateGamma.wgsl

@@ -1,4 +1,6 @@
-
+struct Params{
+    lr: f32
+}
 @group(0) @binding(0)
 var<storage, read> input: array<f32>;
 
@@ -12,7 +14,7 @@ var<storage, read> sumSquares: array<f32>;
 var<storage, read_write> gamma: array<f32>;
 
 @group(0) @binding(4)
-var<storage, read> params: array<f32>;
+var<uniform> params: Params;
 
 var<workgroup> partial: array<f32, 64>;
 
@@ -47,6 +49,6 @@ fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: v
         stride /= 2u;
     }
     if (lid == 0u && dim < D) {
-        gamma[dim] -= params[0] * partial[0];
+        gamma[dim] -= params.lr * partial[0];
     }
 }

package/src/shaders/SB.wgsl

@@ -16,6 +16,11 @@ var<storage, read_write> targetToken: array<f32>;
 
 var<uniform> params: Params;
 
+
+@group(0) @binding(4)
+
+var<storage, read_write> dLogits: array<f32>;
+
 @compute @workgroup_size(256)
 
 fn main(@builtin(global_invocation_id) id: vec3<u32>) {
@@ -29,8 +34,8 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
       return;
     }
     var targetIndex = targetToken[token];
-    probs[index] = probs[index];
+    dLogits[index] = probs[index];
     if(vocabIndex == u32(targetIndex)){
-        probs[index] -= 1;
+        dLogits[index] -= 1;
     }
 }