@imaduddin/fastbrain 1.1.1-patch → 2.0.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/Embedding/index.d.ts +24 -0
- package/dist/Embedding/index.js +93 -0
- package/dist/FFN/index.d.ts +20 -0
- package/dist/FFN/index.js +51 -0
- package/dist/Gelu/index.d.ts +15 -0
- package/dist/Gelu/index.js +58 -0
- package/dist/Linear/index.d.ts +31 -0
- package/dist/Linear/index.js +136 -0
- package/dist/RMSNorm/index.d.ts +16 -0
- package/dist/RMSNorm/index.js +63 -0
- package/dist/Relu/index.d.ts +15 -0
- package/dist/Relu/index.js +58 -0
- package/dist/SelfAttention/index.d.ts +47 -0
- package/dist/SelfAttention/index.js +186 -0
- package/dist/Softmax/index.d.ts +18 -0
- package/dist/Softmax/index.js +96 -0
- package/dist/Tokenizer/index.d.ts +15 -0
- package/dist/Tokenizer/index.js +58 -0
- package/dist/TransformerBlock/index.d.ts +21 -0
- package/dist/TransformerBlock/index.js +43 -0
- package/dist/fileManager.d.ts +6 -0
- package/dist/fileManager.js +34 -0
- package/dist/index.d.ts +10 -0
- package/dist/index.js +10 -0
- package/dist/residualConnection/index.d.ts +9 -0
- package/dist/residualConnection/index.js +25 -0
- package/package.json +14 -3
- package/shaders/add.wgsl +22 -0
- package/shaders/embedding/backward.wgsl +36 -0
- package/shaders/embedding/forward.wgsl +35 -0
- package/shaders/linear/ComputeDinput.wgsl +72 -0
- package/shaders/linear/forward.wgsl +63 -0
- package/shaders/linear/updateBias.wgsl +33 -0
- package/shaders/linear/updateWeights.wgsl +38 -0
- package/shaders/relu/backward.wgsl +43 -0
- package/shaders/relu/forward.wgsl +28 -0
- package/shaders/rmsnorm/backward.wgsl +54 -0
- package/shaders/rmsnorm/forward.wgsl +51 -0
- package/shaders/selfAttention/SB.wgsl +30 -0
- package/shaders/selfAttention/add3.wgsl +27 -0
- package/shaders/selfAttention/matMul.wgsl +63 -0
- package/shaders/selfAttention/qK.wgsl +66 -0
- package/shaders/selfAttention/transpose.wgsl +24 -0
- package/shaders/softmax/backward.wgsl +36 -0
- package/shaders/softmax/crossEntropy.wgsl +36 -0
- package/shaders/softmax/forward.wgsl +68 -0
- package/src/Embedding/index.ts +81 -0
- package/src/FFN/index.ts +47 -0
- package/src/Gelu/index.ts +53 -0
- package/src/Linear/index.ts +122 -0
- package/src/RMSNorm/index.ts +56 -0
- package/src/SelfAttention/index.ts +177 -0
- package/src/Softmax/index.ts +82 -0
- package/src/Tokenizer/index.ts +60 -0
- package/src/TransformerBlock/index.ts +41 -0
- package/src/fileManager.ts +46 -0
- package/src/index.ts +10 -0
- package/src/residualConnection/index.ts +24 -0
- package/tsconfig.json +27 -0
- package/index.js +0 -22
- package/src/Embedding/index.d.ts +0 -21
- package/src/Embedding/index.js +0 -140
- package/src/FFN/index.d.ts +0 -23
- package/src/FFN/index.js +0 -45
- package/src/Linear/index.d.ts +0 -31
- package/src/Linear/index.js +0 -196
- package/src/OperationManager.js +0 -169
- package/src/PositionEmbedding/index.d.ts +0 -30
- package/src/PositionEmbedding/index.js +0 -70
- package/src/RMSNorm/index.d.ts +0 -30
- package/src/RMSNorm/index.js +0 -91
- package/src/Relu/index.d.ts +0 -25
- package/src/Relu/index.js +0 -108
- package/src/SelfAttention/index.d.ts +0 -28
- package/src/SelfAttention/index.js +0 -138
- package/src/Softmax/index.d.ts +0 -26
- package/src/Softmax/index.js +0 -195
- package/src/Tokenizer/index.d.ts +0 -12
- package/src/Tokenizer/index.js +0 -164
- package/src/Tokenizer/vocabs.json +0 -19618
- package/src/TransformerBlock/index.d.ts +0 -28
- package/src/TransformerBlock/index.js +0 -78
- package/src/getVec.js +0 -30
- package/src/gpu.js +0 -0
- package/src/model/index.d.ts +0 -43
- package/src/model/index.js +0 -82
- package/src/residualConnection/index.js +0 -31
- package/src/shaders/CalculateAcc.wgsl +0 -50
- package/src/shaders/CalculateLoss.wgsl +0 -33
- package/src/shaders/ComputeDInput.wgsl +0 -46
- package/src/shaders/EB.wgsl +0 -33
- package/src/shaders/EF.wgsl +0 -25
- package/src/shaders/LB.wgsl +0 -71
- package/src/shaders/LF.wgsl +0 -60
- package/src/shaders/Norm/rmsBackward.wgsl +0 -65
- package/src/shaders/Norm/rmsForward.wgsl +0 -30
- package/src/shaders/Norm/sumSquare.wgsl +0 -38
- package/src/shaders/Norm/updateGamma.wgsl +0 -54
- package/src/shaders/PB.wgsl +0 -25
- package/src/shaders/PF.wgsl +0 -24
- package/src/shaders/RB.wgsl +0 -37
- package/src/shaders/RF.wgsl +0 -33
- package/src/shaders/ReductionFunction.wgsl +0 -28
- package/src/shaders/SB.wgsl +0 -41
- package/src/shaders/SF.wgsl +0 -75
- package/src/shaders/SoftmaxBackward.wgsl +0 -31
- package/src/shaders/add.wgsl +0 -30
- package/src/shaders/attentionShader/CompueFinalForwardOutput.wgsl +0 -29
- package/src/shaders/attentionShader/ComputeDA.wgsl +0 -30
- package/src/shaders/attentionShader/ComputeDV.wgsl +0 -37
- package/src/shaders/attentionShader/ComputeK.wgsl +0 -30
- package/src/shaders/attentionShader/ComputedQ.wgsl +0 -31
- package/src/shaders/attentionShader/attentionSocres.wgsl +0 -40
- package/src/shaders/attentionShader/sumGradient.wgsl +0 -24
- package/src/shaders/calculateTotal.wgsl +0 -26
|
@@ -0,0 +1,82 @@
|
|
|
1
|
+
import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
|
|
2
|
+
import { isBufferOrNumber } from '../fileManager.js'
|
|
3
|
+
|
|
4
|
+
export class Softmax {
|
|
5
|
+
forwardShader: GPUComputePipeline | undefined;
|
|
6
|
+
lossShader: GPUComputePipeline | undefined;
|
|
7
|
+
backwardShader: GPUComputePipeline | undefined;
|
|
8
|
+
output: Buffer | undefined;
|
|
9
|
+
loss: Buffer | undefined;
|
|
10
|
+
gpu: GPUManager | undefined;
|
|
11
|
+
forwardParams: Buffer | undefined;
|
|
12
|
+
gradOutput: Buffer | undefined
|
|
13
|
+
targetBuffer: Buffer | undefined;
|
|
14
|
+
constructor(public device: GPUDevice, public VocabSize: number) {
|
|
15
|
+
this.gpu = new GPUManager(device, {
|
|
16
|
+
shaderDir: './shaders/softmax'
|
|
17
|
+
})
|
|
18
|
+
}
|
|
19
|
+
async forward(logit: Buffer | number[], encoderShared: GPUCommandEncoder | undefined) {
|
|
20
|
+
let logits = await isBufferOrNumber(logit, this.gpu!)
|
|
21
|
+
const tokenCount = logits.buffer.size / (this.VocabSize * 4)
|
|
22
|
+
const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
|
|
23
|
+
if (!this.forwardShader) {
|
|
24
|
+
this.forwardShader = await this.gpu?.getPiplineShader("forward.wgsl")
|
|
25
|
+
}
|
|
26
|
+
if (!this.output || this.output.buffer.size !== logits.buffer.size || this.output.isDestroy()) {
|
|
27
|
+
if (this.output) this.output.destroy()
|
|
28
|
+
this.output = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.VocabSize), "storage")
|
|
29
|
+
}
|
|
30
|
+
if (!this.forwardParams) {
|
|
31
|
+
this.forwardParams = await this.gpu?.createBuffer(new Uint32Array([this.VocabSize, tokenCount]), "uniform")
|
|
32
|
+
} else {
|
|
33
|
+
await this.forwardParams.changeData(new Uint32Array([this.VocabSize, tokenCount]))
|
|
34
|
+
}
|
|
35
|
+
const bindGroup = await this.gpu?.createBindGroup(this.forwardShader!, [logits!, this.output!, this.forwardParams!], 0)
|
|
36
|
+
await this.gpu?.RunPipline(encoder, this.forwardShader!, bindGroup!, tokenCount)
|
|
37
|
+
if (!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
|
|
38
|
+
return this.output
|
|
39
|
+
}
|
|
40
|
+
async backward(probs: Buffer, target: number[], encoderShared: GPUCommandEncoder | undefined) {
|
|
41
|
+
let tokenCount = probs.buffer.size / (4 * this.VocabSize)
|
|
42
|
+
if (!this.backwardShader) {
|
|
43
|
+
this.backwardShader = await this.gpu?.getPiplineShader("backward.wgsl")
|
|
44
|
+
}
|
|
45
|
+
const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
|
|
46
|
+
if (!this.gradOutput || this.gradOutput.buffer.size !== probs.buffer.size || this.gradOutput.isDestroy()) {
|
|
47
|
+
if (this.gradOutput) this.gradOutput.destroy()
|
|
48
|
+
this.gradOutput = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.VocabSize), "storage")
|
|
49
|
+
}
|
|
50
|
+
if (!this.targetBuffer || this.targetBuffer.buffer.size !== tokenCount * Uint32Array.BYTES_PER_ELEMENT) {
|
|
51
|
+
if (this.targetBuffer) this.targetBuffer.destroy()
|
|
52
|
+
this.targetBuffer = await this.gpu?.createBuffer(new Uint32Array(target), "storage")
|
|
53
|
+
} else {
|
|
54
|
+
this.targetBuffer.changeData(new Uint32Array(target))
|
|
55
|
+
}
|
|
56
|
+
const bindGroup = await this.gpu?.createBindGroup(this.backwardShader!, [probs, this.gradOutput!, this.targetBuffer!, this.forwardParams!], 0)
|
|
57
|
+
await this.gpu?.RunPipline(encoder, this.backwardShader!, bindGroup!, Math.ceil(this.VocabSize / 64), tokenCount)
|
|
58
|
+
if (!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
|
|
59
|
+
return this.gradOutput
|
|
60
|
+
}
|
|
61
|
+
async crossEntropy(probs: Buffer, target: number[], encoderShared : GPUCommandEncoder | undefined) {
|
|
62
|
+
let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder()
|
|
63
|
+
let tokenCount = probs.buffer.size / (4 * this.VocabSize)
|
|
64
|
+
if (!this.targetBuffer || this.targetBuffer.buffer.size !== tokenCount * Uint32Array.BYTES_PER_ELEMENT) {
|
|
65
|
+
if (this.targetBuffer) this.targetBuffer.destroy()
|
|
66
|
+
this.targetBuffer = await this.gpu?.createBuffer(new Uint32Array(target), "storage")
|
|
67
|
+
} else {
|
|
68
|
+
this.targetBuffer.changeData(new Uint32Array(target))
|
|
69
|
+
}
|
|
70
|
+
if(!this.loss || this.loss.buffer.size !== tokenCount * 4){
|
|
71
|
+
if(this.loss) this.loss.destroy()
|
|
72
|
+
this.loss = await this.gpu?.createBuffer(new Float32Array(tokenCount), "storage")
|
|
73
|
+
}
|
|
74
|
+
if(!this.lossShader){
|
|
75
|
+
this.lossShader = await this.gpu?.getPiplineShader("crossEntropy.wgsl")
|
|
76
|
+
}
|
|
77
|
+
const bindGroup = await this.gpu?.createBindGroup(this.lossShader!, [probs,this.targetBuffer!, this.loss!, this.forwardParams!], 0)
|
|
78
|
+
await this.gpu?.RunPipline(encoder, this.lossShader!, bindGroup!, Math.ceil(tokenCount / 16))
|
|
79
|
+
if(!encoderShared) await this.gpu?.SumbitEncoder(encoder!)
|
|
80
|
+
return this.loss
|
|
81
|
+
}
|
|
82
|
+
}
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
import { GPUManager } from '@imaduddin/webgpu-manager';
|
|
2
|
+
import fs from 'fs'
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
export class Tokenizer {
|
|
6
|
+
vocabs: string[] = ["hi"]
|
|
7
|
+
vocabSize: number | undefined = 0
|
|
8
|
+
idToString: Record<number, string> = {12: "hi"};
|
|
9
|
+
stringToID: Record<string, number> = {"hi": 12};
|
|
10
|
+
constructor(public configs:{
|
|
11
|
+
vocabFile: string
|
|
12
|
+
}) {
|
|
13
|
+
const isVocabExist = fs.existsSync(configs.vocabFile);
|
|
14
|
+
if(isVocabExist){
|
|
15
|
+
console.log(`[LOG] Tokenizer is Loaded with ${configs.vocabFile}`)
|
|
16
|
+
this.vocabs = JSON.parse(fs.readFileSync(configs.vocabFile , "utf-8"))
|
|
17
|
+
this.vocabSize = this.vocabs?.length;
|
|
18
|
+
this.idToString = Object.fromEntries(this.vocabs?.map((e,idx)=> [idx,e]))
|
|
19
|
+
this.stringToID = Object.fromEntries(this.vocabs?.map((e,idx)=> [e,idx]))
|
|
20
|
+
}else{
|
|
21
|
+
console.log(`[LOG] Tokenizer exited with error 1. No such File or Dir.`)
|
|
22
|
+
}
|
|
23
|
+
}
|
|
24
|
+
encode(text: string){
|
|
25
|
+
const split = text.toLocaleLowerCase().split(/(\s+|[.:?,;!:'-()])/).filter(Boolean)
|
|
26
|
+
const tokens : number[] | []= []
|
|
27
|
+
for(let a = 0; a < split.length; a++){
|
|
28
|
+
const token : number | undefined = this.stringToID[split[a]] as never;
|
|
29
|
+
if(token !== undefined){
|
|
30
|
+
tokens.push(token)
|
|
31
|
+
}else{
|
|
32
|
+
const unkToken = this.stringToID["<unk>"] as never
|
|
33
|
+
tokens.push(unkToken)
|
|
34
|
+
}
|
|
35
|
+
}
|
|
36
|
+
return tokens
|
|
37
|
+
}
|
|
38
|
+
decode(token: number[]){
|
|
39
|
+
let text = ""
|
|
40
|
+
for(let a = 0; a < token.length; a++){
|
|
41
|
+
const string = this.idToString[token[a]];
|
|
42
|
+
if(string !== undefined){
|
|
43
|
+
text += string
|
|
44
|
+
}else{
|
|
45
|
+
text += "<unk>"
|
|
46
|
+
}
|
|
47
|
+
}
|
|
48
|
+
return text
|
|
49
|
+
}
|
|
50
|
+
trainTokenizer(text: string, file: string) {
|
|
51
|
+
let oldArray = fs.existsSync(file) ? JSON.parse(fs.readFileSync(file, 'utf-8')) : []
|
|
52
|
+
const splitedText = text.toLocaleLowerCase().split(/\r\n|\r|\n/).filter((e) => e.trim().length > 0).map((e) => e.replaceAll("<eos>", "<eos>").split(/(\s+|[.:?,;:'!-()])/).filter(Boolean)).flat()
|
|
53
|
+
const uqiuneText = [...new Set([...splitedText, "<unk>", "_m", "_e", "<eos>", ...oldArray])]
|
|
54
|
+
console.log("Operation Completed!")
|
|
55
|
+
console.log("New Vocab Size", uqiuneText.length - oldArray.length);
|
|
56
|
+
fs.writeFileSync(file, JSON.stringify(uqiuneText))
|
|
57
|
+
}
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
}
|
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
import {FFN, RMSNorm, SelfAttention, residualConnection} from '../index.js'
|
|
2
|
+
import { Buffer } from '@imaduddin/webgpu-manager'
|
|
3
|
+
|
|
4
|
+
export class TransformerBlock{
|
|
5
|
+
ffn: FFN | undefined
|
|
6
|
+
attention: SelfAttention | undefined
|
|
7
|
+
norm : RMSNorm | undefined
|
|
8
|
+
norm1 : RMSNorm | undefined
|
|
9
|
+
residualConnection1: residualConnection | undefined;
|
|
10
|
+
residualConnection2: residualConnection | undefined;
|
|
11
|
+
constructor(public device : GPUDevice, public embeddingSize : number, public configs : {
|
|
12
|
+
dirname: string
|
|
13
|
+
}){
|
|
14
|
+
this.ffn = new FFN(device, embeddingSize, configs)
|
|
15
|
+
this.norm = new RMSNorm(device, embeddingSize)
|
|
16
|
+
this.norm1 = new RMSNorm(device, embeddingSize)
|
|
17
|
+
this.attention = new SelfAttention(device, embeddingSize, configs)
|
|
18
|
+
this.residualConnection1 = new residualConnection(device)
|
|
19
|
+
this.residualConnection2 = new residualConnection(device)
|
|
20
|
+
}
|
|
21
|
+
async forward(input : Buffer){
|
|
22
|
+
const norm = await this.norm?.forward(input, undefined);
|
|
23
|
+
const att = await this.attention?.forward(norm!)
|
|
24
|
+
const residual = await this.residualConnection1?.add(att!, input);
|
|
25
|
+
const norm1 = await this.norm1?.forward(residual!, undefined)
|
|
26
|
+
const ff = await this.ffn?.forward(norm1!, undefined);
|
|
27
|
+
return await this.residualConnection2?.add(residual!, ff!);
|
|
28
|
+
}
|
|
29
|
+
async backward(outGradient: Buffer , lr : number){
|
|
30
|
+
const d_FFN = await this.ffn?.backward(outGradient, lr, undefined)
|
|
31
|
+
const d_norm1 = await this.norm1?.backward(d_FFN!, undefined)
|
|
32
|
+
let d_residual2 = await this.residualConnection2?.add(outGradient, d_norm1!);
|
|
33
|
+
const d_att = await this.attention?.backward(d_residual2!, lr)
|
|
34
|
+
const d_norm2 = await this.norm?.backward(d_att!,undefined)
|
|
35
|
+
return await this.residualConnection1?.add(outGradient, d_norm2!)
|
|
36
|
+
}
|
|
37
|
+
async Save(){
|
|
38
|
+
await this.ffn?.Save()
|
|
39
|
+
await this.attention?.Save()
|
|
40
|
+
}
|
|
41
|
+
}
|
|
@@ -0,0 +1,46 @@
|
|
|
1
|
+
import fs from 'fs'
|
|
2
|
+
import path, { join } from 'path'
|
|
3
|
+
import { Buffer as WEBB, GPUManager } from '@imaduddin/webgpu-manager';
|
|
4
|
+
|
|
5
|
+
process.stdout.write("\x1bc")
|
|
6
|
+
|
|
7
|
+
let totalParameter = 0;
|
|
8
|
+
export function GetVectorsForFile(file: string, total: number, isBias: boolean = false, aSize: number, bSize: number) {
|
|
9
|
+
if (!fs.existsSync(file)) return null;
|
|
10
|
+
totalParameter += total;
|
|
11
|
+
|
|
12
|
+
console.log(`Loading File ${join(file)} (${aSize}x${bSize})...`)
|
|
13
|
+
|
|
14
|
+
const rawBuffer = fs.readFileSync(file)
|
|
15
|
+
const vec = new Float32Array(rawBuffer.buffer, rawBuffer.byteOffset, rawBuffer.byteLength / 4)
|
|
16
|
+
if (total <= vec.length) {
|
|
17
|
+
return vec;
|
|
18
|
+
}
|
|
19
|
+
const expanded = new Float32Array(total);
|
|
20
|
+
expanded.set(vec);
|
|
21
|
+
for (let i = vec.length; i < total; i++) {
|
|
22
|
+
expanded[i] = isBias ? 0 : (Math.random() * 2 - 1) * 0.05;
|
|
23
|
+
}
|
|
24
|
+
return expanded;
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
export function SaveFileVector(file: string, vector: Float32Array) {
|
|
29
|
+
fs.writeFileSync(file, Buffer.from(
|
|
30
|
+
vector.buffer,
|
|
31
|
+
vector.byteOffset,
|
|
32
|
+
vector.byteLength
|
|
33
|
+
))
|
|
34
|
+
}
|
|
35
|
+
|
|
36
|
+
export function initRandomVector(total: number) {
|
|
37
|
+
return new Float32Array(total).map((e) => (Math.random() * 2 - 1) * 0.02)
|
|
38
|
+
}
|
|
39
|
+
|
|
40
|
+
export async function isBufferOrNumber(f: WEBB | number[], gpu: GPUManager) {
|
|
41
|
+
return f instanceof WEBB ? f : await gpu.createBuffer(new Float32Array(f), "storage");
|
|
42
|
+
}
|
|
43
|
+
|
|
44
|
+
export function getParameter() {
|
|
45
|
+
console.log(`Total Parameter ${(totalParameter / 1000000).toFixed(2)}M.`)
|
|
46
|
+
}
|
package/src/index.ts
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
export * from './Embedding/index.js'
|
|
2
|
+
export * from './Linear/index.js'
|
|
3
|
+
export * from './Softmax/index.js'
|
|
4
|
+
export * from './Gelu/index.js'
|
|
5
|
+
export * from './FFN/index.js'
|
|
6
|
+
export * from './RMSNorm/index.js'
|
|
7
|
+
export * from './SelfAttention/index.js'
|
|
8
|
+
export * from './Tokenizer/index.js'
|
|
9
|
+
export * from './TransformerBlock/index.js'
|
|
10
|
+
export * from './residualConnection/index.js'
|
|
@@ -0,0 +1,24 @@
|
|
|
1
|
+
import {Buffer,GPUManager} from "@imaduddin/webgpu-manager"
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
export class residualConnection{
|
|
5
|
+
output: Buffer | undefined;
|
|
6
|
+
gpu : GPUManager | undefined;
|
|
7
|
+
addShader: GPUComputePipeline | undefined;
|
|
8
|
+
constructor(public device : GPUDevice){
|
|
9
|
+
this.gpu = new GPUManager(device, {
|
|
10
|
+
shaderDir: "./shaders"
|
|
11
|
+
})
|
|
12
|
+
}
|
|
13
|
+
async add(a: Buffer, b: Buffer){
|
|
14
|
+
const encoder = await this.gpu?.CreateEncoder()
|
|
15
|
+
if(!this.output || this.output.buffer.size !== a.buffer.size / 4){
|
|
16
|
+
if(this.output) this.output.destroy()
|
|
17
|
+
this.output = await this.gpu?.createBuffer(new Float32Array(a.buffer.size / 4), "storage")
|
|
18
|
+
}
|
|
19
|
+
const bindGroup = await this.gpu?.createBindGroup(this.addShader!, [a,b,this.output!], 0)
|
|
20
|
+
await this.gpu?.RunPipline(encoder, this.addShader!, bindGroup!, Math.ceil((a.buffer.size / 4) / 64))
|
|
21
|
+
await this.gpu?.SumbitEncoder(encoder!)
|
|
22
|
+
return this.output;
|
|
23
|
+
}
|
|
24
|
+
}
|
package/tsconfig.json
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
{
|
|
2
|
+
"compilerOptions": {
|
|
3
|
+
"target": "ES2022",
|
|
4
|
+
"module": "NodeNext",
|
|
5
|
+
"moduleResolution": "NodeNext",
|
|
6
|
+
"rootDir": "src",
|
|
7
|
+
"outDir": "dist",
|
|
8
|
+
"strict": true,
|
|
9
|
+
"declaration": true,
|
|
10
|
+
"lib": [
|
|
11
|
+
"ES2022",
|
|
12
|
+
"DOM"
|
|
13
|
+
],
|
|
14
|
+
"types": [
|
|
15
|
+
"node",
|
|
16
|
+
"@webgpu/types"
|
|
17
|
+
],
|
|
18
|
+
"paths": {
|
|
19
|
+
"@/*": [
|
|
20
|
+
"./src/*"
|
|
21
|
+
]
|
|
22
|
+
}
|
|
23
|
+
},
|
|
24
|
+
"include": [
|
|
25
|
+
"src/**/*.ts"
|
|
26
|
+
]
|
|
27
|
+
}
|
package/index.js
DELETED
|
@@ -1,22 +0,0 @@
|
|
|
1
|
-
import { Embedding as E } from "./src/Embedding/index.js";
|
|
2
|
-
import { Linear as L } from "./src/Linear/index.js";
|
|
3
|
-
import { Softmax as S } from "./src/Softmax/index.js";
|
|
4
|
-
import { Relu as r } from "./src/Relu/index.js";
|
|
5
|
-
import { FFN as ffn } from "./src/FFN/index.js";
|
|
6
|
-
import { Tokenizer as to } from './src/Tokenizer/index.js'
|
|
7
|
-
import {SelfAttention as ST} from './src/SelfAttention/index.js'
|
|
8
|
-
import {TransformerBlock as TT} from './src/TransformerBlock/index.js'
|
|
9
|
-
import {PositionEmbedding as PM } from './src/PositionEmbedding/index.js'
|
|
10
|
-
import { RMSNorm as RM } from "./src/RMSNorm/index.js";
|
|
11
|
-
import {model as M} from './src/model/index.js'
|
|
12
|
-
export const Embedding = E
|
|
13
|
-
export const Linear = L
|
|
14
|
-
export const Softmax = S;
|
|
15
|
-
export const RELU = r
|
|
16
|
-
export const FFN = ffn
|
|
17
|
-
export const Tokenizer = to;
|
|
18
|
-
export const SelfAttention = ST
|
|
19
|
-
export const TransformerBlock = TT;
|
|
20
|
-
export const PositionEmbedding = PM;
|
|
21
|
-
export const RMSNorm = RM
|
|
22
|
-
export const model = M
|
package/src/Embedding/index.d.ts
DELETED
|
@@ -1,21 +0,0 @@
|
|
|
1
|
-
export class Embedding {
|
|
2
|
-
constructor(device: GPUBuffer, embedding: number, vocabSize: number, config: {
|
|
3
|
-
cpuReadBack: boolean,
|
|
4
|
-
save: {
|
|
5
|
-
filename: string
|
|
6
|
-
}
|
|
7
|
-
})
|
|
8
|
-
async forward(tokens: number[]): Promise<{
|
|
9
|
-
read: function(): Promise<Float32Array>,
|
|
10
|
-
mainBuffer: GPUBuffer,
|
|
11
|
-
isDestroyed: function(): Boolean,
|
|
12
|
-
destroy: function(): void
|
|
13
|
-
}>
|
|
14
|
-
async Save(): void
|
|
15
|
-
async backward(inputGradient: {
|
|
16
|
-
read: function(): Promise<Float32Array>,
|
|
17
|
-
mainBuffer: GPUBuffer,
|
|
18
|
-
isDestroyed: function(): Boolean,
|
|
19
|
-
destroy: function(): void
|
|
20
|
-
}, lr: number): Promise<void>
|
|
21
|
-
}
|
package/src/Embedding/index.js
DELETED
|
@@ -1,140 +0,0 @@
|
|
|
1
|
-
import { config } from "process";
|
|
2
|
-
|
|
3
|
-
import { SaveVector, getVector } from '../getVec.js'
|
|
4
|
-
import { OperationManager } from "../OperationManager.js";
|
|
5
|
-
|
|
6
|
-
function GenerateVectors(embeddingSize, vocabSize) {
|
|
7
|
-
const vec = new Float32Array(embeddingSize * vocabSize).map(() => (Math.random() * 2 - 1) * 0.02)
|
|
8
|
-
return vec
|
|
9
|
-
}
|
|
10
|
-
|
|
11
|
-
export class Embedding {
|
|
12
|
-
embeddingSize;
|
|
13
|
-
vocabSize;
|
|
14
|
-
gpu;
|
|
15
|
-
firstTrue;
|
|
16
|
-
vectors;
|
|
17
|
-
configs;
|
|
18
|
-
vectorsBuffer;
|
|
19
|
-
forwardPipline;
|
|
20
|
-
constructor(device, embeddingSize, vocabSize, configs) {
|
|
21
|
-
this.configs = configs
|
|
22
|
-
this.embeddingSize = embeddingSize;
|
|
23
|
-
this.vocabSize = vocabSize;
|
|
24
|
-
this.device = device
|
|
25
|
-
let total = embeddingSize * vocabSize;
|
|
26
|
-
if (configs && configs.save) {
|
|
27
|
-
const vec = getVector(configs.save.filename, total);
|
|
28
|
-
if (vec) {
|
|
29
|
-
this.vectors = vec;
|
|
30
|
-
|
|
31
|
-
} else {
|
|
32
|
-
this.vectors = GenerateVectors(embeddingSize, vocabSize)
|
|
33
|
-
}
|
|
34
|
-
|
|
35
|
-
} else {
|
|
36
|
-
this.vectors = GenerateVectors(embeddingSize, vocabSize)
|
|
37
|
-
|
|
38
|
-
}
|
|
39
|
-
this.gpu = new OperationManager(device)
|
|
40
|
-
this.forwardPipline = this.gpu.getShaderPipline("./src/shaders/EF.wgsl", {
|
|
41
|
-
embeddingSize: embeddingSize
|
|
42
|
-
})
|
|
43
|
-
this.backwardPipline = this.gpu.getShaderPipline("./src/shaders/EB.wgsl", {
|
|
44
|
-
embeddingSize: embeddingSize
|
|
45
|
-
})
|
|
46
|
-
SaveVector(this.configs.save.filename, this.vectors)
|
|
47
|
-
}
|
|
48
|
-
async forward(input, encoder2) {
|
|
49
|
-
const encoder = encoder2 ? encoder2 : this.device.createCommandEncoder()
|
|
50
|
-
if (!this.vectorsBuffer) {
|
|
51
|
-
this.vectorsBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize);
|
|
52
|
-
await this.gpu.WriteBuffer(this.vectorsBuffer, this.vectors)
|
|
53
|
-
}
|
|
54
|
-
if (!this.tokenBuffer) this.tokenBuffer = await this.gpu.createBuffer(Uint32Array.BYTES_PER_ELEMENT * input.length);
|
|
55
|
-
if (this.tokenBuffer.size !== Uint32Array.BYTES_PER_ELEMENT * input.length) {
|
|
56
|
-
this.tokenBuffer.destroy()
|
|
57
|
-
this.tokenBuffer = await this.gpu.createBuffer(Uint32Array.BYTES_PER_ELEMENT * input.length);
|
|
58
|
-
}
|
|
59
|
-
|
|
60
|
-
await this.gpu.WriteBuffer(this.tokenBuffer, new Uint32Array(input))
|
|
61
|
-
if (!this.outputBuffer) this.outputBuffer = await this.gpu.ManageBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length)
|
|
62
|
-
if (this.outputBuffer.mainBuffer.size !== Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length || this.outputBuffer.isDestroyed()) {
|
|
63
|
-
this.outputBuffer.destroy()
|
|
64
|
-
this.outputBuffer = await this.gpu.ManageBuffer(Float32Array.BYTES_PER_ELEMENT * this.embeddingSize * input.length);
|
|
65
|
-
}
|
|
66
|
-
const bind = await this.device.createBindGroup({
|
|
67
|
-
layout: await this.forwardPipline.getBindGroupLayout(0),
|
|
68
|
-
entries: [
|
|
69
|
-
{
|
|
70
|
-
binding: 0,
|
|
71
|
-
resource: {
|
|
72
|
-
buffer: this.vectorsBuffer
|
|
73
|
-
}
|
|
74
|
-
},
|
|
75
|
-
{
|
|
76
|
-
binding: 1,
|
|
77
|
-
resource: {
|
|
78
|
-
buffer: this.outputBuffer.mainBuffer
|
|
79
|
-
}
|
|
80
|
-
}, {
|
|
81
|
-
binding: 2,
|
|
82
|
-
resource: {
|
|
83
|
-
buffer: this.tokenBuffer
|
|
84
|
-
}
|
|
85
|
-
}
|
|
86
|
-
]
|
|
87
|
-
})
|
|
88
|
-
await this.gpu.RunPipline(encoder, this.forwardPipline, bind, Math.ceil((this.embeddingSize * input.length) / 256))
|
|
89
|
-
if(!encoder2) await this.device.queue.submit([encoder.finish()])
|
|
90
|
-
return this.outputBuffer
|
|
91
|
-
|
|
92
|
-
}
|
|
93
|
-
async Save() {
|
|
94
|
-
const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.vocabSize * this.embeddingSize, this.vectorsBuffer)
|
|
95
|
-
SaveVector(this.configs.save.filename, vec)
|
|
96
|
-
}
|
|
97
|
-
async backward(inputGradient, lr, en2) {
|
|
98
|
-
const encoder = en2 ? en2 : this.device.createCommandEncoder()
|
|
99
|
-
if (!this.oldParam) {
|
|
100
|
-
this.oldParam = new Float32Array([0, 0]);
|
|
101
|
-
this.param = await this.gpu.createUniFormBuffer(this.oldParam.byteLength);
|
|
102
|
-
}
|
|
103
|
-
if (this.oldParam[0] !== lr) {
|
|
104
|
-
this.oldParam = new Float32Array([lr, 0]);
|
|
105
|
-
this.gpu.WriteBuffer(this.param, this.oldParam)
|
|
106
|
-
}
|
|
107
|
-
const bind = await this.device.createBindGroup({
|
|
108
|
-
layout: await this.backwardPipline.getBindGroupLayout(0),
|
|
109
|
-
entries: [
|
|
110
|
-
{
|
|
111
|
-
binding: 0,
|
|
112
|
-
resource: {
|
|
113
|
-
buffer: this.vectorsBuffer
|
|
114
|
-
}
|
|
115
|
-
},
|
|
116
|
-
{
|
|
117
|
-
binding: 1,
|
|
118
|
-
resource: {
|
|
119
|
-
buffer: inputGradient.mainBuffer
|
|
120
|
-
}
|
|
121
|
-
}, {
|
|
122
|
-
binding: 2,
|
|
123
|
-
resource: {
|
|
124
|
-
buffer: this.tokenBuffer
|
|
125
|
-
}
|
|
126
|
-
|
|
127
|
-
}, {
|
|
128
|
-
binding: 3,
|
|
129
|
-
resource: {
|
|
130
|
-
buffer: this.param
|
|
131
|
-
}
|
|
132
|
-
|
|
133
|
-
}
|
|
134
|
-
]
|
|
135
|
-
})
|
|
136
|
-
await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * (this.tokenBuffer.size / 4)) / 256))
|
|
137
|
-
if(!en2) this.device.queue.submit([encoder.finish()])
|
|
138
|
-
}
|
|
139
|
-
|
|
140
|
-
}
|
package/src/FFN/index.d.ts
DELETED
|
@@ -1,23 +0,0 @@
|
|
|
1
|
-
export class FFN {
|
|
2
|
-
constructor(device: GPUAdapter, embeddingSize: number, configs: {
|
|
3
|
-
dirname: string
|
|
4
|
-
})
|
|
5
|
-
async forward(vec: GPUBuffer): Promise<{
|
|
6
|
-
read: function(): Promise<Float32Array>,
|
|
7
|
-
mainBuffer: GPUBuffer,
|
|
8
|
-
isDestroyed: function(): Boolean,
|
|
9
|
-
destroy: function(): void
|
|
10
|
-
}>
|
|
11
|
-
async backward(inputGradient: {
|
|
12
|
-
read: function(): Promise<Float32Array>,
|
|
13
|
-
mainBuffer: GPUBuffer,
|
|
14
|
-
isDestroyed: function(): Boolean,
|
|
15
|
-
destroy: function(): void
|
|
16
|
-
}, lr: number): Promise<{
|
|
17
|
-
read: function(): Promise<Float32Array>,
|
|
18
|
-
mainBuffer: GPUBuffer,
|
|
19
|
-
isDestroyed: function(): Boolean,
|
|
20
|
-
destroy: function(): void
|
|
21
|
-
}>
|
|
22
|
-
async Save(): Promise<void>
|
|
23
|
-
}
|
package/src/FFN/index.js
DELETED
|
@@ -1,45 +0,0 @@
|
|
|
1
|
-
import { RELU, Linear } from "../../index.js";
|
|
2
|
-
import { join } from 'path'
|
|
3
|
-
import { OperationManager } from "../OperationManager.js";
|
|
4
|
-
|
|
5
|
-
|
|
6
|
-
export class FFN {
|
|
7
|
-
linear1;
|
|
8
|
-
linear2;
|
|
9
|
-
hiddenLayer;
|
|
10
|
-
Op
|
|
11
|
-
constructor(device, embeddingSize, configs) {
|
|
12
|
-
this.Op = new OperationManager(device)
|
|
13
|
-
this.linear1 = new Linear(device, embeddingSize, embeddingSize * 4, {
|
|
14
|
-
save: {
|
|
15
|
-
filename: [join(configs.dirname, "FFN_weights1.bin"), join(configs.dirname, "FFN_bias1.bin")]
|
|
16
|
-
}
|
|
17
|
-
});
|
|
18
|
-
this.hiddenLayer = new RELU(device, embeddingSize * 4);
|
|
19
|
-
this.linear2 = new Linear(device, embeddingSize * 4, embeddingSize, {
|
|
20
|
-
save: {
|
|
21
|
-
filename: [join(configs.dirname, "FFN_weights2.bin"), join(configs.dirname, "FFN_bias2.bin")]
|
|
22
|
-
}
|
|
23
|
-
})
|
|
24
|
-
}
|
|
25
|
-
async forward(vec, encoderMain) {
|
|
26
|
-
const logit1 = await this.linear1.forward(vec, encoderMain)
|
|
27
|
-
const hidden = await this.hiddenLayer.forward(logit1, encoderMain)
|
|
28
|
-
const final = await this.linear2.forward(hidden, encoderMain)
|
|
29
|
-
return final
|
|
30
|
-
}
|
|
31
|
-
async backward(inputGradient, lr, encoder) {
|
|
32
|
-
const logit2 = await this.linear2.backward(inputGradient, lr, encoder)
|
|
33
|
-
const hidden = await this.hiddenLayer.backward(logit2,encoder)
|
|
34
|
-
const final = await this.linear1.backward(hidden, lr, encoder)
|
|
35
|
-
return final
|
|
36
|
-
}
|
|
37
|
-
async Save() {
|
|
38
|
-
await this.linear1.Save()
|
|
39
|
-
await this.linear2.Save()
|
|
40
|
-
}
|
|
41
|
-
async ClearInputCache() {
|
|
42
|
-
this.linear1.ClearInputCache()
|
|
43
|
-
this.linear2.ClearInputCache()
|
|
44
|
-
}
|
|
45
|
-
}
|
package/src/Linear/index.d.ts
DELETED
|
@@ -1,31 +0,0 @@
|
|
|
1
|
-
export class Linear {
|
|
2
|
-
constructor(device: GPUAdapter, embeddingSize: number, vocabSize: number, configs: {
|
|
3
|
-
save: {
|
|
4
|
-
filename: string[]
|
|
5
|
-
}
|
|
6
|
-
})
|
|
7
|
-
async forward(vec: {
|
|
8
|
-
read: function(): Promise<Float32Array>,
|
|
9
|
-
mainBuffer: GPUBuffer,
|
|
10
|
-
isDestroyed: function(): Boolean,
|
|
11
|
-
destroy: function(): void
|
|
12
|
-
}): Promise<{
|
|
13
|
-
read: function(): Promise<Float32Array>,
|
|
14
|
-
mainBuffer: GPUBuffer,
|
|
15
|
-
isDestroyed: function(): Boolean,
|
|
16
|
-
destroy: function(): void
|
|
17
|
-
}>
|
|
18
|
-
async backward(dOutput: {
|
|
19
|
-
read: function(): Promise<Float32Array>,
|
|
20
|
-
mainBuffer: GPUBuffer,
|
|
21
|
-
isDestroyed: function(): Boolean,
|
|
22
|
-
destroy: function(): void
|
|
23
|
-
}, lr: number): Promise<{
|
|
24
|
-
read: function(): Promise<Float32Array>,
|
|
25
|
-
mainBuffer: GPUBuffer,
|
|
26
|
-
isDestroyed: function(): Boolean,
|
|
27
|
-
destroy: function(): void
|
|
28
|
-
}>
|
|
29
|
-
async Save(): Promise<void>
|
|
30
|
-
|
|
31
|
-
}
|