freywar/nn
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add residuals

Browse Source
This commit is contained in:
freywar
2024-07-27 14:39:08 +00:00
parent def3a479be
commit dda7369ee7
17 changed files with 366 additions and 271 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/data/gpu/gpu-struct.ts
+2
View File
@@ -151,6 +151,8 @@ export abstract class GPUStruct<A extends Uint32Array | Int32Array | Float32Arra
this._op = null;
}
public abstract get(): Promise<unknown>;
public buffer(): GPUBuffer {
if (this._dirty === 'front') {
this.write();
src/data/gpu/matrix-batch.ts
+20
View File
@@ -9,6 +9,7 @@ const SHADER_SCALE = fs.readFileSync('shader/bscale.wgsl').toString();
const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString();
const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString();
const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString();
const SHADER_TMUL = fs.readFileSync('shader/tmul.wgsl').toString();
const SHADER_AVG = fs.readFileSync('shader/avg.wgsl').toString();
const SHADER_TOT = fs.readFileSync('shader/tot.wgsl').toString();
@@ -253,6 +254,25 @@ export class MatrixBatch<
return result.calculate('transpose', SHADER_TRANS, [m]);
}
public static tmul<
L extends number,
H extends number,
S extends number,
W extends number,
A extends Uint32Array | Int32Array | Float32Array,
T extends NonNullable<ReturnType<A['at']>>,
>(
l: MatrixBatch<L, H, S, A, T>,
tr: MatrixBatch<L, W, S, A, T>,
result: MatrixBatch<L, H, W, A, T> = new MatrixBatch(l._device, l._array, l.layers, l.height, tr.height),
): MatrixBatch<L, H, W, A, T> {
this.assertUnique(l, tr, result);
this.assertColocated(l, tr, result);
this.assertPackable(1, l, tr, result);
return result.calculate('tmul', SHADER_TMUL, [l, tr]);
}
public static avg<
L extends number,
H extends number,
src/data/gpu/matrix.ts
+19
View File
@@ -8,6 +8,7 @@ import { Vector } from './vector';
const SHADER_ISUM = fs.readFileSync('shader/isum.wgsl').toString();
const SHADER_MUL = fs.readFileSync('shader/mul.wgsl').toString();
const SHADER_TRANS = fs.readFileSync('shader/transpose.wgsl').toString();
const SHADER_TMUL = fs.readFileSync('shader/tmul.wgsl').toString();
const SHADER_UNAVG = fs.readFileSync('shader/unavg.wgsl').toString();
export class Matrix<
@@ -197,6 +198,24 @@ export class Matrix<
return result.calculate('mul', SHADER_MUL, [l, r]);
}
public static tmul<
H extends number,
S extends number,
W extends number,
A extends Uint32Array | Int32Array | Float32Array,
T extends NonNullable<ReturnType<A['at']>>,
>(
l: Matrix<H, S, A, T>,
tr: Matrix<W, S, A, T>,
result: Matrix<H, W, A, T> = new Matrix(l._device, l._array, l.height, tr.height),
): Matrix<H, W, A, T> {
this.assertUnique(l, tr, result);
this.assertColocated(l, tr, result);
this.assertPackable(1, l, tr, result);
return result.calculate('tmul', SHADER_TMUL, [l, tr]);
}
public static transpose<
H extends number,
W extends number,
src/model-gpu.ts
+166 -62
View File
@@ -1,26 +1,27 @@
import fs from 'node:fs';
import { Writable } from 'node:stream';
import { GPUStruct } from './data/gpu/gpu-struct';
import { Matrix } from './data/gpu/matrix';
import { Scalar } from './data/gpu/scalar';
import { Vector } from './data/gpu/vector';
import { Matrix as CPUMatrix, Vector as CPUVector, drm, mat, randseq, slope, stddev, vec } from './data/math';
import { TID, Tokenizer } from './tokenizer';
const SHADER_HOTONES = fs.readFileSync('shader/hotones.wgsl').toString();
const SHADER_ONE_HOTS = fs.readFileSync('shader/onehots.wgsl').toString();
const SHADER_EMBED = fs.readFileSync('shader/embed.wgsl').toString();
const SHADER_ATTW = fs.readFileSync('shader/attw.wgsl').toString();
const SHADER_CENTER = fs.readFileSync('shader/center.wgsl').toString();
const SHADER_NORMALIZE = fs.readFileSync('shader/normalize.wgsl').toString();
const SHADER_TWEIGH = fs.readFileSync('shader/tweigh.wgsl').toString();
const SHADER_SOFTMAX_EXP = fs.readFileSync('shader/softmaxexp.wgsl').toString();
const SHADER_SOFTMAX_NORM = fs.readFileSync('shader/softmaxnorm.wgsl').toString();
const SHADER_DLDO = fs.readFileSync('shader/dldo.wgsl').toString();
const SHADER_OTDLDO = fs.readFileSync('shader/otdldo.wgsl').toString();
const SHADER_TOTDLDOWA = fs.readFileSync('shader/totdldowa.wgsl').toString();
const SHADER_DLDS = fs.readFileSync('shader/dlds.wgsl').toString();
const SHADER_GRAD_TOUT = fs.readFileSync('shader/gradtout.wgsl').toString();
const SHADER_GRAD_TATTW = fs.readFileSync('shader/gradtattw.wgsl').toString();
const SHADER_GRAD_INP = fs.readFileSync('shader/gradinp.wgsl').toString();
const SHADER_SOFTMAX_GRAD = fs.readFileSync('shader/softmaxgrad.wgsl').toString();
const SHADER_LAYERNORM_GRAD = fs.readFileSync('shader/layernormgrad.wgsl').toString();
const SHADER_ARGMAX = fs.readFileSync('shader/argmax.wgsl').toString();
const SHADER_SAMPLES = fs.readFileSync('shader/samples.wgsl').toString();
const LOGGING = false;
export interface ModelInitData<
V extends number,
W extends number,
@@ -126,14 +127,16 @@ export class Model<
readonly itids: Vector<W, Int32Array, TID>;
readonly ilogits: Matrix<W, V, Float32Array>;
readonly e: Matrix<W, D, Float32Array>;
readonly ce: Matrix<W, D, Float32Array>;
readonly ne: Matrix<W, D, Float32Array>;
readonly q: Matrix<W, D, Float32Array>;
readonly k: Matrix<W, D, Float32Array>;
readonly tk: Matrix<D, W, Float32Array>;
readonly qtk: Matrix<W, W, Float32Array>;
readonly v: Matrix<W, D, Float32Array>;
readonly s: Matrix<W, W, Float32Array>;
readonly se: Matrix<W, W, Float32Array>;
readonly wa: Matrix<W, W, Float32Array>;
readonly ar: Matrix<W, D, Float32Array>;
readonly a: Matrix<W, D, Float32Array>;
readonly o: Matrix<W, V, Float32Array>;
readonly temp: Scalar<Float32Array>;
@@ -156,6 +159,12 @@ export class Model<
readonly ttids: Vector<W, Int32Array, TID>;
};
protected async log<A extends Uint32Array | Int32Array | Float32Array>(name: string, s: GPUStruct<A>, sample: number = 4): Promise<void> {
if (LOGGING) {
console.log(name, (await s.get() as number[]).slice(0, sample));
}
}
protected softmax<H extends number, W extends number>(
m: Matrix<H, W, Float32Array>,
e: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
@@ -166,28 +175,75 @@ export class Model<
return r;
}
protected softmax_grad<H extends number, W extends number>(
m: Matrix<H, W, Float32Array>,
dldm: Matrix<H, W, Float32Array>,
result: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
): Matrix<H, W, Float32Array> {
result.calculate('softmaxgrad', SHADER_SOFTMAX_GRAD, [m, dldm]);
return result;
}
protected async layernorm<H extends number, W extends number>(
m: Matrix<H, W, Float32Array>,
c: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
r: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
): Promise<Matrix<H, W, Float32Array>> {
c.calculate('center', SHADER_CENTER, [m], false, 4);
r.calculate('normalize', SHADER_NORMALIZE, [m], false, 4);
await this.log('layernorm c', c);
await this.log('layernorm r', r);
return r;
}
protected layernorm_grad<H extends number, W extends number>(
m: Matrix<H, W, Float32Array>,
dldm: Matrix<H, W, Float32Array>,
result: Matrix<H, W, Float32Array> = new Matrix(this.device, Float32Array, m.height, m.width),
): Matrix<H, W, Float32Array> {
result.calculate('layernormgrad', SHADER_LAYERNORM_GRAD, [m, dldm]);
return result;
}
protected async ilogits(itids: Vector<W, Int32Array, TID> = this.processing.itids): Promise<Matrix<W, V, Float32Array>> {
this.processing.ilogits.calculate('hotones', SHADER_HOTONES, [itids]);
this.processing.ilogits.calculate('onehots', SHADER_ONE_HOTS, [itids]);
await this.log('itids', itids);
await this.log('ilogits', this.processing.ilogits);
return this.processing.ilogits;
}
protected async e(ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits): Promise<Matrix<W, D, Float32Array>> {
this.processing.e.calculate('embed', SHADER_EMBED, [ilogits, this.weights.inp, this.weights.pos]);
await this.log('inp', this.weights.inp);
await this.log('pos', this.weights.pos);
await this.log('e', this.processing.e);
return this.processing.e;
}
protected q(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
this.processing.q.calculate('attw', SHADER_ATTW, [e, this.weights.tqry]);
protected async ne(e: Matrix<W, D, Float32Array> = this.processing.e): Promise<Matrix<W, D, Float32Array>> {
this.layernorm(e, this.processing.ce, this.processing.ne);
await this.log('ne', this.processing.ne);
return this.processing.ne;
}
protected async q(ne: Matrix<W, D, Float32Array> = this.processing.ne): Promise<Matrix<W, D, Float32Array>> {
this.processing.q.calculate('tqry', SHADER_TWEIGH, [ne, this.weights.tqry]);
await this.log('tqry', this.weights.tqry);
await this.log('q', this.processing.q);
return this.processing.q;
}
protected k(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
this.processing.k.calculate('attw', SHADER_ATTW, [e, this.weights.tkey]);
protected async k(ne: Matrix<W, D, Float32Array> = this.processing.ne): Promise<Matrix<W, D, Float32Array>> {
this.processing.k.calculate('tkey', SHADER_TWEIGH, [ne, this.weights.tkey]);
await this.log('tkey', this.weights.tkey);
await this.log('k', this.processing.k);
return this.processing.k;
}
protected v(e: Matrix<W, D, Float32Array> = this.processing.e): Matrix<W, D, Float32Array> {
this.processing.v.calculate('attw', SHADER_ATTW, [e, this.weights.tval]);
protected async v(ne: Matrix<W, D, Float32Array> = this.processing.ne): Promise<Matrix<W, D, Float32Array>> {
this.processing.v.calculate('tval', SHADER_TWEIGH, [ne, this.weights.tval]);
await this.log('tval', this.weights.tval);
await this.log('v', this.processing.v);
return this.processing.v;
}
@@ -195,8 +251,7 @@ export class Model<
q: Matrix<W, D, Float32Array> = this.processing.q,
k: Matrix<W, D, Float32Array> = this.processing.k,
): Matrix<W, W, Float32Array> {
Matrix.transpose(k, this.processing.tk);
Matrix.mul(q, this.processing.tk, this.processing.qtk);
Matrix.tmul(q, k, this.processing.qtk);
Matrix.scale(this.processing.ds, this.processing.qtk, this.processing.s);
return this.processing.s;
}
@@ -207,15 +262,17 @@ export class Model<
}
protected a(
e: Matrix<W, D, Float32Array> = this.processing.e,
wa: Matrix<W, W, Float32Array> = this.processing.wa,
v: Matrix<W, D, Float32Array> = this.processing.v,
): Matrix<W, D, Float32Array> {
Matrix.mul(wa, v, this.processing.a);
Matrix.mul(wa, v, this.processing.ar);
Matrix.sum(e, this.processing.ar, this.processing.a);
return this.processing.a;
}
protected o(a: Matrix<W, D, Float32Array> = this.processing.a): Matrix<W, V, Float32Array> {
this.processing.o.calculate('attw', SHADER_ATTW, [a, this.weights.tout]);
this.processing.o.calculate('tout', SHADER_TWEIGH, [a, this.weights.tout]);
return this.processing.o;
}
@@ -252,7 +309,7 @@ export class Model<
protected async upd(
ilogits: Matrix<W, V, Float32Array> = this.processing.ilogits,
e: Matrix<W, D, Float32Array> = this.processing.e,
ne: Matrix<W, D, Float32Array> = this.processing.ne,
q: Matrix<W, D, Float32Array> = this.processing.q,
k: Matrix<W, D, Float32Array> = this.processing.k,
v: Matrix<W, D, Float32Array> = this.processing.v,
@@ -270,39 +327,80 @@ export class Model<
this.processing.lr,
]);
this.processing.otdldo.calculate('otdldo', SHADER_OTDLDO, [this.processing.dldo, this.weights.tout]);
this.processing.dlds.calculate('dlds', SHADER_DLDS, [
this.processing.ds,
this.softmax_grad(
wa,
v,
this.processing.otdldo,
]);
Matrix.mul(this.processing.dldo, Matrix.mul(this.weights.tout, Matrix.transpose(v))),
this.processing.dlds,
);
Matrix.mul(this.processing.dlds, k, this.processing.dldsk);
Matrix.mul(Matrix.transpose(this.processing.dlds), q, this.processing.dldsq);
Matrix.mul(
Matrix.transpose(ilogits),
Matrix.sum(
Matrix.mul(this.processing.dldo, this.weights.tout),
this.layernorm_grad(
ne,
Matrix.sum(
Matrix.mul(
Matrix.mul(
this.processing.dlds,
Matrix.scale(this.processing.ds, k),
),
this.weights.tqry,
)
,
Matrix.sum(
Matrix.mul(
Matrix.mul(
Matrix.transpose(this.processing.dlds),
Matrix.scale(this.processing.ds, q),
),
this.weights.tkey,
),
Matrix.mul(
Matrix.mul(Matrix.transpose(wa), this.processing.dldo),
Matrix.mul(this.weights.tout, this.weights.tval),
),
),
),
),
),
this.processing.ginp,
);
Matrix.transpose(
Matrix.scale(this.processing.ds, Matrix.mul(
Matrix.transpose(ne),
Matrix.mul(this.processing.dlds, k),
)),
this.processing.gtqry,
);
Matrix.transpose(
Matrix.scale(this.processing.ds, Matrix.mul(
Matrix.transpose(ne),
Matrix.mul(Matrix.transpose(this.processing.dlds), q),
)),
this.processing.gtkey,
);
Matrix.transpose(
Matrix.mul(
Matrix.mul(Matrix.transpose(ne), Matrix.transpose(wa)),
Matrix.mul(this.processing.dldo, this.weights.tout),
),
this.processing.gtval,
);
Matrix.transpose(
Matrix.mul(Matrix.transpose(a), this.processing.dldo),
this.processing.gtout,
);
this.processing.ginp.calculate('gradinp', SHADER_GRAD_INP, [
ilogits,
this.processing.dldsq,
this.processing.dldsk,
this.processing.totdldowa,
this.weights.tqry,
this.weights.tkey,
this.weights.tval,
]);
Matrix.sum(this.processing.ginp, this.weights.inp, this.weights.inp);
this.processing.gtout.calculate('gradtout', SHADER_GRAD_TOUT, [a, this.processing.dldo]);
Matrix.sum(this.processing.gtout, this.weights.tout, this.weights.tout);
this.processing.gtqry.calculate('gradtqry', SHADER_GRAD_TATTW, [e, this.processing.dldsk]);
Matrix.sum(this.processing.gtqry, this.weights.tqry, this.weights.tqry);
this.processing.gtkey.calculate('gradtkey', SHADER_GRAD_TATTW, [e, this.processing.dldsq]);
Matrix.sum(this.processing.gtkey, this.weights.tkey, this.weights.tkey);
this.processing.totdldowa.calculate('totdldowa', SHADER_TOTDLDOWA, [wa, this.processing.otdldo]);
this.processing.gtval.calculate('gradtval', SHADER_GRAD_TATTW, [e, this.processing.totdldowa]);
Matrix.sum(this.processing.gtval, this.weights.tval, this.weights.tval);
Matrix.sum(this.processing.gtout, this.weights.tout, this.weights.tout);
}
public constructor(
@@ -349,15 +447,17 @@ export class Model<
itids: new Vector<W, Int32Array, TID>(this.device, Int32Array, this.ws),
ilogits: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
e: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
ce: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
ne: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
q: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
k: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
tk: new Matrix<D, W, Float32Array>(this.device, Float32Array, this.ds, this.ws),
qtk: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
v: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
s: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
se: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
wa: new Matrix<W, W, Float32Array>(this.device, Float32Array, this.ws, this.ws),
a: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
ar: new Matrix<W, D, Float32Array>(this.device, Float32Array, this.ws, this.ds),
o: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
temp: new Scalar<Float32Array>(this.device, Float32Array, 1),
tempo: new Matrix<W, V, Float32Array>(this.device, Float32Array, this.ws, this.vs),
@@ -458,6 +558,8 @@ export class Model<
const tids = this.tokenizer.tokenize(data);
const batches = tids.length - this.ws - 1;
const seed = tids.slice(tids.length / 2, tids.length / 2 + 100);
let lr = 0.003;
for (let epoch = 0; epoch < epochs; epoch++) {
@@ -471,15 +573,16 @@ export class Model<
const ilogits = await this.ilogits(this.processing.itids);
const e = await this.e(ilogits);
const q = this.q(e);
const k = this.k(e);
const v = this.v(e);
const ne = await this.ne(e);
const q = await this.q(ne);
const k = await this.k(ne);
const v = await this.v(ne);
const s = this.s(q, k);
const wa = await this.wa(s);
const a = this.a(wa, v);
const a = this.a(e, wa, v);
const o = this.o(a);
const ologits = await this.ologits(o);
await this.upd(ilogits, e, q, k, v, wa, a, ologits, this.processing.ttids, lr);
await this.upd(ilogits, ne, q, k, v, wa, a, ologits, this.processing.ttids, lr);
const queued: number = Date.now();
@@ -567,19 +670,19 @@ export class Model<
samples?.write(`Epoch ${epoch + 1}:\n`);
samples?.write('Sample (temp=0.3):\n');
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.3 })).join(''));
samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 0.3 })).join(''));
samples?.write('\nEnd of sample.\n');
samples?.write('Sample (temp=0.7):\n');
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 0.7 })).join(''));
samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 0.7 })).join(''));
samples?.write('\nEnd of sample.\n');
samples?.write('Sample (temp=1):\n');
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1 })).join(''));
samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 1 })).join(''));
samples?.write('\nEnd of sample.\n');
samples?.write('Sample (temp=1.3):\n');
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 1.3 })).join(''));
samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 1.3 })).join(''));
samples?.write('\nEnd of sample.\n\n');
samples?.write('Sample (argmax):\n');
samples?.write(this.tokenizer.detokenize(await this.run(this.tokenizer.tokenize('Hello?'), { temp: 'max' })).join(''));
samples?.write(this.tokenizer.detokenize(await this.run(seed, { temp: 'max' })).join(''));
samples?.write('\nEnd of sample.\n\n');
snapshotted = end;
@@ -606,12 +709,13 @@ export class Model<
await this.processing.ttids.set([...itids.slice(1), 0]);
const ilogits = await this.ilogits(this.processing.itids);
const e = await this.e(ilogits);
const q = this.q(e);
const k = this.k(e);
const v = this.v(e);
const ne = await this.ne(e);
const q = await this.q(ne);
const k = await this.k(ne);
const v = await this.v(ne);
const s = this.s(q, k);
const wa = await this.wa(s);
const a = this.a(wa, v);
const a = this.a(e, wa, v);
const o = this.o(a);
const ologits = await this.ologits(o, temp === 'max' ? 1 : temp);
const out = temp === 'max' ? this.otids(ologits) : this.otids(await this.rnd(), ologits);