p_1_2_atg_anw.livemd

some space to sink 2u - supasimpsink- - bgbg em - fork

System.put_env("XLA_BUILD", "true")

Mix.install([
  {:nx, "~> 0.11.0"},
  {:exla, "~> 0.11.0"},
  {:jason, "~> 1.4"}
])

System.put_env("HSA_OVERRIDE_GFX_VERSION", "11.0.0")
System.put_env("XLA_TARGET", "rocm")

System.put_env("TF_ROCM_FUSION_ENABLE", "0")

System.put_env("XLA_FLAGS", "--xla_gpu_autotune_level=0 --xla_gpu_enable_reassociation_for_converted_ar=false --xla_gpu_graph_enable_concurrent_region=false")

System.put_env("ROCM_DISABLE_CU_MEMSET", "1")

Application.put_env(:exla, :clients,
  rocm: [
    platform: :rocm,
    memory_fraction: 0.75,
    preallocate: false,
    num_replicas: 1
  ]
)

Application.put_env(:exla, :default_client, :rocm)

gpu = {EXLA.Backend, client: :rocm}
Nx.global_default_backend(gpu)
Nx.Defn.default_options(compiler: EXLA)

platforms = EXLA.Client.get_supported_platforms()
if platforms[:rocm] do
  IO.puts("gpu found")
else
  IO.puts("gpu missing")
end

Helpers

defmodule DLM.Tree do
  @moduledoc """
  Recursively maps a function over nested tuples of tensors.
  """
  
  def map(tuple, func) when is_tuple(tuple) do
    tuple
    |> Tuple.to_list()
    |> Enum.map(&map(&1, func))
    |> List.to_tuple()
  end

  # The base case: When it hits a tensor (or a compile-time Expr), apply the function
  def map(tensor, func) do
    func.(tensor)
  end
end

{:module, DLM.Tree, <<70, 79, 82, 49, 0, 0, 8, ...>>, {:map, 2}}

defmodule BabyLM.CurriculumStreamer do
  def build_infinite_stream(directory_path, batch_size) do
    target_shards = Path.wildcard(Path.join(directory_path, "*.bin"))

    if length(target_shards) == 0 do
      raise "CRITICAL: No .bin shards found in directory: #{directory_path}"
    end

    Stream.iterate(1, &(&1 + 1))
    |> Stream.flat_map(fn epoch ->
      :rand.seed(:exsss, {epoch, epoch, epoch})
      shuffled_shards = Enum.shuffle(target_shards)
      IO.puts("\n🔄 Epoch #{epoch} Started | Shuffling #{length(target_shards)} shards...")
      Stream.flat_map(shuffled_shards, fn path ->
        tensor = File.read!(path) |> :erlang.binary_to_term()
          
        total_seqs = Nx.axis_size(tensor, 0)
        num_batches = div(total_seqs, batch_size)

        # 1. Generate all the slice indices
        0..(num_batches - 1)
        # 2. Slice the tensor into actual batch tensors
        |> Enum.map(fn i ->
          Nx.slice_along_axis(tensor, i * batch_size, batch_size, axis: 0)
        end)
        # 3. Shuffle the list of batches
        |> Enum.shuffle()
        end) #|> Enum.shuffle() # shuffle it all together now
    end)
  end
end

{:module, BabyLM.CurriculumStreamer, <<70, 79, 82, 49, 0, 0, 17, ...>>, {:build_infinite_stream, 2}}

defmodule DLM.StatelessMuon do
  import Nx.Defn

  defn newton_schulz(g) do
    frobenius_norm = Nx.LinAlg.norm(g) # defaults to frobenius if type not given
    x = g / frobenius_norm
    
    {final_x, _} = while {curr_x = x, i = 0}, Nx.less(i, 5) do
      x_t_x = Nx.dot(Nx.transpose(curr_x), curr_x)
      next_x = 1.5 * curr_x - 0.5 * Nx.dot(curr_x, x_t_x)
      {next_x, i + 1}
    end
    final_x
  end

  defn step(em, grad, lr) do
    ns_grad = newton_schulz(grad)
    update = Nx.multiply(lr, ns_grad)
    Nx.subtract(em, update)
  end

  defn step_decay(em, grad, lr, wd \\ 1.0e-4) do
    ns_grad = newton_schulz(grad)
    update = Nx.multiply(lr, ns_grad)
    em_decayed = Nx.multiply(em, Nx.subtract(1.0, Nx.multiply(lr, wd)))
    Nx.subtract(em_decayed, update)
  end
end

{:module, DLM.StatelessMuon, <<70, 79, 82, 49, 0, 0, 23, ...>>, true}

defmodule DLM.SparseSinkGD do
  import Nx.Defn

  @doc """
  Applies Sinkhorn Gradient Scaling strictly to active (non-zero) gradients.
  """
  defn normalize_sparse_grad(grad) do
    eps = 1.0e-9
    # 1. Identify which characters actually appeared in this batch
    raw_row_norms = Nx.pow(grad, 2) |> Nx.sum(axes: [-1], keep_axes: true) |> Nx.sqrt()
    active_mask = Nx.greater(raw_row_norms, eps)
    
    # 2. Sinkhorn Alternating Loop
    result = while {i = 0, g = grad, am = active_mask}, Nx.less(i, 9) do
      # --- Row Normalization ---
      row_norms = g |> Nx.pow(2) |> Nx.sum(axes: [-1], keep_axes: true) |> Nx.sqrt()
      safe_row_norms = Nx.select(am, row_norms, 1.0)
      g_row = Nx.divide(g, safe_row_norms)
      
      # --- Column Normalization ---
      col_norms = g_row |> Nx.pow(2) |> Nx.sum(axes: [0], keep_axes: true) |> Nx.sqrt()
      safe_col_norms = Nx.select(Nx.greater(col_norms, eps), col_norms, 1.0)
      g_col = Nx.divide(g_row, safe_col_norms)
      
      {i + 1, g_col, am}
    end
    
    elem(result, 1)
  end

  defn step(embed_matrix, grad, lr) do
    # 1. Scale the raw sparse gradient to preserve boolean mask logic
    scaled_grad = normalize_sparse_grad(grad)
    
    # 2. Apply the pure gradient descent subtraction
    new_embeds = Nx.subtract(embed_matrix, Nx.multiply(scaled_grad, lr))
    
    new_embeds
  end

  defn normalize_sparse_grad_3d(grad) do
    eps = 1.0e-9

    # 1. Identify active tokens -> shape: {128, 1, 1}
    raw_row_norms =
      Nx.pow(grad, 2)
      |> Nx.sum(axes: [-2, -1], keep_axes: true)
      |> Nx.sqrt()

    active_mask = Nx.greater(raw_row_norms, eps)

    # 2. Sinkhorn Alternating Loop — one pass per axis, innermost → outermost
    result =
      while {i = 0, g = grad, am = active_mask}, Nx.less(i, 9) do

        # --- Axis -1 (last) Normalization ---
        row_norms = g |> Nx.pow(2) |> Nx.sum(axes: [-1], keep_axes: true) |> Nx.sqrt()
        
        # FIX: Explicitly broadcast the condition tensor (am) to match row_norms ({128, 128, 1})
        # This allows Nx.select to map elements correctly under loop vectorization.
        am_broadcasted = Nx.broadcast(am, Nx.shape(row_norms))
        safe_row_norms = Nx.select(am_broadcasted, row_norms, 1.0)
        g_ax2 = Nx.divide(g, safe_row_norms)

        # --- Axis -2 (middle) Normalization ---
        mid_norms = g_ax2 |> Nx.pow(2) |> Nx.sum(axes: [-2], keep_axes: true) |> Nx.sqrt()
        safe_mid_norms = Nx.select(Nx.greater(mid_norms, eps), mid_norms, 1.0)
        g_ax1 = Nx.divide(g_ax2, safe_mid_norms)

        # --- Axis 0 (first) Normalization ---
        col_norms = g_ax1 |> Nx.pow(2) |> Nx.sum(axes: [0], keep_axes: true) |> Nx.sqrt()
        safe_col_norms = Nx.select(Nx.greater(col_norms, eps), col_norms, 1.0)
        g_col = Nx.divide(g_ax1, safe_col_norms)

        {i + 1, g_col, am}
      end

    elem(result, 1)
  end

  defn step_3d(embed_matrix, grad, lr) do
    scaled_grad = normalize_sparse_grad(grad)
    Nx.subtract(embed_matrix, Nx.multiply(scaled_grad, lr))
  end
end

{:module, DLM.SparseSinkGD, <<70, 79, 82, 49, 0, 0, 30, ...>>, true}

SSM

defmodule DLM.SvdSSM do
  import Nx.Defn

  defn rms_norm(x, epsilon \\ 1.0e-6) do
    variance = Nx.mean(Nx.pow(x, 2), axes: [-1], keep_axes: true)
    x * Nx.rsqrt(variance + epsilon)
  end

  defn layer_norm(x, gamma, beta, epsilon \\ 1.0e-6) do
    norm = rms_norm(x, epsilon)
    scaled = norm * gamma
    scaled + beta
  end

  defn gelu(x) do
    cdf = Nx.multiply(
      0.7978845608,
      Nx.add(x, Nx.multiply(0.044715, Nx.pow(x, 3)))
    )
    Nx.multiply(Nx.multiply(0.5, x), Nx.add(1.0, Nx.tanh(cdf)))
  end

  defn softmax(t) do
    max_val = Nx.reduce_max(t, axes: [-1], keep_axes: true)
    safe_t = Nx.subtract(t, max_val)
    exps = Nx.exp(safe_t)
    sum_exps = Nx.sum(exps, axes: [-1], keep_axes: true)
    Nx.divide(exps, Nx.add(sum_exps, 1.0e-8))
  end
  
  defn feed_forward_layer(x, um, dm, bi, sc, ga, be) do
    up = Nx.dot(x, um) # up projection
    act = gelu(up) # gelu activation, tanh also works similarly well
    scaled = Nx.multiply(act, sc) # scale projection
    biased = Nx.add(scaled, bi) # add bias to projection
    out = Nx.dot(biased, dm) |> Nx.add(x) # down projection
    layer_norm(out, ga, be) # rms norm with scaler and bias
  end

  defn apply_xsa(x, y) do # exclusive 'space' attention
    y_norm_sq = Nx.add(Nx.sum(Nx.pow(y, 2), axes: [-1], keep_axes: true), 1.0e-6)
    y_unit = Nx.divide(y, Nx.sqrt(y_norm_sq))
    y_ov = Nx.sum(Nx.multiply(x, y_unit), axes: [-1], keep_axes: true)
    Nx.subtract(x, Nx.multiply(y_ov, y_unit))
  end

  defn apply_dot_attention(q, em) do
    r = Nx.axis_size(em, 1)
    scores  = Nx.dot(q, Nx.transpose(em)) 
    scaled  = Nx.divide(scores, Nx.sqrt(r))
    weights = softmax(scaled)
    attended = Nx.dot(weights, em)
    apply_xsa(attended, q) |> Nx.add(q)
  end

  defn apply_attention(q, v) do
    k = Nx.transpose(v, axes: [0, 2, 1])
    r = Nx.axis_size(v, 2)
    qu = Nx.new_axis(q, 1)
    score = Nx.dot(qu, [2], [0], k, [1], [0])
    scaled_score = Nx.divide(score, Nx.sqrt(r))
    soft_score = softmax(scaled_score)
    at = Nx.dot(soft_score, [2], [0], v, [1], [0]) |> Nx.squeeze(axes: [1])
    apply_xsa(at, q) |> Nx.add(q)
  end

  defn forward_step(x_t, h_prev, p) do
    {a1, a2, ab, aq, ag, ae,
     b1, b2, bb, bq, bg, be, 
     c1, c2, cb, cq, cg, ce,
     d1, d2, db, dq, dg, de, emb, gw} = p
    batch_size = Nx.axis_size(x_t, 0)
    em = Nx.take(emb, x_t)
    # process h_prev
    hp_at = apply_attention(h_prev, em)
    h_term = feed_forward_layer(hp_at, a1, a2, ab, aq, ag, ae)
    # process current character 
    batch_indices = Nx.iota({batch_size})
    coords = Nx.stack([batch_indices, x_t], axis: -1)
    u_t = Nx.gather(em, coords)
    u_at = apply_attention(u_t, em)
    b_term = feed_forward_layer(u_at, b1, b2, bb, bq, bg, be)
    # update hidden state
    h_new = Nx.add(b_term, h_term)
    # attend 'pieces' of h_new to vocab
    b_at = apply_attention(b_term, em)
    h_at = apply_attention(h_term, em)
    # predictive ffw layers
    c_term = feed_forward_layer(h_at, c1, c2, cb, cq, cg, ce)
    d_term = feed_forward_layer(b_at, d1, d2, db, dq, dg, de)
    # add results for final logits
    ctxr = Nx.add(c_term, d_term)
    # Learn a scalar gate per character dim
    gate = Nx.sigmoid(Nx.dot(h_new, gw))   # gate_w is a 128x128 learned matrix
    logits = Nx.add(
      Nx.multiply(gate, u_t),
      Nx.multiply(Nx.subtract(1.0, gate), ctxr)
    )
    {logits, h_new}
  end
end

{:module, DLM.SvdSSM, <<70, 79, 82, 49, 0, 0, 52, ...>>, true}

defmodule DLM.SvdOptimizer do
  import Nx.Defn

  defn update(params, grads, step, lr) do
    {a1, a2, ab, aq, ag, ae,
     b1, b2, bb, bq, bg, be, 
     c1, c2, cb, cq, cg, ce,
     d1, d2, db, dq, dg, de, bem, gw}  = params   
    {ga1, ga2, gab, gaq, gag, gae,
     gb1, gb2, gbb, gbq, gbg, gbe, 
     gc1, gc2, gcb, gcq, gcg, gce, 
     gd1, gd2, gdb, gdq, gdg, gde, gbem, ggw}  = grads
    
    new_step = Nx.add(step, 1)

    # optimizer update step
    a1_new = DLM.SparseSinkGD.step(a1, ga1, lr)
    a2_new = DLM.SparseSinkGD.step(a2, ga2, lr)
    ab_new = DLM.SparseSinkGD.step(ab, gab, lr)
    aq_new = DLM.SparseSinkGD.step(aq, gaq, lr)
    ae_new = DLM.SparseSinkGD.step(ae, gae, lr)
    ag_new = DLM.SparseSinkGD.step(ag, gag, lr)

    b1_new = DLM.SparseSinkGD.step(b1, gb1, lr)
    b2_new = DLM.SparseSinkGD.step(b2, gb2, lr)
    bb_new = DLM.SparseSinkGD.step(bb, gbb, lr)
    bq_new = DLM.SparseSinkGD.step(bq, gbq, lr)
    be_new = DLM.SparseSinkGD.step(be, gbe, lr)
    bg_new = DLM.SparseSinkGD.step(bg, gbg, lr)

    c1_new = DLM.SparseSinkGD.step(c1, gc1, lr)
    c2_new = DLM.SparseSinkGD.step(c2, gc2, lr)
    cb_new = DLM.SparseSinkGD.step(cb, gcb, lr)
    cq_new = DLM.SparseSinkGD.step(cq, gcq, lr)
    ce_new = DLM.SparseSinkGD.step(ce, gce, lr)
    cg_new = DLM.SparseSinkGD.step(cg, gcg, lr)

    d1_new = DLM.SparseSinkGD.step(d1, gd1, lr)
    d2_new = DLM.SparseSinkGD.step(d2, gd2, lr)
    db_new = DLM.SparseSinkGD.step(db, gdb, lr)
    dq_new = DLM.SparseSinkGD.step(dq, gdq, lr)
    de_new = DLM.SparseSinkGD.step(de, gde, lr)
    dg_new = DLM.SparseSinkGD.step(dg, gdg, lr)

    gw_new = DLM.SparseSinkGD.step(gw, ggw, lr)

    bem_new = DLM.SparseSinkGD.step_3d(bem, gbem, lr)
    
    
    new_params = {a1_new, a2_new, ab_new, aq_new, ag_new, ae_new,
                  b1_new, b2_new, bb_new, bq_new, bg_new, be_new,
                  c1_new, c2_new, cb_new, cq_new, cg_new, ce_new,
                  d1_new, d2_new, db_new, dq_new, dg_new, de_new,
                  bem_new, gw_new}

    {new_params, new_step}
  end
end

{:module, DLM.SvdOptimizer, <<70, 79, 82, 49, 0, 0, 24, ...>>, true}

defmodule DLM.PhysicsGenesis do
  import Nx.Defn

  @dim 512
  @vocab 128

  defn init(k1) do
    std_voc = Nx.divide(1.0, Nx.sqrt(Nx.tensor(@vocab, type: :f32)))
    std_dim = Nx.divide(1.0, Nx.sqrt(Nx.tensor(@dim, type: :f32)))

    # projection matrices
    {a1, k2} = Nx.Random.normal(k1, 0.0, std_voc, shape: {@vocab, @dim})
    {a2, k3} = Nx.Random.normal(k2, 0.0, std_dim, shape: {@dim, @vocab})
    {b1, k4} = Nx.Random.normal(k3, 0.0, std_voc, shape: {@vocab, @dim})
    {b2, k5} = Nx.Random.normal(k4, 0.0, std_dim, shape: {@dim, @vocab})
    {c1, k6} = Nx.Random.normal(k5, 0.0, std_voc, shape: {@vocab, @dim})
    {c2, k7} = Nx.Random.normal(k6, 0.0, std_dim, shape: {@dim, @vocab})
    {d1, k8} = Nx.Random.normal(k7, 0.0, std_voc, shape: {@vocab, @dim})
    {d2, k9} = Nx.Random.normal(k8, 0.0, std_dim, shape: {@dim, @vocab})
    {bem, k10} = Nx.Random.normal(k9, 0.0, std_voc, shape: {@vocab, @vocab, @vocab})
    {gw, k11} = Nx.Random.normal(k10, 0.0, std_voc, shape: {@vocab, @vocab})

    # projection bias vectors
    ab = Nx.broadcast(0.0, {@dim}) |> Nx.as_type(:f32)
    bb = Nx.broadcast(0.0, {@dim}) |> Nx.as_type(:f32)
    cb = Nx.broadcast(0.0, {@dim}) |> Nx.as_type(:f32)
    db = Nx.broadcast(0.0, {@dim}) |> Nx.as_type(:f32)
    # projection scaling vectors
    aq = Nx.broadcast(1.0, {@dim}) |> Nx.as_type(:f32)
    bq = Nx.broadcast(1.0, {@dim}) |> Nx.as_type(:f32)
    cq = Nx.broadcast(1.0, {@dim}) |> Nx.as_type(:f32)
    dq = Nx.broadcast(1.0, {@dim}) |> Nx.as_type(:f32)
    # rms norm beta and gamma vectors
    ae = Nx.broadcast(0.0, {@vocab}) |> Nx.as_type(:f32)
    ag = Nx.broadcast(1.0, {@vocab}) |> Nx.as_type(:f32)
    be = Nx.broadcast(0.0, {@vocab}) |> Nx.as_type(:f32)
    bg = Nx.broadcast(1.0, {@vocab}) |> Nx.as_type(:f32)
    ce = Nx.broadcast(0.0, {@vocab}) |> Nx.as_type(:f32)
    cg = Nx.broadcast(1.0, {@vocab}) |> Nx.as_type(:f32)
    de = Nx.broadcast(0.0, {@vocab}) |> Nx.as_type(:f32)
    dg = Nx.broadcast(1.0, {@vocab}) |> Nx.as_type(:f32)


    
    trainable_params = {a1, a2, ab, aq, ag, ae,
                        b1, b2, bb, bq, bg, be, 
                        c1, c2, cb, cq, cg, ce,
                        d1, d2, db, dq, dg, de, bem, gw} 
    
    {trainable_params, k11}
  end
end

{:module, DLM.PhysicsGenesis, <<70, 79, 82, 49, 0, 0, 21, ...>>, true}

BLiMP

defmodule DLM.Phase2_Eval do
  import Nx.Defn

  defn step_cross_entropy(logits, targets) do
    vocab_size = Nx.axis_size(logits, -1)
    one_hot = Nx.equal(Nx.iota({vocab_size}), Nx.new_axis(targets, -1))
    
    max_logits = Nx.reduce_max(logits, axes: [-1], keep_axes: true)
    safe_logits = Nx.subtract(logits, max_logits)
    log_probs = Nx.subtract(safe_logits, Nx.log(Nx.sum(Nx.exp(safe_logits), axes: [-1], keep_axes: true)))
    
    Nx.negate(Nx.sum(Nx.multiply(one_hot, log_probs), axes: [-1]))
  end

  defn compute_sentence_loss(input_tokens, target_tokens, h_init, trainable) do
    seq_len = Nx.axis_size(input_tokens, 1) -1
    batch_size = Nx.axis_size(input_tokens, 0)
      
    # Explicitly initialize the batch-sized accumulators
    zero_batch = Nx.broadcast(0.0, {batch_size})

    result =
      while {i = 0, 
             total_loss = zero_batch, 
             valid_tokens = zero_batch,
             h = h_init,
             inf = input_tokens,
             tgf = target_tokens,
             tr = trainable},
            Nx.less(i, seq_len) do
  
        # 1. Grab the single character context and target
        x_t = inf[[.., i]]
        t_t = tgf[[.., i]]
        
        # 2. Forward pass through the reservoir
        {logits, h_new} = DLM.SvdSSM.forward_step(x_t, h, tr)
        
        # 3. Calculate 1-step prediction loss
        step_loss = step_cross_entropy(logits, t_t)
  
        # 4. MASK: 1.0 if valid text, 0.0 if ASCII 1 (SOH / Padding)
        mask = Nx.not_equal(t_t, 1) |> Nx.as_type(:f32)
        masked_loss = Nx.multiply(step_loss, mask)

        {i + 1, 
         Nx.add(total_loss, masked_loss), 
         Nx.add(valid_tokens, mask),
         h_new, inf, tgf, tr}
      end
  
    {_, acc_loss, val_tokens, _, _, _, _} = result

    total_loss = Nx.divide(acc_loss, Nx.add(val_tokens, 1.0e-8))

    total_loss
  end
end

{:module, DLM.Phase2_Eval, <<70, 79, 82, 49, 0, 0, 25, ...>>, true}

gpu = {EXLA.Backend, client: :rocm}
checkpoint_path = "v7667_DLM_512sq_bgbgem1237gw_checkpoint_step_3400.bin"
%{trainable: t_p} = File.read!(checkpoint_path) |> :erlang.binary_to_term()
trainable_gpu = DLM.Tree.map(t_p, &Nx.backend_copy(&1, gpu))

warmup_input  = Nx.broadcast(1, {200, 255}) |> Nx.as_type(:s32) |> Nx.backend_copy(gpu)
warmup_target = Nx.broadcast(1, {200, 255}) |> Nx.as_type(:s32) |> Nx.backend_copy(gpu)
h_init_warm   = Nx.broadcast(0.0, {200, 128}) |> Nx.backend_copy(gpu)

DLM.Phase2_Eval.compute_sentence_loss(warmup_input, warmup_target, h_init_warm, trainable_gpu)

run_blimp = fn path ->
  pairs =
    File.read!(path)
    |> String.split("\n", trim: true)
    |> Enum.map(&Jason.decode!/1)

  total_pairs = length(pairs)
  # 50 for supplimental
  pairs_per_batch = 100 

  {passes, total_loss_diff} =
    pairs
    |> Enum.chunk_every(pairs_per_batch)
    |> Enum.reduce({0, 0.0}, fn chunk, {acc_passes, acc_diff} ->
      
      flat_data =
        Enum.flat_map(chunk, fn pair ->
          good_chars = pair["sentence_good"] |> to_charlist()
          bad_chars  = pair["sentence_bad"]  |> to_charlist()

          good_len = length(good_chars)
          bad_len  = length(bad_chars)

          good_t = Nx.tensor(good_chars, type: :s32)
          bad_t  = Nx.tensor(bad_chars,  type: :s32)

          good_pad = 256 - good_len
          bad_pad  = 256 - bad_len

          # Pad the full sentence to 512
          good_full = Nx.slice_along_axis(good_t, 0, good_len) |> Nx.pad(1, [{0, good_pad, 0}])
          bad_full = Nx.slice_along_axis(bad_t, 0, bad_len) |> Nx.pad(1, [{0, bad_pad, 0}])

          [good_full, bad_full]
        end)

      actual_sentences = length(flat_data)
      batch_full = flat_data |> Nx.stack() |> Nx.backend_copy(gpu)
      
      seq_len = Nx.axis_size(batch_full, 1) - 1
      batch_input  = Nx.slice_along_axis(batch_full, 0, seq_len, axis: 1)
      batch_target = Nx.slice_along_axis(batch_full, 1, seq_len, axis: 1)

      h_init = Nx.broadcast(0.0, {actual_sentences, 128}) |> Nx.backend_copy(gpu)
      
      losses = DLM.Phase2_Eval.compute_sentence_loss(
        batch_input, batch_target, h_init, trainable_gpu
      )

      losses_list = Nx.to_flat_list(losses)

      {chunk_passes, chunk_diff} =
        losses_list
        |> Enum.chunk_every(2)
        |> Enum.reduce({0, 0.0}, fn [g_loss, b_loss], {p_acc, d_acc} ->
          new_p = if g_loss < b_loss, do: p_acc + 1, else: p_acc
          {new_p, d_acc + (b_loss - g_loss)}
        end)

      :erlang.garbage_collect(self())

      {acc_passes + chunk_passes, acc_diff + chunk_diff}
    end)

  accuracy   = (passes / total_pairs) * 100.0
  avg_margin = total_loss_diff / total_pairs
  name       = Path.basename(path, ".jsonl")

  {name, accuracy, avg_margin, passes, total_pairs}
end

# unfiltered - ./data/blimp_unfiltered
# ./data/bbylm_non_blimp_evals/*.jsonl -> ./data/bblm_blimp_26_quick/*.jsonl
dataset_files = Path.wildcard("./data/bblm_blimp_26_quick/*.jsonl") |> Enum.sort()
total_files   = length(dataset_files)

IO.puts("🚀 Running BLiMP batched eval across #{total_files} datasets...\n")

results =
  dataset_files
  |> Enum.with_index(1)
  |> Enum.map(fn {path, idx} ->
    IO.puts("  [#{idx}/#{total_files}] #{Path.basename(path)}...")
    result = run_blimp.(path)
    :erlang.garbage_collect(self())
    result
  end)

sorted = Enum.sort_by(results, fn {_, acc, _, _, _} -> acc end, :desc)

IO.puts("\n")
IO.puts(String.duplicate("=", 70))
IO.puts("🏆 BLIMP BATCH RESULTS — sorted by accuracy")
IO.puts(String.duplicate("=", 70))
IO.puts(String.pad_trailing("Dataset", 48) <> "Accuracy   Margin")
IO.puts(String.duplicate("-", 70))

Enum.each(sorted, fn {name, acc, margin, passes, total} ->
  flag = cond do
    acc >= 60.0 -> "✅"
    acc <= 40.0 -> "❌"
    true        -> "〰️"
  end
  IO.puts("#{flag} #{String.pad_trailing(name, 46)} #{String.pad_leading("#{Float.round(acc, 1)}%", 7)} (#{passes}/#{total})   #{Float.round(margin, 4)}")
end)

IO.puts(String.duplicate("-", 70))

# Overall average
avg_acc = results |> Enum.map(fn {_, acc, _, _, _} -> acc end) |> Enum.sum() |> Kernel./(total_files)
IO.puts("📊 Overall average accuracy: #{Float.round(avg_acc, 2)}%")
IO.puts(String.duplicate("=", 70))

🚀 Running BLiMP batched eval across 66 datasets...

  [1/66] adjunct_island.jsonl...
  [2/66] anaphor_gender_agreement.jsonl...
  [3/66] anaphor_number_agreement.jsonl...
  [4/66] animate_subject_passive.jsonl...
  [5/66] animate_subject_trans.jsonl...
  [6/66] causative.jsonl...
  [7/66] complex_NP_island.jsonl...
  [8/66] coordinate_structure_constraint_complex_left_branch.jsonl...
  [9/66] coordinate_structure_constraint_object_extraction.jsonl...
  [10/66] determiner_noun_agreement_1.jsonl...
  [11/66] determiner_noun_agreement_2.jsonl...
  [12/66] determiner_noun_agreement_irregular_1.jsonl...
  [13/66] determiner_noun_agreement_irregular_2.jsonl...
  [14/66] determiner_noun_agreement_with_adj_2.jsonl...
  [15/66] determiner_noun_agreement_with_adj_irregular_1.jsonl...
  [16/66] determiner_noun_agreement_with_adj_irregular_2.jsonl...
  [17/66] determiner_noun_agreement_with_adjective_1.jsonl...
  [18/66] distractor_agreement_relational_noun.jsonl...
  [19/66] distractor_agreement_relative_clause.jsonl...
  [20/66] ellipsis_n_bar_1.jsonl...
  [21/66] ellipsis_n_bar_2.jsonl...
  [22/66] existential_there_object_raising.jsonl...
  [23/66] existential_there_quantifiers_1.jsonl...
  [24/66] existential_there_quantifiers_2.jsonl...
  [25/66] existential_there_subject_raising.jsonl...
  [26/66] expletive_it_object_raising.jsonl...
  [27/66] inchoative.jsonl...
  [28/66] intransitive.jsonl...
  [29/66] irregular_past_participle_adjectives.jsonl...
  [30/66] irregular_past_participle_verbs.jsonl...
  [31/66] irregular_plural_subject_verb_agreement_1.jsonl...
  [32/66] irregular_plural_subject_verb_agreement_2.jsonl...
  [33/66] left_branch_island_echo_question.jsonl...
  [34/66] left_branch_island_simple_question.jsonl...
  [35/66] matrix_question_npi_licensor_present.jsonl...
  [36/66] npi_present_1.jsonl...
  [37/66] npi_present_2.jsonl...
  [38/66] only_npi_licensor_present.jsonl...
  [39/66] only_npi_scope.jsonl...
  [40/66] passive_1.jsonl...
  [41/66] passive_2.jsonl...
  [42/66] principle_A_c_command.jsonl...
  [43/66] principle_A_case_1.jsonl...
  [44/66] principle_A_case_2.jsonl...
  [45/66] principle_A_domain_1.jsonl...
  [46/66] principle_A_domain_2.jsonl...
  [47/66] principle_A_domain_3.jsonl...
  [48/66] principle_A_reconstruction.jsonl...
  [49/66] regular_plural_subject_verb_agreement_1.jsonl...
  [50/66] regular_plural_subject_verb_agreement_2.jsonl...
  [51/66] sentential_negation_npi_licensor_present.jsonl...
  [52/66] sentential_negation_npi_scope.jsonl...
  [53/66] sentential_subject_island.jsonl...
  [54/66] superlative_quantifiers_1.jsonl...
  [55/66] superlative_quantifiers_2.jsonl...
  [56/66] tough_vs_raising_1.jsonl...
  [57/66] tough_vs_raising_2.jsonl...
  [58/66] transitive.jsonl...
  [59/66] wh_island.jsonl...
  [60/66] wh_questions_object_gap.jsonl...
  [61/66] wh_questions_subject_gap.jsonl...
  [62/66] wh_questions_subject_gap_long_distance.jsonl...
  [63/66] wh_vs_that_no_gap.jsonl...
  [64/66] wh_vs_that_no_gap_long_distance.jsonl...
  [65/66] wh_vs_that_with_gap.jsonl...
  [66/66] wh_vs_that_with_gap_long_distance.jsonl...


======================================================================
🏆 BLIMP BATCH RESULTS — sorted by accuracy
======================================================================
Dataset                                         Accuracy   Margin
----------------------------------------------------------------------
✅ principle_A_case_1                              100.0% (200/200)   0.0865
✅ principle_A_domain_1                            100.0% (200/200)   0.0794
✅ superlative_quantifiers_1                        96.5% (193/200)   0.0597
✅ wh_vs_that_no_gap                                95.5% (191/200)   0.0649
✅ wh_vs_that_no_gap_long_distance                  93.5% (187/200)   0.0473
✅ wh_questions_subject_gap                         89.0% (178/200)   0.0635
✅ wh_island                                        88.5% (177/200)   0.065
✅ wh_questions_subject_gap_long_distance           87.0% (174/200)   0.0427
✅ left_branch_island_echo_question                 78.0% (156/200)   0.069
✅ wh_questions_object_gap                          77.5% (155/200)   0.0434
✅ principle_A_case_2                               75.0% (150/200)   0.027
✅ anaphor_gender_agreement                         70.5% (141/200)   0.0421
✅ superlative_quantifiers_2                        63.5% (127/200)   0.0127
✅ sentential_negation_npi_licensor_present         63.0% (126/200)   0.0208
✅ determiner_noun_agreement_irregular_2            62.0% (124/200)   0.0219
✅ existential_there_quantifiers_2                  61.5% (123/200)   0.0107
✅ principle_A_c_command                            61.0% (122/200)   0.0058
✅ existential_there_subject_raising                60.5% (121/200)   0.0236
〰️ passive_2                                        59.0% (118/200)   0.0165
〰️ animate_subject_passive                          58.5% (117/200)   0.0226
〰️ coordinate_structure_constraint_object_extraction   58.5% (117/200)   0.0147
〰️ irregular_past_participle_verbs                  57.5% (115/200)   0.026
〰️ determiner_noun_agreement_with_adj_irregular_2   56.0% (112/200)   0.0072
〰️ existential_there_object_raising                 55.5% (111/200)   0.0034
〰️ distractor_agreement_relative_clause             55.0% (110/200)   0.0081
〰️ sentential_subject_island                        55.0% (110/200)   0.0025
〰️ transitive                                       55.0% (110/200)   0.01
〰️ distractor_agreement_relational_noun             54.5% (109/200)   0.0078
〰️ ellipsis_n_bar_1                                 54.5% (109/200)   0.002
〰️ determiner_noun_agreement_with_adjective_1       54.0% (108/200)   0.0024
〰️ determiner_noun_agreement_with_adj_2             53.5% (107/200)   -0.0015
〰️ determiner_noun_agreement_1                      53.0% (106/200)   0.0
〰️ principle_A_domain_3                             53.0% (106/200)   0.0034
〰️ anaphor_number_agreement                         52.0% (104/200)   0.0059
〰️ determiner_noun_agreement_with_adj_irregular_1   52.0% (104/200)   0.0044
〰️ tough_vs_raising_2                               52.0% (104/200)   0.0128
〰️ expletive_it_object_raising                      51.0% (102/200)   0.0008
〰️ adjunct_island                                   48.5% (97/200)   -0.0007
〰️ determiner_noun_agreement_2                      48.0% (96/200)   0.007
〰️ tough_vs_raising_1                               48.0% (96/200)   -0.0026
〰️ irregular_plural_subject_verb_agreement_1        47.5% (95/200)   -0.0045
〰️ complex_NP_island                                47.0% (94/200)   -0.0015
〰️ determiner_noun_agreement_irregular_1            47.0% (94/200)   0.0017
〰️ existential_there_quantifiers_1                  47.0% (94/200)   -0.003
〰️ passive_1                                        47.0% (94/200)   -0.0122
〰️ regular_plural_subject_verb_agreement_2          47.0% (94/200)   -0.0054
〰️ irregular_plural_subject_verb_agreement_2        46.0% (92/200)   -0.0142
〰️ sentential_negation_npi_scope                    46.0% (92/200)   -0.0012
〰️ ellipsis_n_bar_2                                 42.5% (85/200)   -0.0078
〰️ principle_A_reconstruction                       42.5% (85/200)   -0.0123
〰️ left_branch_island_simple_question               41.5% (83/200)   -0.017
〰️ causative                                        41.0% (82/200)   -0.018
❌ coordinate_structure_constraint_complex_left_branch   38.5% (77/200)   -0.0086
❌ animate_subject_trans                            38.0% (76/200)   -0.0552
❌ only_npi_scope                                   37.0% (74/200)   -0.0099
❌ principle_A_domain_2                             36.5% (73/200)   -0.0156
❌ inchoative                                       34.5% (69/200)   -0.0572
❌ regular_plural_subject_verb_agreement_1          33.0% (66/200)   -0.0288
❌ irregular_past_participle_adjectives             31.5% (63/200)   -0.0402
❌ intransitive                                     30.0% (60/200)   -0.0722
❌ only_npi_licensor_present                        26.5% (53/200)   -0.0296
❌ npi_present_2                                    11.0% (22/200)   -0.1032
❌ npi_present_1                                     9.0% (18/200)   -0.0904
❌ matrix_question_npi_licensor_present              7.0% (14/200)   -0.1539
❌ wh_vs_that_with_gap                               4.5% (9/200)   -0.0746
❌ wh_vs_that_with_gap_long_distance                 3.0% (6/200)   -0.0519
----------------------------------------------------------------------
📊 Overall average accuracy: 52.86%
======================================================================

:ok

Inference

defmodule DLM.Inference do
  import Nx.Defn

  @allowed_chars ~c" abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789.,:;-'\"()[]"

  def build_muzzle(vocab_size \\ 128, allowed \\ @allowed_chars) do
    muzzle = Nx.broadcast(Nx.tensor(-1.0e10, type: :f32), {vocab_size})
    Enum.reduce(allowed, muzzle, fn char, acc ->
      Nx.indexed_put(acc, Nx.tensor([[char]]), Nx.tensor([0.0]))
    end)
  end

  defn compute_next_logits(x_t, h, tp, muzzle) do
    {logits, h_new} = DLM.SvdSSM.forward_step(x_t, h, tp)
    muzzled = (logits + muzzle) |> Nx.squeeze()
    {muzzled, h_new}
  end

  defn sample_top_k(x_t, h, tp, muzzle, key) do
    {muzzled, h_new} = compute_next_logits(x_t, h, tp, muzzle)
     scaled = muzzled / 0.65
    {top_values, top_indices} = Nx.top_k(scaled, k: 5)
    probs = Nx.exp(top_values) / Nx.sum(Nx.exp(top_values))
    {sampled_tensor, new_key} = Nx.Random.choice(key, top_indices, probs)
    {sampled_tensor[0], h_new, new_key}
  end

  # k = 5 temp = 0.8
  def top_k(params, prompt, muzzle) do
    prompt_chars = String.to_charlist(prompt)
    h_start = Nx.broadcast(0.0, {1, 128}) |> Nx.as_type(:f32)
    resp_init = 0

    prompt_length = length(prompt_chars)
    pad_len = 256 - prompt_length
    padded  = prompt_chars ++ List.duplicate(0, max(0, pad_len))

    Enum.reduce(0..255, {~c"", resp_init, h_start, Nx.Random.key(42)}, fn step, {chars, resp, h_new, key} ->

      x_i = 
        cond do
          step > prompt_length ->
              resp |> Nx.new_axis(-1)
          true -> 
            Enum.at(padded, step) |> Nx.new_axis(-1)
        end
      {muzzled, h_next, new_key} = sample_top_k(x_i, h_new, params, muzzle, key)
      next_token = muzzled |> Nx.to_number()

      new_chars = chars ++ [next_token]
      
      {new_chars, next_token, h_next, new_key}
    end)
    |> elem(0)
    |> List.to_string()
  end
end

IO.puts("✅ DLM.Inference loaded")

✅ DLM.Inference loaded

:ok

defmodule DLM.Probe do
  def run(params) do
    gpu    = Nx.default_backend()
    muzzle = DLM.Inference.build_muzzle() |> Nx.backend_copy(gpu)

    IO.puts("\n" <> String.duplicate("=", 60))
    IO.puts("DLM UNIVERSAL SSM INFERENCE PROBE")
    IO.puts(String.duplicate("=", 60))

    IO.puts("\n📚 IN-DISTRIBUTION")
    IO.puts(String.duplicate("-", 40))
    probe(params, "[W]happy [T]adj [D]",  "Definition of 'happy'", muzzle)
    probe(params, "[W]water [T]noun [D]", "Definition of 'water'", muzzle)
    probe(params, "[W]run [T]verb [D]",   "Definition of 'run'",   muzzle)

    IO.puts("\n🔄 REVERSE FORMAT")
    IO.puts(String.duplicate("-", 40))
    probe(params, "[D]A feeling of great pleasure [W]", "Word for pleasure", muzzle)
    probe(params, "[D]To move swiftly on foot [W]",     "Word for swift",    muzzle)

    IO.puts("\n🔀 NEAR-DISTRIBUTION")
    IO.puts(String.duplicate("-", 40))
    probe(params, "The word 'happy' means",         "Prose definition",  muzzle)
    probe(params, "[W]love [R]Related:",            "Relational format", muzzle)
    probe(params, "[W]cold [T]adj [D]The opposite", "Partial def",       muzzle)

    IO.puts("\n💬 OUT-OF-DISTRIBUTION")
    IO.puts(String.duplicate("-", 40))
    probe(params, "What is the meaning of life?", "Philosophical", muzzle)
    probe(params, "Hello, how are you?",          "Greeting",      muzzle)
    probe(params, "Can you define the word",      "Conversational", muzzle)
    probe(params, "Do state space models dream of recursive sheep?", "Special", muzzle)

    IO.puts("\n🔤 MORPHOLOGY")
    IO.puts(String.duplicate("-", 40))
    probe(params, "[W]Electr", "Electr-", muzzle)
    probe(params, "[W]Mecha",  "Mecha-",  muzzle)
    probe(params, "[W]Inter",  "Inter-",  muzzle)
    probe(params, "[W]Ortho",  "Ortho-",  muzzle)
    probe(params, "[W]Trans",  "Trans-",  muzzle)

    IO.puts("\n" <> String.duplicate("=", 60))
  end

  defp probe(params, prompt, label, muzzle) do
    IO.puts("\n▶ #{label}")
    IO.puts("  Prompt:  #{inspect(prompt)}")

    # Call the V17 Inference
    sampled = DLM.Inference.top_k(params, prompt, muzzle)

    IO.puts("  Top-k:   #{inspect(String.slice(sampled, String.length(prompt)..-1//1))}")

    :erlang.garbage_collect(self())
  end
end

IO.puts("✅ DLM.Probe loaded")

✅ DLM.Probe loaded

:ok

# --- Inference Probe ---
full_state_path = "v7667_DLM_512sq_bgbgem1237gw_checkpoint_step_3400.bin"

if File.exists?(full_state_path) do
  IO.puts("💾 Loading latest V18 Checkpoint...")
  %{trainable: t_p} = File.read!(full_state_path) |> :erlang.binary_to_term()
    params_gpu = DLM.Tree.map(t_p, &Nx.backend_copy(&1, gpu))


  IO.puts("✅ Model Loaded")
  
  # Fire the battery!
  DLM.Probe.run(params_gpu)
else
  IO.puts("❌ No checkpoint found at #{full_state_path}")
end

💾 Loading latest V18 Checkpoint...
✅ Model Loaded

============================================================
DLM UNIVERSAL SSM INFERENCE PROBE
============================================================

📚 IN-DISTRIBUTION
----------------------------------------

▶ Definition of 'happy'
  Prompt:  "[W]happy [T]adj [D]"
  Top-k:   "e shere what we wh ha s ouly se and the aran wan and ancer thoutour the are to sono had and he ther and the the st the sea der ar an are and aid in he tha ain the se hand an the and to to on t and winge songh se she hand and, wand hear t"

▶ Definition of 'water'
  Prompt:  "[W]water [T]noun [D]"
  Top-k:   "ind shead the saringingre anging and he hinghin the anong ane te beathe and whorand ofre so hare was and and the athe toferer and ared and the athe ar the ar and han he and and the to fare the thes ass in she and and whe an tore heallly"

▶ Definition of 'run'
  Prompt:  "[W]run [T]verb [D]"
  Top-k:   "er aid thas tha d hin ar on w io and thin t the ain the has ong tin the the the wo he d he sorora fin the hand the ar he atorongere to the she me me se the to and sen se the me se so hast it he the wat and are wher he t the sind the ain t"

🔄 REVERSE FORMAT
----------------------------------------

▶ Word for pleasure
  Prompt:  "[D]A feeling of great pleasure [W]"
  Top-k:   "andin andoun t he aro t ther whe he the art ane ta dond and an toun har ou he the hou the sould some tomer hond and and ond hing tout the bo s and that and an the and the har the ther ard the an he he arat ther wing the th"

▶ Word for swift
  Prompt:  "[D]To move swiftly on foot [W]"
  Top-k:   "\" the mustingan s t and sher arin in and so the thou ase he the thate an to the and we and and and and and the wonond he sase sathe so an a an thing thin tou the sous at that was the ther arand are hin s hathe hate ar and whe "

🔀 NEAR-DISTRIBUTION
----------------------------------------

▶ Prose definition
  Prompt:  "The word 'happy' means"
  Top-k:   "an an st toroury hou and ass and ard ang ton and shee har he whin the soute ooure sono the cheris the and and whe ar t at and and the an a d ar at the hes and hed ofor ang and the hast whe her ast hou her asastere to ould aid and she "

▶ Relational format
  Prompt:  "[W]love [R]Related:"
  Top-k:   "and ar artonofone mang tit anet the aill sthe some to oute that wes one sonone mas serer and the tar the sould whis the hile sond the she torer to t he ald sould and or and whe has her tore the thes ongere boust ithe sand thing the aid t"

▶ Partial def
  Prompt:  "[W]cold [T]adj [D]The opposite"
  Top-k:   "e boof o te t an toone se ster boor ar and aid and, ald, the are to shat his and hereres ofor s the send and oratonofofofred the sorer soof her hand and shie the but he the torer and and and and he storono t the ally and the t"

💬 OUT-OF-DISTRIBUTION
----------------------------------------

▶ Philosophical
  Prompt:  "What is the meaning of life?"
  Top-k:   " he as seandourd and she shering inge to of a d he ave bun the and hine and ofre and his wat the the so the chisted and and se and hed are she hat he and the sous ar he sher hene sher the whe llle werer the an the te mo ar hat h"

▶ Greeting
  Prompt:  "Hello, how are you?"
  Top-k:   ", aid an or astorend hering t in thing ther ain wing thonond thee wishe hat to ier her ore the when w he mere sond wat tounthingh his we sate wat the soute of he her the hare the healitoro and the hat and as as, s and she somene the wat "

▶ Conversational
  Prompt:  "Can you define the word"
  Top-k:   "ed towing st ain wato ar o he has s ato o o the watl and thee withe sand ofere an t an the woring ther and an the an hore and and hat ar ain a hand ar to the sally and ar the the ange and the ait she hes atinghe has and wat the aid a"

▶ Special
  Prompt:  "Do state space models dream of recursive sheep?"
  Top-k:   "edin the se tome monthe wate was shen what the the cean t t our st shen and our and hin the hile and and the the to t ang the that the as soule to the and hed har and hed has s ar ald the hea he t torere as an"

🔤 MORPHOLOGY
----------------------------------------

▶ Electr-
  Prompt:  "[W]Electr"
  Top-k:   "is as she here hereate hof at teres as oore harar ang to to an tho ond thed his the ther his ofing the heand, the aid the sus the ashas as and the sais so he ce angh was and shar and the ar at s of as and that hous to oul st as in ther hand her an"

▶ Mecha-
  Prompt:  "[W]Mecha"
  Top-k:   "in the saste to and whish at an s aid wa the sherore he mo s and aid and, ald, the are to shat his and hereres ofor s the send and oratonofofofred the sorer soof her hand and shie the but he the torer and and and and he storono t the ally and the t"

▶ Inter-
  Prompt:  "[W]Inter"
  Top-k:   "in the the he are tone torongo te tan wa and hat out oust hout he mare te t o the are the chand wh ho he denet and the and as the shithe the the aid the sout to and she and whis, and ar and and wat s the sonore sun the ala ald dowat hen hand wit wh"

▶ Ortho-
  Prompt:  "[W]Ortho"
  Top-k:   "oror he she sored towa the shoull tis stofof t wa ther wou d and aid and, ald, whe the to f an t ar and the ave the be and to he so tonge and and she shere hare and wat an the torere to the and ont the the s in st in and wat the hand at orer sere a"

▶ Trans-
  Prompt:  "[W]Trans"
  Top-k:   "ind and and an to asst to ongat wit, and wand ther sad as at the sust th hen s the she hissered dowand whit we sout he s ond and theride the the as an her are wer ait lll wand as and as and and shain sas so ad dowat to orango as s ing war and serer"

============================================================

:ok

Train

defmodule DLM.Phase2Trainer do
  import Nx.Defn

  defn step_cross_entropy(logits, targets) do
    vocab_size = Nx.axis_size(logits, -1)
    one_hot = Nx.equal(Nx.iota({vocab_size}), Nx.new_axis(targets, -1))
    
    max_logits = Nx.reduce_max(logits, axes: [-1], keep_axes: true)
    safe_logits = Nx.subtract(logits, max_logits)
    log_probs = Nx.subtract(safe_logits, Nx.log(Nx.sum(Nx.exp(safe_logits), axes: [-1], keep_axes: true)))
    
    # Returns a {batch_size} 1D tensor of individual sequence losses
    Nx.negate(Nx.sum(Nx.multiply(one_hot, log_probs), axes: [-1]))
  end

  defn compute_sequence_loss(input_tokens, params) do
    batch_size = Nx.axis_size(input_tokens, 0)

    h_start = Nx.broadcast(0.0, {batch_size, 128}) |> Nx.as_type(:f32)
    
    zero_scalar = Nx.tensor(0.0, type: :f32)

    result =
      while {i = 0, 
             total_loss = zero_scalar, 
             valid_tokens = zero_scalar,
             h = h_start,
             inf = input_tokens,
             tp = params},
            Nx.less(i, Nx.axis_size(inf, 1)-1) do

        x_t = inf[[.., i]]
        t_t = inf[[.., i+1]]
        
        {logits, h_new} = DLM.SvdSSM.forward_step(x_t, h, tp)
        
        step_loss = step_cross_entropy(logits, t_t)
        mask = Nx.not_equal(t_t, 1) |> Nx.as_type(:f32)
        
        masked_loss_scalar = Nx.sum(Nx.multiply(step_loss, mask))
        valid_tokens_scalar = Nx.sum(mask)

        {i + 1, 
         Nx.add(total_loss, masked_loss_scalar), 
         Nx.add(valid_tokens, valid_tokens_scalar),
         h_new, inf, tp}
      end
  
    {_, acc_loss, val_tokens, _, _, _} = result
    # compress and save final hlong for a long term memory context?

    total_loss = Nx.divide(acc_loss, Nx.add(val_tokens, 1.0e-8))
    total_loss
  end

  defn compute_grad(params, batch_tokens) do
    seq_len = Nx.axis_size(batch_tokens, 1) - 1
    input_tokens  = Nx.slice_along_axis(batch_tokens, 0, seq_len, axis: 1)

    {loss, raw_grads} = 
      value_and_grad(params, fn p -> compute_sequence_loss(input_tokens, p) end)
      
    {loss, raw_grads}
  end
end

{:module, DLM.Phase2Trainer, <<70, 79, 82, 49, 0, 0, 29, ...>>, true}

defmodule DLM.CheckpointManager do
  @moduledoc """
  Scans the directory for the latest checkpoint of a specific run prefix.
  If no checkpoints exist, it triggers a completely fresh initialization.
  """
  def get_resume_state(run_prefix) do
    checkpoints = Path.wildcard("#{run_prefix}_checkpoint_step_*.bin")

    if Enum.empty?(checkpoints) do
      if File.exists?("#{run_prefix}_final.bin") do
        IO.puts("Found completed run: #{run_prefix}_final.bin")
        {"#{run_prefix}_final.bin", :completed}
      else
        IO.puts("No checkpoints found for '#{run_prefix}'. Starting FRESH from Genesis.")
        {:fresh, 0}
      end
    else
      # Extract step numbers using regex and find the maximum
      latest_file = Enum.max_by(checkpoints, fn file ->
        case Regex.run(~r/_step_(\d+)\.bin$/, file) do
          [_, step_str] -> String.to_integer(step_str)
          _ -> -1
        end
      end)

      [_, step_str] = Regex.run(~r/_step_(\d+)\.bin$/, latest_file)
      step = String.to_integer(step_str)
      
      IO.puts("Auto-Resuming '#{run_prefix}' from Step #{step}...")
      {latest_file, step}
    end
  end
end

{:module, DLM.CheckpointManager, <<70, 79, 82, 49, 0, 0, 17, ...>>, {:get_resume_state, 1}}

run_prefix  = "v7667_DLM_512sq_bgbgem1237gw"
# frozen_path = "./data/f_semantic_reservoir_128D.bin"
total_steps = 7500

{load_path, start_step} = DLM.CheckpointManager.get_resume_state(run_prefix)

if start_step == :completed do
  IO.puts("This run has already reached #{total_steps} steps. Exiting.") 
  
else
  batch_size = 128
  # ./data/bbylm_512_fullmix (256 too) bbylm_lg_wiki_512 # try data with sq len of just one word
  data_dir = "./data/bbylm_512_strict" #/babylm_high_sanitize(sq 256), /bbylm_512_strict(sq 512)
  
  IO.puts("Initializing Lazy Curriculum Data Stream from [#{data_dir}]...")
  
  data_stream = 
    BabyLM.CurriculumStreamer.build_infinite_stream(data_dir, batch_size)
    |> Stream.drop(start_step) 

   {starting_trainable, initial_key} =
    if load_path == :fresh do
      IO.puts("Igniting brand new Trainable weights for #{run_prefix}...")
      {ip, ik} = DLM.PhysicsGenesis.init(Nx.Random.key(System.system_time()))
      # %{frozen: {r, _m}} = File.read!(frozen_path) |> :erlang.binary_to_term()
      {ip, ik}
    else
      IO.puts("Loading Trainable Weights from #{load_path}...")
      # Checkpoints now only contain %{trainable: ...}
      %{trainable: t_p} = File.read!(load_path) |> :erlang.binary_to_term()
      # %{frozen: {r, _m}} = File.read!(frozen_path) |> :erlang.binary_to_term()
      {t_p, Nx.Random.key(System.system_time())}
    end

  trainable_gpu = DLM.Tree.map(starting_trainable, &Nx.backend_copy(&1, gpu))
  step_count = Nx.tensor(start_step) |> Nx.backend_copy(gpu)
  
  # ZIP FROM THE START_STEP
  {final_p, _key, _step} =
    Enum.reduce(Stream.zip(start_step..(total_steps - 1), data_stream), 
      {trainable_gpu, initial_key, step_count}, 
      fn {step, cpu_batch}, {cur_tr, cur_key, c_step} ->        
        
      gpu_batch = Nx.backend_copy(cpu_batch, gpu)
     
      base_lr = 7.0e-4
      min_lr = 7.0e-5
      warmup_steps = 1600
      cooldown_start = 5600
      
      current_lr = 
        if step < warmup_steps do
          base_lr * (step / warmup_steps)
        else
          if step < cooldown_start do
            base_lr
          else
            progress = min(max((step - cooldown_start) / (total_steps - cooldown_start), 0.0), 1.0) 
            min_lr + 0.5 * (base_lr - min_lr) * (1.0 + :math.cos(:math.pi() * progress))
          end
        end
      
      lr_tensor = Nx.tensor(current_lr, type: :f32) |> Nx.backend_copy(gpu)

      split_keys = Nx.Random.split(cur_key)
      next_key = split_keys[1]
      
      {loss_tensor, grads} = DLM.Phase2Trainer.compute_grad(cur_tr, gpu_batch)
      loss_val = Nx.to_number(loss_tensor)

      {up_tr, up_step} = DLM.SvdOptimizer.update(cur_tr, grads, c_step, lr_tensor)
  
      # Telemetry
      if rem(step, 100) == 0 do
        {a1, a2, ab, aq, ag, ae,
         b1, b2, bb, bq, bg, be, 
         c1, c2, cb, cq, cg, ce,
         d1, d2, db, dq, dg, de, em, gw}  = up_tr

        a1_norm = a1 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        a2_norm = a2 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        ab_norm = ab |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        aq_norm = aq |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        ag_norm = ag |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        ae_norm = ae |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        b1_norm = b1 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        b2_norm = b2 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        bb_norm = bb |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        bq_norm = bq |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        bg_norm = bg |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        be_norm = be |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        c1_norm = c1 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        c2_norm = c2 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        cb_norm = cb |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        cq_norm = cq |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        cg_norm = cg |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        ce_norm = ce |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        d1_norm = d1 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        d2_norm = d2 |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        db_norm = db |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        dq_norm = dq |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        dg_norm = dg |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        de_norm = de |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
      
        em_norm = em |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        gw_norm = gw |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        
        IO.puts("➡️ Step #{step} | LR: #{Float.round(current_lr, 6)} | Loss: #{Float.round(loss_val, 4)} | em #{em_norm} | gw #{gw_norm}")
        IO.puts("➡️ PN 1: a1 #{a1_norm} |  a2 #{a2_norm} |  ab #{ab_norm} |  aq #{aq_norm} |  ae #{ae_norm} |  ag #{ag_norm}")
        IO.puts("➡️ PN 2: b1 #{b1_norm} |  b2 #{b2_norm} |  bb #{bb_norm} |  bq #{bq_norm} |  be #{be_norm} |  bg #{bg_norm}")
        IO.puts("➡️ PN 3: c1 #{c1_norm} |  c2 #{c2_norm} |  cb #{cb_norm} |  cq #{cq_norm} |  ce #{ce_norm} |  cg #{cg_norm}")
        IO.puts("➡️ PN 4: d1 #{d1_norm} |  d2 #{d2_norm} |  db #{db_norm} |  dq #{dq_norm} |  de #{de_norm} |  dg #{dg_norm}")
        :erlang.garbage_collect()
      end
  
      # Checkpoint Saving (Save both structures!) try every 500 for 64bx512sqln
      if rem(step, 200) == 0 and step > 0 do
        IO.puts("Checkpoint reached! Saving Step #{step}...")
        cpu_tr = DLM.Tree.map(up_tr, &Nx.backend_copy(&1, Nx.BinaryBackend))
        filename = "#{run_prefix}_checkpoint_step_#{step}.bin"
        File.write!(filename, :erlang.term_to_binary(%{trainable: cpu_tr}))
        :erlang.garbage_collect()
      end
  
      {up_tr, next_key, up_step}
    end)

  IO.puts("Saving Final Fourier Reservoir Weights...")
  fin_t = DLM.Tree.map(final_p, &Nx.backend_copy(&1, Nx.BinaryBackend))
  # You could save frozen here too if desired, but since they are deterministic it's optional.
  File.write!("#{run_prefix}_final.bin", :erlang.term_to_binary(%{trainable: fin_t}))
  IO.puts("Training Complete!")
end

Auto-Resuming 'v7667_DLM_512sq_bgbgem1237gw' from Step 800...
Initializing Lazy Curriculum Data Stream from [./data/bbylm_512_strict]...
Loading Trainable Weights from v7667_DLM_512sq_bgbgem1237gw_checkpoint_step_800.bin...

13:03:53.465 [info] XLA service 0x75be240b6730 initialized for platform ROCM (this does not guarantee that XLA will be used). Devices:

13:03:53.467 [info]   StreamExecutor [0]: AMD Radeon RX 7600, AMDGPU ISA version: gfx1100 (Driver: 6.4.43482; Runtime: 6.4.43482; Toolkit: 6.4.43482; DNN: 0.0.0)

13:03:53.467 [info] Using BFC allocator.

13:03:53.467 [info] XLA backend will use up to 6429868032 bytes on device 0 for BFCAllocator.

13:03:53.467 [info] XLA backend will use up to 2143289344 bytes on device 0 for CollectiveBFCAllocator.

🔄 Epoch 1 Started | Shuffling 6 shards...

🔄 Epoch 2 Started | Shuffling 6 shards...

13:03:53.945 [info] Merging Dots in computation: region_27.98.clone
➡️ Step 800 | LR: 3.5e-4 | Loss: 3.0348 | em 128.6124 | gw 11.2801
➡️ PN 1: a1 22.6883 |  a2 11.2873 |  ab 0.1012 |  aq 22.6557 |  ae 0.1129 |  ag 11.3867
➡️ PN 2: b1 22.9387 |  b2 11.4736 |  bb 0.0853 |  bq 22.707 |  be 0.1099 |  bg 11.3898
➡️ PN 3: c1 22.702 |  c2 11.34 |  cb 0.1099 |  cq 22.6652 |  ce 0.1178 |  cg 11.3461
➡️ PN 4: d1 22.754 |  d2 11.421 |  db 0.1143 |  dq 22.6895 |  de 0.1178 |  dg 11.3203
Checkpoint reached! Saving Step 800...
➡️ Step 900 | LR: 3.94e-4 | Loss: 2.6824 | em 128.8973 | gw 11.2856
➡️ PN 1: a1 22.7102 |  a2 11.3027 |  ab 0.1196 |  aq 22.667 |  ae 0.1415 |  ag 11.4138
➡️ PN 2: b1 23.0406 |  b2 11.5156 |  bb 0.0967 |  bq 22.7349 |  be 0.1371 |  bg 11.4173
➡️ PN 3: c1 22.7386 |  c2 11.353 |  cb 0.1347 |  cq 22.682 |  ce 0.146 |  cg 11.3573
➡️ PN 4: d1 22.8273 |  d2 11.4254 |  db 0.1386 |  dq 22.7126 |  de 0.146 |  dg 11.3296
➡️ Step 1000 | LR: 4.38e-4 | Loss: 2.6327 | em 129.2342 | gw 11.2956
➡️ PN 1: a1 22.7353 |  a2 11.321 |  ab 0.1367 |  aq 22.681 |  ae 0.1737 |  ag 11.4468
➡️ PN 2: b1 23.1535 |  b2 11.5617 |  bb 0.1033 |  bq 22.7652 |  be 0.1663 |  bg 11.4492
➡️ PN 3: c1 22.7855 |  c2 11.3726 |  cb 0.1622 |  cq 22.7033 |  ce 0.1767 |  cg 11.3687
➡️ PN 4: d1 22.9123 |  d2 11.4256 |  db 0.1622 |  dq 22.7393 |  de 0.1767 |  dg 11.3416
Checkpoint reached! Saving Step 1000...
➡️ Step 1100 | LR: 4.81e-4 | Loss: 2.6279 | em 129.6303 | gw 11.3101
➡️ PN 1: a1 22.7565 |  a2 11.3374 |  ab 0.1486 |  aq 22.6925 |  ae 0.2064 |  ag 11.485
➡️ PN 2: b1 23.2615 |  b2 11.6059 |  bb 0.1034 |  bq 22.7961 |  be 0.1938 |  bg 11.485
➡️ PN 3: c1 22.8366 |  c2 11.4024 |  cb 0.1867 |  cq 22.7258 |  ce 0.2046 |  cg 11.3775
➡️ PN 4: d1 22.9978 |  d2 11.4331 |  db 0.1796 |  dq 22.7664 |  de 0.2046 |  dg 11.3534
➡️ Step 1200 | LR: 5.25e-4 | Loss: 2.5298 | em 130.1287 | gw 11.3337
➡️ PN 1: a1 22.7727 |  a2 11.35 |  ab 0.1606 |  aq 22.7013 |  ae 0.2439 |  ag 11.5289
➡️ PN 2: b1 23.3644 |  b2 11.6486 |  bb 0.106 |  bq 22.8264 |  be 0.2268 |  bg 11.5249
➡️ PN 3: c1 22.9072 |  c2 11.4353 |  cb 0.218 |  cq 22.7544 |  ce 0.2401 |  cg 11.3909
➡️ PN 4: d1 23.1041 |  d2 11.4488 |  db 0.1995 |  dq 22.7984 |  de 0.2401 |  dg 11.3701
Checkpoint reached! Saving Step 1200...
➡️ Step 1300 | LR: 5.69e-4 | Loss: 2.542 | em 130.6134 | gw 11.3618
➡️ PN 1: a1 22.7856 |  a2 11.3596 |  ab 0.1696 |  aq 22.7069 |  ae 0.28 |  ag 11.5758
➡️ PN 2: b1 23.4418 |  b2 11.6804 |  bb 0.1087 |  bq 22.8498 |  be 0.2569 |  bg 11.5678
➡️ PN 3: c1 22.9767 |  c2 11.4686 |  cb 0.2438 |  cq 22.7828 |  ce 0.2715 |  cg 11.4033
➡️ PN 4: d1 23.1952 |  d2 11.4693 |  db 0.2139 |  dq 22.8282 |  de 0.2715 |  dg 11.3854
➡️ Step 1400 | LR: 6.12e-4 | Loss: 2.5527 | em 131.1047 | gw 11.4042
➡️ PN 1: a1 22.801 |  a2 11.37 |  ab 0.1776 |  aq 22.7121 |  ae 0.3184 |  ag 11.626
➡️ PN 2: b1 23.5025 |  b2 11.7067 |  bb 0.1106 |  bq 22.8685 |  be 0.2859 |  bg 11.6151
➡️ PN 3: c1 23.0446 |  c2 11.4896 |  cb 0.2676 |  cq 22.8102 |  ce 0.3034 |  cg 11.4173
➡️ PN 4: d1 23.2757 |  d2 11.4892 |  db 0.2227 |  dq 22.8562 |  de 0.3034 |  dg 11.4016
Checkpoint reached! Saving Step 1400...

🔄 Epoch 3 Started | Shuffling 6 shards...
➡️ Step 1500 | LR: 6.56e-4 | Loss: 2.4545 | em 131.6313 | gw 11.4615
➡️ PN 1: a1 22.8212 |  a2 11.384 |  ab 0.1837 |  aq 22.7209 |  ae 0.3577 |  ag 11.6809
➡️ PN 2: b1 23.5534 |  b2 11.7282 |  bb 0.1103 |  bq 22.8846 |  be 0.314 |  bg 11.6667
➡️ PN 3: c1 23.1168 |  c2 11.5 |  cb 0.2881 |  cq 22.8384 |  ce 0.3357 |  cg 11.4331
➡️ PN 4: d1 23.3477 |  d2 11.5093 |  db 0.2284 |  dq 22.8818 |  de 0.3357 |  dg 11.419
➡️ Step 1600 | LR: 0.0007 | Loss: 2.4704 | em 132.1776 | gw 11.5313
➡️ PN 1: a1 22.851 |  a2 11.4031 |  ab 0.1898 |  aq 22.7347 |  ae 0.3978 |  ag 11.7401
➡️ PN 2: b1 23.6041 |  b2 11.7465 |  bb 0.1109 |  bq 22.9011 |  be 0.34 |  bg 11.7213
➡️ PN 3: c1 23.1943 |  c2 11.5062 |  cb 0.3032 |  cq 22.867 |  ce 0.3655 |  cg 11.4476
➡️ PN 4: d1 23.4099 |  d2 11.5294 |  db 0.2326 |  dq 22.9029 |  de 0.3655 |  dg 11.4354
Checkpoint reached! Saving Step 1600...
➡️ Step 1700 | LR: 0.0007 | Loss: 2.4093 | em 132.6823 | gw 11.6155
➡️ PN 1: a1 22.8866 |  a2 11.4257 |  ab 0.1948 |  aq 22.751 |  ae 0.4339 |  ag 11.8002
➡️ PN 2: b1 23.6507 |  b2 11.7642 |  bb 0.1125 |  bq 22.917 |  be 0.3624 |  bg 11.7777
➡️ PN 3: c1 23.2691 |  c2 11.518 |  cb 0.3152 |  cq 22.8942 |  ce 0.3921 |  cg 11.4604
➡️ PN 4: d1 23.4655 |  d2 11.5491 |  db 0.2351 |  dq 22.9215 |  de 0.3921 |  dg 11.4506
➡️ Step 1800 | LR: 0.0007 | Loss: 2.5808 | em 133.2253 | gw 11.6988
➡️ PN 1: a1 22.9112 |  a2 11.4434 |  ab 0.205 |  aq 22.7629 |  ae 0.4723 |  ag 11.8583
➡️ PN 2: b1 23.6953 |  b2 11.783 |  bb 0.1123 |  bq 22.9326 |  be 0.3841 |  bg 11.836
➡️ PN 3: c1 23.336 |  c2 11.5472 |  cb 0.3321 |  cq 22.9197 |  ce 0.4172 |  cg 11.47
➡️ PN 4: d1 23.5212 |  d2 11.5738 |  db 0.2364 |  dq 22.941 |  de 0.4172 |  dg 11.4688
Checkpoint reached! Saving Step 1800...
➡️ Step 1900 | LR: 0.0007 | Loss: 2.395 | em 133.7089 | gw 11.7873
➡️ PN 1: a1 22.9383 |  a2 11.4629 |  ab 0.2134 |  aq 22.7778 |  ae 0.5067 |  ag 11.9145
➡️ PN 2: b1 23.7281 |  b2 11.7959 |  bb 0.1117 |  bq 22.9428 |  be 0.4059 |  bg 11.8918
➡️ PN 3: c1 23.3954 |  c2 11.565 |  cb 0.3429 |  cq 22.9424 |  ce 0.4412 |  cg 11.482
➡️ PN 4: d1 23.5668 |  d2 11.5926 |  db 0.2404 |  dq 22.9562 |  de 0.4412 |  dg 11.4855
➡️ Step 2000 | LR: 0.0007 | Loss: 2.3077 | em 134.0441 | gw 11.8875
➡️ PN 1: a1 22.9729 |  a2 11.483 |  ab 0.212 |  aq 22.7952 |  ae 0.5334 |  ag 11.9728
➡️ PN 2: b1 23.7497 |  b2 11.8036 |  bb 0.1104 |  bq 22.947 |  be 0.4215 |  bg 11.9471
➡️ PN 3: c1 23.448 |  c2 11.5724 |  cb 0.3476 |  cq 22.9633 |  ce 0.4619 |  cg 11.4952
➡️ PN 4: d1 23.5988 |  d2 11.606 |  db 0.2411 |  dq 22.9671 |  de 0.4619 |  dg 11.5021
Checkpoint reached! Saving Step 2000...
➡️ Step 2100 | LR: 0.0007 | Loss: 2.4165 | em 134.475 | gw 12.014
➡️ PN 1: a1 23.0063 |  a2 11.5031 |  ab 0.2162 |  aq 22.8104 |  ae 0.5629 |  ag 12.0332
➡️ PN 2: b1 23.7743 |  b2 11.8134 |  bb 0.1096 |  bq 22.9529 |  be 0.4364 |  bg 12.0055
➡️ PN 3: c1 23.5042 |  c2 11.596 |  cb 0.3535 |  cq 22.9847 |  ce 0.4836 |  cg 11.5081
➡️ PN 4: d1 23.6409 |  d2 11.6294 |  db 0.2422 |  dq 22.9799 |  de 0.4836 |  dg 11.5177

🔄 Epoch 4 Started | Shuffling 6 shards...
➡️ Step 2200 | LR: 0.0007 | Loss: 2.3357 | em 134.9102 | gw 12.1409
➡️ PN 1: a1 23.0407 |  a2 11.5234 |  ab 0.222 |  aq 22.8243 |  ae 0.5894 |  ag 12.0918
➡️ PN 2: b1 23.8103 |  b2 11.8293 |  bb 0.1099 |  bq 22.9649 |  be 0.4496 |  bg 12.0612
➡️ PN 3: c1 23.5611 |  c2 11.6096 |  cb 0.3584 |  cq 23.0045 |  ce 0.5058 |  cg 11.5217
➡️ PN 4: d1 23.6822 |  d2 11.6524 |  db 0.2425 |  dq 22.9909 |  de 0.5058 |  dg 11.5328
Checkpoint reached! Saving Step 2200...
➡️ Step 2300 | LR: 0.0007 | Loss: 2.2756 | em 135.2217 | gw 12.2663
➡️ PN 1: a1 23.0713 |  a2 11.5418 |  ab 0.2276 |  aq 22.837 |  ae 0.6143 |  ag 12.1486
➡️ PN 2: b1 23.8374 |  b2 11.8421 |  bb 0.1112 |  bq 22.9732 |  be 0.4605 |  bg 12.1107
➡️ PN 3: c1 23.6034 |  c2 11.6171 |  cb 0.3657 |  cq 23.0188 |  ce 0.5255 |  cg 11.5319
➡️ PN 4: d1 23.7104 |  d2 11.6701 |  db 0.2434 |  dq 22.9982 |  de 0.5255 |  dg 11.5464
➡️ Step 2400 | LR: 0.0007 | Loss: 2.2732 | em 135.529 | gw 12.383
➡️ PN 1: a1 23.1037 |  a2 11.5611 |  ab 0.2317 |  aq 22.8509 |  ae 0.6375 |  ag 12.2082
➡️ PN 2: b1 23.8626 |  b2 11.8532 |  bb 0.1141 |  bq 22.9807 |  be 0.4705 |  bg 12.1624
➡️ PN 3: c1 23.6452 |  c2 11.628 |  cb 0.3698 |  cq 23.0336 |  ce 0.5442 |  cg 11.5431
➡️ PN 4: d1 23.7404 |  d2 11.6911 |  db 0.2435 |  dq 23.006 |  de 0.5442 |  dg 11.5608
Checkpoint reached! Saving Step 2400...
➡️ Step 2500 | LR: 0.0007 | Loss: 2.5296 | em 135.9707 | gw 12.4914
➡️ PN 1: a1 23.1252 |  a2 11.5741 |  ab 0.2399 |  aq 22.8607 |  ae 0.6673 |  ag 12.2652
➡️ PN 2: b1 23.8917 |  b2 11.8631 |  bb 0.1143 |  bq 22.9892 |  be 0.4803 |  bg 12.2178
➡️ PN 3: c1 23.6679 |  c2 11.6508 |  cb 0.3794 |  cq 23.0425 |  ce 0.5613 |  cg 11.5498
➡️ PN 4: d1 23.7746 |  d2 11.7136 |  db 0.2433 |  dq 23.0163 |  de 0.5613 |  dg 11.5754
➡️ Step 2600 | LR: 0.0007 | Loss: 2.4247 | em 136.2832 | gw 12.6191
➡️ PN 1: a1 23.1526 |  a2 11.5908 |  ab 0.2423 |  aq 22.8727 |  ae 0.69 |  ag 12.3223
➡️ PN 2: b1 23.9182 |  b2 11.8739 |  bb 0.1144 |  bq 22.9975 |  be 0.4923 |  bg 12.2694
➡️ PN 3: c1 23.6951 |  c2 11.6543 |  cb 0.3826 |  cq 23.053 |  ce 0.5801 |  cg 11.5623
➡️ PN 4: d1 23.8012 |  d2 11.7327 |  db 0.2449 |  dq 23.0238 |  de 0.5801 |  dg 11.5892
Checkpoint reached! Saving Step 2600...
➡️ Step 2700 | LR: 0.0007 | Loss: 2.298 | em 136.5394 | gw 12.7703
➡️ PN 1: a1 23.183 |  a2 11.609 |  ab 0.2428 |  aq 22.8842 |  ae 0.7083 |  ag 12.3802
➡️ PN 2: b1 23.9488 |  b2 11.8864 |  bb 0.1137 |  bq 23.0075 |  be 0.501 |  bg 12.3211
➡️ PN 3: c1 23.7246 |  c2 11.6587 |  cb 0.3847 |  cq 23.0633 |  ce 0.597 |  cg 11.5733
➡️ PN 4: d1 23.8267 |  d2 11.7519 |  db 0.2447 |  dq 23.031 |  de 0.597 |  dg 11.602
➡️ Step 2800 | LR: 0.0007 | Loss: 2.3799 | em 136.8061 | gw 12.9377
➡️ PN 1: a1 23.2137 |  a2 11.6267 |  ab 0.2438 |  aq 22.8955 |  ae 0.7286 |  ag 12.4379
➡️ PN 2: b1 23.9749 |  b2 11.8977 |  bb 0.1138 |  bq 23.0157 |  be 0.5096 |  bg 12.3723
➡️ PN 3: c1 23.7568 |  c2 11.6713 |  cb 0.3882 |  cq 23.0752 |  ce 0.6139 |  cg 11.5836
➡️ PN 4: d1 23.8519 |  d2 11.7668 |  db 0.2444 |  dq 23.038 |  de 0.6139 |  dg 11.6151
Checkpoint reached! Saving Step 2800...
➡️ Step 2900 | LR: 0.0007 | Loss: 2.2407 | em 137.1291 | gw 13.1012
➡️ PN 1: a1 23.2454 |  a2 11.6448 |  ab 0.2446 |  aq 22.9074 |  ae 0.7472 |  ag 12.4972
➡️ PN 2: b1 24.0015 |  b2 11.9086 |  bb 0.1133 |  bq 23.0238 |  be 0.5153 |  bg 12.423
➡️ PN 3: c1 23.7966 |  c2 11.6853 |  cb 0.3899 |  cq 23.0889 |  ce 0.632 |  cg 11.5944
➡️ PN 4: d1 23.877 |  d2 11.7841 |  db 0.2442 |  dq 23.0439 |  de 0.632 |  dg 11.629

🔄 Epoch 5 Started | Shuffling 6 shards...
➡️ Step 3000 | LR: 0.0007 | Loss: 2.2943 | em 137.4459 | gw 13.2632
➡️ PN 1: a1 23.2788 |  a2 11.6639 |  ab 0.2468 |  aq 22.9191 |  ae 0.7658 |  ag 12.558
➡️ PN 2: b1 24.0292 |  b2 11.9192 |  bb 0.1133 |  bq 23.0318 |  be 0.5205 |  bg 12.4705
➡️ PN 3: c1 23.84 |  c2 11.6989 |  cb 0.3917 |  cq 23.1029 |  ce 0.6498 |  cg 11.6053
➡️ PN 4: d1 23.8985 |  d2 11.8037 |  db 0.2442 |  dq 23.0486 |  de 0.6498 |  dg 11.6426
Checkpoint reached! Saving Step 3000...
➡️ Step 3100 | LR: 0.0007 | Loss: 2.291 | em 137.7198 | gw 13.4129
➡️ PN 1: a1 23.3062 |  a2 11.6813 |  ab 0.2494 |  aq 22.9281 |  ae 0.782 |  ag 12.6174
➡️ PN 2: b1 24.0601 |  b2 11.9316 |  bb 0.1129 |  bq 23.0406 |  be 0.5258 |  bg 12.515
➡️ PN 3: c1 23.877 |  c2 11.7142 |  cb 0.3936 |  cq 23.1141 |  ce 0.665 |  cg 11.6141
➡️ PN 4: d1 23.9164 |  d2 11.8241 |  db 0.2444 |  dq 23.0522 |  de 0.665 |  dg 11.6544
➡️ Step 3200 | LR: 0.0007 | Loss: 2.1878 | em 137.8962 | gw 13.5354
➡️ PN 1: a1 23.3309 |  a2 11.6974 |  ab 0.2506 |  aq 22.9388 |  ae 0.7981 |  ag 12.6754
➡️ PN 2: b1 24.0849 |  b2 11.9415 |  bb 0.1121 |  bq 23.0477 |  be 0.5309 |  bg 12.5566
➡️ PN 3: c1 23.9005 |  c2 11.7204 |  cb 0.396 |  cq 23.1211 |  ce 0.6787 |  cg 11.623
➡️ PN 4: d1 23.9304 |  d2 11.8377 |  db 0.244 |  dq 23.0552 |  de 0.6787 |  dg 11.6665
Checkpoint reached! Saving Step 3200...
➡️ Step 3300 | LR: 0.0007 | Loss: 2.3679 | em 138.0634 | gw 13.6595
➡️ PN 1: a1 23.3593 |  a2 11.7147 |  ab 0.2532 |  aq 22.9494 |  ae 0.8134 |  ag 12.7312
➡️ PN 2: b1 24.1001 |  b2 11.9488 |  bb 0.1108 |  bq 23.0516 |  be 0.536 |  bg 12.5961
➡️ PN 3: c1 23.9234 |  c2 11.7271 |  cb 0.3984 |  cq 23.1279 |  ce 0.6924 |  cg 11.6305
➡️ PN 4: d1 23.9468 |  d2 11.8531 |  db 0.2458 |  dq 23.0589 |  de 0.6924 |  dg 11.6759
➡️ Step 3400 | LR: 0.0007 | Loss: 2.3901 | em 138.2334 | gw 13.8147
➡️ PN 1: a1 23.3853 |  a2 11.7303 |  ab 0.2548 |  aq 22.9588 |  ae 0.829 |  ag 12.7865
➡️ PN 2: b1 24.1138 |  b2 11.956 |  bb 0.1117 |  bq 23.0548 |  be 0.5417 |  bg 12.635
➡️ PN 3: c1 23.9431 |  c2 11.7372 |  cb 0.4002 |  cq 23.1334 |  ce 0.7043 |  cg 11.6369
➡️ PN 4: d1 23.9631 |  d2 11.8707 |  db 0.2473 |  dq 23.0622 |  de 0.7043 |  dg 11.684
Checkpoint reached! Saving Step 3400...
➡️ Step 3500 | LR: 0.0007 | Loss: 2.2954 | em 138.4104 | gw 13.99
➡️ PN 1: a1 23.413 |  a2 11.7471 |  ab 0.2566 |  aq 22.969 |  ae 0.8441 |  ag 12.8433
➡️ PN 2: b1 24.1287 |  b2 11.9621 |  bb 0.1127 |  bq 23.0581 |  be 0.5475 |  bg 12.6715
➡️ PN 3: c1 23.9689 |  c2 11.7456 |  cb 0.4017 |  cq 23.1404 |  ce 0.7171 |  cg 11.6447
➡️ PN 4: d1 23.9817 |  d2 11.8878 |  db 0.2483 |  dq 23.0664 |  de 0.7171 |  dg 11.6937