@mineproness/tensorloop 2.1.68 → 2.1.69

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.69",
   "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,5 @@
+export class RMSNorm{
+    constructor(device : GPUAdapter, embeddingSize : number);
+    async forward(input : GPUBuffer) : Promise<GPUBuffer>
+    async backward(outputGradient : GPUBuffer) : Promise<GPUBuffer>
+}

package/src/RMSNorm/index.js

@@ -0,0 +1,45 @@
+import { OperationManager } from '../OperationManager.js'
+import fs from 'fs'
+
+export class RMSNorm {
+  embeddingSize;
+  reductionShader;
+  forwardShader;
+  backwardShader;
+  op;
+  constructor(device, embeddingSize) {
+    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")
+  }
+  async forward(input) {
+    const outputReduce = await this.op.createBuffer(4 * 1);
+    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])
+    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) {
+    const output = await this.op.createBuffer(outputGradient.size);
+    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input, outputGradient, this.rms, output]);
+    const encoder = await this.op.CreateEncoder()
+    await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size / 4) / 64))
+    await this.op.submitQueue(encoder)
+    if (this.prevOut) this.prevOut.destroy()
+    this.prevOut = outputGradient
+    return output
+  }
+
+}

package/src/TransformerBlock/index.js

@@ -1,47 +1,74 @@
-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';
+
+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);
+        this.rmsnorm2 = new RMSNorm(device, embeddingSize);
+
+        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);
+        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 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();
     }
 }

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,35 @@
+
+@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> gradInput: array<f32>;
+
+var<workgroup> partial: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32) {
+    let D = arrayLength(&input);
+    partial[lid] = gradOutput[lid] * input[lid];
+    workgroupBarrier();
+    var stride = 32u;
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+        workgroupBarrier();
+        stride = stride / 2u;
+    }
+    let dot = partial[0];
+    let r = rms[0];
+    gradInput[lid] = gradOutput[lid] / r - input[lid] * dot / (f32(D) * r * r * r);
+}

package/src/shaders/Norm/rmsForward.wgsl

@@ -0,0 +1,19 @@
+@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>;
+
+@compute @workgroup_size(64)
+
+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;
+}

package/src/shaders/Norm/sumSquare.wgsl

@@ -0,0 +1,37 @@
+
+@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>;
+
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32) {
+    var sum = 0.0;
+    var i = lid;
+    while (i < arrayLength(&input)) {
+        sum += input[i] * input[i];
+        i += 64u;
+    }
+    partial[lid] = sum;
+    workgroupBarrier();
+    var stride = 32u;
+
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+
+        workgroupBarrier();
+        stride = stride / 2u;
+    }
+    if (lid == 0u) {
+        out[0] = partial[0];
+    }
+}