toy 0.8.0

data/lib/toy/train/sampler.rb

@@ -0,0 +1,314 @@
+# Sampler — composable transforms applied to logits before pick.
+#
+# All transforms take a `Mat` (1 × vocab) of logits + a context object
+# and return a `Mat` (1 × vocab) — possibly mutating in place. Final
+# pick is `argmax_or_multinomial` (deterministic at T=0; sampled
+# otherwise).
+#
+# Standard pipeline:
+#   logits = repetition_penalty(logits, ctx, cfg.rep_penalty)
+#   logits = temperature(logits, cfg.temperature)
+#   logits = top_k(logits, cfg.top_k)
+#   logits = top_p(logits, cfg.top_p)
+#   pick   = argmax_or_multinomial(logits, cfg)
+#
+# Context (`SamplerContext`) carries generated IDs (for rep penalty)
+# + RNG seed. Config (`SamplerConfig`) carries per-call hyperparams.
+
+require_relative "../models/transformer"   # for Mat
+
+class SamplerConfig
+  attr_accessor :temperature, :top_k, :top_p, :rep_penalty, :seed, :max_tokens
+
+  def initialize
+    @temperature = 0.0       # 0 = greedy argmax
+    @top_k       = 0         # 0 = disabled
+    @top_p       = 1.0       # 1.0 = disabled
+    @rep_penalty = 1.0       # 1.0 = disabled
+    @seed        = 42
+    @max_tokens  = 64
+  end
+end
+
+class SamplerContext
+  attr_accessor :generated_ids, :prompt_ids, :rng_state
+
+  def initialize(prompt_ids, seed)
+    @generated_ids = [0]
+    @generated_ids.pop
+    @prompt_ids = prompt_ids
+    @rng_state  = seed
+  end
+
+  # xorshift32 — Spinel-safe deterministic RNG. Don't use Kernel#rand;
+  # we want reproducibility across Spinel builds.
+  def next_u32
+    x = @rng_state
+    x = x ^ ((x << 13) & 0xFFFFFFFF)
+    x = x ^ (x >> 17)
+    x = x ^ ((x << 5) & 0xFFFFFFFF)
+    @rng_state = x & 0xFFFFFFFF
+    @rng_state
+  end
+
+  # Float in [0, 1).
+  def next_unit
+    (next_u32.to_f) / 4294967296.0
+  end
+end
+
+module Sampler
+  # In-place divide by temperature. T=0 means "do nothing here, let
+  # argmax_or_multinomial fall through to argmax."
+  def self.temperature(logits, t)
+    if t <= 0.0 || t == 1.0
+      return logits
+    end
+    inv = 1.0 / t
+    n = logits.ncols
+    j = 0
+    while j < n
+      logits.flat[j] = logits.flat[j] * inv
+      j = j + 1
+    end
+    logits
+  end
+
+  # Subtract `rep_penalty` from logits of any token already in the
+  # generated context. Default 1.0 = disabled. The HF convention
+  # DIVIDES positive logits and MULTIPLIES negative ones; we use the
+  # simpler subtract-on-positive variant to keep it Spinel-friendly.
+  # For most fine-tunes a value of 1.05–1.2 is reasonable.
+  def self.repetition_penalty(logits, ctx, p)
+    if p <= 1.0
+      return logits
+    end
+    seen = ctx.generated_ids
+    i = 0
+    while i < seen.length
+      tid = seen[i]
+      if tid >= 0 && tid < logits.ncols
+        v = logits.flat[tid]
+        if v > 0.0
+          logits.flat[tid] = v / p
+        else
+          logits.flat[tid] = v * p
+        end
+      end
+      i = i + 1
+    end
+    logits
+  end
+
+  # Keep top-k logits; mask the rest with -INFINITY (= -1e30 here so
+  # softmax never sees -Inf). k=0 disables.
+  def self.top_k(logits, k)
+    if k <= 0 || k >= logits.ncols
+      return logits
+    end
+    n = logits.ncols
+    # Find the k-th largest by k passes of argmax. O(k*n); fine for
+    # k ≤ ~100 at vocab=150K. For bigger k a real partial-sort would
+    # be better; current sizes don't need it.
+    kept = [0]
+    kept.pop
+    snapshot = [0.0]
+    snapshot.pop
+    j = 0
+    while j < n
+      snapshot.push(logits.flat[j])
+      j = j + 1
+    end
+    pass = 0
+    while pass < k
+      best_i = -1
+      best_v = NEG_INF_SCORE
+      j = 0
+      while j < n
+        v = snapshot[j]
+        if v > best_v
+          best_v = v
+          best_i = j
+        end
+        j = j + 1
+      end
+      if best_i < 0
+        # already all masked
+        return logits
+      end
+      kept.push(best_i)
+      snapshot[best_i] = NEG_INF_SCORE
+      pass = pass + 1
+    end
+    # Build keep-set as a flag array
+    keep = [false]
+    keep.pop
+    j = 0
+    while j < n
+      keep.push(false)
+      j = j + 1
+    end
+    j = 0
+    while j < kept.length
+      keep[kept[j]] = true
+      j = j + 1
+    end
+    j = 0
+    while j < n
+      if !keep[j]
+        logits.flat[j] = NEG_INF_SCORE
+      end
+      j = j + 1
+    end
+    logits
+  end
+
+  # Top-p / nucleus: softmax → cumulative sort → keep smallest set
+  # whose probability mass ≥ p. p>=1 disables.
+  def self.top_p(logits, p)
+    if p >= 1.0 || p <= 0.0
+      return logits
+    end
+    n = logits.ncols
+    # softmax in place into a copy
+    probs = [0.0]
+    probs.pop
+    max_v = NEG_INF_SCORE
+    j = 0
+    while j < n
+      v = logits.flat[j]
+      if v > max_v
+        max_v = v
+      end
+      j = j + 1
+    end
+    sum = 0.0
+    j = 0
+    while j < n
+      e = Math.exp(logits.flat[j] - max_v)
+      probs.push(e)
+      sum = sum + e
+      j = j + 1
+    end
+    inv_sum = 1.0 / sum
+    j = 0
+    while j < n
+      probs[j] = probs[j] * inv_sum
+      j = j + 1
+    end
+    # Sort indices by descending prob (selection-sort with mark; O(n^2)).
+    # Acceptable because vocab ≤ 200K and we typically prune via top_k
+    # first; a partial-sort would help if top_p were used solo at full
+    # vocab.
+    order = [0]
+    order.pop
+    taken = [false]
+    taken.pop
+    j = 0
+    while j < n
+      taken.push(false)
+      j = j + 1
+    end
+    cum = 0.0
+    pass = 0
+    while pass < n
+      best_i = -1
+      best_v = -1.0
+      j = 0
+      while j < n
+        if !taken[j] && probs[j] > best_v
+          best_v = probs[j]
+          best_i = j
+        end
+        j = j + 1
+      end
+      if best_i < 0
+        break
+      end
+      taken[best_i] = true
+      cum = cum + best_v
+      order.push(best_i)
+      if cum >= p
+        break
+      end
+      pass = pass + 1
+    end
+    # Mask anything NOT in `order`.
+    keep = [false]
+    keep.pop
+    j = 0
+    while j < n
+      keep.push(false)
+      j = j + 1
+    end
+    j = 0
+    while j < order.length
+      keep[order[j]] = true
+      j = j + 1
+    end
+    j = 0
+    while j < n
+      if !keep[j]
+        logits.flat[j] = NEG_INF_SCORE
+      end
+      j = j + 1
+    end
+    logits
+  end
+
+  # Final pick. If cfg.temperature <= 0, return argmax (greedy).
+  # Otherwise softmax + multinomial draw using ctx's RNG.
+  def self.pick(logits, cfg, ctx)
+    if cfg.temperature <= 0.0
+      return Sampler.argmax(logits)
+    end
+    Sampler.multinomial(logits, ctx)
+  end
+
+  def self.argmax(logits)
+    n = logits.ncols
+    best_i = 0
+    best_v = logits.flat[0]
+    j = 1
+    while j < n
+      v = logits.flat[j]
+      if v > best_v
+        best_v = v
+        best_i = j
+      end
+      j = j + 1
+    end
+    best_i
+  end
+
+  def self.multinomial(logits, ctx)
+    n = logits.ncols
+    # softmax
+    max_v = NEG_INF_SCORE
+    j = 0
+    while j < n
+      v = logits.flat[j]
+      if v > max_v
+        max_v = v
+      end
+      j = j + 1
+    end
+    sum = 0.0
+    j = 0
+    while j < n
+      sum = sum + Math.exp(logits.flat[j] - max_v)
+      j = j + 1
+    end
+    target = ctx.next_unit * sum
+    cum = 0.0
+    j = 0
+    while j < n
+      cum = cum + Math.exp(logits.flat[j] - max_v)
+      if cum >= target
+        return j
+      end
+      j = j + 1
+    end
+    n - 1
+  end
+end

data/lib/toy/train/toy_chat_template.rb

@@ -0,0 +1,179 @@
+# Per-arch chat templating: list of (role, content) messages → a
+# single String the tokenizer can encode. Mirrors HuggingFace's
+# `tokenizer.chat_template` but hardcoded per arch family (no Jinja
+# evaluator under Spinel).
+#
+# Usage:
+#   text = ToyChatTemplate.apply("chatml",
+#            [["system", "Be brief."], ["user", "Hi"]],
+#            true)   # true = append generation prompt
+#   ids  = tokenizer.encode(text)
+#
+# Families supported:
+#   chatml   — Qwen 2 / Qwen 2.5 / Qwen 3 / SmolLM2 / OLMoE / Mistral-Instruct-v0.3+
+#   llama3   — Llama 3 / Llama 3.1 / Llama 3.2 / Llama 3.3
+#   mistral  — Mistral-Instruct-v0.1 / v0.2 (pre-ChatML adoption)
+#   gemma2   — Gemma 2 (user / model turns)
+#
+# Detection convention (see ToyChatTemplate.detect_family below):
+#   architecture + tokenizer.ggml.model + presence of specific special
+#   tokens in the vocab. Returns "chatml" as the safe default for
+#   unknown llama-family GGUFs.
+#
+# Spinel notes:
+#   - messages is Array<Array<String>> where each inner array is
+#     [role, content]. Symbols not used (Spinel symbol-as-key fights
+#     poly dispatch).
+#   - No defaults / kwargs (per landmine #4). add_generation_prompt
+#     is a regular positional Boolean.
+
+module ToyChatTemplate
+  # Render the 4 supported families. Returns the concatenated string;
+  # tokenization is the caller's job (existing Tokenizer.encode path).
+  def self.apply(family, messages, add_generation_prompt)
+    if family == "chatml"
+      render_chatml(messages, add_generation_prompt)
+    elsif family == "llama3"
+      render_llama3(messages, add_generation_prompt)
+    elsif family == "mistral"
+      render_mistral(messages, add_generation_prompt)
+    elsif family == "gemma2"
+      render_gemma2(messages, add_generation_prompt)
+    else
+      # Default to chatml — the broadest covering format among toy's
+      # supported arches. Caller wanting strict matching should pass
+      # the explicit family string.
+      render_chatml(messages, add_generation_prompt)
+    end
+  end
+
+  # ChatML: <|im_start|>role\ncontent<|im_end|>\n
+  # Used by Qwen 2 / Qwen 2.5 / Qwen 3, SmolLM2, OLMoE, modern Mistral.
+  def self.render_chatml(messages, add_generation_prompt)
+    s = ""
+    i = 0
+    while i < messages.length
+      role    = messages[i][0]
+      content = messages[i][1]
+      s = s + "<|im_start|>" + role + "\n" + content + "<|im_end|>\n"
+      i = i + 1
+    end
+    if add_generation_prompt
+      s = s + "<|im_start|>assistant\n"
+    end
+    s
+  end
+
+  # Llama-3 turn structure:
+  #   <|begin_of_text|><|start_header_id|>role<|end_header_id|>\n\n
+  #   content<|eot_id|>
+  # The system tokenizer auto-prepends <|begin_of_text|> when add_bos
+  # is set; we emit it explicitly so the template is self-contained.
+  def self.render_llama3(messages, add_generation_prompt)
+    s = "<|begin_of_text|>"
+    i = 0
+    while i < messages.length
+      role    = messages[i][0]
+      content = messages[i][1]
+      s = s + "<|start_header_id|>" + role + "<|end_header_id|>\n\n"
+      s = s + content + "<|eot_id|>"
+      i = i + 1
+    end
+    if add_generation_prompt
+      s = s + "<|start_header_id|>assistant<|end_header_id|>\n\n"
+    end
+    s
+  end
+
+  # Mistral pre-ChatML: [INST] user [/INST] assistant</s>
+  # Subtleties (HF v0.1/v0.2):
+  #   - system message gets fused into the FIRST user turn
+  #   - [INST]/[/INST] wraps user; assistant text follows verbatim
+  #   - </s> closes each assistant turn; final has just [INST] user [/INST]
+  def self.render_mistral(messages, add_generation_prompt)
+    s = "<s>"
+    pending_system = ""
+    i = 0
+    while i < messages.length
+      role    = messages[i][0]
+      content = messages[i][1]
+      if role == "system"
+        # Fuse into next user turn — Mistral v0.x has no system role.
+        pending_system = content + "\n\n"
+      elsif role == "user"
+        s = s + "[INST] " + pending_system + content + " [/INST]"
+        pending_system = ""
+      elsif role == "assistant"
+        s = s + " " + content + "</s>"
+      end
+      i = i + 1
+    end
+    # Mistral doesn't need an explicit assistant marker — generation
+    # starts right after the last [/INST]. add_generation_prompt is
+    # effectively the absence-of-trailing-eos state.
+    s
+  end
+
+  # Gemma 2 turn structure:
+  #   <start_of_turn>role\ncontent<end_of_turn>\n
+  # Roles are "user" and "model" (not "assistant"); system messages
+  # have no dedicated role — they're fused into the first user turn,
+  # like Mistral.
+  def self.render_gemma2(messages, add_generation_prompt)
+    s = ""
+    pending_system = ""
+    i = 0
+    while i < messages.length
+      role    = messages[i][0]
+      content = messages[i][1]
+      if role == "system"
+        pending_system = content + "\n\n"
+      elsif role == "user"
+        body = pending_system + content
+        pending_system = ""
+        s = s + "<start_of_turn>user\n" + body + "<end_of_turn>\n"
+      elsif role == "assistant" || role == "model"
+        s = s + "<start_of_turn>model\n" + content + "<end_of_turn>\n"
+      end
+      i = i + 1
+    end
+    if add_generation_prompt
+      s = s + "<start_of_turn>model\n"
+    end
+    s
+  end
+
+  # Heuristic family detection from GGUF metadata. The caller passes
+  # the arch string from general.architecture plus the tokenizer model
+  # kind (tokenizer.ggml.model). Returns a family string suitable for
+  # apply().
+  #
+  # Rules:
+  #   - arch=gemma2 → gemma2
+  #   - arch=llama AND tokenizer.ggml.tokens contains "<|im_start|>" → chatml
+  #     (SmolLM2 + modern Mistral fall here)
+  #   - arch=llama AND tokens contains "<|begin_of_text|>" → llama3
+  #   - arch=llama (otherwise, classical) → mistral
+  #   - arch=qwen2 / qwen3 → chatml
+  #   - other → chatml as the modern default
+  #
+  # `has_im_start` and `has_bot` come from a vocab scan the caller
+  # does (tokenizer's @vocab_inv.has_key? on the marker strings).
+  def self.detect_family(arch, has_im_start, has_bot)
+    if arch == "gemma2"
+      "gemma2"
+    elsif arch == "qwen2" || arch == "qwen3"
+      "chatml"
+    elsif arch == "llama"
+      if has_im_start
+        "chatml"
+      elsif has_bot
+        "llama3"
+      else
+        "mistral"
+      end
+    else
+      "chatml"
+    end
+  end
+end

data/lib/toy/train/toy_drift_grad.rb

@@ -0,0 +1,176 @@
+# tao#drift-grad-sentinels — opt-in producers for `drift` and `grad`
+# events (schema v1, docs/events-schema.md).
+#
+# Two independent toggles, both driven from the caller-side loop:
+#
+#   ToyDriftGrad.params(sess) → Array<:ptr>
+#     enumerate PARAM-flagged tensors in dispatch order
+#
+#   ToyDriftGrad.snapshot_one(sess, t) → Mat
+#     download one PARAM's current values for use as a step-0 baseline
+#
+#   ToyDriftGrad.emit_drift_event(sess, t, snap_mat, step, t_now)
+#     compute (cos_to_init, l2_to_init) against `snap_mat`, emit one
+#     `drift` event
+#
+#   ToyDriftGrad.emit_grad_event(sess, t, step, t_now)
+#     download t's grad, emit one `grad` event with (l2, abs_mean,
+#     nan_count) computed via tnn_scratch_*
+#
+# Both are cheap-when-off — caller-side env guards (TOY_DRIFT_EVERY,
+# TOY_GRAD_SENTINELS) gate the calls entirely, so the unconfigured
+# path costs nothing beyond an env-read at boot.
+#
+# Spinel notes:
+#   - sp_Mat_ptr_array isn't a generated type, so we keep Mat-typed
+#     arrays in main scope (caller manages the snapshot[] array); the
+#     module never accepts or returns Array<Mat>.
+#   - Use `** 0.5` over Math.sqrt where the receiver could be poly-
+#     typed via Spinel's inference cascade (per landmine #12).
+#   - Mat round-trip uses tnn_download_to_f64_array → Mat.flat[i],
+#     same primitive Mat-roundtrip uses elsewhere in toy.
+#   - Param enumeration walks the graph leaves + nodes (same logic as
+#     ToyDescribeFlow.build_index but inlined to avoid cross-module
+#     dependency).
+
+require_relative "../../../vendor/spinel/spinel_kit/lib/spinel_kit/json_builder"
+
+module ToyDriftGrad
+  # Walk session graph + leaves, return ordered array of (ptr, name)
+  # for every PARAM-flagged tensor (flags bit 4 set). The flag is
+  # `ggml_set_param`-marked, so this exactly matches the params the
+  # training graph differentiates over.
+  def self.params(sess)
+    ptrs = [TinyNN.tnn_null_ptr]; ptrs.pop
+    seen = [TinyNN.tnn_null_ptr]; seen.pop
+
+    # Collect every distinct tensor reachable from the graph: nodes +
+    # their src dependencies. Filter to PARAM-flagged.
+    n_nodes = TinyNN.tnn_graph_n_nodes(sess)
+    i = 0
+    while i < n_nodes
+      node = TinyNN.tnn_graph_node(sess, i)
+      # The node itself
+      seen.push(node)
+      # Its srcs (leaves usually carry the PARAM flag)
+      si = 0
+      while si < 10
+        src = TinyNN.tnn_tensor_src(node, si)
+        if src == nil || src == TinyNN.tnn_null_ptr
+          si = 10
+        else
+          # Already seen?
+          dup = false
+          k = 0
+          while k < seen.length
+            if seen[k] == src
+              dup = true
+              k = seen.length
+            else
+              k = k + 1
+            end
+          end
+          if !dup
+            seen.push(src)
+          end
+          si = si + 1
+        end
+      end
+      i = i + 1
+    end
+
+    # Filter by PARAM flag.
+    i = 0
+    while i < seen.length
+      t = seen[i]
+      flags = TinyNN.tnn_tensor_flags(t)
+      if (flags & 4) != 0
+        ptrs.push(t)
+      end
+      i = i + 1
+    end
+    ptrs
+  end
+
+  # Snapshot ONE param's current values into a fresh Mat. Caller
+  # keeps these in a main-scope Array<Mat> (a separate per-step Array
+  # would force Spinel to declare sp_Mat_ptr_array, which it can't).
+  def self.snapshot_one(sess, t)
+    n = TinyNN.tnn_tensor_nelements(t)
+    m = Mat.new(1, n)
+    TinyNN.tnn_download_to_f64_array(sess, t, m.flat, n)
+    m
+  end
+
+  # Emit one `drift` event for tensor t, comparing against snap_mat.
+  # Single-pass cos + L2 loop over the f64 flats.
+  def self.emit_drift_event(sess, t, snap_mat, step, t_now)
+    n = TinyNN.tnn_tensor_nelements(t)
+    cur = Mat.new(1, n)
+    TinyNN.tnn_download_to_f64_array(sess, t, cur.flat, n)
+
+    dot = 0.0
+    sum_sq_s = 0.0
+    sum_sq_c = 0.0
+    sum_sq_diff = 0.0
+    i = 0
+    while i < n
+      sv = snap_mat.flat[i]
+      cv = cur.flat[i]
+      dot = dot + sv * cv
+      sum_sq_s = sum_sq_s + sv * sv
+      sum_sq_c = sum_sq_c + cv * cv
+      d = sv - cv
+      sum_sq_diff = sum_sq_diff + d * d
+      i = i + 1
+    end
+    norm_s = sum_sq_s ** 0.5
+    norm_c = sum_sq_c ** 0.5
+    cos_to_init = 0.0
+    if norm_s > 0.0 && norm_c > 0.0
+      cos_to_init = dot / (norm_s * norm_c)
+    end
+    l2_to_init = sum_sq_diff ** 0.5
+
+    name = TinyNN.tnn_tensor_name(t)
+    ev = SpinelKit::Json::Builder.new
+    ev.add_str("kind",  "drift")
+    ev.add_str("phase", "train")
+    ev.add_num("t",           t_now)
+    ev.add_num("step",        step)
+    ev.add_str("param",       name)
+    ev.add_num("cos_to_init", cos_to_init)
+    ev.add_num("l2_to_init",  l2_to_init)
+    TinyNN.tnn_events_emit(ev.dump)
+  end
+
+  # Emit one `grad` event for tensor t. Call after tnn_compute_backward
+  # but before the next step's reset (so the grad tensor is still live).
+  def self.emit_grad_event(sess, t, step, t_now)
+    g = TinyNN.tnn_tensor_grad(sess, t)
+    if g == nil || g == TinyNN.tnn_null_ptr
+      return
+    end
+    n = TinyNN.tnn_tensor_nelements(g)
+    TinyNN.tnn_download(sess, g)
+    sum_sq   = TinyNN.tnn_scratch_sum_sq_f32(sess, n)
+    sum_abs  = TinyNN.tnn_scratch_sum_abs_f32(sess, n)
+    nan      = TinyNN.tnn_scratch_nan_count_f32(sess, n)
+    l2       = sum_sq ** 0.5
+    abs_mean = sum_abs / n.to_f
+    ne0 = TinyNN.tnn_tensor_ne0(g)
+    ne1 = TinyNN.tnn_tensor_ne1(g)
+    name = TinyNN.tnn_tensor_name(t)
+    ev = SpinelKit::Json::Builder.new
+    ev.add_str("kind",  "grad")
+    ev.add_str("phase", "train")
+    ev.add_num("t",         t_now)
+    ev.add_num("step",      step)
+    ev.add_str("param",     name)
+    ev.add_raw("shape",     "[" + ne0.to_s + "," + ne1.to_s + "]")
+    ev.add_num("l2",        l2)
+    ev.add_num("abs_mean",  abs_mean)
+    ev.add_num("nan_count", nan)
+    TinyNN.tnn_events_emit(ev.dump)
+  end
+end