@imaduddin/fastbrain 1.1.1-patch → 2.0.0

package/src/Linear/index.js

@@ -1,196 +0,0 @@
-
-import fs from 'fs'
-import { buffer } from "stream/consumers";
-import { SaveVector, getVector } from '../getVec.js'
-import { OperationManager } from "../OperationManager.js";
-
-
-
-function GenerateVectors(embeddingSize, vocabSize) {
-    const vec = new Float32Array(embeddingSize * vocabSize).map(() => (Math.random() * 2 - 1) * 0.02)
-    return vec
-}
-
-function GenerateBias(vocabSize) {
-    const vec = new Float32Array(vocabSize)
-    return vec
-}
-
-export class Linear {
-    device;
-    embeddingSize;
-    vocabSize;
-    configs;
-    bias;
-    weights;
-    forwardPipline;
-    configs;
-    input;
-    wBuffer;
-    bBuffer;
-    gpu;
-    paramForward;
-    oldParam;
-    constructor(device, embeddingSize, vocabSize, configs) {
-        this.configs = configs
-        this.embeddingSize = embeddingSize;
-        this.vocabSize = vocabSize;
-        this.configs = configs;
-        this.device = device
-        let total = embeddingSize * vocabSize;
-        if (configs && configs.save) {
-            const w = getVector(configs.save.filename[0], total);
-            const b = getVector(configs.save.filename[1], vocabSize, true)
-            if (!w && !b) {
-                this.bias = GenerateBias(vocabSize)
-                this.weights = GenerateVectors(embeddingSize, vocabSize)
-                SaveVector(configs.save.filename[0], this.weights)
-                SaveVector(configs.save.filename[1], this.bias)
-            } else {
-                this.bias = b;
-                this.weights = w;
-            }
-        } else {
-            this.bias = GenerateBias(vocabSize)
-            this.weights = GenerateVectors(embeddingSize, vocabSize)
-        }
-        this.gpu = new OperationManager(device);
-        this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/LF.wgsl")
-        this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/LB.wgsl")
-        this.dInputBackward = this.gpu.getShaderPipline("./src/shaders/ComputeDInput.wgsl")
-    }
-    async forward(vec, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
-        const tokenCount = vec.mainBuffer.size / (this.embeddingSize * 4)
-        if (!this.wBuffer) {
-
-            this.wBuffer = await this.gpu.createBuffer(this.embeddingSize * this.vocabSize * 4)
-            await this.gpu.WriteBuffer(this.wBuffer, this.weights)
-            this.bBuffer = await this.gpu.createBuffer(this.vocabSize * 4)
-            await this.gpu.WriteBuffer(this.bBuffer, this.bias)
-            this.paramForward = await this.gpu.createUniFormBuffer(4 * 4)
-            this.paramBackward = await this.gpu.createUniFormBuffer(16)
-        }
-        if (!this.oldParam) {
-            this.oldParam = new Uint32Array([this.embeddingSize, this.vocabSize, 0, 0])
-        }
-        if (this.oldParam[2] !== tokenCount) {
-            await this.gpu.WriteBuffer(this.paramForward, new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount, 0]))
-            this.oldParam = new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount, 0]);
-        }
-        if (!this.outputBuffer) this.outputBuffer = await this.gpu.ManageBuffer(4 * this.vocabSize * tokenCount)
-        if (this.outputBuffer.mainBuffer.size !== 4 * this.vocabSize * tokenCount || this.outputBuffer.isDestroyed()) {
-            this.outputBuffer.destroy()
-            this.outputBuffer = await this.gpu.ManageBuffer(4 * this.vocabSize * tokenCount)
-        }
-        const bind = await this.device.createBindGroup({
-            layout: await this.forwardPipline.getBindGroupLayout(0),
-            entries: [
-                {
-                    binding: 0,
-                    resource: this.wBuffer
-                },
-                {
-                    binding: 1,
-                    resource: this.outputBuffer.mainBuffer
-                },
-                {
-                    binding: 2,
-                    resource: vec.mainBuffer
-                },
-                {
-                    binding: 3,
-                    resource: this.bBuffer
-                },
-                {
-                    binding: 4,
-                    resource: this.paramForward
-                }
-            ]
-        })
-        await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil(this.vocabSize / 16), Math.ceil(tokenCount / 4))
-
-        if(!encoderMain) this.device.queue.submit([encoder.finish()])
-        this.input = vec;
-        return this.outputBuffer
-    }
-    async backward(dOutput, learningRate, encoderMain) {
-        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
-        const tokenCount = dOutput.size / (this.vocabSize * 4)
-        const data = new ArrayBuffer(16);
-        const view = new DataView(data);
-        view.setUint32(0, this.vocabSize, true);
-        view.setUint32(4, tokenCount, true);
-        view.setFloat32(8, learningRate, true);
-        view.setUint32(12, this.embeddingSize, true);
-        await this.gpu.WriteBuffer(this.paramBackward, data);
-
-
-        if (!this.gradientBuffer) this.gradientBuffer = await this.gpu.ManageBuffer(this.embeddingSize * tokenCount * 4);
-        if (this.gradientBuffer.size !== 4 * this.embeddingSize * tokenCount || this.gradientBuffer.isDestroyed()) {
-            this.gradientBuffer.destroy()
-            this.gradientBuffer = await this.gpu.ManageBuffer(this.embeddingSize * tokenCount * 4)
-
-        }
-        await this.gradientBuffer.count()
-        const bind = await this.device.createBindGroup({
-            layout: await this.backwardPipline.getBindGroupLayout(0),
-            entries: [
-                {
-                    binding: 0,
-                    resource: this.bBuffer
-                },
-                {
-                    binding: 1,
-                    resource: this.wBuffer
-                },
-                {
-                    binding: 2,
-                    resource: dOutput.mainBuffer
-                },
-                {
-                    binding: 3,
-                    resource: this.input.mainBuffer
-                },
-                {
-                    binding: 4,
-                    resource: this.paramBackward
-                }
-            ]
-        })
-        const bind2 = await this.device.createBindGroup({
-            layout: await this.dInputBackward.getBindGroupLayout(0),
-            entries: [
-                {
-                    binding: 0,
-                    resource: this.wBuffer
-                },
-                {
-                    binding: 1,
-                    resource: dOutput.mainBuffer
-                },
-                {
-                    binding: 3,
-                    resource: this.paramBackward
-                },
-                {
-                    binding: 4,
-                    resource: this.gradientBuffer.mainBuffer
-                }
-            ]
-        })
-        await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 64));
-        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, this.vocabSize);
-        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
-        return this.gradientBuffer
-
-    }
-    async Save() {
-        const b = await this.gpu.readBuffer(this.bBuffer.size, this.bBuffer);
-        const w = await this.gpu.readBuffer(this.wBuffer.size, this.wBuffer);
-        SaveVector(this.configs.save.filename[0], w)
-        SaveVector(this.configs.save.filename[1], b)
-    }
-
-
-}

package/src/OperationManager.js

@@ -1,169 +0,0 @@
-import fs from 'fs'
-import path, { join, dirname, format } from 'path';
-import { fileURLToPath } from 'url';
-
-const __dirname = dirname(fileURLToPath(import.meta.url))
-
-export class OperationManager {
-    device;
-    matrixMatSHhaders;
-    params;
-    addShader;
-    constructor(device) {
-        this.device = device;
-        this.addShader = this.getShaderPipline("./src/shaders/add.wgsl")
-    }
-    async createBuffer(size) {
-        return this.device.createBuffer({
-            size: size,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
-        })
-    }
-    async createUniFormBuffer(size) {
-        return this.device.createBuffer({
-            size,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.UNIFORM
-        })
-    }
-    async readBuffer(size, buffer) {
-        const readBuffer = this.device.createBuffer({
-            size: size,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
-        })
-        const encoder = this.device.createCommandEncoder()
-        encoder.copyBufferToBuffer(buffer, 0, readBuffer, 0, size);
-        this.device.queue.submit([await encoder.finish()])
-        await readBuffer.mapAsync(GPUMapMode.READ);
-        const mapped = readBuffer.getMappedRange();
-        const result = new Float32Array(
-            mapped.slice(0)
-        );
-        readBuffer.unmap()
-        readBuffer.destroy()
-        return result
-    }
-    getShaderPipline(filename, value = { "embedding": 123 }, entryPoint = "main") {
-        let readShaderCode = fs.readFileSync(path.join(__dirname, `../${filename}`), "utf-8");
-        if (value) {
-            for (const a in value) {
-                readShaderCode = readShaderCode.replaceAll(`%${a}%`, value[a])
-            }
-        };
-        const compiledKernel = this.device.createShaderModule({
-            code: readShaderCode
-        });
-        const compiledPipLine = this.device.createComputePipeline({
-            layout: "auto",
-            compute: {
-                module: compiledKernel,
-                entryPoint: entryPoint
-            }
-        });
-        return compiledPipLine;
-    }
-    async WriteBuffer(buffer, data) {
-        this.device.queue.writeBuffer(buffer, 0, data)
-    }
-    async RunPipline(encoder, pipline, binding, dispatchWorkgroups, y = 1, z = 1) {
-        const pass = encoder.beginComputePass();
-        pass.setPipeline(pipline)
-        pass.setBindGroup(0, binding)
-        pass.dispatchWorkgroups(dispatchWorkgroups, y, z)
-        pass.end()
-    }
-
-    async MakeACopyBuffer(encoder2, sourceBuffer) {
-        const encoder = encoder2 ? encoder2 : await this.CreateEncoder()
-        const buffer = this.device.createBuffer({
-            size: sourceBuffer.size,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
-        })
-        encoder.copyBufferToBuffer(sourceBuffer, 0, buffer, 0, sourceBuffer.size);
-        if (!encoder2) await this.submitQueue(encoder)
-        return buffer
-    }
-    CreateEncoder() {
-        return this.device.createCommandEncoder()
-    }
-    createBindGroup(shader, groupNum, ent) {
-
-        return this.device.createBindGroup({
-            layout: shader.getBindGroupLayout(groupNum),
-            entries: ent.map((e, x) => {
-                return {
-                    binding: x,
-                    resource: { buffer: e}
-                }
-            })
-        })
-    }
-    async submitQueue(encoder) {
-        this.device.queue.submit([encoder.finish()])
-    }
-    async add(a, b, size) {
-        if (!this.params) {
-            this.params = await this.createUniFormBuffer(4 * 4)
-
-        }
-        const encoder = await this.CreateEncoder()
-        const output = await this.ManageBuffer(a.size);
-        const bind = await this.createBindGroup(this.addShader, 0, [a, b, output.mainBuffer, this.params]);
-        await this.WriteBuffer(this.params, new Uint32Array([size, 0]))
-        await this.RunPipline(encoder, this.addShader, bind, Math.ceil(size / 64));
-        await this.submitQueue(encoder)
-        // a.destroy()
-        // b.destroy()
-        return output
-    }
-    ManageBuffer(size) {
-        let buffer = this.device.createBuffer({
-            size,
-            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
-        })
-        let isDestroy = false;
-        let ret = {
-            id: Math.floor(Math.random() * 9999),
-            mainBuffer: buffer,
-            reuse: 0,
-            destroy: function () {
-                isDestroy = true
-                this.mainBuffer.destroy()
-            },
-            isDestroyed: function () {
-                return isDestroy
-            },
-            read: async () => {
-                return await this.readBuffer(buffer.size, buffer)
-            },
-            write: async (data) => {
-                if (ret.reuse == 1000) {
-                    await ret.mainBuffer.destroy()
-                    ret.mainBuffer = this.device.createBuffer({
-                        size: ret.size,
-                        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
-                    })
-                    ret.reuse = 0;
-                }
-                await this.WriteBuffer(buffer, data)
-                ret.reuse++;
-
-            },
-            size: size,
-            count: async () => {
-                if (ret.reuse == 1000) {
-                    await ret.mainBuffer.destroy()
-                    ret.mainBuffer = this.device.createBuffer({
-                        size: ret.size,
-                        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
-                    })
-                    ret.reuse = 0;
-                }
-                ret.reuse++;
-
-            },
-
-        }
-        return ret
-    }
-
-}   

package/src/PositionEmbedding/index.d.ts

@@ -1,30 +0,0 @@
-export class PositionEmbedding {
-    constructor(device: GPUBuffer, embedding: number, contextSize: number, config: {
-        save: {
-            filename: 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 Save(): Promise<void>
-    async backward(inputGradient: {
-        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
-    }>
-}

package/src/PositionEmbedding/index.js

@@ -1,70 +0,0 @@
-import { OperationManager } from "../OperationManager.js"
-import { getVector, SaveVector } from "../getVec.js"
-function GenerateVectors(embeddingSize, contextSize) {
-    const vec = new Float32Array(embeddingSize * contextSize).map(() => (Math.random() * 2 - 1) * 0.02)
-    return vec
-}
-
-export class PositionEmbedding {
-    vectors;
-    vecBuffer;
-    gpu;
-    embeddingSize;
-    contextSize;
-    params;
-    forwardShader;
-    backwardShader;
-    configs;
-    constructor(device, embeddingSize, contextSize, configs) {
-        this.embeddingSize = embeddingSize;
-        const total = embeddingSize * contextSize;
-        this.configs = configs;
-        this.contextSize = contextSize;
-        if (configs && configs.save) {
-            const vec = getVector(configs.save.filename, total);
-            if (vec) {
-                this.vectors = vec;
-
-            } else {
-                this.vectors = GenerateVectors(embeddingSize, contextSize)
-                SaveVector(this.configs.save.filename, this.vectors)
-            }
-
-        } else {
-            this.vectors = GenerateVectors(embeddingSize, contextSize)
-
-        }
-        this.gpu = new OperationManager(device)
-        this.forwardShader = this.gpu.getShaderPipline("./src/shaders/PF.wgsl")
-        this.backwardShader = this.gpu.getShaderPipline("./src/shaders/PB.wgsl")
-    }
-    async forward(input, en2) {
-        const encoder =  en2 ? en2 : this.gpu.CreateEncoder()
-        if (!this.params) this.params = await this.gpu.createUniFormBuffer(4 * 4)
-        if (!this.vecBuffer) {
-            this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
-            await this.gpu.WriteBuffer(this.vecBuffer, this.vectors)
-        }
-        if(!this.output) this.output = await this.gpu.ManageBuffer(input.size);
-        if(this.output.size !== input.size || this.output.isDestroyed()){
-            this.output.destroy()
-            this.output = await this.gpu.ManageBuffer(input.size)
-        }
-        const size = Math.ceil((input.size / 4) / 64);
-        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input.mainBuffer, this.output.mainBuffer])
-        await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
-        if(!en2) await this.gpu.submitQueue(encoder);
-        return this.output
-    }
-    async backward(outputGradient, lr, en2){
-        const encoder = en2 ? en2 : await this.gpu.CreateEncoder()
-        await this.gpu.WriteBuffer(this.params, new Float32Array([lr,0]))
-        const bindGroup = await this.gpu.createBindGroup(this.backwardShader, 0, [this.vecBuffer, outputGradient.mainBuffer, this.params])
-        await this.gpu.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((outputGradient.size / 4) / 64))
-        if(!en2) await this.gpu.submitQueue(encoder);
-    }
-    async Save() {
-        const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize, this.vecBuffer)
-        SaveVector(this.configs.save.filename, vec)
-    }
-}

package/src/RMSNorm/index.d.ts

@@ -1,30 +0,0 @@
-export class RMSNorm {
-    constructor(device: GPUAdapter, embeddingSize: number, {
-        save: {
-            filename: 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(outputGradient: {
-        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 Save(): Promise<void>
-}

package/src/RMSNorm/index.js

@@ -1,91 +0,0 @@
-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;
-  rmsReduction;
-  output;
-  gradientOutput;
-  input;
-  device
-  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
-    })
-    this.device = device;
-    if (configs) {
-      const gammaFile = getVector(configs.save.filename, embeddingSize);
-      if (!gammaFile) {
-        this.gamma = new Float32Array(embeddingSize).fill(1);
-        SaveVector(configs.save.filename, this.gamma)
-      } else {
-        this.gamma = gammaFile
-      }
-    } else {
-      this.gamma = new Float32Array(embeddingSize).fill(1)
-    }
-  }
-  async forward(vec, en) {
-    if (!this.gammabuffer) {
-      this.gammabuffer = await this.op.createBuffer(this.embeddingSize * 4);
-      await this.op.WriteBuffer(this.gammabuffer, this.gamma);
-    }
-    const context = vec.size / (this.embeddingSize * 4)
-    if (!this.rmsReduction) this.rmsReduction = await this.op.createBuffer(context * 4);
-    if (!this.output) this.output = await this.op.ManageBuffer(vec.size)
-    if (this.rmsReduction.size !== context * 4) {
-      this.rmsReduction.destroy()
-      this.rmsReduction = await this.op.createBuffer(4 * context)
-    }
-    if (this.output) this.output.destroy();
-    this.output = await this.op.ManageBuffer(vec.size)
-    const encoder = en ? en : await this.op.CreateEncoder()
-    const bind1 = await this.op.createBindGroup(this.reductionShader, 0, [vec.mainBuffer, this.rmsReduction])
-    const bind2 = await this.op.createBindGroup(this.forwardShader, 0, [this.rmsReduction, vec.mainBuffer, this.output.mainBuffer, this.gammabuffer])
-    await this.op.RunPipline(encoder, this.reductionShader, bind1, context)
-    await this.op.RunPipline(encoder, this.forwardShader, bind2, Math.ceil((this.embeddingSize * context) / 64))
-    if(!en) await this.op.submitQueue(encoder);
-    this.input = vec;
-    return this.output
-
-  }
-  async backward(outGradient, lr, en) {
-    const contextSize = outGradient.size / (this.embeddingSize * 4);
-    if (this.gradientOutput) this.gradientOutput.destroy()
-    this.gradientOutput = await this.op.ManageBuffer(this.embeddingSize * 4 * contextSize)
-    if (!this.params) this.params = await this.op.createUniFormBuffer(4 * 2);
-    await this.op.WriteBuffer(this.params, new Float32Array([lr, 0]))
-    const encoder = en ? en : this.op.CreateEncoder()
-    const bind1 = await this.op.createBindGroup(this.backwardShader, 0, [this.input.mainBuffer, outGradient.mainBuffer, this.rmsReduction, this.gammabuffer, this.gradientOutput.mainBuffer])
-    await this.op.RunPipline(encoder, this.backwardShader, bind1, contextSize);
-    const bind2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input.mainBuffer, outGradient.mainBuffer, this.rmsReduction, this.gammabuffer, this.params])
-    await this.op.RunPipline(encoder, this.updateGamma, bind2, this.embeddingSize);
-    if(!en) this.op.submitQueue(encoder);
-    return this.gradientOutput
-  }
-  async Save() {
-    const gpuGamma = await this.op.readBuffer(4 * this.embeddingSize, this.gammabuffer);
-    SaveVector(this.configs.save.filename, gpuGamma)
-  }
-
-}

package/src/Relu/index.d.ts

@@ -1,25 +0,0 @@
-export class Relu {
-    constructor(device: GPUAdapter, embeddingSize: number);
-    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(outGradinet: {
-        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
-    }>
-}