@mineproness/tensorloop 2.1.69 → 2.1.70

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mineproness/tensorloop",
-  "version": "2.1.69",
+  "version": "2.1.70",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called Tensorloop.",
   "keywords": [
     "ai",

package/src/RMSNorm/index.d.ts

@@ -1,5 +1,10 @@
 export class RMSNorm{
-    constructor(device : GPUAdapter, embeddingSize : number);
+    constructor(device : GPUAdapter, embeddingSize : number , {
+        save:{
+            filename: string
+        }
+    });
     async forward(input : GPUBuffer) : Promise<GPUBuffer>
     async backward(outputGradient : GPUBuffer) : Promise<GPUBuffer>
+    async Save() : Promise<void>
 }

package/src/RMSNorm/index.js

@@ -1,5 +1,6 @@
 import { OperationManager } from '../OperationManager.js'
 import fs from 'fs'
+import { SaveVector, getVector } from '../getVec.js'
 
 export class RMSNorm {
   embeddingSize;
@@ -7,20 +8,49 @@ export class RMSNorm {
   forwardShader;
   backwardShader;
   op;
-  constructor(device, embeddingSize) {
+  gamma;
+  gammabuffer;
+  configs;
+  constructor(device, embeddingSize, configs) {
+    this.configs = configs
     this.op = new OperationManager(device);
     this.embeddingSize = embeddingSize;
-    this.reductionShader = this.op.getShaderPipline("./src/shaders/Norm/sumSquare.wgsl")
-    this.forwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsForward.wgsl")
-    this.backwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsBackward.wgsl")
+    this.reductionShader = this.op.getShaderPipline("./src/shaders/Norm/sumSquare.wgsl", {
+      embeddingSize: embeddingSize,
+    })
+    this.forwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsForward.wgsl", {
+      embeddingSize: embeddingSize
+    })
+    this.backwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsBackward.wgsl", {
+      embeddingSize: embeddingSize
+    })
+    this.updateGamma = this.op.getShaderPipline("./src/shaders/Norm/updateGamma.wgsl", {
+      embeddingSize: embeddingSize
+    })
+    if (configs) {
+      const ga = getVector(configs.save.filename, embeddingSize);
+      if (!ga) {
+        this.gamma = new Float32Array(embeddingSize).fill(1);
+        SaveVector(configs.save.filename, this.gamma)
+      } else {
+        this.gamma = ga
+      }
+    } else {
+      this.gamma = new Float32Array(embeddingSize).fill(1)
+    }
   }
   async forward(input) {
-    const outputReduce = await this.op.createBuffer(4 * 1);
+    if (!this.gammabuffer) {
+      this.gammabuffer = await this.op.createBuffer(4 * this.embeddingSize);
+      await this.op.WriteBuffer(this.gammabuffer, this.gamma);
+    }
+    let context = input.size / (this.embeddingSize * 4)
+    const outputReduce = await this.op.createBuffer(4 * context);
     const output = await this.op.createBuffer(input.size);
     const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input, outputReduce])
     const encoder = await this.op.CreateEncoder()
     await this.op.RunPipline(encoder, this.reductionShader, bindGroup, Math.ceil((input.size / 4) / 64))
-    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [outputReduce, input, output])
+    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [outputReduce, input, output, this.gammabuffer])
     await this.op.RunPipline(encoder, this.forwardShader, bindGroup1, Math.ceil((input.size / 4) / 64));
     await this.op.submitQueue(encoder);
     if (this.rms) this.rms.destroy()
@@ -31,15 +61,24 @@ export class RMSNorm {
     this.outputPre = output
     return output
   }
-  async backward(outputGradient) {
+  async backward(outputGradient, lr) {
     const output = await this.op.createBuffer(outputGradient.size);
-    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input, outputGradient, this.rms, output]);
+    let params = await this.op.createBuffer(4 * 2)
+    await this.op.WriteBuffer(params, new Float32Array([lr, 0]))
+    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input, outputGradient, this.rms, this.gammabuffer, output]);
+    const bindGroup2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input, outputGradient, this.rms, this.gammabuffer, params]);
     const encoder = await this.op.CreateEncoder()
     await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size / 4) / 64))
+    await this.op.RunPipline(encoder, this.updateGamma, bindGroup2, this.embeddingSize)
     await this.op.submitQueue(encoder)
     if (this.prevOut) this.prevOut.destroy()
     this.prevOut = outputGradient
+    params.destroy()
     return output
   }
+  async Save(){
+    const read = await this.op.readBuffer(4 * this.embeddingSize, this.gammabuffer);
+    SaveVector(this.configs.save.filename, read)
+  }
 
 }

package/src/TransformerBlock/index.js

@@ -4,6 +4,7 @@ import { FFN } from '../FFN/index.js';
 import { existsSync, mkdirSync } from 'fs';
 import { RMSNorm } from '../../index.js';
 import { OperationManager } from '../OperationManager.js';
+import { join } from 'path';
 
 export class TransformerBlock {
     attention;
@@ -16,8 +17,16 @@ export class TransformerBlock {
     constructor(device, embeddingSize, configs) {
         this.size = embeddingSize;
 
-        this.rmsnorm = new RMSNorm(device, embeddingSize);
-        this.rmsnorm2 = new RMSNorm(device, embeddingSize);
+        this.rmsnorm = new RMSNorm(device, embeddingSize, {
+            save: {
+                filename: join(configs.dirname, "LAYERNORM-1.bin")
+            }
+        });
+        this.rmsnorm2 = new RMSNorm(device, embeddingSize, {
+            save: {
+                filename: join(configs.dirname, "LAYERNORM-2.bin")
+            }
+        });
 
         if (!existsSync(configs.dirname)) {
             mkdirSync(configs.dirname, { recursive: true });
@@ -48,10 +57,10 @@ export class TransformerBlock {
     async backward(outGradient, lr) {
         const count = outGradient.size / 4;
         const dNorm2 = await this.ffn.backward(outGradient, lr);
-        const dHFromFFN = await this.rmsnorm2.backward(dNorm2);
+        const dHFromFFN = await this.rmsnorm2.backward(dNorm2 , lr);
         const dH = await this.gpu.add(outGradient, dHFromFFN, count);
         const dNorm1 = await this.attention.backward(dH, lr);
-        const dInputFromAttention = await this.rmsnorm.backward(dNorm1);
+        const dInputFromAttention = await this.rmsnorm.backward(dNorm1, lr);
         const dInput = await this.gpu.add(dH, dInputFromAttention, count);
         dNorm2.destroy();
         dHFromFFN.destroy();
@@ -70,5 +79,7 @@ export class TransformerBlock {
     async Save() {
         await this.ffn.Save();
         await this.attention.Save();
+        await tthis.rmsnorm.Save()
+        await this.rmsnorm2.Save()
     }
 }

package/src/shaders/Norm/rmsBackward.wgsl

@@ -1,4 +1,3 @@
-
 @group(0) @binding(0)
 var<storage, read> input: array<f32>;
 
@@ -9,16 +8,25 @@ var<storage, read> gradOutput: array<f32>;
 var<storage, read> rms: array<f32>;
 
 @group(0) @binding(3)
+var<storage, read> gamma: array<f32>;
+
+@group(0) @binding(4)
 var<storage, read_write> gradInput: array<f32>;
 
-var<workgroup> partial: array<f32, 64>;
+var<workgroup> partial: array<f32, %embeddingSize%>;
 
-@compute @workgroup_size(64)
-fn main(@builtin(local_invocation_index) lid: u32) {
-    let D = arrayLength(&input);
-    partial[lid] = gradOutput[lid] * input[lid];
+@compute @workgroup_size(%embeddingSize%)
+fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
+    let D = %embeddingSize%u;
+    let token = wid.x;
+    let idx = token * D + lid;
+    let x = input[idx];
+    let dy = gradOutput[idx];
+    let g = gamma[lid];
+    let r = rms[token];
+    partial[lid] = dy * g * x;
     workgroupBarrier();
-    var stride = 32u;
+    var stride = D / 2u;
     loop {
         if (stride == 0u) {
             break;
@@ -27,9 +35,8 @@ fn main(@builtin(local_invocation_index) lid: u32) {
             partial[lid] += partial[lid + stride];
         }
         workgroupBarrier();
-        stride = stride / 2u;
+        stride /= 2u;
     }
     let dot = partial[0];
-    let r = rms[0];
-    gradInput[lid] = gradOutput[lid] / r - input[lid] * dot / (f32(D) * r * r * r);
+    gradInput[idx] = (dy * g) / r - x * dot / (f32(D) * r * r * r);
 }

package/src/shaders/Norm/rmsForward.wgsl

@@ -1,19 +1,30 @@
 @group(0) @binding(0)
-
-var <storage, read> reduce: array<f32>;
+var<storage, read> reduce: array<f32>;
 
 @group(0) @binding(1)
-
-var <storage, read> vector: array<f32>;
+var<storage, read> vector: array<f32>;
 
 @group(0) @binding(2)
+var<storage, read_write> output: array<f32>;
 
-var <storage, read_write> output : array<f32>;
+@group(0) @binding(3)
+var<storage, read> gamma: array<f32>;
 
 @compute @workgroup_size(64)
-
-fn main(@builtin(global_invocation_id) id : vec3<u32>){
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
     let x = id.x;
-    let rms = sqrt(reduce[0u] / f32(arrayLength(&vector)) + 1e-6);
-    output[x] = vector[x] / rms;
+
+    if (x >= arrayLength(&vector)) {
+        return;
+    }
+
+    let dim = %embeddingSize%u;
+    let token = x / dim;
+    let d = x % dim;
+
+    let rms = sqrt(
+        reduce[token] / f32(dim) + 1e-6
+    );
+
+    output[x] = (vector[x] / rms) * gamma[d];
 }

package/src/shaders/Norm/sumSquare.wgsl

@@ -1,37 +1,35 @@
-
 @group(0) @binding(0)
 var<storage, read> input: array<f32>;
 
 @group(0) @binding(1)
 var<storage, read_write> out: array<f32>;
 
-var<workgroup> partial: array<f32, 64>;
 
+var<workgroup> partial: array<f32, %embeddingSize%>;
+
+const embeddingSize : u32 = %embeddingSize%u;
 @compute @workgroup_size(64)
-fn main(@builtin(local_invocation_index) lid: u32) {
+fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
+    let dim = embeddingSize;
+    let index = wid.x * dim + lid;
     var sum = 0.0;
-    var i = lid;
-    while (i < arrayLength(&input)) {
-        sum += input[i] * input[i];
-        i += 64u;
+    if (lid < dim && index < arrayLength(&input)) {
+        sum = input[index] * input[index];
     }
     partial[lid] = sum;
     workgroupBarrier();
     var stride = 32u;
-
     loop {
         if (stride == 0u) {
             break;
         }
-
         if (lid < stride) {
             partial[lid] += partial[lid + stride];
         }
-
         workgroupBarrier();
-        stride = stride / 2u;
+        stride /= 2u;
     }
     if (lid == 0u) {
-        out[0] = partial[0];
+        out[wid.x] = partial[0];
     }
 }

package/src/shaders/Norm/updateGamma.wgsl

@@ -0,0 +1,61 @@
+@group(0) @binding(0)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read> gradOutput: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read> rms: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read_write> gamma: array<f32>;
+
+@group(0) @binding(4)
+var<storage, read> params: array<f32>;
+
+var<workgroup> partial: array<f32, %embeddingSize%>;
+
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
+    let D = %embeddingSize%u;
+    let dim = wid.x;
+
+    if (dim >= D) {
+        return;
+    }
+
+    let T = arrayLength(&rms);
+    let lr = params[0];
+
+    var sum = 0.0;
+    var token = lid;
+
+    while (token < T) {
+        let idx = token * D + dim;
+
+        sum += gradOutput[idx] * input[idx] / rms[token];
+
+        token += 64u;
+    }
+
+    partial[lid] = sum;
+    workgroupBarrier();
+
+    var stride = 32u;
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+
+        workgroupBarrier();
+        stride /= 2u;
+    }
+    if (lid == 0u) {
+        let gradGamma = partial[0];
+        gamma[dim] -= lr * gradGamma;
+    }
+}