|
| 1 | +defmodule Emily.Conformance.CompilerNativeTest do |
| 2 | + @moduledoc """ |
| 3 | + CM5 — the **no-fallback** gate: real Bumblebee model forwards compile |
| 4 | + through the native single-NIF `Emily.Compiler` (`native: true`) and |
| 5 | + match the Evaluator-on-`Emily.Backend` path, with **zero** fallback to |
| 6 | + `Nx.BinaryBackend`. |
| 7 | +
|
| 8 | + The native compiler lowers the whole `Nx.Defn.Expr` to one program and |
| 9 | + **raises** on any op it can't lower (no silent fallback by design), so a |
| 10 | + forward that completes proves full native op coverage for that model. |
| 11 | + We additionally fail on any `[:emily, :fallback, *]` telemetry — the |
| 12 | + Backend-level BinaryBackend fallback — so an op that silently round-trips |
| 13 | + to the host is caught too. |
| 14 | +
|
| 15 | + Gated `:conformance` (downloads ~3 MB tiny-random HF fixtures); run with |
| 16 | + `mix test --only conformance`. |
| 17 | + """ |
| 18 | + use ExUnit.Case, async: false |
| 19 | + |
| 20 | + @moduletag :conformance |
| 21 | + @moduletag timeout: 600_000 |
| 22 | + |
| 23 | + setup do |
| 24 | + prev = Nx.default_backend() |
| 25 | + Nx.global_default_backend(Emily.Backend) |
| 26 | + on_exit(fn -> Nx.global_default_backend(prev) end) |
| 27 | + :ok |
| 28 | + end |
| 29 | + |
| 30 | + # Run `model`'s forward both ways and assert every output leaf agrees, |
| 31 | + # while asserting no Backend fallback fires on the native path. |
| 32 | + defp assert_native_matches(model, params, inputs) do |
| 33 | + {_init, native_predict} = Axon.build(model, compiler: Emily.Compiler, native: true) |
| 34 | + {_init, eval_predict} = Axon.build(model, compiler: Emily.Compiler) |
| 35 | + |
| 36 | + {native, fallbacks} = with_fallback_count(fn -> native_predict.(params, inputs) end) |
| 37 | + eval = eval_predict.(params, inputs) |
| 38 | + |
| 39 | + assert fallbacks == 0, |
| 40 | + "native compile path triggered #{fallbacks} Backend BinaryBackend fallback(s)" |
| 41 | + |
| 42 | + compare_outputs(native, eval) |
| 43 | + end |
| 44 | + |
| 45 | + # Count [:emily, :fallback, *] telemetry events during `fun`. |
| 46 | + defp with_fallback_count(fun) do |
| 47 | + ref = make_ref() |
| 48 | + me = self() |
| 49 | + id = {__MODULE__, ref} |
| 50 | + |
| 51 | + :telemetry.attach_many( |
| 52 | + id, |
| 53 | + [[:emily, :fallback, :start], [:emily, :fallback, :stop], [:emily, :fallback, :exception]], |
| 54 | + &__MODULE__.handle_fallback/4, |
| 55 | + {me, ref} |
| 56 | + ) |
| 57 | + |
| 58 | + result = fun.() |
| 59 | + :telemetry.detach(id) |
| 60 | + |
| 61 | + count = drain(ref, 0) |
| 62 | + {result, count} |
| 63 | + end |
| 64 | + |
| 65 | + @doc false |
| 66 | + def handle_fallback(_event, _measure, _meta, {pid, ref}), do: send(pid, {ref, :fallback}) |
| 67 | + |
| 68 | + defp drain(ref, n) do |
| 69 | + receive do |
| 70 | + {^ref, :fallback} -> drain(ref, n + 1) |
| 71 | + after |
| 72 | + 0 -> n |
| 73 | + end |
| 74 | + end |
| 75 | + |
| 76 | + defp compare_outputs(%Nx.Tensor{} = native, %Nx.Tensor{} = eval) do |
| 77 | + assert native.shape == eval.shape |
| 78 | + assert native.type == eval.type |
| 79 | + # Same MLX kernels in the same order => exact; allow a tiny tolerance |
| 80 | + # against fp reassociation in the lazy-graph eval. |
| 81 | + n = Nx.to_flat_list(native) |
| 82 | + e = Nx.to_flat_list(eval) |
| 83 | + |
| 84 | + assert Enum.zip(n, e) |> Enum.all?(fn {a, b} -> abs(a - b) <= 1.0e-4 + 1.0e-4 * abs(b) end), |
| 85 | + "native vs evaluator outputs diverge beyond tolerance" |
| 86 | + end |
| 87 | + |
| 88 | + defp compare_outputs(native, eval) when is_map(native) and not is_struct(native) do |
| 89 | + for {k, nv} <- native, Map.has_key?(eval, k) do |
| 90 | + compare_outputs(nv, Map.fetch!(eval, k)) |
| 91 | + end |
| 92 | + end |
| 93 | + |
| 94 | + # Axon.None placeholders, tuples, and other non-tensor leaves: skip. |
| 95 | + defp compare_outputs(_native, _eval), do: :ok |
| 96 | + |
| 97 | + test "tiny DistilBERT base forward: single-NIF native == evaluator, no fallback" do |
| 98 | + {:ok, %{model: model, params: params, spec: _spec}} = |
| 99 | + Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertModel"}) |
| 100 | + |
| 101 | + inputs = %{ |
| 102 | + "input_ids" => Nx.tensor([[1, 5, 7, 2, 3, 9]], backend: Emily.Backend), |
| 103 | + "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1]], backend: Emily.Backend) |
| 104 | + } |
| 105 | + |
| 106 | + assert_native_matches(model, params, inputs) |
| 107 | + end |
| 108 | + |
| 109 | + test "tiny DistilBERT for-masked-LM forward: native == evaluator, no fallback" do |
| 110 | + {:ok, %{model: model, params: params}} = |
| 111 | + Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertForMaskedLM"}) |
| 112 | + |
| 113 | + inputs = %{ |
| 114 | + "input_ids" => Nx.tensor([[1, 5, 7, 2]], backend: Emily.Backend), |
| 115 | + "attention_mask" => Nx.tensor([[1, 1, 1, 1]], backend: Emily.Backend) |
| 116 | + } |
| 117 | + |
| 118 | + assert_native_matches(model, params, inputs) |
| 119 | + end |
| 120 | + |
| 121 | + test "tiny ViT base forward (conv patch embed): native == evaluator, no fallback" do |
| 122 | + {:ok, %{model: model, params: params}} = |
| 123 | + Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-ViTModel"}, architecture: :base) |
| 124 | + |
| 125 | + inputs = %{ |
| 126 | + "pixel_values" => Nx.broadcast(Nx.tensor(0.1, backend: Emily.Backend), {1, 30, 30, 3}) |
| 127 | + } |
| 128 | + |
| 129 | + assert_native_matches(model, params, inputs) |
| 130 | + end |
| 131 | +end |
0 commit comments