p_2_1_wembi.livemd

some space to sink 2u - world model - yeet - fork - trudat - yoyo - 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 MOE.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, MOE.Tree, <<70, 79, 82, 49, 0, 0, 8, ...>>, {:map, 2}}

defmodule MOE.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, MOE.CurriculumStreamer, <<70, 79, 82, 49, 0, 0, 17, ...>>, {:build_infinite_stream, 2}}

defmodule MOE.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, 6) 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, 6) 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, MOE.SparseSinkGD, <<70, 79, 82, 49, 0, 0, 29, ...>>, true}

SSM

defmodule MOE.SSM 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, skip connection
    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_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)
  end

  defn ot_routing(affinity) do
    max_a = Nx.reduce_max(affinity, axes: [-1], keep_axes: true)
    p_init = Nx.exp(Nx.subtract(affinity, max_a))
    # p_init = Nx.exp(affinity)

    result = 
      while {i = 0, p = p_init}, Nx.less(i, 4) do
        # row norm
        row_sums = Nx.sum(p, axes: [-1], keep_axes: true)
        p_row = Nx.divide(p, Nx.add(row_sums, 1.0e-8))
        # col norm
        col_sums = Nx.sum(p_row, axes: [-2], keep_axes: true)
        p_col = Nx.divide(p_row, Nx.add(col_sums, 1.0e-8))

        {i + 1, p_col}
      end

    elem(result, 1)
  end

  defn forward_step(x_t, h_prev, p) do
    {emb, dec, key} = p
    em = Nx.take(emb, x_t) # get embedding for current character
    u_t = Nx.take(key, x_t) |> gelu() |> rms_norm() # get key for compression
    hp_ot = ot_routing(h_prev) # 'focus' previous hidden state
    hp_sc = Nx.multiply(h_prev, Nx.sigmoid(dec)) # learned decay gate
    h_new = Nx.add(hp_sc, em) # update hidden state
    # compress em to 'character', activate, normalize
    uqc = apply_attention(u_t, em) # attend to current em
    uqh = apply_attention(u_t, hp_ot) # attend to 'focused' h_prev
    logits = Nx.add(uqc, uqh) |> Nx.add(u_t) # combine with 'skip connection'
    
    {logits, h_new}
  end

  # defn forward_step(x_t, h_prev, p) do
  #   {emb} = p
  #   # process current character 
  #   batch_size = Nx.axis_size(x_t, 0)
  #   batch_indices = Nx.iota({batch_size})
  #   em = Nx.take(emb, x_t)
  #   h_new = Nx.add(h_prev, em)
  #   coords = Nx.stack([batch_indices, x_t], axis: -1)
  #   u_t = Nx.gather(h_new, coords) |> gelu() |> rms_norm()
  #   u_x = apply_attention(u_t, em) |> Nx.add(u_t)
  #   u_h = apply_attention(u_h, h_new) |> Nx.add(u_h)
  #   # get context
  #   logits = Nx.add(u_x, u_h) 
    
  #   {logits, h_new}
  # end
end

{:module, MOE.SSM, <<70, 79, 82, 49, 0, 0, 50, ...>>, true}

defmodule MOE.Optimizer do
  import Nx.Defn

  defn update(params, grads, step, lr) do
    {emb, dec, key}  = params   
    {gemb, gdec, gkey}  = grads
    
    new_step = Nx.add(step, 1)

    # optimizer update step
    emb_new = MOE.SparseSinkGD.step_3d(emb, gemb, lr)
    dec_new = MOE.SparseSinkGD.step(dec, gdec, lr)
    key_new = MOE.SparseSinkGD.step(key, gkey, lr)
    
    
    new_params = {emb_new, dec_new, key_new}

    {new_params, new_step}
  end
end

{:module, MOE.Optimizer, <<70, 79, 82, 49, 0, 0, 13, ...>>, true}

defmodule MOE.Parameter do
  import Nx.Defn

  @vocab 128

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

    # projection matrices
    {emb, k2} = Nx.Random.normal(k1, 0.0, std_voc, shape: {@vocab, @vocab, @vocab})
    {key, k3} = Nx.Random.normal(k2, 0.0, std_voc, shape: {@vocab, @vocab})

    dec = Nx.broadcast(0.5, {@vocab}) |> Nx.as_type(:f32)

    trainable_params = {emb, dec, key} 
    
    {trainable_params, k3}
  end
end

{:module, MOE.Parameter, <<70, 79, 82, 49, 0, 0, 12, ...>>, true}

BLiMP

defmodule MOE.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} = MOE.SSM.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, MOE.Phase2_Eval, <<70, 79, 82, 49, 0, 0, 25, ...>>, true}

gpu = {EXLA.Backend, client: :rocm}
checkpoint_path = "v67_MOE_512sq_yolo_em2b4u_finmb_checkpoint_step_2200.bin"
%{trainable: t_p} = File.read!(checkpoint_path) |> :erlang.binary_to_term()
trainable_gpu = MOE.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, 128}) |> Nx.backend_copy(gpu)

MOE.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, 128}) |> Nx.backend_copy(gpu)
      
      losses = MOE.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
----------------------------------------------------------------------
✅ existential_there_quantifiers_2                 100.0% (200/200)   0.0608
✅ principle_A_case_1                              100.0% (200/200)   0.1427
✅ principle_A_domain_1                            100.0% (200/200)   0.108
✅ superlative_quantifiers_1                       100.0% (200/200)   0.0611
✅ wh_questions_subject_gap                        100.0% (200/200)   0.0613
✅ wh_questions_subject_gap_long_distance          100.0% (200/200)   0.0371
✅ wh_vs_that_no_gap                               100.0% (200/200)   0.0532
✅ wh_vs_that_no_gap_long_distance                 100.0% (200/200)   0.0389
✅ wh_questions_object_gap                          93.0% (186/200)   0.0391
✅ principle_A_case_2                               91.0% (182/200)   0.0405
✅ left_branch_island_echo_question                 87.0% (174/200)   0.0335
✅ irregular_past_participle_verbs                  80.5% (161/200)   0.0333
✅ wh_island                                        71.5% (143/200)   0.0117
✅ anaphor_gender_agreement                         70.5% (141/200)   0.0533
✅ principle_A_c_command                            70.0% (140/200)   0.0403
✅ determiner_noun_agreement_irregular_2            64.0% (128/200)   0.0322
✅ npi_present_2                                    62.5% (125/200)   0.0133
✅ determiner_noun_agreement_with_adj_irregular_2   61.5% (123/200)   0.0143
✅ superlative_quantifiers_2                        61.0% (122/200)   0.0215
✅ animate_subject_passive                          60.5% (121/200)   0.0207
✅ determiner_noun_agreement_with_adj_irregular_1   60.5% (121/200)   0.0067
〰️ ellipsis_n_bar_1                                 57.5% (115/200)   0.0064
〰️ tough_vs_raising_2                               57.5% (115/200)   0.022
〰️ distractor_agreement_relative_clause             55.0% (110/200)   0.0066
〰️ existential_there_subject_raising                54.5% (109/200)   0.0057
〰️ determiner_noun_agreement_irregular_1            54.0% (108/200)   0.004
〰️ causative                                        53.0% (106/200)   0.0054
〰️ distractor_agreement_relational_noun             52.0% (104/200)   0.0053
〰️ determiner_noun_agreement_2                      51.5% (103/200)   0.0095
〰️ transitive                                       51.5% (103/200)   0.008
〰️ determiner_noun_agreement_with_adjective_1       51.0% (102/200)   0.0044
〰️ regular_plural_subject_verb_agreement_2          51.0% (102/200)   0.0023
〰️ anaphor_number_agreement                         49.5% (99/200)   -0.0093
〰️ complex_NP_island                                49.5% (99/200)   0.0009
〰️ determiner_noun_agreement_with_adj_2             49.0% (98/200)   -0.002
〰️ existential_there_object_raising                 49.0% (98/200)   -0.0011
〰️ determiner_noun_agreement_1                      48.5% (97/200)   0.0
〰️ expletive_it_object_raising                      48.5% (97/200)   0.0018
〰️ adjunct_island                                   47.5% (95/200)   0.0078
〰️ tough_vs_raising_1                               47.5% (95/200)   -0.015
〰️ ellipsis_n_bar_2                                 47.0% (94/200)   -0.0031
〰️ coordinate_structure_constraint_object_extraction   46.0% (92/200)   -0.0038
〰️ irregular_plural_subject_verb_agreement_1        45.5% (91/200)   -0.001
〰️ npi_present_1                                    45.5% (91/200)   -0.0067
〰️ irregular_plural_subject_verb_agreement_2        44.0% (88/200)   -0.0223
〰️ passive_2                                        44.0% (88/200)   -0.0186
〰️ principle_A_domain_3                             44.0% (88/200)   -0.0034
〰️ inchoative                                       43.5% (87/200)   -0.027
〰️ sentential_negation_npi_licensor_present         41.0% (82/200)   0.0015
❌ intransitive                                     40.0% (80/200)   -0.0342
❌ left_branch_island_simple_question               40.0% (80/200)   -0.0101
❌ passive_1                                        39.0% (78/200)   -0.0248
❌ regular_plural_subject_verb_agreement_1          39.0% (78/200)   -0.0303
❌ principle_A_domain_2                             35.5% (71/200)   -0.0262
❌ principle_A_reconstruction                       34.5% (69/200)   -0.0176
❌ animate_subject_trans                            29.0% (58/200)   -0.0872
❌ existential_there_quantifiers_1                  26.0% (52/200)   -0.0465
❌ irregular_past_participle_adjectives             23.0% (46/200)   -0.0214
❌ sentential_subject_island                        18.5% (37/200)   -0.0124
❌ coordinate_structure_constraint_complex_left_branch   12.0% (24/200)   -0.0
❌ sentential_negation_npi_scope                     6.0% (12/200)   -0.0
❌ matrix_question_npi_licensor_present              3.5% (7/200)   -0.0824
❌ only_npi_scope                                    1.0% (2/200)   -0.0201
❌ only_npi_licensor_present                         0.0% (0/200)   -0.1067
❌ wh_vs_that_with_gap                               0.0% (0/200)   -0.0592
❌ wh_vs_that_with_gap_long_distance                 0.0% (0/200)   -0.0422
----------------------------------------------------------------------
📊 Overall average accuracy: 52.4%
======================================================================

:ok

Inference

defmodule MOE.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} = MOE.SSM.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, 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 MOE.Probe do
  def run(params) do
    gpu    = Nx.default_backend()
    muzzle = MOE.Inference.build_muzzle() |> Nx.backend_copy(gpu)

    IO.puts("\n" <> String.duplicate("=", 60))
    IO.puts("MOE 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 = MOE.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("✅ MOE.Probe loaded")

✅ DLM.Probe loaded

:ok

# --- Inference Probe ---
full_state_path = "v7667_DLM_512sq_yolo_em2b4u_p3d6expa2fx_checkpoint_step_4600.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 = MOE.Tree.map(t_p, &Nx.backend_copy(&1, gpu))


  IO.puts("✅ Model Loaded")
  
  # Fire the battery!
  MOE.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:   "No s wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Definition of 'water'
  Prompt:  "[W]water [T]noun [D]"
  Top-k:   "Dend thilind we a th th thanof thend s t he s t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Definition of 'run'
  Prompt:  "[W]run [T]verb [D]"
  Top-k:   "Nou s wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

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

▶ Word for pleasure
  Prompt:  "[D]A feeling of great pleasure [W]"
  Top-k:   "No he he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Word for swift
  Prompt:  "[D]To move swiftly on foot [W]"
  Top-k:   "Nor we he he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

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

▶ Prose definition
  Prompt:  "The word 'happy' means"
  Top-k:   "pere h we st atin he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Relational format
  Prompt:  "[W]love [R]Related:"
  Top-k:   "No s wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Partial def
  Prompt:  "[W]cold [T]adj [D]The opposite"
  Top-k:   "Nor we he he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

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

▶ Philosophical
  Prompt:  "What is the meaning of life?"
  Top-k:   "pathas atin he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Greeting
  Prompt:  "Hello, how are you?"
  Top-k:   "No s wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Conversational
  Prompt:  "Can you define the word"
  Top-k:   "Sthed. aner atin he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Special
  Prompt:  "Do state space models dream of recursive sheep?"
  Top-k:   "Dof thend s t he s t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

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

▶ Electr-
  Prompt:  "[W]Electr"
  Top-k:   "Dind the as tous tear tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Mecha-
  Prompt:  "[W]Mecha"
  Top-k:   "Sore whe t wane wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Inter-
  Prompt:  "[W]Inter"
  Top-k:   "Sore whe t wane wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Ortho-
  Prompt:  "[W]Ortho"
  Top-k:   "Sore whe t wane wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

▶ Trans-
  Prompt:  "[W]Trans"
  Top-k:   "Sore whe t wane wanofo tonghean he he wa he t s win a ato t the th are te tine athe an tounge whato the asereand there the a t t angofis the that s ange s t ang an sthathe t areng t thand and whaner athe and he woure s torer the s ang hend t asthe "

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

:ok

Train

defmodule MOE.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, 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,
             inf = input_tokens,
             h = h_start,
             tp = params},
            Nx.less(i, Nx.axis_size(inf, 1)-1) do

        x_t = inf[[.., i]]
        t_t = inf[[.., i+1]]
        
        {logits, h_new} = MOE.SSM.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),
         inf, h_new, 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, MOE.Phase2Trainer, <<70, 79, 82, 49, 0, 0, 29, ...>>, true}

defmodule MOE.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, MOE.CheckpointManager, <<70, 79, 82, 49, 0, 0, 17, ...>>, {:get_resume_state, 1}}

run_prefix  = "v67_MOE_512sq_yolo_em2b4u_finmb"
# frozen_path = "./data/f_semantic_reservoir_512D.bin"
total_steps = 7301

{load_path, start_step} = MOE.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 = 
    MOE.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} = MOE.Parameter.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()
      {t_p, Nx.Random.key(System.system_time())}
    end

  trainable_gpu = MOE.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 = 800
      cooldown_start = 6600
      
      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} = MOE.Phase2Trainer.compute_grad(cur_tr, gpu_batch)
      loss_val = Nx.to_number(loss_tensor)

      {up_tr, up_step} = MOE.Optimizer.update(cur_tr, grads, c_step, lr_tensor)
  
      # Telemetry
      if rem(step, 100) == 0 do
        {emb, dec, key} = up_tr

        emb_norm = emb |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        dec_norm = dec |> Nx.pow(2) |> Nx.sum() |> Nx.sqrt() |> Nx.to_number() |> Float.round(4)
        key_norm = key |> 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)} | emb #{emb_norm} | dec #{dec_norm} | key #{key_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 = MOE.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 = MOE.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

No checkpoints found for 'v67_MOE_512sq_yolo_em2b4u_finmb'. Starting FRESH from Genesis.
Initializing Lazy Curriculum Data Stream from [./data/bbylm_512_strict]...
Igniting brand new Trainable weights for v67_MOE_512sq_yolo_em2b4u_finmb...

14:42:45.058 [info] XLA service 0x7368dc17f180 initialized for platform ROCM (this does not guarantee that XLA will be used). Devices:

14:42:45.067 [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)

14:42:45.067 [info] Using BFC allocator.

14:42:45.067 [info] XLA backend will use up to 6429868032 bytes on device 0 for BFCAllocator.

14:42:45.067 [info] XLA backend will use up to 2143289344 bytes on device 0 for CollectiveBFCAllocator.

🔄 Epoch 1 Started | Shuffling 6 shards...
➡️ Step 0 | LR: 0.0 | Loss: 5.4445 | emb 128.0558 | dec 5.6569 | key 11.291
➡️ Step 100 | LR: 8.7e-5 | Loss: 5.4512 | emb 128.0567 | dec 5.6571 | key 11.2904
➡️ Step 200 | LR: 1.75e-4 | Loss: 5.3982 | emb 128.0697 | dec 5.6578 | key 11.2879
Checkpoint reached! Saving Step 200...
➡️ Step 300 | LR: 2.62e-4 | Loss: 5.3389 | emb 128.127 | dec 5.6593 | key 11.2818
➡️ Step 400 | LR: 3.5e-4 | Loss: 5.2714 | emb 128.2798 | dec 5.6614 | key 11.2698
Checkpoint reached! Saving Step 400...
➡️ Step 500 | LR: 4.38e-4 | Loss: 5.1545 | emb 128.6007 | dec 5.6653 | key 11.2469
➡️ Step 600 | LR: 5.25e-4 | Loss: 5.019 | emb 129.0987 | dec 5.6736 | key 11.2169
Checkpoint reached! Saving Step 600...
➡️ Step 700 | LR: 6.12e-4 | Loss: 4.7583 | emb 129.9321 | dec 5.6834 | key 11.1697

🔄 Epoch 2 Started | Shuffling 6 shards...
➡️ Step 800 | LR: 0.0007 | Loss: 4.6208 | emb 131.1795 | dec 5.6933 | key 11.0994
Checkpoint reached! Saving Step 800...
➡️ Step 900 | LR: 0.0007 | Loss: 4.398 | emb 132.7678 | dec 5.7048 | key 11.0195
➡️ Step 1000 | LR: 0.0007 | Loss: 4.172 | emb 134.5133 | dec 5.7174 | key 10.94
Checkpoint reached! Saving Step 1000...
➡️ Step 1100 | LR: 0.0007 | Loss: 3.949 | emb 136.5118 | dec 5.7295 | key 10.8486
➡️ Step 1200 | LR: 0.0007 | Loss: 3.76 | emb 138.7041 | dec 5.7429 | key 10.7532
Checkpoint reached! Saving Step 1200...
➡️ Step 1300 | LR: 0.0007 | Loss: 3.6448 | emb 140.8889 | dec 5.7575 | key 10.6599
➡️ Step 1400 | LR: 0.0007 | Loss: 3.491 | emb 143.0412 | dec 5.7708 | key 10.562
Checkpoint reached! Saving Step 1400...

🔄 Epoch 3 Started | Shuffling 6 shards...
➡️ Step 1500 | LR: 0.0007 | Loss: 3.3006 | emb 145.1632 | dec 5.7794 | key 10.456
➡️ Step 1600 | LR: 0.0007 | Loss: 3.2203 | emb 147.2269 | dec 5.7845 | key 10.3444
Checkpoint reached! Saving Step 1600...
➡️ Step 1700 | LR: 0.0007 | Loss: 3.1176 | emb 149.2014 | dec 5.7886 | key 10.2232
➡️ Step 1800 | LR: 0.0007 | Loss: 3.2085 | emb 151.1498 | dec 5.795 | key 10.1314
Checkpoint reached! Saving Step 1800...
➡️ Step 1900 | LR: 0.0007 | Loss: 3.0736 | emb 152.897 | dec 5.7994 | key 10.046
➡️ Step 2000 | LR: 0.0007 | Loss: 2.9582 | emb 154.3849 | dec 5.8018 | key 9.9579
Checkpoint reached! Saving Step 2000...
➡️ Step 2100 | LR: 0.0007 | Loss: 2.913 | emb 155.9341 | dec 5.8063 | key 9.832

🔄 Epoch 4 Started | Shuffling 6 shards...
➡️ Step 2200 | LR: 0.0007 | Loss: 2.8184 | emb 157.4671 | dec 5.8125 | key 9.7109
Checkpoint reached! Saving Step 2200...