@imaduddin/fastbrain 1.0.0 → 1.0.2

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@imaduddin/fastbrain",
-  "version": "1.0.0",
+  "version": "1.0.2",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called fastBrain.",
   "keywords": [
     "ai",
@@ -21,7 +21,7 @@
     "url": "https://github.com/imaduddin/fastBrain.git"
   },
   "dependencies": {
-    "webgpu": "^0.6.0"
+    "webgpu": "^0.6.1"
   },
   "allowScripts": {
     "webgpu@0.6.0": true

package/src/Embedding/index.js

@@ -63,7 +63,6 @@ export class Embedding {
             this.outputBuffer.destroy()
             this.outputBuffer = await this.gpu.ManageBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length);
         }
-        await this.outputBuffer.count()
         const bind = await this.device.createBindGroup({
             layout: await this.forwardPipline.getBindGroupLayout(0),
             entries: [
@@ -87,7 +86,7 @@ export class Embedding {
             ]
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
-        if(!encoderMain) this.device.queue.submit([encoder.finish()])
+        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
         return this.outputBuffer
 
     }

package/src/Linear/index.js

@@ -83,7 +83,6 @@ export class Linear {
             this.outputBuffer.destroy()
             this.outputBuffer = await this.gpu.ManageBuffer(4 * this.vocabSize * tokenCount)
         }
-        await this.outputBuffer.count()
         const bind = await this.device.createBindGroup({
             layout: await this.forwardPipline.getBindGroupLayout(0),
             entries: [

package/src/RMSNorm/index.js

@@ -11,6 +11,11 @@ export class RMSNorm {
   gamma;
   gammabuffer;
   configs;
+  rmsReduction;
+  output;
+  gradientOutput;
+  input;
+  device
   constructor(device, embeddingSize, configs) {
     this.configs = configs
     this.op = new OperationManager(device);
@@ -27,70 +32,64 @@ export class RMSNorm {
     this.updateGamma = this.op.getShaderPipline("./src/shaders/Norm/updateGamma.wgsl", {
       embeddingSize: embeddingSize
     })
+    this.device = device;
     if (configs) {
-      const ga = getVector(configs.save.filename, embeddingSize);
-      if (!ga) {
+      const gammaFile = getVector(configs.save.filename, embeddingSize);
+      if (!gammaFile) {
         this.gamma = new Float32Array(embeddingSize).fill(1);
         SaveVector(configs.save.filename, this.gamma)
       } else {
-        this.gamma = ga
+        this.gamma = gammaFile
       }
     } else {
       this.gamma = new Float32Array(embeddingSize).fill(1)
     }
   }
-  async forward(input) {
-    this.input = input;
+  async forward(vec, encoderMain) {
     if (!this.gammabuffer) {
-      this.gammabuffer = await this.op.createBuffer(4 * this.embeddingSize);
+      this.gammabuffer = await this.op.createBuffer(this.embeddingSize * 4);
       await this.op.WriteBuffer(this.gammabuffer, this.gamma);
     }
-    let context = input.size / (this.embeddingSize * 4)
-    if (!this.rms) this.rms = await this.op.createBuffer(4 * context);
-    if (this.rms.size !== 4 * context) {
-      this.rms.destroy()
-      this.rms = await this.op.createBuffer(4 * context);
+    const context = vec.size / (this.embeddingSize * 4)
+    if (!this.rmsReduction) this.rmsReduction = await this.op.createBuffer(context * 4);
+    if (!this.output) this.output = await this.op.ManageBuffer(vec.size)
+    if (this.rmsReduction.size !== context * 4) {
+      this.rmsReduction.destroy()
+      this.rmsReduction = await this.op.createBuffer(4 * context)
     }
-    if (!this.output) this.output = await this.op.ManageBuffer(input.size);
-    if (this.output.size !== 4 * this.embeddingSize * context || this.output.isDestroyed() || this.output.size == 0) {
+    if (this.output.size !== vec.size) {
       this.output.destroy()
-      this.output = await this.op.ManageBuffer(input.size)
+      this.output = await this.op.ManageBuffer(vec.size)
     }
-    await this.output.count()
-    const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input.mainBuffer, this.rms])
-    const encoder = await this.op.CreateEncoder()
-    await this.op.RunPipline(encoder, this.reductionShader, bindGroup, context)
-    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [this.rms, input.mainBuffer, this.output.mainBuffer, this.gammabuffer])
-    await this.op.RunPipline(encoder, this.forwardShader, bindGroup1, Math.ceil((input.size / 4) / 64));
-    await this.op.submitQueue(encoder);
+    const encoder = encoderMain ? encoderMain : 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);
+    this.input = vec;
     return this.output
+
   }
-  async backward(outputGradient, lr) {
-    if (!this.input || !this.input.mainBuffer || this.input.isDestroyed?.()) {
-      throw new Error("RMSNorm.backward: input buffer is gone");
-    }
-    if (!this.rms || this.rms.isDestroyed?.()) {
-      throw new Error("RMSNorm.backward: rms buffer is gone");
-    }
-    if (!this.gradOut) this.gradOut = await this.op.ManageBuffer(outputGradient.size);
-    if (this.gradOut.size !== outputGradient.size || this.gradOut.isDestroyed() || this.gradOut.size == 0) {
-      this.gradOut.destroy()
-      this.gradOut = await this.op.ManageBuffer(outputGradient.size)
+  async backward(outGradient, lr, encoderMain) {
+    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)
     }
-    await this.gradOut.count()
-    if (!this.params) this.params = await this.op.createBuffer(4 * 2)
+    if (!this.params) this.params = await this.op.createBuffer(4 * 2);
     await this.op.WriteBuffer(this.params, new Float32Array([lr, 0]))
-    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input.mainBuffer, outputGradient.mainBuffer, this.rms, this.gammabuffer, this.gradOut.mainBuffer]);
-    const bindGroup2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input.mainBuffer, outputGradient.mainBuffer, this.rms, this.gammabuffer, this.params]);
-    const encoder = await this.op.CreateEncoder()
-    await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size) / (this.embeddingSize * 4)))
-    await this.op.RunPipline(encoder, this.updateGamma, bindGroup2, Math.ceil(this.embeddingSize / 64))
-    await this.op.submitQueue(encoder)
-    return this.gradOut
+    const encoder = encoderMain ? encoderMain : 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);
+    return this.gradientOutput
   }
   async Save() {
-    const read = await this.op.readBuffer(4 * this.embeddingSize, this.gammabuffer);
-    SaveVector(this.configs.save.filename, read)
+    const gpuGamma = await this.op.readBuffer(4 * this.embeddingSize, this.gammabuffer);
+    SaveVector(this.configs.save.filename, gpuGamma)
   }
 
 }

package/src/Relu/index.js

@@ -27,7 +27,6 @@ export class Relu {
             this.output.destroy()
             this.output = await this.gpu.ManageBuffer(input.mainBuffer.size);
         }
-        await this.output.count()
         if (this.oldParam[1] !== tokenCount) {
             this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
             await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
@@ -71,7 +70,6 @@ export class Relu {
             this.gradOut.destroy();
             this.gradOut = await this.gpu.ManageBuffer(inputGradient.size)
         }
-        await this.gradOut.count()
         if (this.oldParam[1] !== tokenCount) {
             this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
             await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)

package/src/Softmax/index.js

@@ -42,7 +42,6 @@ export class Softmax {
             this.output.destroy()
             this.output = await this.gpu.ManageBuffer(logits.size)
         }
-        await this.output.count()
         const bind = await this.device.createBindGroup({
             layout: await this.forwardPipline.getBindGroupLayout(0),
             entries: [

package/src/TransformerBlock/index.js

@@ -45,10 +45,10 @@ export class TransformerBlock {
 
     async forward(input, encoderMain) {
         const count = input.size / 4;
-        const norm1 = await this.rmsnorm.forward(input);
+        const norm1 = await this.rmsnorm.forward(input, encoderMain);
         const att = await this.attention.forward(norm1, encoderMain)
         const residual1 = await this.rFor1.add(att, input, encoderMain)
-        const norm2 = await this.rmsnorm2.forward(residual1)
+        const norm2 = await this.rmsnorm2.forward(residual1, encoderMain)
         const f = await this.ffn.forward(norm2, encoderMain)
         const residual2 = await this.rFor2.add(f, residual1, encoderMain)
         return residual2;