@imaduddin/fastbrain 0.0.0-stage → 1.0.1

package/src/Embedding/index.d.ts

@@ -0,0 +1,21 @@
+export class Embedding {
+    constructor(device: GPUBuffer, embedding: number, vocabSize: number, config: {
+        cpuReadBack: boolean,
+        save: {
+            filename: string
+        }
+    })
+    async forward(tokens: number[]): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async Save(): void
+    async backward(inputGradient: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, lr: number): Promise<void>
+}

package/src/Embedding/index.js

@@ -0,0 +1,140 @@
+import { config } from "process";
+
+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
+}
+
+export class Embedding {
+    embeddingSize;
+    vocabSize;
+    gpu;
+    firstTrue;
+    vectors;
+    configs;
+    vectorsBuffer;
+    forwardPipline;
+    constructor(device, embeddingSize, vocabSize, configs) {
+        this.configs = configs
+        this.embeddingSize = embeddingSize;
+        this.vocabSize = vocabSize;
+        this.device = device
+        let total = embeddingSize * vocabSize;
+        if (configs && configs.save) {
+            const vec = getVector(configs.save.filename, total);
+            if (vec) {
+                this.vectors = vec;
+
+            } else {
+                this.vectors = GenerateVectors(embeddingSize, vocabSize)
+            }
+
+        } else {
+            this.vectors = GenerateVectors(embeddingSize, vocabSize)
+
+        }
+        this.gpu = new OperationManager(device)
+        this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/EF.wgsl", {
+            embeddingSize: embeddingSize
+        })
+        this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/EB.wgsl", {
+            embeddingSize: embeddingSize
+        })
+        SaveVector(this.configs.save.filename, this.vectors)
+    }
+    async forward(input, encoderMain) {
+        const encoder = encoderMain ? encoderMain :  this.device.createCommandEncoder()
+        if (!this.vectorsBuffer) {
+            this.vectorsBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize);
+            await this.gpu.WriteBuffer(this.vectorsBuffer, this.vectors)
+        }
+        if (!this.tokenBuffer) this.tokenBuffer = await this.gpu.createBuffer(Uint32Array.BYTES_PER_ELEMENT * input.length);
+        if (this.tokenBuffer.size !== Uint32Array.BYTES_PER_ELEMENT * input.length) {
+            this.tokenBuffer.destroy()
+            this.tokenBuffer = await this.gpu.createBuffer(Uint32Array.BYTES_PER_ELEMENT * input.length);
+        }
+        
+        await this.gpu.WriteBuffer(this.tokenBuffer, new Uint32Array(input))
+        if (!this.outputBuffer) this.outputBuffer = await this.gpu.ManageBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length)
+            if (this.outputBuffer.mainBuffer.size !== Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length || this.outputBuffer.isDestroyed()) {
+            this.outputBuffer.destroy()
+            this.outputBuffer = await this.gpu.ManageBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length);
+        }
+        const bind = await this.device.createBindGroup({
+            layout: await this.forwardPipline.getBindGroupLayout(0),
+            entries: [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: this.vectorsBuffer
+                    }
+                },
+                {
+                    binding: 1,
+                    resource: {
+                        buffer: this.outputBuffer.mainBuffer
+                    }
+                }, {
+                    binding: 2,
+                    resource: {
+                        buffer: this.tokenBuffer
+                    }
+                }
+            ]
+        })
+        await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
+        if(!encoderMain) await this.device.queue.submit([encoder.finish()])
+        return this.outputBuffer
+
+    }
+    async Save() {
+        const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize, this.vectorsBuffer)
+        SaveVector(this.configs.save.filename, vec)
+    }
+    async backward(inputGradient, lr, encoderMain) {
+        const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder()
+        if (!this.oldParam) {
+            this.oldParam = new Float32Array([0, 0]);
+            this.param = await this.gpu.createUniFormBuffer(this.oldParam.byteLength);
+        }
+        if (this.oldParam[0] !== lr) {
+            this.oldParam = new Float32Array([lr, 0]);
+            this.gpu.WriteBuffer(this.param, this.oldParam)
+        }
+        const bind = await this.device.createBindGroup({
+            layout: await this.backwardPipline.getBindGroupLayout(0),
+            entries: [
+                {
+                    binding: 0,
+                    resource: {
+                        buffer: this.vectorsBuffer
+                    }
+                },
+                {
+                    binding: 1,
+                    resource: {
+                        buffer: inputGradient.mainBuffer
+                    }
+                }, {
+                    binding: 2,
+                    resource: {
+                        buffer: this.tokenBuffer
+                    }
+
+                }, {
+                    binding: 3,
+                    resource: {
+                        buffer: this.param
+                    }
+
+                }
+            ]
+        })
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * (this.tokenBuffer.size / 4)) / 256))
+        if(!encoderMain) this.device.queue.submit([encoder.finish()])
+    }
+
+}

package/src/FFN/index.d.ts

@@ -0,0 +1,23 @@
+export class FFN {
+    constructor(device: GPUAdapter, embeddingSize: number, configs: {
+        dirname: string
+    })
+    async forward(vec: GPUBuffer): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): 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
+    }>
+    async Save(): Promise<void>
+}

package/src/FFN/index.js

@@ -0,0 +1,45 @@
+import { RELU, Linear } from "../../index.js";
+import { join } from 'path'
+import { OperationManager } from "../OperationManager.js";
+
+
+export class FFN {
+    linear1;
+    linear2;
+    hiddenLayer;
+    Op
+    constructor(device, embeddingSize, configs) {
+        this.Op = new OperationManager(device)
+        this.linear1 = new Linear(device, embeddingSize, embeddingSize * 4, {
+            save: {
+                filename: [join(configs.dirname, "FFN_weights1.bin"), join(configs.dirname, "FFN_bias1.bin")]
+            }
+        });
+        this.hiddenLayer = new RELU(device, embeddingSize * 4);
+        this.linear2 = new Linear(device, embeddingSize * 4, embeddingSize, {
+            save: {
+                filename: [join(configs.dirname, "FFN_weights2.bin"), join(configs.dirname, "FFN_bias2.bin")]
+            }
+        })
+    }
+    async forward(vec, encoderMain) {
+        const logit1 = await this.linear1.forward(vec, encoderMain)
+        const hidden = await this.hiddenLayer.forward(logit1, encoderMain)
+        const final = await this.linear2.forward(hidden, encoderMain)
+        return final
+    }
+    async backward(inputGradient, lr, encoderMain) {
+        const logit2 = await this.linear2.backward(inputGradient, lr, encoderMain)
+        const hidden = await this.hiddenLayer.backward(logit2, encoderMain)
+        const final = await this.linear1.backward(hidden, lr, encoderMain)
+        return final
+    }
+    async Save() {
+        await this.linear1.Save()
+        await this.linear2.Save()
+    }
+    async ClearInputCache() {
+        this.linear1.ClearInputCache()
+        this.linear2.ClearInputCache()
+    }
+}

package/src/Linear/index.d.ts

@@ -0,0 +1,31 @@
+export class Linear {
+    constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number, configs: {
+        save: {
+            filename: string[]
+        }
+    })
+    async forward(vec: {
+        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(dOutput: {
+        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/Linear/index.js

@@ -0,0 +1,197 @@
+
+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
+                }
+            ]
+        })
+        // console.log((await this.gpu.readBuffer(this.input.size, this.input))[0])
+        await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 256));
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil(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

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

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