@imaduddin/fastbrain 1.1.1 → 2.0.0

package/src/TransformerBlock/index.d.ts

@@ -1,28 +0,0 @@
-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

@@ -1,78 +0,0 @@
-
-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, encoderMain);
-        const att = await this.attention.forward(norm1, encoderMain)
-        const residual1 = await this.rFor1.add(att, input, encoderMain)
-        const norm2 = await this.rmsnorm2.forward(residual1, encoderMain)
-        const f = await this.ffn.forward(norm2, encoderMain)
-        const residual2 = await this.rFor2.add(f, residual1, encoderMain)
-        return residual2;
-    }
-
-    async backward(outGradient, lr, encoder) {
-        const dAtt = await this.ffn.backward(outGradient, lr, encoder);
-        const dNorm2 = await this.rmsnorm2.backward(dAtt, lr, encoder)
-        const dResidual1 = await this.rback2.add(outGradient, dNorm2, encoder)
-        const dFF = await this.attention.backward(dResidual1, lr, encoder)
-        const dNorm = await this.rmsnorm.backward(dFF, lr, encoder)
-        const dResidual2 = await this.rback1.add(dNorm, dResidual1 ,encoder)
-        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

@@ -1,30 +0,0 @@
-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/model/index.d.ts

@@ -1,43 +0,0 @@
-import { Embedding } from '../Embedding/index.ts'
-import { Linear } from '../Linear/index.ts'
-import { PositionEmbedding } from '../PositionEmbedding/index.ts'
-import { TransformerBlock } from '../TransformerBlock/index.ts'
-import { Softmax } from '../Softmax/index.ts'
-import { Tokenizer } from '../Tokenizer/index.ts'
-
-
-
-export class model {
-    tokenizer: Tokenizer
-    embedding: Embedding
-    linear: Linear
-    PosEmbedding: PositionEmbedding
-    blocks: TransformerBlock[]
-    softmax: Softmax
-    embeddingSize: number
-    loopSize: number
-    contextSize: number
-    constructor(device: GPUAdapter, configs: {
-        embeddingSize: number,
-        contextSize: number,
-        blockSize: number,
-        dirname: string,
-        vocabFile: string
-    })
-    getLoss : Function
-    getAcc: Function
-    clearLoss : Function
-    async forward(input: number[]): Promise<{
-        read: function(): Promise<Float32Array>,
-        mainBuffer: GPUBuffer,
-        isDestroyed: function(): Boolean,
-        destroy: function(): void
-    }>
-    async backward(x: {
-        read: function(): Promise<Float32Array>,
-        mainBuffer: GPUBuffer,
-        isDestroyed: function(): Boolean,
-        destroy: function(): void
-    }, target: number[], lr : number)
-    async Save() : Promise<{ loss: Function, acc: Function }>
-}

package/src/model/index.js

@@ -1,83 +0,0 @@
-import { Embedding } from "../Embedding/index.js"
-import { Linear } from "../Linear/index.js"
-import { PositionEmbedding } from "../PositionEmbedding/index.js"
-import { Softmax } from "../Softmax/index.js"
-import { Tokenizer } from "../Tokenizer/index.js"
-import { join } from 'path'
-import { TransformerBlock } from "../TransformerBlock/index.js"
-import { OperationManager } from '../OperationManager.js'
-export class model {
-    tokenizer
-    embedding
-    linear
-    PosEmbedding
-    blocks = []
-    softmax
-    embeddingSize
-    contextSize
-    op;
-    clearLoss;
-    getLoss;
-    getAcc;
-    constructor(device, configs) {
-        this.contextSize = configs.contextSize
-        this.embeddingSize = configs.embeddingSize;
-        this.tokenizer = new Tokenizer(configs.vocabFile)
-        this.op = new OperationManager(device)
-        this.embedding = new Embedding(device, this.embeddingSize, this.tokenizer.vocabSize, {
-            save: {
-                filename: join(configs.dirname, "vec.bin")
-            }
-        })
-        this.softmax = new Softmax(device, this.embeddingSize, this.tokenizer.vocabSize);
-        this.PosEmbedding = new PositionEmbedding(device, this.embeddingSize, configs.contextSize, {
-            save: {
-                filename: join(configs.dirname, "posVec.bin")
-            }
-        })
-        this.linear = new Linear(device, this.embeddingSize, this.tokenizer.vocabSize, {
-            save: {
-                filename: [join(configs.dirname, "w.bin"), join(configs.dirname, "b.bin")]
-            }
-        })
-        for (let a = 0; a < configs.blockSize; a++) {
-            this.blocks[a] = new TransformerBlock(device, this.embeddingSize, {
-                dirname: join(configs.dirname, `block-${a + 1}`)
-            })
-        }
-    }
-    async forward(input) {
-        if (input.length > this.contextSize) throw new Error("The input get large then ContextSize")
-        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)
-        }
-
-        let logits = await this.linear.forward(vec)
-        let probs = await this.softmax.forward(logits)
-
-        return probs
-    }
-    async backward(x, target, lr) {
-        const { acc, loss } = await this.softmax.crossEntropy(x, target)
-        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)
-        };
-        await this.embedding.backward(dInput, lr)
-        await this.PosEmbedding.backward(dInput, lr)
-        await this.op.submitQueue(encoder)
-        return { acc, loss }
-
-    }
-    async Save() {
-        await this.PosEmbedding.Save()
-        await this.embedding.Save()
-        await this.linear.Save()
-        for (let i = this.blocks.length - 1; i >= 0; i--) {
-            await this.blocks[i].Save()
-        };
-    }
-}

package/src/residualConnection/index.js

@@ -1,31 +0,0 @@
-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, en) {
-        if (!this.params) {
-            this.params = await this.gpu.createUniFormBuffer(4 * 4)
-        }
-        const encoder = en ? en : 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);
-        }
-        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(!en) await this.gpu.submitQueue(encoder)
-        return this.output
-    }
-}

package/src/shaders/CalculateAcc.wgsl

@@ -1,50 +0,0 @@
-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

@@ -1,33 +0,0 @@
-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;
-    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

@@ -1,46 +0,0 @@
-
-struct Params {
-    vocabSize: u32,
-    tokenCount: u32,
-    lr: f32,
-    inputSize: u32,
-}
-
-@group(0) @binding(0)
-var<storage, read> 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 >= total) {
-        return;
-    }
-
-    let inputSize = params.inputSize;
-    let vocabSize = params.vocabSize;
-
-    let token = index / inputSize;
-    let inputIndex = index % inputSize;
-
-    let weightBase = inputIndex;
-    let gradBase = token * vocabSize;
-    var sum = 0.0;
-
-    for (var i = 0u; i < vocabSize; i++) {
-        let wIndex = i * inputSize + weightBase;
-        let gIndex = gradBase + i;
-        sum += outGradient[gIndex] * weight[wIndex];
-    }
-    dInput[index] = sum / f32(params.tokenCount);
-}

package/src/shaders/EB.wgsl

@@ -1,33 +0,0 @@
-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

@@ -1,25 +0,0 @@
-@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

@@ -1,71 +0,0 @@
-
-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;
-
-const WG_SIZE: u32 = 256u;
-
-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;
-    if (outputIndex >= vocab) {
-        return;
-    }
-    let gradBase = outputIndex;
-    let weightBase = outputIndex * inputSize;
-    for (var j = tid; j < inputSize; j += WG_SIZE) {
-        var grad = 0.0;
-        for (var t = 0u; t < tokens; t++) {
-            let inputIndex = t * inputSize + j;
-            let outputGradIndex = t * vocab + gradBase;
-            grad += input[inputIndex] * outGradient[outputGradIndex];
-        }
-
-        if (tokens > 0u) {
-            let idx = weightBase + j;
-            weight[idx] -= params.lr * (grad / f32(params.tokenCount));
-        }
-    }
-
-    var biasGrad = 0.0;
-    for (var t = tid; t < tokens; t += WG_SIZE) {
-        biasGrad += outGradient[t * vocab + gradBase];
-    }
-
-    sharedBias[tid] = biasGrad;
-    workgroupBarrier();
-    for (var stride = WG_SIZE / 2u; 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(params.tokenCount));
-    }
-}

package/src/shaders/LF.wgsl

@@ -1,60 +0,0 @@
-
-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

@@ -1,65 +0,0 @@
-
-@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;
-    }
-}