Merge pull request #180 from ausimian/feat/expr-compiler-training-convergence

Native training convergence tests (#174)

Author: ausimian

RELEASE.md

@@ -68,7 +68,15 @@
   the eager NIFs and the compiled replay share one implementation. A
   small-CNN **training step** (conv + maxpool forward and backward, grad,
   SGD) now lowers fully native under `native_fallback: :raise`, producing a
-  loss bit-identical to the evaluator.
+  loss bit-identical to the evaluator. Native training is now
+  **convergence**-tested, not just verified-lowering: handwritten CNN
+  (30 SGD steps) and MLP (50 steps) trajectories match the op-by-op
+  evaluator bit-for-bit and a `BinaryBackend` oracle to f32 tolerance, and
+  full **Axon** training drives native end-to-end — `Axon.Loop.run`
+  forwards `native: true`/`native_fallback:` to the defn jit, so a LeNet
+  CNN and a dense MLP train on real MNIST entirely through the single-NIF
+  path (forward, categorical-cross-entropy, backward, Adam) and reach the
+  same >97% / >96% accuracy as the evaluator (`:training_full`).
 
 - **`Bumblebee.Text.generation` compiles fully native — greedy and sampling.**
   The headline result: an end-to-end Bumblebee generation (the transformer

test/emily/training/cnn_native_curve_test.exs

@@ -0,0 +1,134 @@
+defmodule Emily.Training.CnnNativeCurveTest do
+  @moduledoc """
+  Native single-NIF CNN training convergence (issue #174).
+
+  The training analogue of the conformance native lanes
+  (`Emily.Conformance.CompilerNativeTest`): a full conv + maxpool
+  training step — forward, backward, grad, and SGD update — is driven
+  through `compiler: Emily.Compiler, native: true, native_fallback:
+  :raise` for 30 steps and the per-step loss trajectory is checked
+  against two references.
+
+  Why this exists. `cnn_curve_test.exs` already curve-matches the
+  handwritten CNN, but only in **eval** mode (`Emily.Compiler` walking
+  the Expr op-by-op via the Evaluator). Every other `training/*` test
+  is eval-only too, so CNN training was verified-lowering (the
+  `compiler_equivalence_test.exs` op gates) but never **convergence**-
+  tested under the single-NIF replay. This closes that gap.
+
+  Three lanes, same deterministic init and data:
+
+    * **native** — `native: true, native_fallback: :raise`. The
+      `:raise` makes this a no-fallback gate: if any op in the
+      forward+backward+grad+SGD step fails to lower (the maxpool
+      backward lands on `window_scatter_max` every step; the conv
+      backward flips the kernel with `reverse`), the run raises here
+      instead of silently degrading to the evaluator.
+    * **eval** — `Emily.Compiler` op-by-op. Same MLX kernels in the
+      same order as the native replay, so the two track **bit-
+      identically** through training. A 1e-6 bar asserts the single-
+      NIF lowering reproduces op-by-op exactly across 30 SGD updates.
+    * **binary** — `Nx.Defn.Evaluator` on `Nx.BinaryBackend`, the
+      non-MLX convergence oracle. Looser bar (1e-2 rtol, as in
+      `cnn_curve_test.exs`) absorbs f32 reduction-order drift between
+      MLX's parallel reductions and BinaryBackend's sequential ones.
+
+  No Axon — the handwritten path keeps the failure surface tiny (see
+  `cnn_curve_test.exs`). The Axon CNN canary stays in
+  `mnist_cnn_full_test.exs` (`:training_full`).
+  """
+
+  use ExUnit.Case, async: true
+
+  alias Emily.TrainingHelper, as: TH
+  import TH, only: [close?: 4, flunk_trajectory: 5]
+
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+  @eval [compiler: Emily.Compiler]
+
+  @input_shape {1, 10, 10}
+  @batch 4
+  @classes 3
+  @steps 30
+  @lr_val 0.05
+
+  test "per-step CNN loss trajectory matches under native single-NIF compile" do
+    # Native single-NIF lane — the system under test. `native_fallback:
+    # :raise` proves full native coverage of the training step.
+    params_native = TH.init_cnn(@input_shape, @classes, 0, Emily.Backend)
+    {x_native, y_native} = TH.cnn_batch({@batch, 10, 10}, @classes, Emily.Backend)
+    lr_native = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend)
+
+    losses_native =
+      TH.run_steps(
+        &TH.cnn_step_with_loss/4,
+        params_native,
+        [x_native, y_native, lr_native],
+        @steps,
+        @native
+      )
+
+    # Op-by-op Emily eval lane — same MLX kernels, isolates single-NIF
+    # lowering bugs from backend numerics.
+    params_eval = TH.init_cnn(@input_shape, @classes, 0, Emily.Backend)
+    {x_eval, y_eval} = TH.cnn_batch({@batch, 10, 10}, @classes, Emily.Backend)
+    lr_eval = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend)
+
+    losses_eval =
+      TH.run_steps(
+        &TH.cnn_step_with_loss/4,
+        params_eval,
+        [x_eval, y_eval, lr_eval],
+        @steps,
+        @eval
+      )
+
+    # BinaryBackend oracle — the non-MLX convergence reference.
+    params_bin = TH.init_cnn(@input_shape, @classes, 0, Nx.BinaryBackend)
+    {x_bin, y_bin} = TH.cnn_batch({@batch, 10, 10}, @classes, Nx.BinaryBackend)
+    lr_bin = Nx.tensor(@lr_val, type: {:f, 32}, backend: Nx.BinaryBackend)
+
+    losses_bin =
+      TH.run_steps(
+        &TH.cnn_step_with_loss/4,
+        params_bin,
+        [x_bin, y_bin, lr_bin],
+        @steps,
+        Nx.Defn.Evaluator
+      )
+
+    assert length(losses_native) == @steps
+    assert length(losses_eval) == @steps
+    assert length(losses_bin) == @steps
+
+    # 1. Single-NIF native == op-by-op eval. Both are MLX in the same
+    #    order, so they track bit-identically; the tight bar makes a
+    #    divergent native trajectory a hard failure.
+    for {{ln, le}, i} <- Enum.zip(losses_native, losses_eval) |> Enum.with_index() do
+      close?(ln, le, 1.0e-6, 1.0e-6) ||
+        flunk_trajectory(i, ln, le, losses_native, losses_eval)
+    end
+
+    # 2. Native trajectory matches the BinaryBackend oracle within the
+    #    CNN tolerance — same bar as cnn_curve_test.exs.
+    for {{ln, lb}, i} <- Enum.zip(losses_native, losses_bin) |> Enum.with_index() do
+      close?(ln, lb, 1.0e-4, 1.0e-2) ||
+        flunk_trajectory(i, ln, lb, losses_native, losses_bin)
+    end
+
+    # 3. Convergence — the native loss actually decreased over the run.
+    assert List.first(losses_native) > List.last(losses_native),
+           "native loss did not decrease: first=#{List.first(losses_native)} " <>
+             "last=#{List.last(losses_native)}"
+
+    # 4. Final loss agrees with the oracle (convergence correctness:
+    #    catches a run where per-step drift averaged out but the
+    #    optimizer ended up somewhere wrong).
+    ln_final = List.last(losses_native)
+    lb_final = List.last(losses_bin)
+
+    assert close?(ln_final, lb_final, 1.0e-4, 1.0e-2),
+           "final loss divergence: native=#{ln_final} bin=#{lb_final} " <>
+             "reldiff=#{abs(ln_final - lb_final) / abs(lb_final)}"
+  end
+end

test/emily/training/mlp_native_curve_test.exs

@@ -0,0 +1,113 @@
+defmodule Emily.Training.MlpNativeCurveTest do
+  @moduledoc """
+  Native single-NIF MLP training convergence (issue #174).
+
+  The dense + SGD companion to `cnn_native_curve_test.exs`: a 2-layer
+  ReLU MLP training step — forward, backward, grad, SGD update — is
+  driven through `compiler: Emily.Compiler, native: true,
+  native_fallback: :raise` for 50 steps and the per-step loss
+  trajectory is checked against two references.
+
+  This closes the matmul-dominated half of the training-coverage gap
+  the issue calls out: `mlp_curve_test.exs` already curve-matches this
+  MLP, but only in eval mode. Here the same step replays through the
+  single NIF, with `:raise` proving the dense forward/backward and SGD
+  update lower with **zero** fallback.
+
+  Three lanes, same deterministic init and data:
+
+    * **native** — single-NIF replay, no-fallback gate.
+    * **eval** — `Emily.Compiler` op-by-op; bit-identical to native
+      (same MLX kernels, same order), asserted at a 1e-6 bar.
+    * **binary** — `Nx.Defn.Evaluator` on `Nx.BinaryBackend`, the
+      non-MLX oracle. The MLP is matmul-dominated, so the bar matches
+      `mlp_curve_test.exs` (1e-3 per-step rtol, 1e-4 final).
+  """
+
+  use ExUnit.Case, async: true
+
+  alias Emily.TrainingHelper, as: TH
+  import TH, only: [close?: 4, flunk_trajectory: 5]
+
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+  @eval [compiler: Emily.Compiler]
+
+  @dims {4, 8, 3}
+  @batch_shape {16, 4, 3}
+  @steps 50
+  @lr_val 0.5
+
+  test "per-step MLP loss trajectory matches under native single-NIF compile" do
+    # Native single-NIF lane — the system under test.
+    params_native = TH.init_mlp(@dims, 0, Emily.Backend)
+    {x_native, y_native} = TH.mlp_batch(@batch_shape, Emily.Backend)
+    lr_native = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend)
+
+    losses_native =
+      TH.run_steps(
+        &TH.mlp_step_with_loss/4,
+        params_native,
+        [x_native, y_native, lr_native],
+        @steps,
+        @native
+      )
+
+    # Op-by-op Emily eval lane — same MLX kernels.
+    params_eval = TH.init_mlp(@dims, 0, Emily.Backend)
+    {x_eval, y_eval} = TH.mlp_batch(@batch_shape, Emily.Backend)
+    lr_eval = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend)
+
+    losses_eval =
+      TH.run_steps(
+        &TH.mlp_step_with_loss/4,
+        params_eval,
+        [x_eval, y_eval, lr_eval],
+        @steps,
+        @eval
+      )
+
+    # BinaryBackend oracle.
+    params_bin = TH.init_mlp(@dims, 0, Nx.BinaryBackend)
+    {x_bin, y_bin} = TH.mlp_batch(@batch_shape, Nx.BinaryBackend)
+    lr_bin = Nx.tensor(@lr_val, type: {:f, 32}, backend: Nx.BinaryBackend)
+
+    losses_bin =
+      TH.run_steps(
+        &TH.mlp_step_with_loss/4,
+        params_bin,
+        [x_bin, y_bin, lr_bin],
+        @steps,
+        Nx.Defn.Evaluator
+      )
+
+    assert length(losses_native) == @steps
+    assert length(losses_eval) == @steps
+    assert length(losses_bin) == @steps
+
+    # 1. Single-NIF native == op-by-op eval (bit-identical MLX path).
+    for {{ln, le}, i} <- Enum.zip(losses_native, losses_eval) |> Enum.with_index() do
+      close?(ln, le, 1.0e-6, 1.0e-6) ||
+        flunk_trajectory(i, ln, le, losses_native, losses_eval)
+    end
+
+    # 2. Native trajectory matches the BinaryBackend oracle — same bar
+    #    as mlp_curve_test.exs.
+    for {{ln, lb}, i} <- Enum.zip(losses_native, losses_bin) |> Enum.with_index() do
+      close?(ln, lb, 1.0e-4, 1.0e-3) ||
+        flunk_trajectory(i, ln, lb, losses_native, losses_bin)
+    end
+
+    # 3. Convergence — the native loss actually decreased over the run.
+    assert List.first(losses_native) > List.last(losses_native),
+           "native loss did not decrease: first=#{List.first(losses_native)} " <>
+             "last=#{List.last(losses_native)}"
+
+    # 4. Final loss agrees with the oracle (convergence correctness).
+    ln_final = List.last(losses_native)
+    lb_final = List.last(losses_bin)
+
+    assert close?(ln_final, lb_final, 1.0e-5, 1.0e-4),
+           "final loss divergence: native=#{ln_final} bin=#{lb_final} " <>
+             "reldiff=#{abs(ln_final - lb_final) / abs(lb_final)}"
+  end
+end

test/emily/training/mnist_cnn_native_full_test.exs

@@ -0,0 +1,84 @@
+defmodule Emily.Training.MnistCnnNativeFullTest do
+  @moduledoc """
+  MNIST CNN convergence under the **native single-NIF compiler**
+  (issue #174, `:training_full`).
+
+  The native-lane analogue of `mnist_cnn_full_test.exs`: the same
+  LeNet-style Axon CNN trains on real MNIST, but the whole training
+  step compiles through `compiler: Emily.Compiler, native: true,
+  native_fallback: :raise` instead of the op-by-op evaluator.
+
+  Why `native_fallback: :raise` makes this self-proving. Once
+  `native: true` reaches `Emily.Compiler`, the run is binary: the Expr
+  lowers to one program and replays in a single NIF, or an op it can't
+  lower raises (`Emily.Compiler.build_native/4` reraises the lowering
+  `ArgumentError` under `:raise` — it can never silently degrade to the
+  evaluator, that path only exists under `:eval`). So a training run
+  that *completes* proves the entire step — forward (conv, ReLU,
+  maxpool), categorical-cross-entropy loss, the backward
+  (`window_scatter_max` for the maxpool grad, `reverse` for the conv
+  kernel flip), and the Adam update — lowered fully native with zero
+  fallback. `Axon.Loop.run` forwards the per-call jit options (it pops
+  only `:jit_compile?`/`:force_garbage_collection?`), so the options
+  genuinely reach the compiler.
+
+  Native replay is bit-identical to the evaluator (same MLX kernels in
+  the same order), so the accuracy bar matches the eval canary exactly
+  (>97%). The point isn't a different number — it's that training
+  reaches it through the single-NIF path.
+
+  Opt-in — `mix test --only training_full` (downloads MNIST, multi-
+  minute training).
+  """
+
+  use ExUnit.Case, async: true
+
+  alias Emily.MnistHelper
+
+  @moduletag :training_full
+  @moduletag capture_log: true
+  @moduletag timeout: 600_000
+
+  setup do
+    Nx.default_backend(Emily.Backend)
+    :ok
+  end
+
+  @batch_size 64
+  @epochs 5
+  @target_accuracy 0.97
+
+  # Strict no-fallback native lane. `:raise` makes a completed run a
+  # proof of full native lowering (see the moduledoc).
+  @native [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+
+  test "Axon CNN reaches >#{trunc(@target_accuracy * 100)}% accuracy via the native single-NIF compiler" do
+    {train_batches, test_images, test_labels} = MnistHelper.load_mnist(@batch_size, :cnn)
+
+    # Channels-last (Axon default) — MnistHelper produces {N, 28, 28, 1}.
+    model =
+      Axon.input("input", shape: {nil, 28, 28, 1})
+      |> Axon.conv(8, kernel_size: {3, 3}, activation: :relu)
+      |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2])
+      |> Axon.conv(16, kernel_size: {3, 3}, activation: :relu)
+      |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2])
+      |> Axon.flatten()
+      |> Axon.dense(64, activation: :relu)
+      |> Axon.dense(10, activation: :softmax)
+
+    # The whole training loop (init + step) compiles native: `Axon.Loop.run`
+    # forwards `native:`/`native_fallback:` to the defn jit. Under `:raise`,
+    # reaching the end proves every op lowered — no silent eval fallback.
+    trained_state =
+      model
+      |> Axon.Loop.trainer(:categorical_cross_entropy, :adam)
+      |> Axon.Loop.run(train_batches, %{}, [epochs: @epochs] ++ @native)
+
+    # Evaluate through the native path too, so the accuracy that gates the
+    # test is itself produced by the single-NIF forward.
+    accuracy = MnistHelper.evaluate(model, trained_state, test_images, test_labels, @native)
+
+    assert accuracy >= @target_accuracy,
+           "native MNIST CNN accuracy #{Float.round(accuracy, 4)} below target #{@target_accuracy}"
+  end
+end