@imaduddin/fastbrain 1.1.1 → 2.0.1

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.
Files changed (115) hide show
  1. package/dist/Embedding/index.d.ts +24 -0
  2. package/dist/Embedding/index.js +93 -0
  3. package/dist/FFN/index.d.ts +20 -0
  4. package/dist/FFN/index.js +51 -0
  5. package/dist/Gelu/index.d.ts +15 -0
  6. package/dist/Gelu/index.js +58 -0
  7. package/dist/Linear/index.d.ts +31 -0
  8. package/dist/Linear/index.js +136 -0
  9. package/dist/RMSNorm/index.d.ts +16 -0
  10. package/dist/RMSNorm/index.js +63 -0
  11. package/dist/Relu/index.d.ts +15 -0
  12. package/dist/Relu/index.js +58 -0
  13. package/dist/SelfAttention/index.d.ts +47 -0
  14. package/dist/SelfAttention/index.js +186 -0
  15. package/dist/Softmax/index.d.ts +18 -0
  16. package/dist/Softmax/index.js +96 -0
  17. package/dist/Tokenizer/index.d.ts +15 -0
  18. package/dist/Tokenizer/index.js +58 -0
  19. package/dist/TransformerBlock/index.d.ts +21 -0
  20. package/dist/TransformerBlock/index.js +43 -0
  21. package/dist/fileManager.d.ts +6 -0
  22. package/dist/fileManager.js +34 -0
  23. package/dist/index.d.ts +10 -0
  24. package/dist/index.js +10 -0
  25. package/dist/residualConnection/index.d.ts +9 -0
  26. package/dist/residualConnection/index.js +25 -0
  27. package/package.json +20 -3
  28. package/shaders/add.wgsl +22 -0
  29. package/shaders/embedding/backward.wgsl +36 -0
  30. package/shaders/embedding/forward.wgsl +35 -0
  31. package/shaders/linear/ComputeDinput.wgsl +72 -0
  32. package/shaders/linear/forward.wgsl +63 -0
  33. package/shaders/linear/updateBias.wgsl +33 -0
  34. package/shaders/linear/updateWeights.wgsl +38 -0
  35. package/shaders/relu/backward.wgsl +43 -0
  36. package/shaders/relu/forward.wgsl +28 -0
  37. package/shaders/rmsnorm/backward.wgsl +54 -0
  38. package/shaders/rmsnorm/forward.wgsl +51 -0
  39. package/shaders/selfAttention/SB.wgsl +30 -0
  40. package/shaders/selfAttention/add3.wgsl +27 -0
  41. package/shaders/selfAttention/matMul.wgsl +63 -0
  42. package/shaders/selfAttention/qK.wgsl +66 -0
  43. package/shaders/selfAttention/transpose.wgsl +24 -0
  44. package/shaders/softmax/backward.wgsl +36 -0
  45. package/shaders/softmax/crossEntropy.wgsl +36 -0
  46. package/shaders/softmax/forward.wgsl +68 -0
  47. package/src/Embedding/index.ts +81 -0
  48. package/src/FFN/index.ts +47 -0
  49. package/src/Gelu/index.ts +53 -0
  50. package/src/Linear/index.ts +122 -0
  51. package/src/RMSNorm/index.ts +56 -0
  52. package/src/SelfAttention/index.ts +177 -0
  53. package/src/Softmax/index.ts +82 -0
  54. package/src/Tokenizer/index.ts +60 -0
  55. package/src/TransformerBlock/index.ts +41 -0
  56. package/src/fileManager.ts +46 -0
  57. package/src/index.ts +10 -0
  58. package/src/residualConnection/index.ts +24 -0
  59. package/tsconfig.json +27 -0
  60. package/index.js +0 -22
  61. package/src/Embedding/index.d.ts +0 -21
  62. package/src/Embedding/index.js +0 -140
  63. package/src/FFN/index.d.ts +0 -23
  64. package/src/FFN/index.js +0 -45
  65. package/src/Linear/index.d.ts +0 -31
  66. package/src/Linear/index.js +0 -196
  67. package/src/OperationManager.js +0 -169
  68. package/src/PositionEmbedding/index.d.ts +0 -30
  69. package/src/PositionEmbedding/index.js +0 -70
  70. package/src/RMSNorm/index.d.ts +0 -30
  71. package/src/RMSNorm/index.js +0 -91
  72. package/src/Relu/index.d.ts +0 -25
  73. package/src/Relu/index.js +0 -108
  74. package/src/SelfAttention/index.d.ts +0 -28
  75. package/src/SelfAttention/index.js +0 -138
  76. package/src/Softmax/index.d.ts +0 -26
  77. package/src/Softmax/index.js +0 -195
  78. package/src/Tokenizer/index.d.ts +0 -12
  79. package/src/Tokenizer/index.js +0 -164
  80. package/src/Tokenizer/vocabs.json +0 -19618
  81. package/src/TransformerBlock/index.d.ts +0 -28
  82. package/src/TransformerBlock/index.js +0 -78
  83. package/src/getVec.js +0 -30
  84. package/src/gpu.js +0 -0
  85. package/src/model/index.d.ts +0 -43
  86. package/src/model/index.js +0 -83
  87. package/src/residualConnection/index.js +0 -31
  88. package/src/shaders/CalculateAcc.wgsl +0 -50
  89. package/src/shaders/CalculateLoss.wgsl +0 -33
  90. package/src/shaders/ComputeDInput.wgsl +0 -46
  91. package/src/shaders/EB.wgsl +0 -33
  92. package/src/shaders/EF.wgsl +0 -25
  93. package/src/shaders/LB.wgsl +0 -71
  94. package/src/shaders/LF.wgsl +0 -60
  95. package/src/shaders/Norm/rmsBackward.wgsl +0 -65
  96. package/src/shaders/Norm/rmsForward.wgsl +0 -30
  97. package/src/shaders/Norm/sumSquare.wgsl +0 -38
  98. package/src/shaders/Norm/updateGamma.wgsl +0 -54
  99. package/src/shaders/PB.wgsl +0 -25
  100. package/src/shaders/PF.wgsl +0 -24
  101. package/src/shaders/RB.wgsl +0 -37
  102. package/src/shaders/RF.wgsl +0 -33
  103. package/src/shaders/ReductionFunction.wgsl +0 -28
  104. package/src/shaders/SB.wgsl +0 -41
  105. package/src/shaders/SF.wgsl +0 -75
  106. package/src/shaders/SoftmaxBackward.wgsl +0 -31
  107. package/src/shaders/add.wgsl +0 -30
  108. package/src/shaders/attentionShader/CompueFinalForwardOutput.wgsl +0 -29
  109. package/src/shaders/attentionShader/ComputeDA.wgsl +0 -30
  110. package/src/shaders/attentionShader/ComputeDV.wgsl +0 -37
  111. package/src/shaders/attentionShader/ComputeK.wgsl +0 -30
  112. package/src/shaders/attentionShader/ComputedQ.wgsl +0 -31
  113. package/src/shaders/attentionShader/attentionSocres.wgsl +0 -40
  114. package/src/shaders/attentionShader/sumGradient.wgsl +0 -24
  115. 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
@@ -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
- }
@@ -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
- }
@@ -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
- }
@@ -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
- }