Tiny LM: the code

These are the files that do most of the work. The comments explain why each part works the way it does. If anything is unclear, or you would do it differently, tell me. I will also send you the full code for free, with nothing to sign up for. Ask for the full code

// Tiny LM's forward pass in plain TypeScript, mirroring train/tinylm/model.py.
// Sums run in float64 and are rounded to float32 once, which keeps us within
// PyTorch's own rounding error (tests/parity.test.ts).

export interface TensorInfo {
  name: string;
  shape: number[];
  offset: number; // in float32 elements, not bytes
  length: number; // in float32 elements
}

export interface ModelConfig {
  vocab_size: number;
  n_layer: number;
  n_head: number;
  n_embd: number;
  block_size: number;
  dropout?: number;
}

export interface Manifest {
  config: ModelConfig;
  dtype: string;
  params?: number;
  total_floats: number;
  tensors: TensorInfo[];
}

interface BlockWeights {
  ln1w: Float32Array; ln1b: Float32Array;
  attnW: Float32Array; attnB: Float32Array; // c_attn: [3C, C] and [3C]
  projW: Float32Array; projB: Float32Array; // attn.c_proj: [C, C] and [C]
  ln2w: Float32Array; ln2b: Float32Array;
  fcW: Float32Array; fcB: Float32Array; // mlp.c_fc: [4C, C] and [4C]
  mprojW: Float32Array; mprojB: Float32Array; // mlp.c_proj: [C, 4C] and [C]
}

/** Keys and values of earlier tokens, so step() only computes the new one. */
export interface KVCache {
  k: Float32Array[]; // per layer, block_size * C
  v: Float32Array[]; // per layer, block_size * C
  length: number; // how many tokens are stored
}

const LN_EPS = 1e-5;

// Biased variance (divide by C), as PyTorch does.
function layerNorm(
  x: Float32Array, xOff: number, w: Float32Array, b: Float32Array,
  out: Float32Array, outOff: number, C: number,
): void {
  let mean = 0;
  for (let i = 0; i < C; i++) mean += x[xOff + i]!;
  mean /= C;
  let variance = 0;
  for (let i = 0; i < C; i++) {
    const d = x[xOff + i]! - mean;
    variance += d * d;
  }
  variance /= C;
  const inv = 1 / Math.sqrt(variance + LN_EPS);
  for (let i = 0; i < C; i++) out[outOff + i] = (x[xOff + i]! - mean) * inv * w[i]! + b[i]!;
}

// out = W x + b, with W stored [nOut, nIn] row-major.
function linear(
  W: Float32Array, b: Float32Array | null, x: Float32Array, xOff: number,
  out: Float32Array, outOff: number, nOut: number, nIn: number,
): void {
  for (let o = 0; o < nOut; o++) {
    let sum = b ? b[o]! : 0;
    const row = o * nIn;
    for (let i = 0; i < nIn; i++) sum += W[row + i]! * x[xOff + i]!;
    out[outOff + o] = sum;
  }
}

// The tanh approximation, used on both sides because JS has no erf.
function gelu(z: number): number {
  return 0.5 * z * (1 + Math.tanh(0.7978845608028654 * (z + 0.044715 * z * z * z))); // 0.79788... = sqrt(2/pi)
}

// Subtracting the max keeps exp from overflowing.
function softmaxInPlace(a: Float64Array, n: number): void {
  let max = -Infinity;
  for (let i = 0; i < n; i++) if (a[i]! > max) max = a[i]!;
  let sum = 0;
  for (let i = 0; i < n; i++) {
    a[i] = Math.exp(a[i]! - max);
    sum += a[i]!;
  }
  for (let i = 0; i < n; i++) a[i] = a[i]! / sum;
}

export class TinyLM {
  readonly config: ModelConfig;
  private readonly wte: Float32Array; // [vocab, C] word vectors, also reused as the output head
  private readonly wpe: Float32Array; // [block_size, C] position vectors
  private readonly blocks: BlockWeights[];
  private readonly lnfW: Float32Array;
  private readonly lnfB: Float32Array;

  private constructor(manifest: Manifest, weights: ArrayBuffer) {
    const cfg = manifest.config;
    this.config = cfg;
    const byName = new Map(manifest.tensors.map((t) => [t.name, t]));
    // A view onto the buffer, no copy. weights.bin is little-endian, as browsers are.
    const get = (name: string, shape: number[]): Float32Array => {
      const t = byName.get(name);
      if (!t) throw new Error(`weights manifest has no tensor "${name}"`);
      if (t.shape.length !== shape.length || t.shape.some((d, i) => d !== shape[i])) {
        throw new Error(`tensor "${name}" has shape [${t.shape}] but the config needs [${shape}]`);
      }
      const n = shape.reduce((a, b) => a * b, 1);
      if (t.length !== n) throw new Error(`tensor "${name}" has length ${t.length}, expected ${n}`);
      return new Float32Array(weights, t.offset * 4, t.length);
    };

    const C = cfg.n_embd;
    this.wte = get("wte.weight", [cfg.vocab_size, C]);
    this.wpe = get("wpe.weight", [cfg.block_size, C]);
    this.blocks = [];
    for (let i = 0; i < cfg.n_layer; i++) {
      const p = `h.${i}.`;
      this.blocks.push({
        ln1w: get(p + "ln_1.weight", [C]), ln1b: get(p + "ln_1.bias", [C]),
        attnW: get(p + "attn.c_attn.weight", [3 * C, C]), attnB: get(p + "attn.c_attn.bias", [3 * C]),
        projW: get(p + "attn.c_proj.weight", [C, C]), projB: get(p + "attn.c_proj.bias", [C]),
        ln2w: get(p + "ln_2.weight", [C]), ln2b: get(p + "ln_2.bias", [C]),
        fcW: get(p + "mlp.c_fc.weight", [4 * C, C]), fcB: get(p + "mlp.c_fc.bias", [4 * C]),
        mprojW: get(p + "mlp.c_proj.weight", [C, 4 * C]), mprojB: get(p + "mlp.c_proj.bias", [C]),
      });
    }
    this.lnfW = get("ln_f.weight", [C]);
    this.lnfB = get("ln_f.bias", [C]);
  }

  static fromBuffers(manifest: Manifest, weights: ArrayBuffer): TinyLM {
    if (manifest.dtype !== "float32") throw new Error(`unsupported weight dtype "${manifest.dtype}"`);
    const expected = manifest.total_floats * 4;
    if (weights.byteLength !== expected) {
      throw new Error(`weights.bin is ${weights.byteLength} bytes but the manifest expects ${expected} bytes`);
    }
    const cfg = manifest.config;
    if (cfg.n_embd % cfg.n_head !== 0) throw new Error("n_embd must be divisible by n_head");
    return new TinyLM(manifest, weights);
  }

  static async load(base: string, fetchFn: typeof fetch = fetch): Promise<TinyLM> {
    const mRes = await fetchFn(`${base}/manifest.json`);
    if (!mRes.ok) throw new Error(`could not load ${base}/manifest.json (HTTP ${mRes.status})`);
    const manifest = (await mRes.json()) as Manifest;
    const wRes = await fetchFn(`${base}/weights.bin`);
    if (!wRes.ok) throw new Error(`could not load ${base}/weights.bin (HTTP ${wRes.status})`);
    return TinyLM.fromBuffers(manifest, await wRes.arrayBuffer());
  }

  /** Logits for every position (T x V) and every head's attention weights (T x T). */
  forward(ids: number[]): { logits: Float32Array; attn: Float32Array[][] } {
    const { vocab_size: V, n_layer, n_head, n_embd: C, block_size } = this.config;
    const T = ids.length;
    if (T === 0) throw new Error("forward needs at least one token");
    if (T > block_size) throw new Error(`sequence length ${T} exceeds block_size ${block_size}`);
    this.checkIds(ids);

    const hd = C / n_head;
    const x = new Float32Array(T * C);

    for (let t = 0; t < T; t++) {
      const w = ids[t]! * C;
      const p = t * C;
      for (let i = 0; i < C; i++) x[t * C + i] = this.wte[w + i]! + this.wpe[p + i]!;
    }

    const normed = new Float32Array(T * C);
    const qkv = new Float32Array(T * 3 * C);
    const mixed = new Float32Array(T * C);
    const projOut = new Float32Array(C);
    const hidden = new Float32Array(4 * C);
    const mlpOut = new Float32Array(C);
    const scores = new Float64Array(T);
    const attn: Float32Array[][] = [];

    for (let l = 0; l < n_layer; l++) {
      const blk = this.blocks[l]!;

      for (let t = 0; t < T; t++) {
        layerNorm(x, t * C, blk.ln1w, blk.ln1b, normed, t * C, C);
        linear(blk.attnW, blk.attnB, normed, t * C, qkv, t * 3 * C, 3 * C, C);
      }

      const layerAttn: Float32Array[] = [];
      for (let h = 0; h < n_head; h++) layerAttn.push(new Float32Array(T * T));
      const scale = 1 / Math.sqrt(hd);

      for (let h = 0; h < n_head; h++) {
        const weightsOut = layerAttn[h]!;
        const hOff = h * hd;
        for (let i = 0; i < T; i++) {
          // The causal mask is this loop bound: token i only sees j <= i.
          const q = i * 3 * C + hOff;
          for (let j = 0; j <= i; j++) {
            const k = j * 3 * C + C + hOff;
            let dot = 0;
            for (let d = 0; d < hd; d++) dot += qkv[q + d]! * qkv[k + d]!;
            scores[j] = dot * scale;
          }
          softmaxInPlace(scores, i + 1);
          for (let j = 0; j <= i; j++) weightsOut[i * T + j] = scores[j]!;

          for (let d = 0; d < hd; d++) {
            let sum = 0;
            for (let j = 0; j <= i; j++) sum += scores[j]! * qkv[j * 3 * C + 2 * C + hOff + d]!;
            mixed[i * C + hOff + d] = sum;
          }
        }
      }
      attn.push(layerAttn);

      // Residual: the block adds to x rather than replacing it.
      for (let t = 0; t < T; t++) {
        linear(blk.projW, blk.projB, mixed, t * C, projOut, 0, C, C);
        for (let i = 0; i < C; i++) x[t * C + i] = x[t * C + i]! + projOut[i]!;
      }

      for (let t = 0; t < T; t++) {
        layerNorm(x, t * C, blk.ln2w, blk.ln2b, normed, t * C, C);
        linear(blk.fcW, blk.fcB, normed, t * C, hidden, 0, 4 * C, C);
        for (let i = 0; i < 4 * C; i++) hidden[i] = gelu(hidden[i]!);
        linear(blk.mprojW, blk.mprojB, hidden, 0, mlpOut, 0, C, 4 * C);
        for (let i = 0; i < C; i++) x[t * C + i] = x[t * C + i]! + mlpOut[i]!;
      }
    }

    // Tied weights: the output head is the word table.
    const logits = new Float32Array(T * V);
    for (let t = 0; t < T; t++) {
      layerNorm(x, t * C, this.lnfW, this.lnfB, normed, t * C, C);
      linear(this.wte, null, normed, t * C, logits, t * V, V, C);
    }
    return { logits, attn };
  }

  newCache(): KVCache {
    const { n_layer, block_size, n_embd } = this.config;
    return {
      k: Array.from({ length: n_layer }, () => new Float32Array(block_size * n_embd)),
      v: Array.from({ length: n_layer }, () => new Float32Array(block_size * n_embd)),
      length: 0,
    };
  }

  /** One new token in, next-token logits out. Earlier keys and values never change, so they come from the cache. */
  step(cache: KVCache, id: number): Float32Array {
    const { vocab_size: V, n_layer, n_head, n_embd: C, block_size } = this.config;
    const pos = cache.length;
    if (pos >= block_size) throw new Error(`KV cache is full (${block_size} tokens)`);
    this.checkIds([id]);
    const hd = C / n_head;
    const scale = 1 / Math.sqrt(hd);

    const x = new Float32Array(C);
    for (let i = 0; i < C; i++) x[i] = this.wte[id * C + i]! + this.wpe[pos * C + i]!;

    const normed = new Float32Array(C);
    const qkv = new Float32Array(3 * C);
    const mixed = new Float32Array(C);
    const projOut = new Float32Array(C);
    const hidden = new Float32Array(4 * C);
    const mlpOut = new Float32Array(C);
    const scores = new Float64Array(pos + 1);

    for (let l = 0; l < n_layer; l++) {
      const blk = this.blocks[l]!;
      const K = cache.k[l]!;
      const Vv = cache.v[l]!;

      layerNorm(x, 0, blk.ln1w, blk.ln1b, normed, 0, C);
      linear(blk.attnW, blk.attnB, normed, 0, qkv, 0, 3 * C, C);
      for (let i = 0; i < C; i++) {
        K[pos * C + i] = qkv[C + i]!;
        Vv[pos * C + i] = qkv[2 * C + i]!;
      }

      for (let h = 0; h < n_head; h++) {
        const hOff = h * hd;
        for (let j = 0; j <= pos; j++) {
          let dot = 0;
          for (let d = 0; d < hd; d++) dot += qkv[hOff + d]! * K[j * C + hOff + d]!;
          scores[j] = dot * scale;
        }
        softmaxInPlace(scores, pos + 1);
        for (let d = 0; d < hd; d++) {
          let sum = 0;
          for (let j = 0; j <= pos; j++) sum += scores[j]! * Vv[j * C + hOff + d]!;
          mixed[hOff + d] = sum;
        }
      }

      linear(blk.projW, blk.projB, mixed, 0, projOut, 0, C, C);
      for (let i = 0; i < C; i++) x[i] = x[i]! + projOut[i]!;

      layerNorm(x, 0, blk.ln2w, blk.ln2b, normed, 0, C);
      linear(blk.fcW, blk.fcB, normed, 0, hidden, 0, 4 * C, C);
      for (let i = 0; i < 4 * C; i++) hidden[i] = gelu(hidden[i]!);
      linear(blk.mprojW, blk.mprojB, hidden, 0, mlpOut, 0, C, 4 * C);
      for (let i = 0; i < C; i++) x[i] = x[i]! + mlpOut[i]!;
    }

    cache.length = pos + 1;
    layerNorm(x, 0, this.lnfW, this.lnfB, normed, 0, C);
    const logits = new Float32Array(V);
    linear(this.wte, null, normed, 0, logits, 0, V, C);
    return logits;
  }

  private checkIds(ids: number[]): void {
    const V = this.config.vocab_size;
    for (const id of ids) {
      if (!Number.isInteger(id) || id < 0 || id >= V) throw new Error(`token id ${id} is outside 0..${V - 1}`);
    }
  }
}

// Byte-level BPE tokeniser, a port of train/tinylm/bpe.py.
// Ids 0-255 are bytes; merge i made id 256 + i. Merges never cross a chunk boundary.

import { normaliseText } from "./normalise";

export interface TokenizerJSON {
  pattern: string;
  pattern_flags: string;
  merges: [number, number][];
  vocab_size: number;
  newline_id: number;
}

// Left to right, non-overlapping: [a, a, a] with (a, a) gives [newId, a].
function mergePair(ids: number[], a: number, b: number, newId: number): number[] {
  const out: number[] = [];
  let i = 0;
  while (i < ids.length) {
    if (i + 1 < ids.length && ids[i] === a && ids[i + 1] === b) {
      out.push(newId);
      i += 2;
    } else {
      out.push(ids[i]);
      i += 1;
    }
  }
  return out;
}

export class Tokenizer {
  readonly newlineId: number;
  readonly vocabSize: number;
  private readonly pattern: string;
  private readonly merges: [number, number][];
  // "a,b" -> merge rank; lower is applied first.
  private readonly rank = new Map<string, number>();
  private readonly bytes: Uint8Array[] = [];
  // Words repeat a lot, so chunk encodings are cached.
  private readonly cache = new Map<string, number[]>();
  private readonly utf8 = new TextDecoder();

  private constructor(j: TokenizerJSON) {
    this.pattern = j.pattern;
    this.merges = j.merges;
    this.newlineId = j.newline_id;
    this.vocabSize = 256 + j.merges.length;
    j.merges.forEach(([a, b], i) => this.rank.set(`${a},${b}`, i));
    for (let i = 0; i < 256; i++) this.bytes.push(Uint8Array.of(i));
    for (const [a, b] of j.merges) {
      const joined = new Uint8Array(this.bytes[a].length + this.bytes[b].length);
      joined.set(this.bytes[a], 0);
      joined.set(this.bytes[b], this.bytes[a].length);
      this.bytes.push(joined);
    }
  }

  static fromJSON(j: TokenizerJSON): Tokenizer {
    return new Tokenizer(j);
  }

  private encodeChunk(chunk: string): number[] {
    const cached = this.cache.get(chunk);
    if (cached) return cached;
    let ids = Array.from(new TextEncoder().encode(chunk));
    while (ids.length >= 2) {
      let bestRank = Infinity;
      let bestIndex = -1;
      for (let i = 0; i + 1 < ids.length; i++) {
        const r = this.rank.get(`${ids[i]},${ids[i + 1]}`);
        if (r !== undefined && r < bestRank) {
          bestRank = r;
          bestIndex = i;
        }
      }
      if (bestIndex < 0) break;
      ids = mergePair(ids, ids[bestIndex], ids[bestIndex + 1], 256 + bestRank);
    }
    this.cache.set(chunk, ids);
    return ids;
  }

  // Normalises first (Python does not), since visitors type anything. The result is
  // ASCII, so the regex matches Python's re.ASCII behaviour.
  encode(text: string): number[] {
    const clean = normaliseText(text);
    const out: number[] = [];
    for (const m of clean.matchAll(new RegExp(this.pattern, "g"))) {
      for (const id of this.encodeChunk(m[0])) out.push(id);
    }
    return out;
  }

  tokenBytes(id: number): Uint8Array {
    const b = this.bytes[id];
    if (!b) throw new RangeError(`token id ${id} is outside the vocabulary`);
    return b;
  }

  // Bad UTF-8 shows as U+FFFD.
  tokenText(id: number): string {
    return this.utf8.decode(this.tokenBytes(id));
  }

  decode(ids: number[]): string {
    const parts = ids.map((id) => this.tokenBytes(id));
    const out = new Uint8Array(parts.reduce((n, p) => n + p.length, 0));
    let at = 0;
    for (const p of parts) {
      out.set(p, at);
      at += p.length;
    }
    return this.utf8.decode(out);
  }
}

Ask for the full codeBack to the case study