@imaduddin/fastbrain 1.1.0 → 1.1.1

package/package.json

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

package/src/model/index.js

@@ -47,30 +47,27 @@ export class model {
         }
     }
     async forward(input) {
-        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)
+        let vec = await this.embedding.forward(input)
+        vec = await this.PosEmbedding.forward(vec)
         for (let i = 0; i < this.blocks.length; i++) {
-            vec = await this.blocks[i].forward(vec, encoder)
+            vec = await this.blocks[i].forward(vec)
         }
 
-        let logits = await this.linear.forward(vec, encoder)
-        let probs = await this.softmax.forward(logits, encoder)
-        await this.op.submitQueue(encoder)
+        let logits = await this.linear.forward(vec)
+        let probs = await this.softmax.forward(logits)
 
         return probs
     }
     async backward(x, target, lr) {
-        const encoder = await this.op.CreateEncoder()
         const { acc, loss } = await this.softmax.crossEntropy(x, target)
-        let dLogits = await this.softmax.backward(x, target, encoder)
-        let dInput = await this.linear.backward(dLogits, lr, encoder)
+        let dLogits = await this.softmax.backward(x, target)
+        let dInput = await this.linear.backward(dLogits, lr)
         for (let i = this.blocks.length - 1; i >= 0; i--) {
-            dInput = await this.blocks[i].backward(dInput, lr, encoder)
+            dInput = await this.blocks[i].backward(dInput, lr)
         };
-        await this.embedding.backward(dInput, lr, encoder)
-        await this.PosEmbedding.backward(dInput, lr, encoder)
+        await this.embedding.backward(dInput, lr)
+        await this.PosEmbedding.backward(dInput, lr)
         await this.op.submitQueue(encoder)
         return { acc, loss }
 

package/src/shaders/CalculateLoss.wgsl

@@ -24,6 +24,9 @@ var <uniform> params: Params;
 
 fn main(@builtin(global_invocation_id) id : vec3<u32>) {
     let index = id.x;
+    if(index >= params.vocabSize){
+        return;
+    }
     var targetIndex = targetToken[index];
     let probalitils = probs[index * params.vocabSize + u32(targetIndex)];
     loss[index] = -log(probalitils + 1e-15);

package/src/shaders/ComputeDInput.wgsl

@@ -42,5 +42,5 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
         let gIndex = gradBase + i;
         sum += outGradient[gIndex] * weight[wIndex];
     }
-    dInput[index] = sum;
+    dInput[index] = sum / f32(params.tokenCount);
 }

package/src/shaders/LB.wgsl

@@ -47,7 +47,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
 
         if (tokens > 0u) {
             let idx = weightBase + j;
-            weight[idx] -= params.lr * grad;
+            weight[idx] -= params.lr * (grad / f32(params.tokenCount));
         }
     }
 
@@ -66,6 +66,6 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
     }
 
     if (tid == 0u && tokens > 0u) {
-        bias[outputIndex] -= params.lr * sharedBias[0];
+        bias[outputIndex] -= params.lr * (sharedBias[0] / f32(params.tokenCount));
     }
 }