@mineproness/tensorloop 2.1.55 → 2.1.56

package/src/TransformerBlock/index.d.ts

@@ -1,10 +1,9 @@
 export class TransformerBlock{
-    constructor(embeddingSize: number, configs:{
-        dirname: string,
-        activation: "relu" | "gelu"
+    constructor(device: GPUAdapter, embeddingSize: number, configs: {
+        dirname: string
     })
-    Save(): void
-    clearInputCache() : void
-    forward(input: Float32Array[]) : Float32Array[]
-    backward(dOutput: Float32Array[], LearningRate: number) : Float32Array[]
+    async forward(input : GPUBuffer) : Promise<GPUBuffer>
+    async backward(outInput : GPUBuffer, lr : number) : Promise<GPUBuffer>
+    async ClearInputCache() : Promise<void>
+    async Save() : Promise<void>
 }

package/src/TransformerBlock/index.js

@@ -1,30 +1,31 @@
-import { SelfAttention } from '../SelfAttention/index.js'
-import { FFN } from '../FFN/index.js'
-import fs from 'fs'
-export class TransformerBlock {
+import {SelfAttention} from '../SelfAttention/index.js'
+import {FFN} from '../FFN/index.js'
+import {existsSync, mkdirSync} from 'fs'
+
+export class TransformerBlock{
+    attention;
     ffn;
-    selfAttention;
-    constructor(embeddingSize, configs) {
-        if(!fs.existsSync(configs.dirname)) fs.mkdirSync(configs.dirname)
-        this.ffn = new FFN(embeddingSize, configs.activation, configs)
-        this.selfAttention = new SelfAttention(embeddingSize, configs)
+    constructor(device, embeddingSize, configs){
+        if(!existsSync(configs.dirname)) mkdirSync(configs.dirname);
+        this.attention = new SelfAttention(device, embeddingSize, configs)
+        this.ffn = new FFN(device, embeddingSize, configs)
     }
-    forward(input) {
-        const attention = this.selfAttention.forward(input)
-        return this.ffn.forward(attention)
+    async forward(input){
+        const xAtt = await this.attention.forward(input);
+        const xFF = await this.ffn.forward(xAtt);
+        return xFF
     }
-    backward(dOutput, learningRate) {
-        const dAttention = this.ffn.backward(dOutput, learningRate)
-        return this.selfAttention.backward(dAttention, learningRate)
+    async backward(outGradient, lr){
+        const dAtt = await this.ffn.backward(outGradient, lr);
+        const dFF = await   this.attention.backward(dAtt, lr)
+        return dFF
     }
-    clearInputCache() {
-        this.ffn.ClearInputCache()
-        this.selfAttention.key.ClearInputCache()
-        this.selfAttention.query.ClearInputCache()
-        this.selfAttention.value.ClearInputCache()
+    async ClearInputCache(){
+        await this.ffn.ClearInputCache()
+        await this.attention.ClearInputCache()
     }
-    Save() {
-        this.selfAttention.Save()
-        this.ffn.Save()
+    async Save(){
+        await this.ffn.Save()
+        await this.attention.Save()
     }
 }

package/src/getVec.js

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

package/src/shaders/CalculateAcc.wgsl

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

package/src/shaders/CalculateLoss.wgsl

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

package/src/shaders/ComputeDInput.wgsl

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

package/src/shaders/EB.wgsl

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

package/src/shaders/EF.wgsl

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

package/src/shaders/LB.wgsl

@@ -0,0 +1,57 @@
+struct Params {
+    vocabSize: f32,
+    tokenCount: f32,
+    lr: f32,
+    inputSize: f32
+}
+
+@group(0) @binding(0)
+
+var<storage, read_write> bias: array<f32>;
+
+@group(0) @binding(1)
+
+var<storage, read_write> weight: array<f32>;
+
+@group(0) @binding(2)
+
+var<storage, read_write> outGradient: array<f32>;
+
+@group(0) @binding(3)
+
+var<storage, read_write> input: array<f32>;
+
+@group(0) @binding(4)
+
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+
+    let totalWeights = u32(params.inputSize * params.vocabSize);
+    let token = index / u32(params.inputSize);
+    if (index >= totalWeights) {
+        return;
+    }
+    let outputIndex = index / u32(params.inputSize);
+    let inputIndex = index % u32(params.inputSize);
+    var dWeights = 0.0;
+    for (var token = 0u; token < u32(params.tokenCount); token++) {
+        let inputValue = input[token * u32(params.inputSize) + inputIndex];
+        let gradient = outGradient[token * u32(params.vocabSize) + outputIndex];
+        let dW = inputValue * gradient;
+        dWeights += f32(dW);
+    }
+    let vClap = weight[index] - params.lr * max(- 3.0, min(3.0, dWeights));
+    weight[index] = vClap;
+    if (inputIndex == 0u) {
+        var dBias = 0.0;
+        for (var token = 0u; token < u32(params.tokenCount); token++) {
+            dBias += f32(outGradient[token * u32(params.vocabSize) + outputIndex]);
+        }
+        bias[outputIndex] = bias[outputIndex] - params.lr * max(- 3.0, min(3.0, dBias));
+
+    }
+}

package/src/shaders/LF.wgsl

@@ -0,0 +1,44 @@
+ struct Params {
+    inputSize: u32,
+    outputSize: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+
+var<storage, read_write> weights: array<f32>;
+
+@group(0) @binding(1)
+
+var<storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+
+var<storage, read_write> input: array<f32>;
+
+@group(0) @binding(3)
+
+var<storage, read_write> bias: array<f32>;
+
+@group(0) @binding(4)
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.tokenCount * params.outputSize;
+    if (index >= total) {
+        // output[index] = f32(total);
+        return;
+    }
+    let token = index / params.outputSize;
+    let outIndex = index % params.outputSize;
+    var sum = bias[outIndex];
+    for (var i = 0u; i < params.inputSize; i++) {
+        let inputIndex = token * params.inputSize + i;
+        let weightIndex = outIndex * params.inputSize + i;
+        sum += input[inputIndex] * weights[weightIndex];
+    }
+    output[index] = sum;
+}

package/src/shaders/RB.wgsl

@@ -0,0 +1,37 @@
+struct Param {
+    embeddingSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> input: vec2<f32>;
+
+@group(0) @binding(1)
+
+var<uniform> params: Param;
+
+@group(0) @binding(2)
+
+var<storage, read_write> output: vec2<f32>;
+
+@group(0) @binding(3)
+
+var<storage, read_write> inputGradient: vec2<f32>;
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.embeddingSize * params.tokenCount;
+    if (index >= total) {
+        return;
+    }
+    let x = input[index];
+    let grad = inputGradient[index];
+    if (x > 0) {
+        output[index] = grad;
+    }
+    else {
+        output[index] = grad * 0.02;
+    }
+}

package/src/shaders/RF.wgsl

@@ -0,0 +1,33 @@
+struct Param {
+    embeddingSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> input: vec2<f32>;
+
+@group(0) @binding(1)
+
+var<uniform> params: Param;
+
+@group(0) @binding(2)
+
+var<storage, read_write> output: vec2<f32>;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.embeddingSize * params.tokenCount;
+    if (index >= total) {
+        return;
+    }
+    let x = input[index];
+    if (x > 0) {
+        output[index] = x;
+    }
+    else {
+        output[index] = x * 0.02;
+    }
+}

package/src/shaders/ReductionFunction.wgsl

@@ -0,0 +1,28 @@
+struct Parmas {
+    vocabSize: u32,
+    count: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> data: array<f32>;
+
+@group(0) @binding(1)
+
+var<storage, read_write> result: array<f32>;
+
+@group(0) @binding(2)
+
+var<uniform> params: Parmas;
+
+@compute @workgroup_size(1)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    var sum = 0.0;
+    for (var i = 0u; i < params.count; i++) {
+        sum += data[i];
+    }
+    sum /= f32(params.count);
+    result[index] = sum;
+}

package/src/shaders/SB.wgsl

@@ -0,0 +1,36 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32
+}
+
+@group(0) @binding(0)
+
+var<storage, read> probs: array<f32>;
+
+@group(0) @binding(1)
+
+var<storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+
+var<storage, read_write> targetToken: array<f32>;
+
+@group(0) @binding(3)
+
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let token = index / params.vocabSize;
+    let vocabIndex = index % params.vocabSize;
+    let start = token * params.vocabSize;
+    let end = start + params.vocabSize;
+    var targetIndex = targetToken[token];
+    let total: u32 = params.tokenCount * params.vocabSize;
+    output[index] = probs[index];
+    if(vocabIndex == u32(targetIndex)){
+        output[index] -= 1;
+    }
+}

package/src/shaders/SF.wgsl

@@ -0,0 +1,65 @@
+struct Params {
+    vocabSize: u32,
+    tokenCount: u32,
+}
+
+@group(0) @binding(0)
+var<storage, read> logits: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> output: array<f32>;
+
+@group(0) @binding(2)
+var<uniform> params: Params;
+
+var<workgroup> sharedMax: array<f32, 256>;
+var<workgroup> sharedSum: array<f32, 256>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>,) {
+    let token = wg.x;
+    let tid = lid.x;
+    if (token >= params.tokenCount) {
+        return;
+    }
+    let start = token * params.vocabSize;
+    var localMax = - 3.402823e38;
+    var i = tid;
+    while (i < params.vocabSize) {
+        localMax = max(localMax, logits[start + i]);
+        i += 256;
+    }
+    sharedMax[tid] = localMax;
+    workgroupBarrier();
+    var stride = 128u;
+    while (stride > 0u) {
+        if (tid < stride) {
+            sharedMax[tid] = max(sharedMax[tid], sharedMax[tid + stride]);
+        }
+        workgroupBarrier();
+        stride /= 2u;
+    }
+    let maxLogit = sharedMax[0];
+    var localSum = 0.0;
+    i = tid;
+    while (i < params.vocabSize) {
+        localSum += exp(logits[start + i] - maxLogit);
+        i += 256;
+    }
+    sharedSum[tid] = localSum;
+    workgroupBarrier();
+    stride = 128u;
+    while (stride > 0u) {
+        if (tid < stride) {
+            sharedSum[tid] += sharedSum[tid + stride];
+        }
+        workgroupBarrier();
+        stride /= 2u;
+    }
+    let sum = sharedSum[0];
+    i = tid;
+    while (i < params.vocabSize) {
+        output[start + i] = exp(logits[start + i] - maxLogit) / sum;
+        i += 256;
+    }
+}

package/src/shaders/SoftmaxBackward.wgsl

@@ -0,0 +1,31 @@
+struct Params {
+    tokens: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> P: array<f32>;
+@group(0) @binding(1)
+var<storage, read> dP: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> output: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total = params.tokens * params.tokens;
+    if (index >= total) {
+        return;
+    }
+    let row = index / params.tokens;
+    var dot = 0.0;
+    for (var i = 0u; i < params.tokens; i++) {
+        let pos = row * params.tokens + i;
+        dot += P[pos] * dP[pos];
+    }
+    let p = P[index];
+    let dp = dP[index];
+    output[index] = clamp(p * (dp - dot), -1.0, 1.0);
+}

package/src/shaders/attentionShader/CompueFinalForwardOutput.wgsl

@@ -0,0 +1,29 @@
+struct Params {
+    tokens: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> weights: array<f32>;
+@group(0) @binding(1)
+var<storage, read> value: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> output: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let index = id.x;
+    let total =params.tokens * params.embeddingSize;
+    if (index >= total) {
+        return;
+    }
+    let row = index / params.embeddingSize;
+    let dim = index % params.embeddingSize;
+    var sum = 0.0;
+    for (var j = 0u; j < params.tokens; j++) {
+        let weight = weights[row * params.tokens + j];
+        let v =value[j * params.embeddingSize + dim];
+        sum += weight * v;
+    }
+    output[index] = sum;
+}

package/src/shaders/attentionShader/ComputeDA.wgsl

@@ -0,0 +1,30 @@
+struct Params {
+    contextSize: u32,
+    embeddingSize: u32,
+}
+@group(0) @binding(0)
+var<storage, read> dOutput: array<f32>;
+@group(0) @binding(1)
+var<storage, read> V: array<f32>;
+@group(0) @binding(2)
+var<storage, read_write> dA: array<f32>;
+@group(0) @binding(3)
+var<uniform> params: Params;
+@compute @workgroup_size(16, 16)
+fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+    let row = id.x;
+    let col = id.y;
+    let T = params.contextSize;
+    let D = params.embeddingSize;
+    if (row >= T || col >= T) {
+        return;
+    }
+    var sum = 0.0;
+    for (var k = 0u; k < D; k++) {
+        let dOutputIndex = row * D + k;
+        let vIndex = col * D + k;
+        sum += dOutput[dOutputIndex] * V[vIndex];
+    }
+    let outputIndex = row * T + col;
+    dA[outputIndex] = sum;
+}