@genai-fi/nanogpt 1.4.1 → 1.5.0

package/dist/ops/webgpu/normRMS.js

@@ -7,12 +7,12 @@ import s from "./normRMS32_program.js";
 import c from "./normRMS16_program.js";
 //#region lib/ops/webgpu/normRMS.ts
 function l(t) {
-	let { x: l, gamma: u } = t.inputs, d = t.backend, f = o(d), p = n(l), m = u ? n(u) : !1, h = p || m, g = !h || p ? l : r(l), _ = u && (!h || m ? u : r(u)), v = _ ? [g, _] : [g], y = i(v, -1), b = h ? new c(f, y, !!u) : new s(f, y, !!u);
-	if (_ && e(_.shape, [g.shape[g.shape.length - 1]], "Error in RMSNorm: "), l.shape.length !== 3) throw Error(`rmsNormGPU: input rank ${l.shape.length} not supported, only rank 3 is supported`);
-	if (y.inSize !== g.shape[g.shape.length - 1]) throw Error(`rmsNormGPU: reduction size ${y.inSize} does not match expected size ${g.shape[g.shape.length - 1]}`);
-	if (y.batchSize !== l.shape[0] * l.shape[1]) throw Error(`rmsNormGPU: batch size ${y.batchSize} does not match expected size ${l.shape[0] * l.shape[1]}`);
-	let x = a(b, v, d);
-	return h && !p && g.dispose(), h && !m && _ && _.dispose(), x;
+	let { x: l, gamma: u } = t.inputs, d = t.backend, f = t.attrs?.epsilon ?? 1e-8, p = o(d), m = n(l), h = u ? n(u) : !1, g = m || h, _ = !g || m ? l : r(l), v = u && (!g || h ? u : r(u)), y = v ? [_, v] : [_], b = i(y, -1), x = g ? new c(p, b, !!u, f) : new s(p, b, !!u, f);
+	if (v && e(v.shape, [_.shape[_.shape.length - 1]], "Error in RMSNorm: "), l.shape.length !== 3) throw Error(`rmsNormGPU: input rank ${l.shape.length} not supported, only rank 3 is supported`);
+	if (b.inSize !== _.shape[_.shape.length - 1]) throw Error(`rmsNormGPU: reduction size ${b.inSize} does not match expected size ${_.shape[_.shape.length - 1]}`);
+	if (b.batchSize !== l.shape[0] * l.shape[1]) throw Error(`rmsNormGPU: batch size ${b.batchSize} does not match expected size ${l.shape[0] * l.shape[1]}`);
+	let S = a(x, y, d);
+	return g && !m && _.dispose(), g && !h && v && v.dispose(), S;
 }
 t({
 	kernelName: "RMSNorm",

package/dist/ops/webgpu/normRMS16_program.d.ts

@@ -2,8 +2,9 @@ import { backend_util } from '@tensorflow/tfjs-core';
 import { ReduceProgram } from './utils/reductions';
 import { DeviceInformation } from './utils/deviceInfo';
 export default class RMSProgram16 extends ReduceProgram {
+    private epsilon;
     private hasGamma;
-    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, hasGamma?: boolean);
+    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, hasGamma?: boolean, epsilon?: number);
     getPreprocessSnippet(): string;
     getPostprocessSnippet(): string;
     getWriteSnippet(): string;

package/dist/ops/webgpu/normRMS16_program.js

@@ -1,18 +1,19 @@
 import { ReduceProgram as e } from "./utils/reductions.js";
 //#region lib/ops/webgpu/normRMS16_program.ts
 var t = class extends e {
+	epsilon;
 	hasGamma;
-	constructor(e, t, n = !1) {
+	constructor(e, t, n = !1, r = 1e-8) {
 		super(e, t, {
 			reductionOp: "mean",
 			elementwise: !0
-		}, !0), this.shaderKey = "RMSNorm16", n && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = n;
+		}, !0), this.epsilon = r, this.shaderKey = `RMSNorm16_${n ? "Gamma" : "NoGamma"}_${this.epsilon}`, n && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = n;
 	}
 	getPreprocessSnippet() {
 		return "candidate = candidate * candidate;";
 	}
 	getPostprocessSnippet() {
-		return "bestValue = inverseSqrt(bestValue + 1e-8);";
+		return `bestValue = inverseSqrt(bestValue + ${this.epsilon});`;
 	}
 	getWriteSnippet() {
 		return `

package/dist/ops/webgpu/normRMS32_program.d.ts

@@ -2,8 +2,9 @@ import { backend_util } from '@tensorflow/tfjs-core';
 import { ReduceProgram } from './utils/reductions';
 import { DeviceInformation } from './utils/deviceInfo';
 export default class RMSProgram32 extends ReduceProgram {
+    private epsilon;
     private hasGamma;
-    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, hasGamma?: boolean);
+    constructor(deviceInfo: DeviceInformation, reduceInfo: backend_util.ReduceInfo, hasGamma?: boolean, epsilon?: number);
     protected getPreprocessSnippet(): string;
     protected getPostprocessSnippet(): string;
     protected getWriteSnippet(): string;

package/dist/ops/webgpu/normRMS32_program.js

@@ -1,18 +1,19 @@
 import { ReduceProgram as e } from "./utils/reductions.js";
 //#region lib/ops/webgpu/normRMS32_program.ts
 var t = class extends e {
+	epsilon;
 	hasGamma;
-	constructor(e, t, n = !1) {
+	constructor(e, t, n = !1, r = 1e-8) {
 		super(e, t, {
 			reductionOp: "mean",
 			elementwise: !0
-		}, !1), this.shaderKey = "RMSNorm32", n && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = n;
+		}, !1), this.epsilon = r, this.shaderKey = `RMSNorm32_${n ? "Gamma" : "NoGamma"}_${this.epsilon}`, n && this.variableNames.push("gamma"), this.variableComponents = [1, 1], this.hasGamma = n;
 	}
 	getPreprocessSnippet() {
 		return "candidate = candidate * candidate;";
 	}
 	getPostprocessSnippet() {
-		return "bestValue = inverseSqrt(bestValue + 1e-8);";
+		return `bestValue = inverseSqrt(bestValue + ${this.epsilon});`;
 	}
 	getWriteSnippet() {
 		return `

package/dist/ops/webgpu/qkv.js

@@ -1,6 +1,6 @@
 import { Fs as e, Ii as t } from "../../dist-Da20xy8E.js";
 import { isPackedTensor as n } from "../../utilities/packed.js";
-import { t as r } from "../../matMul16-DGtlziMn.js";
+import { t as r } from "../../matMul16-CO1QhsRt.js";
 import { slice16 as i } from "../slice16.js";
 //#region lib/ops/webgpu/qkv.ts
 function a(t) {

package/dist/ops/webgpu/rope.js

@@ -20,13 +20,14 @@ var o = class {
 	];
 	size = !0;
 	uniforms = "pastLen: i32";
-	constructor(e, t, n, a) {
-		this.shaderKey = `Rope_${a}`, this.outputShape = [
+	interleaved;
+	constructor(e, t, n, a, o = !0) {
+		this.shaderKey = `Rope_${a}_${o ? "interleaved" : "split"}`, this.outputShape = [
 			e,
 			t,
 			n,
 			a
-		], this.dispatchLayout = r(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize);
+		], this.dispatchLayout = r(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize), this.interleaved = o;
 	}
 	getUserCode() {
 		let e = this.outputShape[3];
@@ -49,19 +50,19 @@ var o = class {
                     d;
 
                 if (d < rotaryDim) {
-                    let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d / 2;
+                    let half = rotaryDim / 2;
+                    let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] +
+                        ${this.interleaved ? "d / 2" : "d % half"};
                     let cos = cos[idx];
                     let sin = sin[idx];
 
                     let ownX = x[xIdx] * cos;
                     var evenOdd = 0.0;
 
-                    if (d % 2 == 0) {
-                        // even index
-                        evenOdd = -x[xIdx + 1];
+                    if (${this.interleaved ? "d % 2 == 0" : "d < half"}) {
+                        evenOdd = -x[${this.interleaved ? "xIdx + 1" : "xIdx + half"}];
                     } else {
-                        // odd index
-                        evenOdd = x[xIdx - 1];
+                        evenOdd = x[${this.interleaved ? "xIdx - 1" : "xIdx - half"}];
                     }
 
                     outVal = fma(evenOdd, sin, ownX);
@@ -93,7 +94,7 @@ var o = class {
 	size = !0;
 	uniforms = "pastLen: i32";
 	constructor(e, t, n, a) {
-		this.shaderKey = `Rope_${a}`, this.outputShape = [
+		this.shaderKey = `Rope_${a}_interleaved_packed`, this.outputShape = [
 			e,
 			t,
 			n,
@@ -117,21 +118,14 @@ var o = class {
                     t * uniforms.outShapeStrides[2] +
                     d;
 
-                let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d;
-                let cos = cos[idx];
-                let sin = sin[idx];
-
                 let xPair = unpack2x16float(u32(x[xIdx]));
-                let ownX = vec2<f32>(xPair.x * cos, xPair.y * cos);
-
-                let evenOdd = vec2<f32>(
-                    -xPair.y,
-                    xPair.x
-                );
+                let idx = (t + uniforms.pastLen) * uniforms.cosShape[1] + d;
+                let cosVal = cos[idx];
+                let sinVal = sin[idx];
 
                 outVal = vec2<f32>(
-                    fma(evenOdd.x, sin, ownX.x),
-                    fma(evenOdd.y, sin, ownX.y)
+                    fma(-xPair.y, sinVal, xPair.x * cosVal),
+                    fma(xPair.x, sinVal, xPair.y * cosVal)
                 );
 
                 result[index] = i32(pack2x16float(outVal));
@@ -139,25 +133,96 @@ var o = class {
         }
         `;
 	}
+}, c = class {
+	variableNames = [
+		"x",
+		"sin",
+		"cos"
+	];
+	outputShape;
+	shaderKey = "Rope";
+	dispatchLayout;
+	dispatch;
+	workPerThread = 2;
+	workgroupSize = [
+		64,
+		1,
+		1
+	];
+	size = !0;
+	uniforms = "pastLen: i32";
+	constructor(e, t, n, a) {
+		this.shaderKey = `Rope_${a}_split_packed`, this.outputShape = [
+			e,
+			t,
+			n,
+			a / 2
+		], this.dispatchLayout = r(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+			this.workPerThread,
+			1,
+			1
+		]);
+	}
+	getUserCode() {
+		return `
+        ${a("index")} {
+            let halfPacked = uniforms.outShape[3] / 2;
+            for (var i = 0; i < ${this.workPerThread}; i = i + 1) {
+                let flatIndex = index * ${this.workPerThread} + i;
+                if (flatIndex < uniforms.size) {
+                    let coords = getCoordsFromIndex(flatIndex); // [b, h, t, dPacked]
+                    let d = coords[3];
+
+                    if (d < halfPacked) {
+                        let t = coords[2];
+                        let xLow = unpack2x16float(u32(x[flatIndex]));
+                        let xHigh = unpack2x16float(u32(x[flatIndex + halfPacked]));
+
+                        let scalar0 = d * 2;
+                        let scalar1 = scalar0 + 1;
+                        let cacheOffset = (t + uniforms.pastLen) * uniforms.cosShape[1];
+
+                        let cos0 = cos[cacheOffset + scalar0];
+                        let sin0 = sin[cacheOffset + scalar0];
+                        let cos1 = cos[cacheOffset + scalar1];
+                        let sin1 = sin[cacheOffset + scalar1];
+
+                        let outLow = vec2<f32>(
+                            fma(-xHigh.x, sin0, xLow.x * cos0),
+                            fma(-xHigh.y, sin1, xLow.y * cos1)
+                        );
+                        let outHigh = vec2<f32>(
+                            fma(xLow.x, sin0, xHigh.x * cos0),
+                            fma(xLow.y, sin1, xHigh.y * cos1)
+                        );
+
+                        result[flatIndex] = i32(pack2x16float(outLow));
+                        result[flatIndex + halfPacked] = i32(pack2x16float(outHigh));
+                    }
+                }
+            }
+        }
+        `;
+	}
 };
-function c(t) {
-	let { x: r } = t.inputs, { pastLen: i, negSin: a, ropeCache: c } = t.attrs, l = t.backend, u = n(r), d = r.shape[0], f = r.shape[1], p = r.shape[2], m = u ? r.shape[3] * 2 : r.shape[3], h = a ? c.getNegSin() : c.getSin(), g = c.getCos();
-	if (e(h.shape, g.shape, "Error in Rope: "), h.shape[0] < p + i) throw Error(`Sin tensor shape ${h.shape} is not compatible with seqLength ${p} and pastLen ${i}.`);
-	if (h.shape[1] * 2 < m) throw Error(`Sin tensor shape ${h.shape} is not compatible with feature dimension ${m}.`);
-	if (h.shape.length !== 3) throw Error(`Sin tensor must be 3-dimensional, but got shape ${h.shape}.`);
-	let _ = u ? new s(d, f, p, m) : new o(d, f, p, m), v = [{
+function l(t) {
+	let { x: r } = t.inputs, { pastLen: i, negSin: a, ropeCache: l, interleaved: u } = t.attrs, d = t.backend, f = n(r), p = r.shape[0], m = r.shape[1], h = r.shape[2], g = f ? r.shape[3] * 2 : r.shape[3], _ = a ? l.getNegSin() : l.getSin(), v = l.getCos();
+	if (e(_.shape, v.shape, "Error in Rope: "), _.shape[0] < h + i) throw Error(`Sin tensor shape ${_.shape} is not compatible with seqLength ${h} and pastLen ${i}.`);
+	if (_.shape[1] * 2 < g) throw Error(`Sin tensor shape ${_.shape} is not compatible with feature dimension ${g}.`);
+	if (_.shape.length !== 3) throw Error(`Sin tensor must be 3-dimensional, but got shape ${_.shape}.`);
+	let y = f ? u ? new s(p, m, h, g) : new c(p, m, h, g) : new o(p, m, h, g, u), b = [{
 		type: "int32",
 		data: [i]
-	}], y = u ? "packedF16" : r.dtype;
-	return l.runWebGPUProgram(_, [
+	}], x = f ? "packedF16" : r.dtype;
+	return d.runWebGPUProgram(y, [
 		r,
-		h,
-		g
-	], y, v);
+		_,
+		v
+	], x, b);
 }
 t({
 	kernelName: "Rope",
 	backendName: "webgpu",
-	kernelFunc: c
+	kernelFunc: l
 });
 //#endregion

package/dist/ops/webgpu/silu.d.ts

@@ -0,0 +1,14 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export declare class SiluProgram implements WebGPUProgram {
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    variableNames: string[];
+    workgroupSize: [number, number, number];
+    size: boolean;
+    constructor(outputShape: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/silu.js

@@ -0,0 +1,118 @@
+import { Ii as e } from "../../dist-Da20xy8E.js";
+import { isPackedTensor as t } from "../../utilities/packed.js";
+import { c as n, i as r } from "../../webgpu_util-DfODg6XI.js";
+import { s as i } from "../../webgpu_program-B0PokuRJ.js";
+//#region lib/ops/webgpu/silu.ts
+var a = class {
+	outputShape;
+	shaderKey;
+	dispatchLayout;
+	dispatch;
+	variableNames = ["A"];
+	workgroupSize;
+	size = !0;
+	constructor(e) {
+		this.workgroupSize = [
+			128,
+			1,
+			1
+		], this.outputShape = e, this.dispatchLayout = n(this.outputShape), this.dispatch = r(this.dispatchLayout, this.outputShape, this.workgroupSize), this.shaderKey = "unary_silu";
+	}
+	getUserCode() {
+		return `
+            fn unaryOperation(x : f32) -> f32 {
+                return x / (1.0 + exp(-x));
+            }
+            ${i("index")} {
+                if (index < uniforms.size) {
+                let a = getAByOutputIndex(index);
+                setOutputAtIndex(index, unaryOperation(a));
+                }
+            }
+      `;
+	}
+};
+function o(e) {
+	let { x: t } = e.inputs, n = e.backend, r = new a(t.shape);
+	return n.runWebGPUProgram(r, [t], "float32");
+}
+e({
+	kernelName: "Silu",
+	backendName: "webgpu",
+	kernelFunc: o
+});
+var s = class {
+	variableNames = ["dy", "x"];
+	outputShape;
+	shaderKey = "SiluGrad16";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		128,
+		1,
+		1
+	];
+	size = !0;
+	constructor(e) {
+		this.outputShape = e, this.dispatchLayout = n(this.outputShape), this.dispatch = r(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		return `
+            fn activationGrad(dy: f32, X: f32) -> f32 {
+                let sigmoid = 1.0 / (1.0 + exp(-X));
+                return dy * (sigmoid + X * sigmoid * (1.0 - sigmoid));
+            }
+            ${i("index")} {
+                if (index < uniforms.size) {
+                    let X  = unpack2x16float(u32(x[index]));
+                    let DY = unpack2x16float(u32(dy[index]));
+                    let dsilu = vec2<f32>(
+                        activationGrad(DY.x, X.x),
+                        activationGrad(DY.y, X.y)
+                    );
+                    result[index] = i32(pack2x16float(dsilu));
+                }
+            }`;
+	}
+}, c = class {
+	variableNames = ["dy", "x"];
+	outputShape;
+	shaderKey = "SiluGrad32";
+	dispatchLayout;
+	dispatch;
+	workgroupSize = [
+		128,
+		1,
+		1
+	];
+	size = !0;
+	constructor(e) {
+		this.outputShape = e, this.dispatchLayout = n(this.outputShape), this.dispatch = r(this.dispatchLayout, this.outputShape, this.workgroupSize);
+	}
+	getUserCode() {
+		return `
+            fn activationGrad(dy: f32, X: f32) -> f32 {
+                let sigmoid = 1.0 / (1.0 + exp(-X));
+                return dy * (sigmoid + X * sigmoid * (1.0 - sigmoid));
+            }
+            ${i("index")} {
+                if (index < uniforms.size) {
+                    let X  = getXByOutputIndex(index);
+                    let DY = getDyByOutputIndex(index);
+                    let dsilu = activationGrad(DY, X);
+                    setOutputAtIndex(index, dsilu);
+                }
+            }`;
+	}
+};
+function l(e) {
+	let { dy: n, x: r } = e.inputs, i = e.backend, a = t(n), o = a ? new s(r.shape) : new c(r.shape);
+	return i.runWebGPUProgram(o, [n, r], a ? "packedF16" : "float32");
+}
+e({
+	kernelName: "SiluGrad",
+	backendName: "webgpu",
+	kernelFunc: l
+});
+//#endregion
+export { a as SiluProgram };

package/dist/{rope-C8EAiSGY.js → rope-BPu_uooz.js}

@@ -7,18 +7,19 @@ var n = {
 	inputsToSave: [],
 	outputsToSave: [],
 	gradFunc: (e, t, n) => {
-		let { ropeCache: i } = n, a = r(e, i, 0, !0);
-		return { x: () => a };
+		let { ropeCache: i, interleaved: a } = n, o = r(e, i, 0, !0, a ?? !0);
+		return { x: () => o };
 	}
 };
 e(n);
 //#endregion
 //#region lib/ops/rope.ts
-function r(e, n, r, i = !1) {
+function r(e, n, r, i = !1, a = !0) {
 	return n.ensureRopeCache(e.shape[2] + r), t().runKernel("Rope", { x: e }, {
 		pastLen: r,
 		negSin: i,
-		ropeCache: n
+		ropeCache: n,
+		interleaved: a
 	});
 }
 //#endregion

package/dist/silu-B6bvCgJ7.js

@@ -0,0 +1,25 @@
+import { Fi as e, ii as t } from "./dist-Da20xy8E.js";
+//#region lib/ops/grads/silu.ts
+var n = {
+	kernelName: "Silu",
+	inputsToSave: ["x"],
+	outputsToSave: [],
+	gradFunc: (e, t) => {
+		let [n] = t;
+		return { x: () => i(e, n) };
+	}
+};
+e(n);
+//#endregion
+//#region lib/ops/silu.ts
+function r(e) {
+	return t().runKernel("Silu", { x: e });
+}
+function i(e, n) {
+	return t().runKernel("SiluGrad", {
+		dy: e,
+		x: n
+	});
+}
+//#endregion
+export { r as n, n as r, i as t };

package/dist/training/Evaluator.js

@@ -15,7 +15,10 @@ var r = class {
 	}
 	async calculateBatchLoss(r, i, a, o) {
 		let [s, c] = e(() => {
-			let e = this.model.forward({ training: !1 }, r), s = n(e, i, o, a), c = t(e, i);
+			let e = this.model.forward({
+				training: !1,
+				mixedPrecision: !0
+			}, r), s = n(e, i, o, a), c = t(e, i);
 			return e.dispose(), [s, c];
 		}), l = await s.array(), u = await c.array(), d = l, f = u;
 		return c.dispose(), s.dispose(), Array.isArray(d) ? d.map((e) => ({

package/dist/utilities/maxHeap.d.ts

@@ -0,0 +1,16 @@
+export default class MaxHeap<T> {
+    private readonly fn;
+    private readonly heap;
+    constructor(fn: (a: T) => number);
+    get length(): number;
+    get size(): number;
+    get values(): T[];
+    isEmpty(): boolean;
+    peek(): T | undefined;
+    push(item: T): void;
+    pop(): T | undefined;
+    clear(): void;
+    toArray(): T[];
+    private siftUp;
+    private siftDown;
+}

package/dist/utilities/maxHeap.js

@@ -0,0 +1,55 @@
+//#region lib/utilities/maxHeap.ts
+var e = class {
+	fn;
+	heap = [];
+	constructor(e) {
+		this.fn = e;
+	}
+	get length() {
+		return this.heap.length;
+	}
+	get size() {
+		return this.heap.length;
+	}
+	get values() {
+		return [...this.heap];
+	}
+	isEmpty() {
+		return this.heap.length === 0;
+	}
+	peek() {
+		return this.heap[0];
+	}
+	push(e) {
+		this.heap.push(e), this.siftUp(this.heap.length - 1);
+	}
+	pop() {
+		if (this.heap.length === 0) return;
+		let e = this.heap[0], t = this.heap.pop();
+		return this.heap.length > 0 && t !== void 0 && (this.heap[0] = t, this.siftDown(0)), e;
+	}
+	clear() {
+		this.heap.length = 0;
+	}
+	toArray() {
+		return [...this.heap];
+	}
+	siftUp(e) {
+		let t = e;
+		for (; t > 0;) {
+			let e = Math.floor((t - 1) / 2);
+			if (this.fn(this.heap[t]) <= this.fn(this.heap[e])) break;
+			[this.heap[t], this.heap[e]] = [this.heap[e], this.heap[t]], t = e;
+		}
+	}
+	siftDown(e) {
+		let t = e;
+		for (;;) {
+			let e = t * 2 + 1, n = t * 2 + 2, r = t;
+			if (e < this.heap.length && this.fn(this.heap[e]) > this.fn(this.heap[r]) && (r = e), n < this.heap.length && this.fn(this.heap[n]) > this.fn(this.heap[r]) && (r = n), r === t) break;
+			[this.heap[t], this.heap[r]] = [this.heap[r], this.heap[t]], t = r;
+		}
+	}
+};
+//#endregion
+export { e as default };

package/dist/utilities/parameters.js

@@ -1,27 +1,43 @@
+import { swigluHidden as e } from "../layers/MLP.js";
 //#region lib/utilities/parameters.ts
-var e = 4;
-function t(e) {
-	return 4 * e.nEmbed * e.nEmbed + (e.mlpFactor * e.nEmbed * e.nEmbed + e.nEmbed * e.mlpFactor * e.nEmbed);
+var t = 4;
+function n(t) {
+	let n;
+	switch (t.modelType) {
+		case "GenAI_NanoGPT_v1":
+		case "GenAI_NanoGPT_v2":
+			n = e(t.nEmbed, t.mlpFactor) * t.nEmbed * 2;
+			break;
+		case "llama":
+		case "qwen3":
+			n = Math.max(256, Math.ceil(t.nEmbed * (8 / 3) / 256) * 256) * t.nEmbed * 3;
+			break;
+		default: n = 0;
+	}
+	return n;
 }
-function n(e) {
-	let t = e.vocabSize * e.nEmbed, n = e.nLayer * (4 * e.nEmbed * e.nEmbed), r = e.nLayer * (e.mlpFactor * e.nEmbed * e.nEmbed + e.nEmbed * e.mlpFactor * e.nEmbed);
-	return t + n + r;
+function r(e) {
+	return 4 * e.nEmbed * e.nEmbed + n(e);
 }
-function r(t) {
-	return n(t) * e;
+function i(e) {
+	let t = e.modelType === "llama" ? e.vocabSize * e.nEmbed * 2 : e.vocabSize * e.nEmbed, r = e.nLayer * (4 * e.nEmbed * e.nEmbed), i = n(e) * e.nLayer;
+	return t + r + i;
 }
-function i(t, n) {
-	let i = r(t), a = i * 2, o = i * 2, s = n * t.blockSize * t.nEmbed * e;
-	return i + a + o + s;
+function a(e) {
+	return i(e) * t;
 }
-function a(e, t) {
+function o(e, n) {
+	let r = a(e), i = r * 2, o = r * 2, s = n * e.blockSize * e.nEmbed * t;
+	return r + i + o + s;
+}
+function s(e, t) {
 	return {
-		numParams: n(e),
-		modelMemoryMB: r(e) / (1024 * 1024),
-		trainingMemoryMB: i(e, t) / (1024 * 1024)
+		numParams: i(e),
+		modelMemoryMB: a(e) / (1024 * 1024),
+		trainingMemoryMB: o(e, t) / (1024 * 1024)
 	};
 }
-function o(e) {
+function c(e) {
 	if (e.nEmbed % e.nHead !== 0) throw Error("nEmbed_divisible_nHead");
 	if (e.blockSize <= 0) throw Error("blockSize_positive");
 	if (e.vocabSize <= 0) throw Error("vocabSize_positive");
@@ -35,4 +51,4 @@ function o(e) {
 	if (!Number.isInteger(e.vocabSize)) throw Error("vocabSize_integer");
 }
 //#endregion
-export { t as estimateLayerParameters, r as estimateMemoryUsage, n as estimateParameterCount, a as estimateResources, i as estimateTrainingMemoryUsage, o as validateConfig };
+export { r as estimateLayerParameters, a as estimateMemoryUsage, i as estimateParameterCount, s as estimateResources, o as estimateTrainingMemoryUsage, c as validateConfig };

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "1.4.1",
+    "version": "1.5.0",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",