@imaduddin/fastbrain 0.0.0-stage → 1.0.0

package/src/TransformerBlock/index.d.ts

@@ -0,0 +1,28 @@
+export class TransformerBlock{
+    constructor(device: GPUAdapter, embeddingSize: number, configs: {
+        dirname: string
+    })
+    async forward(input : {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }) : Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async backward(outInput : {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, lr : number) : Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async Save() : Promise<void>
+}

package/src/TransformerBlock/index.js

@@ -0,0 +1,78 @@
+
+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 path, { join } from 'path';
+import { residualConnection } from '../residualConnection/index.js'
+
+export class TransformerBlock {
+    attention;
+    ffn;
+    gpu;
+    size;
+    rmsnorm;
+    rmsnorm2;
+
+    constructor(device, embeddingSize, configs) {
+        this.size = embeddingSize;
+        if (!existsSync(configs.dirname)) {
+            mkdirSync(configs.dirname, { recursive: true });
+        }
+        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")
+            }
+        });
+
+
+        this.attention = new SelfAttention(
+            device, embeddingSize, configs
+        );
+        this.rFor1 = new residualConnection(device)
+        this.rFor2 = new residualConnection(device)
+        this.rback1 = new residualConnection(device)
+        this.rback2 = new residualConnection(device)
+        this.ffn = new FFN(device, embeddingSize, configs);
+        this.gpu = new OperationManager(device);
+    }
+
+    async forward(input, encoderMain) {
+        const count = input.size / 4;
+        const norm1 = await this.rmsnorm.forward(input);
+        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 f = await this.ffn.forward(norm2, encoderMain)
+        const residual2 = await this.rFor2.add(f, residual1, encoderMain)
+        return residual2;
+    }
+
+    async backward(outGradient, lr, encoderMain) {
+        const dAtt = await this.ffn.backward(outGradient, lr, encoderMain);
+        const dNorm2 = await this.rmsnorm2.backward(dAtt, lr)
+        const dResidual1 = await this.rback2.add(outGradient, dNorm2, encoderMain)
+        const dFF = await this.attention.backward(dResidual1, lr, encoderMain)
+        const dNorm = await this.rmsnorm.backward(dFF, lr)
+        const dResidual2 = await this.rback1.add(dNorm, dResidual1, encoderMain)
+        return dResidual2;
+    }
+
+    async ClearInputCache() {
+        await this.ffn.ClearInputCache();
+        await this.attention.ClearInputCache();
+    }
+
+    async Save() {
+        await this.ffn.Save();
+        await this.attention.Save();
+        await this.rmsnorm.Save()
+        await this.rmsnorm2.Save()
+    }
+}

package/src/getVec.js

@@ -0,0 +1,30 @@
+import fs from 'fs'
+
+
+export function getVector(filename, size, bias) {
+    if (!fs.existsSync(filename)) {
+        return null;
+    }
+    const buffer = fs.readFileSync(filename);
+    const r = new Float32Array(
+        buffer.buffer,
+        buffer.byteOffset,
+        buffer.byteLength / 4
+    );
+    if (r.length >= size) {
+        return r;
+    }
+    const output = new Float32Array(size);
+    output.set(r);
+    if (!bias) {
+        for (let i = r.length; i < size; i++) {
+            output[i] = (Math.random() * 2 - 1) * 0.02;
+        }
+    }
+    return output;
+}
+
+export function SaveVector(filename, vectors) {
+    const buffer = Buffer.from(vectors.buffer)
+    fs.writeFileSync(filename, buffer)
+}

package/src/residualConnection/index.js

@@ -0,0 +1,32 @@
+import { OperationManager } from "../OperationManager.js";
+
+export class residualConnection {
+    gpu;
+    addShader;
+    a;
+    b;
+    output;
+    constructor(device) {
+        this.gpu = new OperationManager(device)
+        this.addShader = this.gpu.getShaderPipline("./src/shaders/add.wgsl")
+    }
+    async add(a, b, encoderMain) {
+        if (!this.params) {
+            this.params = await this.gpu.createUniFormBuffer(4 * 4)
+        }
+        const encoder = encoderMain ? encoderMain : await this.gpu.CreateEncoder()
+        if (!this.output) {
+            this.output = await this.gpu.ManageBuffer(a.size);
+        }
+        if (this.output.size !== a.size || this.output.size == 0 || this.output.isDestroyed()) {
+            this.output.destroy()
+            this.output = await this.gpu.ManageBuffer(a.size);
+        }
+        await this.output.count()
+        const bind = await this.gpu.createBindGroup(this.addShader, 0, [a.mainBuffer, b.mainBuffer, this.output.mainBuffer, this.params]);
+        await this.gpu.WriteBuffer(this.params, new Uint32Array([(a.size  / 4), 0]))
+        await this.gpu.RunPipline(encoder, this.addShader, bind, Math.ceil((a.size / 4) / 64));
+        if(!encoderMain) await this.gpu.submitQueue(encoder)
+        return this.output
+    }
+}

package/src/shaders/CalculateAcc.wgsl

@@ -0,0 +1,50 @@
+struct Params{
+    vocabSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var <uniform> params : Params;
+
+
+@group(0) @binding(1)
+
+var <storage, read_write> probs : array<f32>;
+
+@group(0) @binding(3)
+
+var <storage, read_write> tToken : array<u32>;
+
+@group(0) @binding(2)
+
+
+var <storage, read_write> result : array<f32>;
+
+
+@compute @workgroup_size(1)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+
+    var correct = 0.0;
+    for(var a = 0u; a < params.tokenCount; a++){
+        var bestVal = - 3.402823e38;
+        var maxToken = 0u;
+        var start = a * params.vocabSize;
+        let end = start + params.vocabSize;
+        for(var t = start; t < end; t++){
+            let val = probs[t];
+            if(bestVal < val){
+                bestVal = val;
+                maxToken = t - start;
+            }
+        }
+        if(maxToken == tToken[a]){
+            correct += 1.0;
+        }
+
+    }
+    correct /= f32(params.tokenCount);
+    result[x] = correct;
+}

package/src/shaders/CalculateLoss.wgsl

@@ -0,0 +1,30 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+
+var <storage, read> probs : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> targetToken : array<u32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> loss : array<f32>;
+
+@group(0) @binding(3)
+
+var <uniform> params: Params;
+
+
+@compute @workgroup_size(4)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>) {
+    let index = id.x;
+    var targetIndex = targetToken[index];
+    let probalitils = probs[index * params.vocabSize + u32(targetIndex)];
+    loss[index] = -log(probalitils + 1e-15);
+}

package/src/shaders/ComputeDInput.wgsl

@@ -0,0 +1,32 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32,
+    lr: f32,
+    inputSize: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read_write> weight: array<f32>;
+@group(0) @binding(1)
+var<storage, read> outGradient: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@group(0) @binding(4)
+var<storage, read_write> dInput: array<f32>;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.tokenCount * params.inputSize;
+    if (index >= u32(total)) {
+        return;
+    }
+    let token =index / u32(params.inputSize);
+    let inputIndex = index % u32(params.inputSize);
+    var gradInput = 0.0;
+    for (var i = 0u; i < u32(params.vocabSize); i++) {
+        let w = weight[i * u32(params.inputSize) + inputIndex];
+        let grad = outGradient[token * u32(params.vocabSize) + i];
+        gradInput += grad * w;
+    }
+    dInput[index] = gradInput;
+}

package/src/shaders/EB.wgsl

@@ -0,0 +1,33 @@
+struct Params {
+    lr: f32
+}
+
+@group(0) @binding(0)
+
+var<storage, read_write> embeddings: array<f32>;
+
+@group(0) @binding(1)
+
+var<storage, read_write> outputGradient: array<f32>;
+
+@group(0) @binding(2)
+
+var<storage, read_write> tokens: array<u32>;
+
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+const embeddingSize = %embeddingSize%u;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let dim = id.x % embeddingSize;
+    if (id.x >= arrayLength(&outputGradient)) {
+        return;
+    }
+    let tokenIndex = id.x / embeddingSize;
+    let tokenID = tokens[tokenIndex];
+    let embeddingIDx = tokenID * embeddingSize + dim;
+    embeddings[embeddingIDx] -= params.lr * outputGradient[id.x];
+}

package/src/shaders/EF.wgsl

@@ -0,0 +1,25 @@
+@group(0) @binding(0)
+
+var <storage, read_write> embeddings: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> tokens: array<u32>;
+
+const embeddingSize = %embeddingSize%u;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let dim = id.x % embeddingSize;
+    let tokenIndex = id.x / embeddingSize;
+
+    let tokenID = tokens[tokenIndex];
+
+    let embeddingIDx = tokenID * embeddingSize + dim;
+    output[id.x] = embeddings[embeddingIDx];
+}

package/src/shaders/LB.wgsl

@@ -0,0 +1,75 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32,
+    lr: f32,
+    inputSize: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read_write> bias: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> weight: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read> outGradient: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(4)
+var<uniform> params: Params;
+
+var<workgroup> sharedBias: array<f32, 256>;
+
+@compute @workgroup_size(256)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>
+) {
+    let outputIndex = wid.x;
+    let tid = lid.x;
+
+    let vocab = params.vocabSize;
+    let tokens = params.tokenCount;
+    let inputSize = params.inputSize;
+
+    for (
+        var inputIndex = tid;
+        inputIndex < inputSize;
+        inputIndex += 256u
+    ) {
+        var grad = 0.0;
+
+        for (var t = 0u; t < tokens; t++) {
+            grad += input[t * inputSize + inputIndex]
+                  * outGradient[t * vocab + outputIndex];
+        }
+
+        if (tokens > 0u) {
+            weight[outputIndex * inputSize + inputIndex] -=
+                params.lr * grad / f32(tokens);
+        }
+    }
+
+    var biasGrad = 0.0;
+
+    for (var t = tid; t < tokens; t += 256u) {
+        biasGrad += outGradient[t * vocab + outputIndex];
+    }
+
+    sharedBias[tid] = biasGrad;
+    workgroupBarrier();
+
+    for (var stride = 128u; stride > 0u; stride >>= 1u) {
+        if (tid < stride) {
+            sharedBias[tid] += sharedBias[tid + stride];
+        }
+        workgroupBarrier();
+    }
+
+    if (tid == 0u && tokens > 0u) {
+        bias[outputIndex] -=
+            params.lr * sharedBias[0] / f32(tokens);
+    }
+}

package/src/shaders/LF.wgsl

@@ -0,0 +1,60 @@
+
+struct Params {
+    inputSize: u32,
+    outputSize: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> weights: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read> bias: array<f32>;
+
+@group(0) @binding(4)
+var<uniform> params: Params;
+
+const TILE_OUT: u32 = 16u;
+const TILE_TOKEN: u32 = 4u;
+const TILE_K: u32 = 16u;
+
+var<workgroup> tileInput: array<f32, 64>;
+var<workgroup> tileWeights: array<f32, 256>;
+
+@compute @workgroup_size(64)
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+    let outFeature = wid.x * TILE_OUT + lid.x % TILE_OUT;
+    let token = wid.y * TILE_TOKEN + lid.x / TILE_OUT;
+    let localOut = lid.x % TILE_OUT;
+    let localToken = lid.x / TILE_OUT;
+    var sum = 0.0;
+    for (var baseK = 0u; baseK < params.inputSize; baseK += TILE_K) {
+        let loadIndex = lid.x;
+        let inputToken = wid.y * TILE_TOKEN + loadIndex / TILE_K;
+        let inputK = baseK + loadIndex % TILE_K;
+        if (loadIndex < TILE_TOKEN * TILE_K && inputToken < params.tokenCount && inputK < params.inputSize) {
+            tileInput[loadIndex] = input[inputToken * params.inputSize + inputK];
+        }
+        let weightFeature = wid.x * TILE_OUT + loadIndex / TILE_K;
+        let weightK = baseK + loadIndex % TILE_K;
+        if (loadIndex < TILE_OUT * TILE_K && weightFeature < params.outputSize &&  weightK < params.inputSize) {
+            tileWeights[loadIndex] = weights[weightFeature * params.inputSize + weightK];
+        }
+        workgroupBarrier();
+        for (var k = 0u; k < TILE_K; k++) {
+            let a = tileInput[localToken * TILE_K + k];
+            let b = tileWeights[localOut * TILE_K + k];
+            sum += a * b;
+        }
+        workgroupBarrier();
+    }
+    if (token < params.tokenCount && outFeature < params.outputSize) {
+        output[token * params.outputSize + outFeature] = sum + bias[outFeature];
+    }
+}

package/src/shaders/Norm/rmsBackward.wgsl

@@ -0,0 +1,65 @@
+
+@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> sumSquares: 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>;
+
+const D: u32 = %embeddingSize%u;
+const EPSILON: f32 = 1e-6;
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
+    let token = wid.x;
+    if (token >= arrayLength(&sumSquares)) {
+        return;
+    }
+    let base = token * D;
+    let r = sqrt(sumSquares[token] / f32(D) + EPSILON);
+    var localSum = 0.0;
+    var d = lid;
+    while (d < D) {
+        let idx = base + d;
+        localSum += gradOutput[idx] * gamma[d] * input[idx];
+        d += 64u;
+    }
+
+    partial[lid] = localSum;
+    workgroupBarrier();
+    var stride = 32u;
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+
+        workgroupBarrier();
+        stride /= 2u;
+    }
+    let dot = partial[0];
+    d = lid;
+    while (d < D) {
+        let idx = base + d;
+        let x = input[idx];
+        let dy = gradOutput[idx];
+        let g = gamma[d];
+
+        gradInput[idx] =
+            (dy * g) / r
+            - (x * dot) / (f32(D) * r * r * r);
+
+        d += 64u;
+    }
+}

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,38 @@
+
+@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>;
+const D: u32 = %embeddingSize%u;
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
+    let base = wid.x * D;
+    var sum = 0.0;
+    var d = lid;
+    while (d < D) {
+        let idx = base + d;
+        if (idx < arrayLength(&input)) {
+            sum += input[idx] * input[idx];
+        }
+        d += 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) {
+        out[wid.x] = partial[0];
+    }
+}

package/src/shaders/Norm/updateGamma.wgsl

@@ -0,0 +1,52 @@
+
+@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> sumSquares: 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, 64>;
+
+const D: u32 = %embeddingSize%u;
+const EPSILON: f32 = 1e-6;
+
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
+    let dim = wid.x;
+    let T = arrayLength(&sumSquares);
+    var sum = 0.0;
+    if (dim < D) {
+        var token = lid;
+        while (token < T) {
+            let idx = token * D + dim;
+            let rms = sqrt(sumSquares[token] / f32(D) + EPSILON);
+            sum += gradOutput[idx] * input[idx] / rms;
+            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 && dim < D) {
+        gamma[dim] -= params[0] * partial[0];
+    }
+}

package/src/shaders/PB.wgsl

@@ -0,0 +1,25 @@
+struct Params{
+    lr : f32
+}
+
+@group(0) @binding(0)
+var <storage, read_write> posVec: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> outputGradient : array<f32>;
+
+
+@group(0) @binding(2)
+
+var <uniform> params: Params;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    if(x >= arrayLength(&outputGradient)){
+        return;
+    }
+    posVec[x] -= params.lr * outputGradient[x];
+}