@mineproness/tensorloop 2.1.68 → 2.1.70

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 { Model as m } from "./src/Model/index.js";
 import {PositionEmbedding as PM } from './src/PositionEmbedding/index.js'
+import { RMSNorm as RM } from "./src/RMSNorm/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 Model = m;
-export const PositionEmbedding = PM;
+export const PositionEmbedding = PM;
+export const RMSNorm = RM

package/package.json

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

package/src/Model/index.js

@@ -1,4 +1,4 @@
-import { Embedding, Linear, Softmax, TransformerBlock, Tokenizer, PositionEmbedding } from '../../index.js'
+import { Embedding, Linear, Softmax, TransformerBlock, Tokenizer, PositionEmbedding, RMSNorm } from '../../index.js'
 import { join } from 'path'
 import fs from 'fs'
 import { OperationManager } from '../OperationManager.js'
@@ -16,6 +16,7 @@ export class Model {
     tokenizerName;
     op;
     tokens;
+    rmsnorm
     constructor(device, embeddingSize = 64, configs) {
         this.device = device;
         this.dimSize = embeddingSize;

package/src/RMSNorm/index.d.ts

@@ -0,0 +1,10 @@
+export class RMSNorm{
+    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

@@ -0,0 +1,84 @@
+import { OperationManager } from '../OperationManager.js'
+import fs from 'fs'
+import { SaveVector, getVector } from '../getVec.js'
+
+export class RMSNorm {
+  embeddingSize;
+  reductionShader;
+  forwardShader;
+  backwardShader;
+  op;
+  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", {
+      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) {
+    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, 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()
+    this.rms = outputReduce
+    if (this.input) this.input.destroy()
+    this.input = input;
+    if (this.outputPre) this.outputPre.destroy()
+    this.outputPre = output
+    return output
+  }
+  async backward(outputGradient, lr) {
+    const output = await this.op.createBuffer(outputGradient.size);
+    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

@@ -1,47 +1,85 @@
-import {SelfAttention} from '../SelfAttention/index.js'
-import {FFN} from '../FFN/index.js'
-import {existsSync, mkdirSync} from 'fs'
-import {OperationManager} from '../OperationManager.js'
-export class TransformerBlock{
+
+import { SelfAttention } from '../SelfAttention/index.js';
+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;
     ffn;
     gpu;
     size;
-    constructor(device, embeddingSize, configs){
+    rmsnorm;
+    rmsnorm2;
+
+    constructor(device, embeddingSize, configs) {
         this.size = embeddingSize;
-        if(!existsSync(configs.dirname)) mkdirSync(configs.dirname);
-        this.attention = new SelfAttention(device, embeddingSize, configs)
-        this.ffn = new FFN(device, embeddingSize, configs)
-        this.gpu = new OperationManager(device)
+
+        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 });
+        }
+
+        this.attention = new SelfAttention(
+            device, embeddingSize, configs
+        );
+
+        this.ffn = new FFN(device, embeddingSize, configs);
+        this.gpu = new OperationManager(device);
     }
-    async forward(input){
-        let context = input.size / (this.size * 4)
-        const xAtt = await this.attention.forward(input);
-        const re1 = await this.gpu.add(input, xAtt, this.size * context)
-        const xFF = await this.ffn.forward(re1);
-        const re2 = await this.gpu.add(re1, xFF, this.size * context)
-        xAtt.destroy()
-        xFF.destroy()
-        return re2
+
+    async forward(input) {
+        const count = input.size / 4;
+        const norm1 = await this.rmsnorm.forward(input);
+        const att = await this.attention.forward(norm1);
+        const h = await this.gpu.add(input, att, count);
+        const norm2 = await this.rmsnorm2.forward(h);
+        const f = await this.ffn.forward(norm2);
+        const y = await this.gpu.add(h, f, count);
+        norm1.destroy();
+        att.destroy();
+        f.destroy();
+        return y;
     }
-    async backward(outGradient, lr){
-        const context = outGradient.size / (this.size * 4)
-        const dAtt = await this.ffn.backward(outGradient, lr);
-        const re1 = await this.gpu.add(outGradient, dAtt, this.size * context)
-        const dFF = await this.attention.backward(re1, lr)
-        const re2 = await this.gpu.add(re1, dFF , context)
-        outGradient.destroy()
-        re1.destroy()
-        dFF.destroy()
-        dAtt.destroy()
-        return re2
+
+    async backward(outGradient, lr) {
+        const count = outGradient.size / 4;
+        const dNorm2 = await this.ffn.backward(outGradient, lr);
+        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, lr);
+        const dInput = await this.gpu.add(dH, dInputFromAttention, count);
+        dNorm2.destroy();
+        dHFromFFN.destroy();
+        dH.destroy();
+        dNorm1.destroy();
+        dInputFromAttention.destroy();
+        dInputFromAttention.destroy();
+        return dInput;
     }
-    async ClearInputCache(){
-        await this.ffn.ClearInputCache()
-        await this.attention.ClearInputCache()
+
+    async ClearInputCache() {
+        await this.ffn.ClearInputCache();
+        await this.attention.ClearInputCache();
     }
-    async Save(){
-        await this.ffn.Save()
-        await this.attention.Save()
+
+    async Save() {
+        await this.ffn.Save();
+        await this.attention.Save();
+        await tthis.rmsnorm.Save()
+        await this.rmsnorm2.Save()
     }
 }

package/src/shaders/LB.wgsl

@@ -42,14 +42,14 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
         let inputValue = input[token * u32(params.inputSize) + inputIndex];
         let gradient = outGradient[token * u32(params.vocabSize) + outputIndex];
         let dW = inputValue * gradient;
-        dWeights += f32(max(-3.0, min(3.0, dW)));
+        dWeights += dW;
     }
     let vClap = weight[index] - params.lr * dWeights;
     weight[index] = vClap;
     if (inputIndex == 0u) {
         var dBias = 0.0;
         for (var token = 0u; token < u32(params.tokenCount); token++) {
-            dBias += f32(max(-3.0, min(3.0, outGradient[token * u32(params.vocabSize) + outputIndex])));
+            dBias += outGradient[token * u32(params.vocabSize) + outputIndex];
         }
         bias[outputIndex] = bias[outputIndex] - params.lr * dBias;
 

package/src/shaders/Norm/rmsBackward.wgsl

@@ -0,0 +1,42 @@
+@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> gamma: array<f32>;
+
+@group(0) @binding(4)
+var<storage, read_write> gradInput: array<f32>;
+
+var<workgroup> partial: array<f32, %embeddingSize%>;
+
+@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 = D / 2u;
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+        workgroupBarrier();
+        stride /= 2u;
+    }
+    let dot = partial[0];
+    gradInput[idx] = (dy * g) / r - x * dot / (f32(D) * r * r * r);
+}

package/src/shaders/Norm/rmsForward.wgsl

@@ -0,0 +1,30 @@
+@group(0) @binding(0)
+var<storage, read> reduce: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read> vector: array<f32>;
+
+@group(0) @binding(2)
+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>) {
+    let x = id.x;
+
+    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

@@ -0,0 +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, %embeddingSize%>;
+
+const embeddingSize : u32 = %embeddingSize%u;
+@compute @workgroup_size(64)
+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;
+    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 /= 2u;
+    }
+    if (lid == 0u) {
+        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;
+    }
+}