@mineproness/tensorloop 2.1.81 → 2.2.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mineproness/tensorloop",
-  "version": "2.1.81",
+  "version": "2.2.1",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called Tensorloop.",
   "keywords": [
     "ai",

package/src/Embedding/index.d.ts

@@ -1,11 +1,21 @@
-export class Embedding{
-    constructor(device: GPUBuffer ,embedding : number, vocabSize: number, config:{
+export class Embedding {
+    constructor(device: GPUBuffer, embedding: number, vocabSize: number, config: {
         cpuReadBack: boolean,
         save: {
             filename: string
         }
     })
-    async forward(tokens: number[]) : Promise<GPUBuffer>
-    async Save() : void
-    async backward(tokens: number[], inputGradient: GPUBuffer, lr: number) : Promise<void>
+    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

@@ -51,9 +51,18 @@ export class Embedding {
             this.vectorsBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize);
             await this.gpu.WriteBuffer(this.vectorsBuffer, this.vectors)
         }
-        const tokenBuffer = await this.gpu.createBuffer(Uint32Array.BYTES_PER_ELEMENT * input.length);
-        await this.gpu.WriteBuffer(tokenBuffer, new Uint32Array(input))
-        const outputBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length);
+        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: [
@@ -66,37 +75,30 @@ export class Embedding {
                 {
                     binding: 1,
                     resource: {
-                        buffer: outputBuffer
+                        buffer: this.outputBuffer.mainBuffer
                     }
                 }, {
                     binding: 2,
                     resource: {
-                        buffer: tokenBuffer
+                        buffer: this.tokenBuffer
                     }
                 }
             ]
         })
-        await this.gpu.RunPipline(encoder,this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
+        await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
         this.device.queue.submit([encoder.finish()])
-        tokenBuffer.destroy()
-        if (this.configs.cpuReadBack) {
-            const output = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length, outputBuffer);
-            outputBuffer.destroy()
-            return output
-        }
-        return outputBuffer
+        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(token, inputGradient, lr) {
+    async backward(inputGradient, lr) {
         const encoder = this.device.createCommandEncoder()
         if (!this.oldParam) {
             this.oldParam = new Float32Array([0, 0]);
             this.param = await this.gpu.createUniFormBuffer(this.oldParam.byteLength);
-            this.tokenBuffer = await this.gpu.createBuffer(4 * token.length)
         }
         if (this.oldParam[0] !== lr) {
             this.oldParam = new Float32Array([lr, 0]);
@@ -114,7 +116,7 @@ export class Embedding {
                 {
                     binding: 1,
                     resource: {
-                        buffer: inputGradient
+                        buffer: inputGradient.mainBuffer
                     }
                 }, {
                     binding: 2,
@@ -131,10 +133,18 @@ export class Embedding {
                 }
             ]
         })
-        await this.gpu.WriteBuffer(this.tokenBuffer, new Uint32Array(token));
-        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * token.length) / 256))
+        await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * (this.tokenBuffer.size / 4)) / 256))
         this.device.queue.submit([encoder.finish()])
-        if(this.prevGra) this.prevGra.destroy()
-        this.prevGra = inputGradient;
+    }
+    destroy() {
+        this.vectorsBuffer?.destroy();
+        this.tokenBuffer?.destroy();
+        this.outputBuffer?.destroy();
+        this.param?.destroy();
+
+        this.vectorsBuffer = null;
+        this.tokenBuffer = null;
+        this.outputBuffer = null;
+        this.param = null;
     }
 }

package/src/FFN/index.js

@@ -29,8 +29,6 @@ export class FFN {
         return final
     }
     async backward(inputGradient, lr) {
-        if(this.prevOUTGR)  await this.prevOUTGR.destroy()
-        this.prevOUTGR = inputGradient;
         const logit2 = await this.linear2.backward(inputGradient, lr)
         const hidden = await this.hiddenLayer.backward(logit2)
         const final = await this.linear1.backward(hidden, lr)

package/src/Linear/index.d.ts

@@ -1,12 +1,31 @@
 export class Linear {
     constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number, configs: {
-        cpuReadBack: boolean,
         save: {
             filename: string[]
         }
     })
-    async forward(vec: GPUBuffer): Promise<GPUBuffer | Float32Array>
-    async backward(dOutput: GPUBuffer, lr: number): Promise<GPUBuffer>
+    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>
-    async ClearInputCache() : Promise<void>
+    async ClearInputCache(): Promise<void>
 }

package/src/Linear/index.js

@@ -63,7 +63,7 @@ export class Linear {
     }
     async forward(vec) {
         const encoder = this.device.createCommandEncoder()
-        const tokenCount = vec.size / (this.embeddingSize * 4)
+        const tokenCount = vec.mainBuffer.size / (this.embeddingSize * 4)
         if (!this.wBuffer) {
 
             this.wBuffer = await this.gpu.createBuffer(this.embeddingSize * this.vocabSize * 4)
@@ -80,7 +80,12 @@ export class Linear {
             await this.gpu.WriteBuffer(this.paramForward, new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount, 0]))
             this.oldParam = new Uint32Array([this.embeddingSize, this.vocabSize, tokenCount, 0]);
         }
-        const outputBuffer = await this.gpu.createBuffer(4 * this.vocabSize * tokenCount)
+        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: [
@@ -90,11 +95,11 @@ export class Linear {
                 },
                 {
                     binding: 1,
-                    resource: outputBuffer
+                    resource: this.outputBuffer.mainBuffer
                 },
                 {
                     binding: 2,
-                    resource: vec
+                    resource: vec.mainBuffer
                 },
                 {
                     binding: 3,
@@ -107,14 +112,12 @@ export class Linear {
             ]
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.vocabSize * tokenCount) / 256))
-        if (this.input) this.input.destroy()
-        this.input = await this.gpu.MakeACopyBuffer(encoder, vec);
+        
         this.device.queue.submit([encoder.finish()])
-        return outputBuffer
+        this.input = vec;
+        return this.outputBuffer
     }
     async backward(dOutput, learningRate) {
-        if(this.pdOutput) await this.pdOutput.destroy();
-         this.pdOutput = dOutput
         const encoder = this.device.createCommandEncoder()
         const tokenCount = dOutput.size / (this.vocabSize * 4)
         if (!this.oldParamb) {
@@ -124,7 +127,12 @@ export class Linear {
             await this.gpu.WriteBuffer(this.paramBackward, new Float32Array([this.vocabSize, tokenCount, learningRate, this.embeddingSize, 0]));
             this.oldParamb = new Float32Array([this.vocabSize, tokenCount, learningRate, this.embeddingSize, 0]);
         }
-        const outputBuffer = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4);
+        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)
+
+        }
         const bind = await this.device.createBindGroup({
             layout: await this.backwardPipline.getBindGroupLayout(0),
             entries: [
@@ -138,11 +146,11 @@ export class Linear {
                 },
                 {
                     binding: 2,
-                    resource: dOutput
+                    resource: dOutput.mainBuffer
                 },
                 {
                     binding: 3,
-                    resource: this.input
+                    resource: this.input.mainBuffer
                 },
                 {
                     binding: 4,
@@ -159,7 +167,7 @@ export class Linear {
                 },
                 {
                     binding: 1,
-                    resource: dOutput
+                    resource: dOutput.mainBuffer
                 },
                 {
                     binding: 3,
@@ -167,7 +175,7 @@ export class Linear {
                 },
                 {
                     binding: 4,
-                    resource: outputBuffer
+                    resource: this.gradientBuffer.mainBuffer
                 }
             ]
         })
@@ -175,8 +183,7 @@ export class Linear {
         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 * tokenCount) / 256));
         await this.device.queue.submit([encoder.finish()])
-        this.input.destroy()
-        return outputBuffer
+        return this.gradientBuffer
 
     }
     async Save() {
@@ -186,7 +193,7 @@ export class Linear {
         SaveVector(this.configs.save.filename[1], b)
     }
     async ClearInputCache() {
-        if (this.input) await this.input.destroy();
+        
     }
 
 }

package/src/OperationManager.js

@@ -15,7 +15,7 @@ export class OperationManager {
     }
     async createBuffer(size) {
         return this.device.createBuffer({
-            size,
+            size: size,
             usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
         })
     }
@@ -106,8 +106,8 @@ export class OperationManager {
 
         }
         const encoder = await this.CreateEncoder()
-        const output = await this.createBuffer(a.size);
-        const bind = await this.createBindGroup(this.addShader, 0, [a, b, output, this.params]);
+        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)
@@ -115,5 +115,27 @@ export class OperationManager {
         // b.destroy()
         return output
     }
+    async ManageBuffer(size){
+        const buffer = await this.device.createBuffer({
+            size,
+            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
+        })
+        let isDestroy = false;
+        return {
+            mainBuffer: buffer,
+            destroy: function(){
+                isDestroy = true
+                this.mainBuffer.destroy()
+            },
+            isDestroyed: function(){
+                return isDestroy
+            },
+            read: async ()=>{
+                return await  this.readBuffer(buffer.size, buffer)
+            },
+            size: size
+
+        }
+    }
 
 }   

package/src/PositionEmbedding/index.d.ts

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

@@ -45,21 +45,16 @@ export class PositionEmbedding {
             this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
             await this.gpu.WriteBuffer(this.vecBuffer, this.vectors)
         }
-        const output = await this.gpu.createBuffer(input.size);
         const size = Math.ceil((input.size / 4) / 64);
-        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input, output])
-        if (this.prevVec) this.prevVec.destroy();
-        this.prevVec = input;
+        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input.mainBuffer])
         await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
         await this.gpu.submitQueue(encoder);
-        return output
+        return input
     }
     async backward(outputGradient, lr){
         const encoder = 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, this.params])
-        if(this.prevGradient) await this.prevGradient.destroy()
-        this.prevGradient = outputGradient;
+        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))
         await this.gpu.submitQueue(encoder);
     }

package/src/RMSNorm/index.d.ts

@@ -1,10 +1,30 @@
-export class RMSNorm{
-    constructor(device : GPUAdapter, embeddingSize : number , {
-        save:{
+export class RMSNorm {
+    constructor(device: GPUAdapter, embeddingSize: number, {
+        save: {
             filename: string
         }
     });
-    async forward(input : GPUBuffer) : Promise<GPUBuffer>
-    async backward(outputGradient : GPUBuffer) : Promise<GPUBuffer>
-    async Save() : Promise<void>
+    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

@@ -40,43 +40,47 @@ export class RMSNorm {
     }
   }
   async forward(input) {
+    this.input = input;
     if (!this.gammabuffer) {
       this.gammabuffer = await this.op.createBuffer(4 * this.embeddingSize);
       await this.op.WriteBuffer(this.gammabuffer, this.gamma);
     }
     let context = input.size / (this.embeddingSize * 4)
-    const outputReduce = await this.op.createBuffer(4 * context);
-    const output = await this.op.createBuffer(input.size);
-    const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input, outputReduce])
+    if (!this.rms) this.rms = await this.op.createBuffer(4 * context);
+    if (this.rms.size !== 4 * context * this.embeddingSize) {
+      this.rms.destroy()
+      this.rms = await this.op.createBuffer(4 * context);
+    }
+    if (!this.output) this.output = await this.op.ManageBuffer(input.size);
+    if (this.output.size !== 4 * context || this.output.isDestroyed()) {
+      this.output.destroy()
+      this.output = await this.op.ManageBuffer(input.size)
+    }
+    const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input.mainBuffer, this.rms])
     const encoder = await this.op.CreateEncoder()
     await this.op.RunPipline(encoder, this.reductionShader, bindGroup, context)
-    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [outputReduce, input, output, this.gammabuffer])
+    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [this.rms, input.mainBuffer, this.output.mainBuffer, this.gammabuffer])
     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
+    return this.output
   }
   async backward(outputGradient, lr) {
-    const output = await this.op.createBuffer(outputGradient.size);
-    let params = await this.op.createBuffer(4 * 2)
-    await this.op.WriteBuffer(params, new Float32Array([lr, 0]))
-    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input, outputGradient, this.rms, this.gammabuffer, output]);
-    const bindGroup2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input, outputGradient, this.rms, this.gammabuffer, params]);
+    if(!this.gradOut) this.gradOut = await this.op.ManageBuffer(outputGradient.size);
+    if(this.gradOut.size !== outputGradient.size || this.gradOut.isDestroyed()){
+      this.gradOut.destroy()
+      this.gradOut = await this.op.ManageBuffer(outputGradient.size)
+    }
+    if(!this.params) this.params = await this.op.createBuffer(4 * 2)
+    await this.op.WriteBuffer(this.params, new Float32Array([lr, 0]))
+    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input.mainBuffer, outputGradient.mainBuffer, this.rms, this.gammabuffer, this.gradOut.mainBuffer]);
+    const bindGroup2 = await this.op.createBindGroup(this.updateGamma, 0, [this.input.mainBuffer, outputGradient.mainBuffer, this.rms, this.gammabuffer, this.params]);
     const encoder = await this.op.CreateEncoder()
     await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size) / (this.embeddingSize * 4)))
     await this.op.RunPipline(encoder, this.updateGamma, bindGroup2, Math.ceil(this.embeddingSize))
     await this.op.submitQueue(encoder)
-    if (this.prevOut) this.prevOut.destroy()
-    this.prevOut = outputGradient
-    params.destroy()
-    return output
+    return this.gradOut
   }
-  async Save(){
+  async Save() {
     const read = await this.op.readBuffer(4 * this.embeddingSize, this.gammabuffer);
     SaveVector(this.configs.save.filename, read)
   }

package/src/Relu/index.d.ts

@@ -1,5 +1,25 @@
-export class Relu{
-    constructor(device : GPUAdapter, embeddingSize : number);
-    async forward(input: GPUBuffer) : Promise<GPUBuffer>
-    async backward(outGradinet: GPUBuffer) : Promise<GPUBuffer>
+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
+    }>
 }

package/src/Relu/index.js

@@ -17,12 +17,16 @@ export class Relu {
     }
     async forward(input) {
         const encoder = this.device.createCommandEncoder()
-        const tokenCount = input.size / (this.embeddingSize * 4)
+        const tokenCount = input.mainBuffer.size / (this.embeddingSize * 4)
         if (!this.oldParam) {
             this.oldParam = new Uint32Array([0, 0, 0]);
             this.paramBuffer = await this.gpu.createUniFormBuffer(4 * 3);
         }
-        const output = await this.gpu.createBuffer(input.size)
+        if (!this.output) this.output = await this.gpu.ManageBuffer(input.mainBuffer.size);
+        if (this.output.mainBuffer.size !== input.mainBuffer.size || this.output.isDestroyed()) {
+            this.output.destroy()
+            this.output = await this.gpu.ManageBuffer(input.mainBuffer.size);
+        }
         if (this.oldParam[1] !== tokenCount) {
             this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
             await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
@@ -33,7 +37,7 @@ export class Relu {
                 {
                     binding: 0,
                     resource: {
-                        buffer: input
+                        buffer: input.mainBuffer
                     }
                 },
                 {
@@ -44,16 +48,15 @@ export class Relu {
                 }, {
                     binding: 2,
                     resource: {
-                        buffer: output
+                        buffer: this.output.mainBuffer
                     }
                 }
             ]
         })
         await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
         this.device.queue.submit([encoder.finish()])
-        if(this.input) this.input.destroy()
         this.input = input
-        return output
+        return this.output
     }
     async backward(inputGradient) {
         const encoder = this.device.createCommandEncoder()
@@ -62,7 +65,11 @@ export class Relu {
             this.oldParam = new Uint32Array([0, 0, 0]);
             this.paramBuffer = await this.gpu.createUniFormBuffer(4 * 3);
         }
-        const output = await this.gpu.createBuffer(inputGradient.size)
+        if(!this.gradOut) this.gradOut = await this.gpu.ManageBuffer(inputGradient.size);
+        if(this.gradOut.size !== inputGradient.size || this.gradOut.isDestroyed()){
+            this.gradOut.destroy();
+            this.gradOut = await this.gpu.ManageBuffer(inputGradient.size)
+        }
         if (this.oldParam[1] !== tokenCount) {
             this.oldParam = new Uint32Array([this.embeddingSize, tokenCount, 0])
             await this.gpu.WriteBuffer(this.paramBuffer, this.oldParam)
@@ -73,7 +80,7 @@ export class Relu {
                 {
                     binding: 0,
                     resource: {
-                        buffer: this.input
+                        buffer: this.input.mainBuffer
                     }
                 },
                 {
@@ -84,20 +91,18 @@ export class Relu {
                 }, {
                     binding: 2,
                     resource: {
-                        buffer: output
+                        buffer: this.gradOut.mainBuffer
                     }
                 }, {
                     binding: 3,
                     resource: {
-                        buffer: inputGradient
+                        buffer: inputGradient.mainBuffer
                     }
                 }
             ]
         })
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((tokenCount * this.embeddingSize) / 256))
         await this.device.queue.submit([encoder.finish()])
-        this.input.destroy()
-        inputGradient.destroy()
-        return output
+        return this.gradOut
     }
 }

package/src/SelfAttention/index.d.ts

@@ -1,9 +1,14 @@
-export class SelfAttention{
-    constructor(device : GPUAdapter,embeddingSize: number, configs : {
+export class SelfAttention {
+    constructor(device: GPUAdapter, embeddingSize: number, configs: {
         dirname: string
     })
-    async forward(vec : GPUBuffer) : Promise<GPUBuffer>
-    async backward(dO : GPUBuffer, lr : number) : Promise<GPUBuffer>
-    async ClearInputCache() : Promise<void>
-    async Save() : Promise<void>
+    async forward(vec: GPUBuffer): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async backward(dO: GPUBuffer, lr: number): Promise<GPUBuffer>
+    async ClearInputCache(): Promise<void>
+    async Save(): Promise<void>
 } 

package/src/SelfAttention/index.js

@@ -39,82 +39,102 @@ export class SelfAttention {
         this.sumShader = this.gpu.getShaderPipline("./src/shaders/attentionShader/sumGradient.wgsl")
     }
     async forward(vec) {
-        if (this.input) await this.input.destroy()
-        this.input = vec;
-        const tokenCount = vec.size / (this.embeddingSize * 4)
+        this.input = vec.mainBuffer;
+        const tokenCount = vec.mainBuffer.size / (this.embeddingSize * 4)
         if (!this.asParams) this.asParams = await this.gpu.createUniFormBuffer(4 * 3);
         if (!this.softmaxParams) this.softmaxParams = await this.gpu.createUniFormBuffer(4 * 3);
         const encoder = await this.gpu.CreateEncoder()
-        const scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
+        if (!this.scroes) {
+            this.scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
+        }
+        if (this.scroes.size !== tokenCount * tokenCount * 4) {
+            this.scroes.destroy()
+            this.scroes = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
+        }
         await this.gpu.WriteBuffer(this.asParams, new Uint32Array([tokenCount, this.embeddingSize, 0]))
         await this.gpu.WriteBuffer(this.softmaxParams, new Uint32Array([tokenCount, tokenCount, 0]))
-        if (this.query) this.query.destroy()
-        if (this.key) this.key.destroy()
-        if (this.value) this.value.destroy()
         this.query = await this.q.forward(vec)
         this.key = await this.k.forward(vec)
         this.value = await this.v.forward(vec)
-        const bindAttention = await this.gpu.createBindGroup(this.attentionscoresShader, 0, [this.query, this.key, scroes, this.asParams])
+        const bindAttention = await this.gpu.createBindGroup(this.attentionscoresShader, 0, [this.query.mainBuffer, this.key.mainBuffer, this.scroes, this.asParams])
         await this.gpu.RunPipline(encoder, this.attentionscoresShader, bindAttention, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
         if (this.weights) this.weights.destroy()
         this.weights = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
-        const bindSF = await this.gpu.createBindGroup(this.softmaxForward, 0, [scroes, this.weights, this.softmaxParams]);
+        const bindSF = await this.gpu.createBindGroup(this.softmaxForward, 0, [this.scroes, this.weights, this.softmaxParams]);
         await this.gpu.RunPipline(encoder, this.softmaxForward, bindSF, tokenCount)
-        const finalOutput = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4)
-        const outputBind = await this.gpu.createBindGroup(this.outputForward, 0, [this.weights, this.value, finalOutput, this.asParams]);
+        if (!this.finalOutput) {
+            this.finalOutput = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
+        }
+        if (this.finalOutput.mainBuffer.size !== tokenCount * this.embeddingSize * 4 || this.finalOutput.isDestroyed()) {
+            this.finalOutput = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
+        }
+        const outputBind = await this.gpu.createBindGroup(this.outputForward, 0, [this.weights, this.value.mainBuffer, this.finalOutput.mainBuffer, this.asParams]);
         await this.gpu.RunPipline(encoder, this.outputForward, outputBind, Math.ceil((tokenCount * this.embeddingSize) / 256))
         await this.gpu.submitQueue(encoder)
-        scroes.destroy()
-        return finalOutput
+        return this.finalOutput
     }
     async backward(dO, lr) {
         const tokenCount = dO.size / (this.embeddingSize * 4);
         const encoder = await this.gpu.CreateEncoder();
-        const dV = await this.gpu.createBuffer(dO.size);
-        const DvBind = await this.gpu.createBindGroup(this.dVShader, 0, [this.weights, dO, dV, this.asParams]);
+        if (!this.dV) this.dV = await this.gpu.createBuffer(dO.size);
+        if (!this.dA) this.dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        if (!this.dS) this.dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        if (!this.dQ) this.dQ = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
+        if (!this.dK) this.dK = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4)
+        if (this.dV.size !== dO.size) {
+            this.dV.destroy()
+            this.dV = await this.gpu.createBuffer(dO.size);
+        }
+        if (this.dA.size !== tokenCount * tokenCount * 4) {
+            this.dA.destroy()
+            this.dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        }
+        if (this.dS.size !== tokenCount * tokenCount * 4) {
+            this.dS.destroy()
+            this.dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
+        }
+        if (this.dQ.size !== tokenCount * this.embeddingSize * 4) {
+            this.dQ.destroy()
+            this.dQ = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
+        }
+        if (this.dK.size !== tokenCount * this.embeddingSize * 4) {
+            this.dK.destroy()
+            this.dK = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
+        }
+        const DvBind = await this.gpu.createBindGroup(this.dVShader, 0, [this.weights, dO.mainBuffer, this.dV, this.asParams]);
         await this.gpu.RunPipline(encoder, this.dVShader, DvBind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16));
-        const dA = await this.gpu.createBuffer(tokenCount * tokenCount * 4)
-        const DaBind = await this.gpu.createBindGroup(this.dAshader, 0, [dO, this.value, dA, this.asParams]);
+        const DaBind = await this.gpu.createBindGroup(this.dAshader, 0, [dO.mainBuffer, this.value.mainBuffer, this.dA, this.asParams]);
         await this.gpu.RunPipline(encoder, this.dAshader, DaBind, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16))
-        const dS = await this.gpu.createBuffer(tokenCount * tokenCount * 4);
-        const dSBind = await this.gpu.createBindGroup(this.softmaxBackward, 0, [this.weights, dA, dS, this.asParams])
+        const dSBind = await this.gpu.createBindGroup(this.softmaxBackward, 0, [this.weights, this.dA, this.dS, this.asParams])
         await this.gpu.RunPipline(encoder, this.softmaxBackward, dSBind, Math.ceil((tokenCount * tokenCount) / 256))
-        const dQ = await this.gpu.createBuffer(tokenCount * this.embeddingSize * 4);
-        const dQbind = await this.gpu.createBindGroup(this.dQshader, 0, [dS, this.key, dQ, this.asParams]);
+        const dQbind = await this.gpu.createBindGroup(this.dQshader, 0, [this.dS, this.key.mainBuffer, this.dQ, this.asParams]);
         await this.gpu.RunPipline(encoder, this.dQshader, dQbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
-        const dK = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4)
-        const dKbind = await this.gpu.createBindGroup(this.dKShader, 0, [dS, this.query, dK, this.asParams]);
+        const dKbind = await this.gpu.createBindGroup(this.dKShader, 0, [this.dS, this.query.mainBuffer, this.dK, this.asParams]);
         await this.gpu.RunPipline(encoder, this.dKShader, dKbind, Math.ceil(tokenCount / 16), Math.ceil(this.embeddingSize / 16))
         await this.gpu.submitQueue(encoder);
-        const dq = await this.q.backward(dQ, lr)
-        const dk = await this.k.backward(dK, lr)
-        const dv = await this.v.backward(dV, lr)
-        const output = await this.gpu.createBuffer(this.embeddingSize * tokenCount * 4);
+        const dq = await this.q.backward({ mainBuffer: this.dQ, size: this.dQ.size}, lr)
+        const dk = await this.k.backward({ mainBuffer: this.dK, size: this.dK.size}, lr)
+        const dv = await this.v.backward({ mainBuffer: this.dV, size: this.dV.size}, lr)
+        if(!this.gradOut) this.gradOut = await this.gpu.ManageBuffer(tokenCount * this.embeddingSize * 4)
+        if(this.gradOut.size !== tokenCount * 4 * this.embeddingSize || this.gradOut.isDestroyed()){
+            this.gradOut.destroy()
+            this.gradOut = this.gpu.ManageBuffer(tokenCount * 4 * this.embeddingSize)
+        }
         const encode = this.gpu.CreateEncoder()
-        const outputBind = await this.gpu.createBindGroup(this.sumShader, 0 , [dq, dk, dv, output, this.asParams]);
+        const outputBind = await this.gpu.createBindGroup(this.sumShader, 0, [dq.mainBuffer, dk.mainBuffer, dv.mainBuffer, this.gradOut.mainBuffer, this.asParams]);
         await this.gpu.RunPipline(encode, this.sumShader, outputBind, Math.ceil((tokenCount * this.embeddingSize / 256)));
         await this.gpu.submitQueue(encode)
-        dS.destroy()
-        dA.destroy()
-        dq.destroy()
-        dk.destroy()
-        dv.destroy()
-        dK.destroy()
-        dQ.destroy()
-        dV.destroy()
-        if(this.prevdO) await this.prevdO.destroy()
-        this.prevdO = dO
-        return output
-        
+        return this.gradOut
+
     }
-    async ClearInputCache(){
-       await this.q.ClearInputCache()
-       await this.k.ClearInputCache()
-       await this.v.ClearInputCache()
+    async ClearInputCache() {
+        await this.q.ClearInputCache()
+        await this.k.ClearInputCache()
+        await this.v.ClearInputCache()
     }
-    async Save(){
-       await this.q.Save()
-       await this.k.Save()
-       await this.v.Save()
+    async Save() {
+        await this.q.Save()
+        await this.k.Save()
+        await this.v.Save()
     }
 }

package/src/Softmax/index.d.ts

@@ -1,6 +1,26 @@
-export class Softmax{
+export class Softmax {
     constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number);
-    async forward(logit: GPUBuffer, isCpuReadBack) : Promise<Float32Array | GPUBuffer>;
-    async backward(probs : GPUBuffer, target : number[]) : Promise<GPUBuffer>
-    async crossEntropy(probs: GPUBuffer, target: number[]) : Promise<{loss: Float32Array, acc: Float32Array}>
+    async forward(logit: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, isCpuReadBack): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>;
+    async backward(probs: {
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }, target: number[]): Promise<{
+        read: function(): Promise<Float32Array>,
+        mainBuffer: GPUBuffer,
+        isDestroyed: function(): Boolean,
+        destroy: function(): void
+    }>
+    async crossEntropy(probs: GPUBuffer, target: number[]): Promise<{ loss: Float32Array, acc: Float32Array }>
 }

package/src/Softmax/index.js

@@ -33,20 +33,24 @@ export class Softmax {
         }
 
         if (this.oldParam[1] !== tokenCount) { await this.gpu.WriteBuffer(this.paramsForward, new Uint32Array([this.vocabSize, tokenCount, 0])); this.oldParam = new Uint32Array([this.vocabSize, tokenCount, 0]); }
-        const output = await this.gpu.createBuffer(logits.size);
+        if (!this.output) this.output = await this.gpu.ManageBuffer(logits.size);
+        if (this.output.size !== 4 * tokenCount * this.vocabSize || this.output.isDestroyed()) {
+            this.output.destroy()
+            this.output = await this.gpu.ManageBuffer(logits.size)
+        }
         const bind = await this.device.createBindGroup({
             layout: await this.forwardPipline.getBindGroupLayout(0),
             entries: [
                 {
                     binding: 0,
                     resource: {
-                        buffer: logits
+                        buffer: logits.mainBuffer
                     }
                 },
                 {
                     binding: 1,
                     resource: {
-                        buffer: output
+                        buffer: this.output.mainBuffer
                     }
                 }, {
                     binding: 2,
@@ -61,12 +65,7 @@ export class Softmax {
             await this.device.queue.submit([encoder.finish()])
             logits.destroy()
         }
-        if (cpuReadBack) {
-            const outputCPU = await this.gpu.readBuffer(output.size, output);
-            logits.destroy()
-            return outputCPU
-        }
-        return output
+        return this.output
     }
     async backward(probs, target) {
         const encoder = this.device.createCommandEncoder()
@@ -81,27 +80,24 @@ export class Softmax {
             await this.gpu.WriteBuffer(this.paramsForward, new Uint32Array([this.vocabSize, tokenCount, 0]));
             this.oldParam = new Uint32Array([this.vocabSize, tokenCount, 0]);
         }
-        const targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
-        await this.gpu.WriteBuffer(targetBuffer, new Float32Array(target))
-        const output = await this.gpu.createBuffer(probs.size)
+        if (!this.targetBuffer) this.targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT);
+        if (this.targetBuffer.size !== target.length * 4) {
+            this.targetBuffer.destroy()
+            this.targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
+        }
+        await this.gpu.WriteBuffer(this.targetBuffer, new Float32Array(target))
         const bind = await this.device.createBindGroup({
             layout: await this.backwardPipline.getBindGroupLayout(0),
             entries: [
                 {
                     binding: 0,
                     resource: {
-                        buffer: probs
-                    }
-                },
-                {
-                    binding: 1,
-                    resource: {
-                        buffer: output
+                        buffer: probs.mainBuffer
                     }
                 }, {
                     binding: 2,
                     resource: {
-                        buffer: targetBuffer
+                        buffer: this.targetBuffer
                     }
                 }, {
                     binding: 3,
@@ -113,9 +109,7 @@ export class Softmax {
         })
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, tokenCount)
         this.device.queue.submit([encoder.finish()])
-        probs.destroy()
-        targetBuffer.destroy()
-        return output
+        return probs
     }
     async crossEntropy(probs, target) {
         const encoder = this.device.createCommandEncoder()
@@ -131,7 +125,7 @@ export class Softmax {
                 {
                     binding: 0,
                     resource: {
-                        buffer: probs
+                        buffer: probs.mainBuffer
                     }
                 },
                 {
@@ -169,7 +163,7 @@ export class Softmax {
                 {
                     binding: 1,
                     resource: {
-                        buffer: probs
+                        buffer: probs.mainBuffer
                     }
                 }, {
                     binding: 3,
@@ -192,6 +186,7 @@ export class Softmax {
         const lossCpu = await this.gpu.readBuffer(outputLoss.size, outputLoss)
         const accCPU = await this.gpu.readBuffer(acc.size, acc)
         acc.destroy()
+        outputLoss.destroy()
         return { loss: lossCpu[0], acc: accCPU[0] }
     }
 

package/src/TransformerBlock/index.js

@@ -43,32 +43,18 @@ export class TransformerBlock {
     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;
+        const att = await this.attention.forward(norm1)
+        const f = await this.ffn.forward(att)
+        const norm2 = await this.rmsnorm2.forward(f)
+        return norm2;
     }
 
     async backward(outGradient, lr) {
-        const count = outGradient.size / 4;
-        const dNorm2 = await this.ffn.backward(outGradient, lr);
-        const dHFromFFN = await this.rmsnorm2.backward(dNorm2 , lr);
-        const dH = await this.gpu.add(outGradient, dHFromFFN, count);
-        const dNorm1 = await this.attention.backward(dH, lr);
-        const dInputFromAttention = await this.rmsnorm.backward(dNorm1, lr);
-        const dInput = await this.gpu.add(dH, dInputFromAttention, count);
-        dNorm2.destroy();
-        dHFromFFN.destroy();
-        dH.destroy();
-        dNorm1.destroy();
-        dInputFromAttention.destroy();
-        dInputFromAttention.destroy();
-        return dInput;
+        const dNorm2 = await this.rmsnorm2.backward(outGradient, lr)
+        const dAtt = await this.ffn.backward(dNorm2, lr);
+        const dFF = await this.attention.backward(dAtt, lr)
+        const dNorm = await this.rmsnorm.backward(dFF, lr)
+        return dNorm;
     }
 
     async ClearInputCache() {

package/src/shaders/LB.wgsl

@@ -44,14 +44,14 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
         let dW = inputValue * gradient;
         dWeights += dW;
     }
-    let vClap = weight[index] - params.lr * clamp(dWeights, -20.0, 20.0);
+    let vClap = weight[index] - params.lr * clamp(dWeights, -8.0, 8.0);
     weight[index] = vClap;
     if (inputIndex == 0u) {
         var dBias = 0.0;
         for (var token = 0u; token < u32(params.tokenCount); token++) {
             dBias += outGradient[token * u32(params.vocabSize) + outputIndex];
         }
-        bias[outputIndex] = bias[outputIndex] - params.lr * clamp(dBias, -20.0, 20.0);
+        bias[outputIndex] = bias[outputIndex] - params.lr * clamp(dBias, -8.0, 8.0);
 
     }
 }

package/src/shaders/PF.wgsl

@@ -3,11 +3,9 @@ var <storage, read> posVec: array<f32>;
 
 @group(0) @binding(1)
 
-var <storage, read> vecs : array<f32>;
+var <storage, read_write> vecs : array<f32>;
 
-@group(0) @binding(2)
 
-var <storage, read_write> output: array<f32>;
 
 
 @compute @workgroup_size(64)
@@ -17,5 +15,5 @@ fn main(@builtin(global_invocation_id) id : vec3<u32>){
     if(x >= arrayLength(&vecs)){
         return;
     }
-    output[x] = vecs[x] + posVec[x];
+    vecs[x] = vecs[x] + posVec[x];
 }

package/src/shaders/RB.wgsl

@@ -32,6 +32,6 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
         output[index] = grad;
     }
     else {
-        output[index] = 0;
+        output[index] = grad * 0.02;
     }
 }

package/src/shaders/RF.wgsl

@@ -28,6 +28,6 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
         output[index] = x;
     }
     else {
-        output[index] = 0;
+        output[index] = x * 0.02;
     }
 }

package/src/shaders/SB.wgsl

@@ -5,11 +5,8 @@ struct Params {
 
 @group(0) @binding(0)
 
-var<storage, read> probs: array<f32>;
+var<storage, read_write> probs: array<f32>;
 
-@group(0) @binding(1)
-
-var<storage, read_write> output: array<f32>;
 
 @group(0) @binding(2)
 
@@ -32,8 +29,8 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
       return;
     }
     var targetIndex = targetToken[token];
-    output[index] = probs[index];
+    probs[index] = probs[index];
     if(vocabIndex == u32(targetIndex)){
-        output[index] -= 1;
+        probs[index] -= 1;
     }
 }

package/src/Model/index.d.ts

@@ -1,15 +0,0 @@
-export class Model {
-    async forward(input: Number[], buffer: boolean): Promise<Float32Array | GPUBuffer>;
-    async backward(output: GPUBuffer, target: number[], lr: number): Promise<{ acc: number, loss: number }>
-    TrainTokenizer(dataset: string): void;
-    constructor(device :GPUAdapter, embeddingSize: number, configs: {
-        tokenizerName: string
-        blockSize: number,
-        dirname: string,
-        contextSize: number,
-        gpuName: string
-
-    })
-    async Save() : Promise<void>
-    async CreateContext()
-}

package/src/Model/index.js

@@ -1,95 +0,0 @@
-import { Embedding, Linear, Softmax, TransformerBlock, Tokenizer, PositionEmbedding, RMSNorm } from '../../index.js'
-import { join } from 'path'
-import fs from 'fs'
-import { OperationManager } from '../OperationManager.js'
-export class Model {
-    embedding;
-    posEmbedding;
-    linear;
-    softmax;
-    blocks = [];
-    blockSize;
-    dimSize;
-    tokenizer;
-    device;
-    configs;
-    tokenizerName;
-    op;
-    tokens;
-    rmsnorm
-    constructor(device, embeddingSize = 64, configs) {
-        this.device = device;
-        this.dimSize = embeddingSize;
-        this.blockSize = configs.blockSize;
-        this.configs = configs
-        this.tokenizer = new Tokenizer(configs.tokenizerName)
-        this.tokenizerName = configs.tokenizerName
-        this.op = new OperationManager(device)
-    }
-    async CreateContext() {
-        this.embedding = new Embedding(this.device, this.dimSize, this.tokenizer.vocabSize, {
-            save: {
-                filename: join(this.configs.dirname, "./vec.bin")
-            }
-        })
-        this.linear = new Linear(this.device, this.dimSize, this.tokenizer.vocabSize, {
-            save: {
-                filename: [join(this.configs.dirname, "weights.bin"), join(this.configs.dirname, "bias.bin")]
-            }
-        })
-        this.softmax = new Softmax(this.device, this.dimSize, this.tokenizer.vocabSize);
-        for (let a = 0; a < this.blockSize; a++) {
-            console.log(`Load the Transfomer Block ${a + 1}...`)
-            if (!fs.existsSync(join(this.configs.dirname, `block-${a}`))) fs.mkdirSync(join(this.configs.dirname, `block-${a}`));
-            this.blocks[a] = new TransformerBlock(this.device, this.dimSize, {
-                dirname: join(this.configs.dirname, `block-${a}`)
-            })
-        }
-        this.posEmbedding = new PositionEmbedding(this.device, this.dimSize, this.configs.contextSize, {
-            save: {
-                filename: join(this.configs.dirname, "posVec.bin")
-            }
-        })
-    }
-    TrainTokenizer(text = "", mergeSize = 5) {
-        this.tokenizer.trainTokenizer(text, mergeSize, this.tokenizerName)
-    }
-    async ClearInputCache() {
-        await this.linear.ClearInputCache()
-        for (let a = 0; a < this.blockSize; a++) {
-            await this.blocks[a].ClearInputCache()
-        }
-    }
-    async forward(input, readBack) {
-        if (this.configs.contextSize < input.length) throw new Error("The Input is larger then contextSize")
-        this.tokens = input;
-        await this.ClearInputCache()
-        let vec = await this.embedding.forward(input)
-        vec = await this.posEmbedding.forward(vec)
-        for (let a = 0; a < this.blockSize; a++) {
-            vec = await this.blocks[a].forward(vec)
-        }
-        let logits = await this.linear.forward(vec);
-        let probs = await this.softmax.forward(logits, readBack)
-        return probs
-    }
-    async backward(outputGradient, target, lr = 0.02) {
-        const loss = await this.softmax.crossEntropy(outputGradient, target)
-        const dLogits = await this.softmax.backward(outputGradient, target);
-        let DInput = await this.linear.backward(dLogits, lr);
-        for (let a = this.blockSize - 1; a >= 0; a--) {
-            DInput = await this.blocks[a].backward(DInput, lr);
-        }
-        this.embedding.backward(this.tokens, DInput, lr);
-        this.posEmbedding.backward(DInput, lr)
-        return loss
-    }
-    async Save() {
-        await this.embedding.Save()
-        await this.linear.Save()
-        for (let a = this.blockSize - 1; a >= 0; a--) {
-            await this.blocks[a].Save();
-        }
-        await this.posEmbedding.Save()
-    }
-}