Add residuals

shader/center.wgsl

@@ -0,0 +1,24 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> centered: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  centered.size = m.size;
+
+  let vc =  m.size.x / 4;
+  let ro = (global_id.z * m.size.y + global_id.y) * vc;
+
+  var s = vec4(0f);
+  for (var x = 0u; x < vc; x++) {
+    s += m.data[ro + x];
+  }
+
+  centered.data[ro + global_id.x] = m.data[ro + global_id.x] - vec4(s.x + s.y + s.z + s.w) / f32(m.size.x);
+}

shader/dlds.wgsl

@@ -1,48 +0,0 @@
-struct Scalar {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> ds: Scalar;
-@group(0) @binding(1)
-var<storage, read> wa: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read> v: MatrixBatch;
-@group(0) @binding(3)
-var<storage, read> otdldo: MatrixBatch;
-@group(0) @binding(4)
-var<storage, read_write> dlds: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  dlds.size = wa.size;
-
-  var s = 0f;
-  var li = (global_id.z * v.size.y + global_id.x) * v.size.x;
-  var ri = global_id.z * otdldo.size.x * otdldo.size.y + global_id.y;
-
-  for (var i = 0u; i < v.size.x; i++) {
-    var cs = 0f;
-    var lj = global_id.z * wa.size.y * wa.size.x + global_id.y * wa.size.x;
-    var mj = global_id.z * v.size.x * v.size.y + i;
-    for (var j = 0u; j < wa.size.x; j++) {
-      cs += wa.data[lj] * v.data[mj];
-      lj++;
-      mj += v.size.x;
-    }
-
-    s += (v.data[li] - cs) * otdldo.data[ri];
-
-    li++;
-    ri += otdldo.size.x;
-  }
-
-  let resi = (global_id.z * dlds.size.y + global_id.y) * dlds.size.x + global_id.x;
-  dlds.data[resi] = ds.data[0] * wa.data[resi] * s;
-}

shader/gradinp.wgsl

@@ -1,47 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> ilogits: MatrixBatch;
-@group(0) @binding(0)
-var<storage, read> tilogits: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> dldsq: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read> dldsk: MatrixBatch;
-@group(0) @binding(3)
-var<storage, read> totdldowa: MatrixBatch;
-@group(0) @binding(4)
-var<storage, read> tqry: MatrixBatch;
-@group(0) @binding(5)
-var<storage, read> tkey: MatrixBatch;
-@group(0) @binding(6)
-var<storage, read> tval: MatrixBatch;
-@group(0) @binding(7)
-var<storage, read_write> ginp: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  ginp.size.x = tval.size.x;
-  ginp.size.y = ilogits.size.x;
-  ginp.size.z = ilogits.size.z;
-
-  var ts = 0f;
-  var li = global_id.z * ilogits.size.y * ilogits.size.x + global_id.y;
-  for (var i = 0u; i < ilogits.size.y; i++) {
-    var s = 0f;
-    var lj = global_id.z * totdldowa.size.x * totdldowa.size.y + i * totdldowa.size.x;
-    var rj = global_id.z * tval.size.x * tval.size.y + global_id.x;
-    for (var j = 0u; j < totdldowa.size.x; j++) {
-      s += totdldowa.data[lj] * tval.data[rj] + dldsq.data[lj] * tkey.data[rj] + dldsk.data[lj] * tqry.data[rj];
-      lj++;
-      rj += tval.size.x;
-    }
-    ts += ilogits.data[li] * s;
-    li += ilogits.size.x;
-  }
-
-  ginp.data[global_id.z * ginp.size.y * ginp.size.x + global_id.y * ginp.size.x + global_id.x] = ts;
-}

shader/gradtattw.wgsl

@@ -1,28 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> e: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> dlds: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> gtattw: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  gtattw.size.z = e.size.z;
-  gtattw.size.y = e.size.x;
-  gtattw.size.x = e.size.x;
-
-  var s = 0f;
-  var li = global_id.z * e.size.x * e.size.y + global_id.x;
-  var ri = global_id.z * dlds.size.x * dlds.size.y + global_id.y;
-  for (var i = 0u; i < e.size.y; i++) {
-    s += e.data[li] * dlds.data[ri];
-    li += e.size.x;
-    ri += dlds.size.x;
-  }
-  gtattw.data[global_id.z * gtattw.size.y * gtattw.size.x + global_id.y * gtattw.size.x + global_id.x] = s;
-}

shader/gradtout.wgsl

@@ -1,28 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> a: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> dldo: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> gtout: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  gtout.size = a.size;
-  gtout.size.y = dldo.size.x;
-
-  var s = 0f;
-  var ai = global_id.z * a.size.y * a.size.x + global_id.x;
-  var di = global_id.z * dldo.size.x * dldo.size.y + global_id.y;
-  for (var i = 0u; i < a.size.y; i++) {
-    s += a.data[ai] * dldo.data[di];
-    ai += a.size.x;
-    di += dldo.size.x;
-  }
-
-  gtout.data[global_id.z * gtout.size.y * gtout.size.x + global_id.y * gtout.size.x + global_id.x] = s;
-}

shader/layernormgrad.wgsl

@@ -0,0 +1,51 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0) var<storage, read> m: MatrixBatch;
+@group(0) @binding(1) var<storage, read> dldm: MatrixBatch;
+@group(0) @binding(2) var<storage, read_write> grad: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  grad.size = m.size;
+
+  let vc = m.size.x / 4;
+  let ro = (global_id.z * m.size.y + global_id.y) * vc;
+
+  var sum = vec4(0f);
+  var sumsq = vec4(0f);
+  for (var i = 0u; i < vc; i++) {
+    let v = m.data[ro + i];
+    sum += v;
+    sumsq += v * v;
+  }
+
+  let mean = (sum.x + sum.y + sum.z + sum.w) / f32(m.size.x);
+  let variance = (sumsq.x + sumsq.y + sumsq.z + sumsq.w) / f32(m.size.x) - mean * mean;
+  let inv_std = inverseSqrt((sumsq.x + sumsq.y + sumsq.z + sumsq.w) / f32(m.size.x) - mean * mean + 1e-5);
+
+  // Recompute x̂
+  var sum_grad = vec4(0.0);
+  var sum_grad_xhat = vec4(0.0);
+
+  for (var i = 0u; i < vc; i++) {
+    let v = m.data[ro + i];
+    let g = dldm.data[ro + i];
+    let xhat = (v - vec4(mean)) * vec4(inv_std);
+    sum_grad += g;
+    sum_grad_xhat += g * xhat;
+  }
+
+  let mg = (sum_grad.x + sum_grad.y + sum_grad.z + sum_grad.w) / f32(m.size.x);
+  let mgxhat = (sum_grad_xhat.x + sum_grad_xhat.y + sum_grad_xhat.z + sum_grad_xhat.w) / f32(m.size.x);
+
+  for (var i = 0u; i < vc; i++) {
+    let v = m.data[ro + i];
+    let g = dldm.data[ro + i];
+    let xhat = (v - vec4(mean)) * vec4(inv_std);
+    let dx = vec4(inv_std) * (g - vec4(mg) - xhat * vec4(mgxhat));
+    grad.data[ro + i] = dx;
+  }
+}

shader/normalize.wgsl

@@ -0,0 +1,24 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> c: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = c.size;
+
+  let vc = c.size.x / 4;
+  let ro = (global_id.z * c.size.y + global_id.y) * vc;
+
+  var s = vec4(0f);
+  for (var x = 0u; x < vc; x++) {
+    s += c.data[ro + x] * c.data[ro + x];
+  }
+
+  result.data[ro + global_id.x] = c.data[ro + global_id.x] * vec4(inverseSqrt((s.x + s.y + s.z + s.w) / f32(c.size.x) + 1e-5f));
+}

shader/otdldo.wgsl

@@ -1,29 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> dldo: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> tout: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> otdldo: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  otdldo.size.x = dldo.size.y;
-  otdldo.size.y = tout.size.x;
-  otdldo.size.z = dldo.size.z;
-
-  var s = 0f;
-  var li = (global_id.z * dldo.size.y + global_id.x) * dldo.size.x;
-  var ri = global_id.z * tout.size.x * tout.size.y + global_id.y;
-  for (var i = 0u; i < dldo.size.x; i++) {
-    s += dldo.data[li] * tout.data[ri];
-    li++;
-    ri += tout.size.x;
-  }
-
-  otdldo.data[(global_id.z * otdldo.size.y + global_id.y) * otdldo.size.x + global_id.x] = s;
-}

shader/softmaxgrad.wgsl

@@ -0,0 +1,26 @@
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<f32>>,
+}
+
+@group(0) @binding(0)
+var<storage, read> m: MatrixBatch;
+@group(0) @binding(1)
+var<storage, read> dldm: MatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> grad: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  grad.size = m.size;
+
+  let vc = m.size.x / 4;
+  let ro = (global_id.z * m.size.y + global_id.y) * vc;
+
+  var s = 0f;
+  for (var x = 0u; x < vc; x++) {
+    s += dot(m.data[ro + x], dldm.data[ro + x]);
+  }
+
+  grad.data[ro + global_id.x] = m.data[ro + global_id.x] * (dldm.data[ro + global_id.x] - vec4(s));
+}

shader/tmul.wgsl

@@ -0,0 +1,34 @@
+alias number = f32;
+
+struct PackedMatrixBatch {
+  size: vec4<u32>,
+  data: array<vec4<number>>,
+}
+
+struct MatrixBatch {
+  size: vec4<u32>,
+  data: array<number>,
+}
+
+@group(0) @binding(0)
+var<storage, read> l: PackedMatrixBatch;
+@group(0) @binding(1)
+var<storage, read> tr: PackedMatrixBatch;
+@group(0) @binding(2)
+var<storage, read_write> result: MatrixBatch;
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  result.size = l.size;
+  result.size.x = tr.size.y;
+
+  var s = 0f;
+  let vc = l.size.x / 4;
+  let lo = (global_id.z * l.size.y + global_id.y) * vc;
+  let ro = (global_id.z * tr.size.y + global_id.x) * vc;
+  for (var i = 0u; i < vc; i++) {
+    s += dot(l.data[lo + i], tr.data[ro + i]);
+  }
+
+  result.data[global_id.z * result.size.y * result.size.x + global_id.y * result.size.x + global_id.x] = s;
+}

shader/totdldowa.wgsl

@@ -1,29 +0,0 @@
-struct MatrixBatch {
-  size: vec4<u32>,
-  data: array<f32>,
-}
-
-@group(0) @binding(0)
-var<storage, read> wa: MatrixBatch;
-@group(0) @binding(1)
-var<storage, read> otdldo: MatrixBatch;
-@group(0) @binding(2)
-var<storage, read_write> totdldowa: MatrixBatch;
-
-@compute @workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  totdldowa.size.x = otdldo.size.y;
-  totdldowa.size.y = wa.size.x;
-  totdldowa.size.z = otdldo.size.z;
-
-  var s = 0f;
-  var li = global_id.z * otdldo.size.x * otdldo.size.y + global_id.x * otdldo.size.x;
-  var ri = global_id.z * wa.size.x * wa.size.y + global_id.y;
-  for (var i = 0u; i < otdldo.size.x; i++) {
-    s += otdldo.data[li] * wa.data[ri];
-    li++;
-    ri += wa.size.x;
-  }
-
-  totdldowa.data[(global_id.z * totdldowa.size.y + global_id.y) * totdldowa.size.x + global_id.x] = s;
-}

src/data/gpu/gpu-struct.ts

@@ -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

@@ -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

@@ -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

@@ -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(
+      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,
+    );
 
-    Matrix.mul(this.processing.dlds, k, this.processing.dldsk);
-    Matrix.mul(Matrix.transpose(this.processing.dlds), q, this.processing.dldsq);
-
-    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);