@mineproness/tensorloop 2.2.1 → 2.2.3
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/index.js +0 -2
- package/package.json +1 -1
- package/src/Softmax/index.js +7 -71
package/index.js
CHANGED
|
@@ -6,7 +6,6 @@ import { FFN as ffn } from "./src/FFN/index.js";
|
|
|
6
6
|
import { Tokenizer as to } from './src/Tokenizer/index.js'
|
|
7
7
|
import {SelfAttention as ST} from './src/SelfAttention/index.js'
|
|
8
8
|
import {TransformerBlock as TT} from './src/TransformerBlock/index.js'
|
|
9
|
-
import { Model as m } from "./src/Model/index.js";
|
|
10
9
|
import {PositionEmbedding as PM } from './src/PositionEmbedding/index.js'
|
|
11
10
|
import { RMSNorm as RM } from "./src/RMSNorm/index.js";
|
|
12
11
|
export const Embedding = E
|
|
@@ -17,6 +16,5 @@ export const FFN = ffn
|
|
|
17
16
|
export const Tokenizer = to;
|
|
18
17
|
export const SelfAttention = ST
|
|
19
18
|
export const TransformerBlock = TT;
|
|
20
|
-
export const Model = m;
|
|
21
19
|
export const PositionEmbedding = PM;
|
|
22
20
|
export const RMSNorm = RM
|
package/package.json
CHANGED
package/src/Softmax/index.js
CHANGED
|
@@ -116,78 +116,14 @@ export class Softmax {
|
|
|
116
116
|
|
|
117
117
|
const tokenCount = probs.size / (this.vocabSize * 4)
|
|
118
118
|
const targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
|
|
119
|
-
|
|
120
|
-
|
|
119
|
+
if(!this.loss) this.loss = await this.gpu.createBuffer(4 * tokenCount)
|
|
120
|
+
if(!this.acc) this.acc = await this.gpu.createBuffer(4 * 1)
|
|
121
121
|
await this.gpu.WriteBuffer(targetBuffer, new Uint32Array(target))
|
|
122
|
-
const
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
resource: {
|
|
128
|
-
buffer: probs.mainBuffer
|
|
129
|
-
}
|
|
130
|
-
},
|
|
131
|
-
{
|
|
132
|
-
binding: 1,
|
|
133
|
-
resource: {
|
|
134
|
-
buffer: targetBuffer
|
|
135
|
-
}
|
|
136
|
-
}, {
|
|
137
|
-
binding: 2,
|
|
138
|
-
resource: {
|
|
139
|
-
buffer: loss
|
|
140
|
-
}
|
|
141
|
-
}, {
|
|
142
|
-
binding: 3,
|
|
143
|
-
resource: {
|
|
144
|
-
buffer: this.paramsForward
|
|
145
|
-
}
|
|
146
|
-
|
|
147
|
-
},
|
|
148
|
-
]
|
|
149
|
-
})
|
|
150
|
-
await this.gpu.RunPipline(encoder, this.lossPipline, bind, Math.ceil((tokenCount / 4)))
|
|
151
|
-
const outputLoss = await this.gpu.createBuffer(4)
|
|
152
|
-
const reduceBind = await this.gpu.createBindGroup(this.reduceFn, 0, [loss, outputLoss, this.paramsForward])
|
|
153
|
-
await this.gpu.RunPipline(encoder, this.reduceFn, reduceBind, 1);
|
|
154
|
-
const bind2 = await this.device.createBindGroup({
|
|
155
|
-
layout: await this.accPipline.getBindGroupLayout(0),
|
|
156
|
-
entries: [
|
|
157
|
-
{
|
|
158
|
-
binding: 0,
|
|
159
|
-
resource: {
|
|
160
|
-
buffer: this.paramsForward
|
|
161
|
-
}
|
|
162
|
-
},
|
|
163
|
-
{
|
|
164
|
-
binding: 1,
|
|
165
|
-
resource: {
|
|
166
|
-
buffer: probs.mainBuffer
|
|
167
|
-
}
|
|
168
|
-
}, {
|
|
169
|
-
binding: 3,
|
|
170
|
-
resource: {
|
|
171
|
-
buffer: targetBuffer
|
|
172
|
-
}
|
|
173
|
-
},
|
|
174
|
-
{
|
|
175
|
-
binding: 2,
|
|
176
|
-
resource: {
|
|
177
|
-
buffer: acc
|
|
178
|
-
}
|
|
179
|
-
}
|
|
180
|
-
]
|
|
181
|
-
})
|
|
182
|
-
await this.gpu.RunPipline(encoder, this.accPipline, bind2, 1)
|
|
183
|
-
this.device.queue.submit([encoder.finish()])
|
|
184
|
-
targetBuffer.destroy()
|
|
185
|
-
loss.destroy()
|
|
186
|
-
const lossCpu = await this.gpu.readBuffer(outputLoss.size, outputLoss)
|
|
187
|
-
const accCPU = await this.gpu.readBuffer(acc.size, acc)
|
|
188
|
-
acc.destroy()
|
|
189
|
-
outputLoss.destroy()
|
|
190
|
-
return { loss: lossCpu[0], acc: accCPU[0] }
|
|
122
|
+
const bindGroup1 = await this.gpu.createBindGroup(this.lossPipline, 0, [probs.mainBuffer, targetBuffer, this.loss])
|
|
123
|
+
const bindGroup2 = await this.gpu.createBindGroup(this.accPipline, 0 , [this.paramsForward, probs.mainBuffer, targetBuffer, this.acc])
|
|
124
|
+
await this.gpu.RunPipline(encoder, this.lossPipline, bindGroup1, Math.ceil((tokenCount / 4)))
|
|
125
|
+
await this.gpu.RunPipline(encoder, this.accPipline, bindGroup2, 1)
|
|
126
|
+
return {loss: await this.gpu.readBuffer(this.loss.size, this.loss), acc: (await this.gpu.readBuffer(this.acc.size, this.acc))[0]}
|
|
191
127
|
}
|
|
192
128
|
|
|
193
129
|
}
|