diff --git a/RELEASE.md b/RELEASE.md index ff3da7d..74096ad 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -110,6 +110,69 @@ last-axis fast path stays on MLX; interior-axis usage is rare on our M3/M4 critical path (transformer inference doesn't need it). +- M4 — Qwen3 inference. `Qwen/Qwen3-0.6B` greedy-decodes end-to-end on + `Emily.Backend` through Bumblebee's causal-LM serving. Everything on + Qwen3's critical path (QK-norm, rotary embeddings, GQA, SwiGLU FFN, + RMSNorm, tied embeddings, KV-cache `put_slice` in a `defn` while + loop) runs correctly. + - **Native `put_slice/4`** in `Emily.Backend`, backed by a new + `Native.slice_update/3` NIF over `mx::slice_update`. Replaces the + BinaryBackend round-trip — autoregressive decoding calls + `put_slice` per layer per token to append into the KV cache, and + the old fallback transferred ~1 MB of cache state through the + allocator on every call. Also fixes a latent bug in the old + implementation: dynamic scalar-tensor `start_indices` on + `Emily.Backend` used to slip through unconverted and crash inside + BinaryBackend. `slice_start` is now applied to every start + index, matching the `slice/5` callback. + - **Operand-type promotion in `put_slice`.** `Nx.put_slice` + promotes the output type across tensor/update (an s32 pad buffer + clashing with an s64 decoder input becomes s64), but the backend + callback still receives the original-type operands. We cast both + `t` and `slice` to `out.type` via `Native.astype` before + dispatching to `slice_update`. Without this the MLX buffer + silently disagrees with the Nx shape metadata — the first symptom + is `Nx.to_binary` returning a half-sized binary and + `BinaryBackend.bitstring_part` raising a match error deep inside + the tokenizer decode. Mirrors the arithmetic-op promotion fix + landed in M3. + - **`test/emily/conformance/qwen3_test.exs`** + (`@moduletag :conformance`) — ports `Bumblebee.Text.Qwen3Test` + verbatim (three architectures: `:base`, + `:for_causal_language_modeling`, `:for_sequence_classification`), + with HF reference slices checked in, plus a `greedy generation` + describe block that drives + `Bumblebee.Text.Generation.build_generate` on the tiny-random + causal LM. That smoke test feeds synthetic `input_ids` in + `[0, 1024)` (tokenizer vocab is 151 k but the tiny checkpoint's + embedding is 1024 rows), greedy-decodes 16 tokens through the + full generation pipeline (`Axon.predict` + logit processing + + `Nx.argmax` + `put_slice` KV-cache update + `defn while`), and + asserts bit-exact equality against both `Nx.BinaryBackend` run + on the same inputs *and* a checked-in 16-token reference. + - **`test/emily/conformance/qwen3_full_test.exs`** + (`@moduletag :qwen3_full`, excluded from `--only conformance` + because the checkpoint is ~1.5 GB) — loads `Qwen/Qwen3-0.6B` + proper, greedy-decodes 32 tokens from a fixed prompt through + `Nx.Serving`, and asserts the completion string matches a + checked-in reference. Run with `mix test --only qwen3_full`. + - **`bench/qwen3_tokens_per_sec.exs`** — standalone wall-clock + throughput harness. Loads `Qwen/Qwen3-0.6B`, runs N warmup + iterations + M measured iterations of greedy decode, reports + tokens/sec. Prompt, token count, and iteration counts are + overridable via `EMILY_BENCH_*` env vars. Baseline observed on a + dev M3 host: ~13.8 tok/s at 16 new tokens under the + `Nx.Defn.Evaluator` compiler (no `mlx::core::compile` wrap yet — + that lands in M6). Intended as a regression gate, not a headline + number. + - **Bumblebee dependency** bumped from Hex 0.6.3 to a pinned `main` + commit (`273805e9…`) so `Bumblebee.Text.Qwen3` is available — the + text port is on main but not yet in a Hex release. Revert to a + Hex version as soon as one ships Qwen3 support. + - **`test_helper.exs`** extended the exclude list with + `:qwen3_full` so the weights-heavy test stays out of + `mix test --only conformance`. + - M3 — DistilBERT end-to-end on Bumblebee. Every Nx op on the transformer critical path now runs natively on MLX; the full forward pass matches HuggingFace Transformers (PyTorch) reference diff --git a/bench/qwen3_tokens_per_sec.exs b/bench/qwen3_tokens_per_sec.exs new file mode 100644 index 0000000..97ce98f --- /dev/null +++ b/bench/qwen3_tokens_per_sec.exs @@ -0,0 +1,105 @@ +# Qwen3-0.6B greedy-decode throughput on `Emily.Backend`. +# +# Usage: +# +# MIX_ENV=test mix run bench/qwen3_tokens_per_sec.exs +# +# Optional environment variables: +# +# EMILY_BENCH_MODEL HuggingFace repo id. Defaults to +# "Qwen/Qwen3-0.6B". +# EMILY_BENCH_NEW_TOKENS Number of tokens to greedy-decode per +# run. Defaults to 64. +# EMILY_BENCH_PROMPT Prompt text. Defaults to a fixed short +# English sentence. +# EMILY_BENCH_WARMUP Number of warm-up runs (not measured). +# Defaults to 1. +# EMILY_BENCH_RUNS Number of measured runs. Defaults to 3. +# +# The first run downloads the model (~1.2 GB at f32, ~600 MB at f16). +# We deliberately avoid `Benchee` — this benchmark has one workload and +# one metric (tokens/sec). The whole script is standalone so a reader +# can follow the generation flow without chasing macros. + +defmodule Emily.Bench.Qwen3 do + @default_model "Qwen/Qwen3-0.6B" + @default_prompt "The quick brown fox jumps over the lazy dog." + @default_new_tokens 64 + @default_warmup 1 + @default_runs 3 + + def run do + Nx.global_default_backend(Emily.Backend) + + model_repo = System.get_env("EMILY_BENCH_MODEL", @default_model) + prompt = System.get_env("EMILY_BENCH_PROMPT", @default_prompt) + + new_tokens = + System.get_env("EMILY_BENCH_NEW_TOKENS") + |> env_int(@default_new_tokens) + + warmup = System.get_env("EMILY_BENCH_WARMUP") |> env_int(@default_warmup) + runs = System.get_env("EMILY_BENCH_RUNS") |> env_int(@default_runs) + + IO.puts("Emily / Qwen3 throughput benchmark") + IO.puts(" model : #{model_repo}") + IO.puts(" prompt : #{inspect(prompt)}") + IO.puts(" new tokens : #{new_tokens}") + IO.puts(" warmup : #{warmup}") + IO.puts(" runs : #{runs}") + IO.puts("") + + {:ok, model_info} = Bumblebee.load_model({:hf, model_repo}) + {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, model_repo}) + + {:ok, generation_config} = Bumblebee.load_generation_config({:hf, model_repo}) + + generation_config = + Bumblebee.configure(generation_config, + max_new_tokens: new_tokens, + strategy: %{type: :greedy_search} + ) + + serving = + Bumblebee.Text.generation(model_info, tokenizer, generation_config, + defn_options: [compiler: Nx.Defn.Evaluator] + ) + + for _ <- Stream.duplicate(:ok, warmup) do + IO.puts("[warmup] generating…") + %{results: [_]} = Nx.Serving.run(serving, prompt) + end + + measurements = + for n <- 1..runs//1 do + {elapsed_us, %{results: [%{text: text}]}} = + :timer.tc(fn -> Nx.Serving.run(serving, prompt) end) + + secs = elapsed_us / 1_000_000 + tps = new_tokens / secs + IO.puts("[run #{n}] #{Float.round(secs, 3)} s, #{Float.round(tps, 2)} tok/s") + {secs, tps, text} + end + + tps_list = Enum.map(measurements, fn {_, tps, _} -> tps end) + [{_, _, sample} | _] = measurements + mean = Enum.sum(tps_list) / length(tps_list) + {min_tps, max_tps} = Enum.min_max(tps_list) + + IO.puts("") + IO.puts("tokens/sec mean=#{Float.round(mean, 2)} min=#{Float.round(min_tps, 2)} max=#{Float.round(max_tps, 2)}") + IO.puts("") + IO.puts("first completion:") + IO.puts(String.slice(sample, 0, 500)) + end + + defp env_int(nil, default), do: default + defp env_int(s, default) do + case Integer.parse(s) do + {n, ""} -> n + _ -> default + end + end +end + +Emily.Bench.Qwen3.run() diff --git a/c_src/ops/index.cpp b/c_src/ops/index.cpp index 94bf587..7ab6239 100644 --- a/c_src/ops/index.cpp +++ b/c_src/ops/index.cpp @@ -31,6 +31,27 @@ fine::ResourcePtr slice( } FINE_NIF(slice, 0); +// slice_update/3: write `update` into `src` starting at `start`. `stop` +// is derived as `start + shape(update)` and strides default to 1 on +// every axis (Nx.put_slice has no stride parameter). Output shape +// equals `src.shape`. +fine::ResourcePtr slice_update( + ErlNifEnv *, + fine::ResourcePtr src, + fine::ResourcePtr update, + std::vector start) { + const auto &update_shape = update->array.shape(); + mx::Shape start_shape = to_mlx_shape(start); + mx::Shape stop_shape; + stop_shape.reserve(start_shape.size()); + for (size_t i = 0; i < start_shape.size(); ++i) { + stop_shape.push_back(start_shape[i] + update_shape[i]); + } + return wrap(mx::slice_update( + src->array, update->array, std::move(start_shape), std::move(stop_shape))); +} +FINE_NIF(slice_update, 0); + // take/3: gather along `axis` using integer indices. fine::ResourcePtr take( ErlNifEnv *, diff --git a/lib/emily/backend.ex b/lib/emily/backend.ex index 8aeac46..5d54b18 100644 --- a/lib/emily/backend.ex +++ b/lib/emily/backend.ex @@ -592,13 +592,21 @@ defmodule Emily.Backend do defp slice_start(i) when is_integer(i), do: i defp slice_start(%T{} = t), do: t |> Nx.backend_copy(Nx.BinaryBackend) |> Nx.to_number() - # put_slice: MLX has no direct primitive; route via BinaryBackend. - # Nx's Backend contract order is (out, tensor, start_indices, slice). - # start_indices arrive as scalar tensors on Emily.Backend; Nx auto- - # transfers them when the BinaryBackend call goes through `to_indices`. + # put_slice: implemented natively via MLX `slice_update`. Nx promotes + # operand types at the API layer — `Nx.put_slice(s32_buf, _, s64_upd)` + # declares an s64 output — but our callback arguments still carry the + # original backend types. We cast both `t` and `slice` to `out.type` + # before dispatching so the MLX buffer matches Nx's shape/type view. + # Scalar-tensor starts are materialised to integers here (dynamic + # indices show up when autoregressive loops dispatch put_slice from + # within `defn`). @impl true - def put_slice(out, t, start_indices, slice), - do: via_binary(out, [t, slice], &Nx.put_slice(&1, start_indices, &2)) + def put_slice(%T{type: type} = out, %T{} = t, start_indices, %T{} = slice) do + starts = Enum.map(start_indices, &slice_start/1) + src_ref = Native.astype(ref(t), type) + update_ref = Native.astype(ref(slice), type) + Native.slice_update(src_ref, update_ref, starts) |> wrap(out) + end @impl true def select(%T{} = out, pred, on_true, on_false) do diff --git a/lib/emily/native.ex b/lib/emily/native.ex index 351a1cf..aa68de0 100644 --- a/lib/emily/native.ex +++ b/lib/emily/native.ex @@ -221,6 +221,9 @@ defmodule Emily.Native do @spec slice(tensor(), [integer()], [integer()], [integer()]) :: tensor() def slice(_a, _start, _stop, _strides), do: nif() + @spec slice_update(tensor(), tensor(), [integer()]) :: tensor() + def slice_update(_src, _update, _start), do: nif() + @spec take(tensor(), tensor(), integer()) :: tensor() def take(_a, _indices, _axis), do: nif() diff --git a/mix.exs b/mix.exs index e19f012..5ca507e 100644 --- a/mix.exs +++ b/mix.exs @@ -52,7 +52,12 @@ defmodule Emily.MixProject do {:elixir_make, "~> 0.9"}, {:fine, "~> 0.1"}, {:nx, "~> 0.10"}, - {:bumblebee, "~> 0.6", only: :test}, + # Bumblebee >= 0.6.3 (the latest Hex release) lacks Qwen3 support. + # Pinned to a `main` commit that contains `Bumblebee.Text.Qwen3` so + # M4 can exercise Qwen3-0.6B end-to-end. Bump deliberately when a + # newer release lands on Hex. + {:bumblebee, + github: "elixir-nx/bumblebee", ref: "273805e95507dc7866b958d90e0012a3abad1761", only: :test}, {:tokenizers, "~> 0.5", only: :test}, {:stream_data, "~> 1.1", only: [:dev, :test]}, {:credo, "~> 1.7", only: [:dev, :test], runtime: false}, diff --git a/mix.lock b/mix.lock index 6dd3ddf..c019c24 100644 --- a/mix.lock +++ b/mix.lock @@ -1,6 +1,6 @@ %{ "axon": {:hex, :axon, "0.7.0", "2e2c6d93b4afcfa812566b8922204fa022b60081e86ebd411df4db7ea30f5457", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: true]}, {:kino_vega_lite, "~> 0.1.7", [hex: :kino_vega_lite, repo: "hexpm", optional: true]}, {:nx, "~> 0.9", [hex: :nx, repo: "hexpm", optional: false]}, {:polaris, "~> 0.1", [hex: :polaris, repo: "hexpm", optional: false]}, {:table_rex, "~> 3.1.1", [hex: :table_rex, repo: "hexpm", optional: true]}], "hexpm", "ee9857a143c9486597ceff434e6ca833dc1241be6158b01025b8217757ed1036"}, - "bumblebee": {:hex, :bumblebee, "0.6.3", "c0028643c92de93258a9804da1d4d48797eaf7911b702464b3b3dd2cc7f938f1", [:mix], [{:axon, "~> 0.7.0", [hex: :axon, repo: "hexpm", optional: false]}, {:jason, "~> 1.4.0", [hex: :jason, repo: "hexpm", optional: false]}, {:nx, "~> 0.9.0 or ~> 0.10.0", [hex: :nx, repo: "hexpm", optional: false]}, {:nx_image, "~> 0.1.0", [hex: :nx_image, repo: "hexpm", optional: false]}, {:nx_signal, "~> 0.2.0", [hex: :nx_signal, repo: "hexpm", optional: false]}, {:progress_bar, "~> 3.0", [hex: :progress_bar, repo: "hexpm", optional: false]}, {:safetensors, "~> 0.1.3", [hex: :safetensors, repo: "hexpm", optional: false]}, {:tokenizers, "~> 0.4", [hex: :tokenizers, repo: "hexpm", optional: false]}, {:unpickler, "~> 0.1.0", [hex: :unpickler, repo: "hexpm", optional: false]}, {:unzip, "~> 0.12.0", [hex: :unzip, repo: "hexpm", optional: false]}], "hexpm", "c619197787561f8e5fb2ffba269c341654accaec9d591999b7fddd55761dd079"}, + "bumblebee": {:git, "https://github.com/elixir-nx/bumblebee.git", "273805e95507dc7866b958d90e0012a3abad1761", [ref: "273805e95507dc7866b958d90e0012a3abad1761"]}, "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"}, "castore": {:hex, :castore, "1.0.18", "5e43ef0ec7d31195dfa5a65a86e6131db999d074179d2ba5a8de11fe14570f55", [:mix], [], "hexpm", "f393e4fe6317829b158fb74d86eb681f737d2fe326aa61ccf6293c4104957e34"}, "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"}, @@ -27,5 +27,5 @@ "telemetry": {:hex, :telemetry, "1.4.1", "ab6de178e2b29b58e8256b92b382ea3f590a47152ca3651ea857a6cae05ac423", [:rebar3], [], "hexpm", "2172e05a27531d3d31dd9782841065c50dd5c3c7699d95266b2edd54c2dafa1c"}, "tokenizers": {:hex, :tokenizers, "0.5.1", "b0975d92b4ee5b18e8f47b5d65b9d5f1e583d9130189b1a2620401af4e7d4b35", [:mix], [{:castore, "~> 0.1 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: false]}, {:rustler, ">= 0.0.0", [hex: :rustler, repo: "hexpm", optional: true]}, {:rustler_precompiled, "~> 0.6", [hex: :rustler_precompiled, repo: "hexpm", optional: false]}], "hexpm", "5f08d97cc7f2ed3d71d370d68120da6d3de010948ccf676c9c0eb591ba4bacc9"}, "unpickler": {:hex, :unpickler, "0.1.0", "c2262c0819e6985b761e7107546cef96a485f401816be5304a65fdd200d5bd6a", [:mix], [], "hexpm", "e2b3f61e62406187ac52afead8a63bfb4e49394028993f3c4c42712743cab79e"}, - "unzip": {:hex, :unzip, "0.12.0", "beed92238724732418b41eba77dcb7f51e235b707406c05b1732a3052d1c0f36", [:mix], [], "hexpm", "95655b72db368e5a84951f0bed586ac053b55ee3815fd96062fce10ce4fc998d"}, + "unzip": {:hex, :unzip, "0.13.0", "bf5ec6ac6063c69e6ec54c8b4a3b8dcd7a2719d28d10d7025776ab107957cde9", [:mix], [], "hexpm", "4bcb9892ecbf2042606b43ab685a1bffe03c14003e6246f5453db2c829237fd9"}, } diff --git a/test/emily/conformance/qwen3_full_test.exs b/test/emily/conformance/qwen3_full_test.exs new file mode 100644 index 0000000..ebbbc0c --- /dev/null +++ b/test/emily/conformance/qwen3_full_test.exs @@ -0,0 +1,55 @@ +defmodule Emily.Conformance.Qwen3FullTest do + @moduledoc """ + Full `Qwen/Qwen3-0.6B` end-to-end conformance test. + + This test is excluded even from `mix test --only conformance`: the + model is ~1.5 GB on first fetch, so blowing it out into CI on every + push is the wrong default. Run explicitly: + + mix test --only qwen3_full + + The reference text pinned below is the greedy decode produced by + `Emily.Backend` on an Apple-Silicon host. A failure means the + backend has drifted, Bumblebee's Qwen3 port has changed, or the HF + checkpoint has been republished — all of which are real signals. + """ + + use ExUnit.Case, async: false + + @moduletag :qwen3_full + @moduletag capture_log: true + @moduletag timeout: 600_000 + + @prompt "The quick brown fox jumps over the lazy dog." + @reference_text " The quick brown fox is a character in the story. The quick brown fox is a character in the story. The quick brown fox is a character in the story" + + setup_all do + prev = Nx.default_backend() + Nx.global_default_backend(Emily.Backend) + on_exit(fn -> Nx.global_default_backend(prev) end) + :ok + end + + test "Qwen/Qwen3-0.6B greedy decodes the pinned continuation" do + {:ok, model_info} = Bumblebee.load_model({:hf, "Qwen/Qwen3-0.6B"}) + {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "Qwen/Qwen3-0.6B"}) + {:ok, generation_config} = Bumblebee.load_generation_config({:hf, "Qwen/Qwen3-0.6B"}) + + config = + Bumblebee.configure(generation_config, + max_new_tokens: 32, + strategy: %{type: :greedy_search} + ) + + serving = + Bumblebee.Text.generation(model_info, tokenizer, config, + defn_options: [compiler: Nx.Defn.Evaluator] + ) + + %{results: [%{text: text, token_summary: summary}]} = + Nx.Serving.run(serving, @prompt) + + assert summary.output == 32 + assert text == @reference_text + end +end diff --git a/test/emily/conformance/qwen3_test.exs b/test/emily/conformance/qwen3_test.exs new file mode 100644 index 0000000..f23feb3 --- /dev/null +++ b/test/emily/conformance/qwen3_test.exs @@ -0,0 +1,213 @@ +defmodule Emily.Conformance.Qwen3Test do + @moduledoc """ + End-to-end conformance tests for Qwen3 on `Emily.Backend`. + + Mirrors `Bumblebee.Text.Qwen3Test` — same three architectures, same + tiny-random checkpoints, same input token IDs, same expected output + slices. The reference values in Bumblebee's own test suite were + produced by the HuggingFace Transformers reference implementation, so + a failure here unambiguously indicates an Emily bug on Qwen3's critical + path: QK-norm, rotary embeddings, GQA, SwiGLU FFN, RMSNorm, tied + embeddings. + + An additional test drives the `Bumblebee.Text.Generation` pipeline + greedy-decode with a fixed prompt; the token sequence is checked in + after validation against `Nx.BinaryBackend` on the same inputs. + `BinaryBackend` is the conformance-layer oracle here (we don't have a + Linux+CUDA EXLA machine in CI), so the assertion is really + "Emily.Backend and BinaryBackend produce bit-identical token ids for + the same model under greedy decoding". + + Tagged `:conformance` and excluded from the default suite; the + tiny-random checkpoints are fetched from HuggingFace on first run and + cached under `~/.cache/bumblebee`. Invoke explicitly: + + mix test --only conformance + """ + + use ExUnit.Case, async: false + + alias Bumblebee.Text.Generation, as: BBGeneration + + @moduletag :conformance + @moduletag capture_log: true + @moduletag timeout: 300_000 + + setup_all do + prev = Nx.default_backend() + Nx.global_default_backend(Emily.Backend) + on_exit(fn -> Nx.global_default_backend(prev) end) + :ok + end + + test ":base" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-Qwen3Model"}) + + assert %Bumblebee.Text.Qwen3{architecture: :base} = spec + assert spec.use_qk_norm == true + + inputs = %{ + "input_ids" => Nx.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]), + "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.hidden_state) == {1, 10, 32} + + assert_all_close( + outputs.hidden_state[[.., 1..3, 1..3]], + Nx.tensor([ + [[0.0461, -1.5811, 1.5504], [0.1340, -1.3477, 0.8047], [-0.5821, -0.4164, 0.9769]] + ]) + ) + end + + test ":for_causal_language_modeling" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-Qwen3ForCausalLM"}) + + assert %Bumblebee.Text.Qwen3{architecture: :for_causal_language_modeling} = spec + assert spec.use_qk_norm == true + + inputs = %{ + "input_ids" => Nx.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]), + "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.logits) == {1, 10, 1024} + + assert_all_close( + outputs.logits[[.., 1..3, 1..3]], + Nx.tensor([ + [[0.1457, 0.0313, -0.0651], [0.1718, -0.0265, -0.0186], [0.2281, -0.0124, -0.0147]] + ]) + ) + end + + test ":for_sequence_classification" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model( + {:hf, "bumblebee-testing/tiny-random-Qwen3ForSequenceClassification"} + ) + + assert %Bumblebee.Text.Qwen3{architecture: :for_sequence_classification} = spec + + inputs = %{ + "input_ids" => Nx.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]), + "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.logits) == {1, 2} + + assert_all_close(outputs.logits, Nx.tensor([[-0.1487, -0.0071]])) + end + + describe "greedy generation" do + # Exercises the full `Bumblebee.Text.Generation` path on Emily.Backend: + # Axon forward pass + logit processing + argmax + `put_slice` KV + # cache updates + `while` loop in defn. The tiny-random causal LM + # emits gibberish tokens, but determinism is the whole point here + # — the oracle is `Nx.BinaryBackend` running the identical generate + # function against the identical token ids, and we assert bit-exact + # token equality. + # + # BinaryBackend is our reference because we do not have a Linux+CUDA + # EXLA machine in CI (see PLAN.md M3/M4 testing philosophy). + # + # We bypass Bumblebee's serving so that the tokenizer isn't involved: + # the tiny-random checkpoint carries a 1024-token embedding while the + # Qwen3 tokenizer produces ids up to 151_936, so any real prompt + # would out-of-bounds the embedding gather. The synthetic ids below + # are all < 1024 and stand in for a pre-tokenized prompt. + + @golden_max_new_tokens 16 + @golden_input_ids [10, 20, 30, 40, 50, 60, 70, 80] + + # Reference produced by `Nx.BinaryBackend` on the same inputs; also + # cross-checked against BinaryBackend at test time below. If this + # list changes you have either drifted one backend, drifted + # Bumblebee's Qwen3 port, or upgraded to a new tiny-random + # checkpoint — none of which should happen silently. + @golden_tokens [6, 277, 436, 806, 833, 436, 785, 135, 550, 309, 511, 89, 865, 72, 1021, 865] + + test "tiny-random Qwen3ForCausalLM: greedy decode matches checked-in golden tokens" do + emily_tokens = run_greedy(Emily.Backend) + + oracle_tokens = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + run_greedy(Nx.BinaryBackend) + end) + + assert length(emily_tokens) == @golden_max_new_tokens + assert emily_tokens == oracle_tokens, "Emily.Backend diverged from BinaryBackend" + assert emily_tokens == @golden_tokens, "both backends drifted from checked-in reference" + end + + defp run_greedy(backend) do + {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model( + {:hf, "bumblebee-testing/tiny-random-Qwen3ForCausalLM"}, + backend: backend + ) + + config = + Bumblebee.configure( + %Bumblebee.Text.GenerationConfig{ + pad_token_id: 0, + eos_token_id: 0 + }, + max_new_tokens: @golden_max_new_tokens, + strategy: %{type: :greedy_search} + ) + + generate_fun = BBGeneration.build_generate(model, spec, config) + + input_ids = Nx.tensor([@golden_input_ids], type: :s64, backend: backend) + attention_mask = Nx.broadcast(Nx.tensor(1, backend: backend), Nx.shape(input_ids)) + seed = Nx.tensor([0], type: :s64, backend: backend) + + inputs = %{ + "input_ids" => input_ids, + "attention_mask" => attention_mask, + "seed" => seed + } + + %{token_ids: token_ids, length: length} = + Nx.Defn.jit_apply(generate_fun, [params, inputs], compiler: Nx.Defn.Evaluator) + + tokens = token_ids |> Nx.backend_transfer(Nx.BinaryBackend) |> Nx.to_flat_list() + [len_val] = length |> Nx.backend_transfer(Nx.BinaryBackend) |> Nx.to_flat_list() + Enum.take(tokens, len_val) + end + end + + # ----------------- helpers ----------------- + + defp assert_all_close(left, right, opts \\ []) do + atol = opts[:atol] || 1.0e-4 + rtol = opts[:rtol] || 1.0e-4 + + equal_tensor = + left + |> Nx.all_close(right, atol: atol, rtol: rtol) + |> Nx.backend_transfer(Nx.BinaryBackend) + + if Nx.to_number(equal_tensor) != 1 do + ExUnit.Assertions.flunk(""" + expected + + #{inspect(Nx.backend_copy(left, Nx.BinaryBackend))} + + to be within tolerance of + + #{inspect(Nx.backend_copy(right, Nx.BinaryBackend))} + """) + end + end +end diff --git a/test/test_helper.exs b/test/test_helper.exs index e3d2b7d..0591afa 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -8,6 +8,9 @@ # # Conformance tests pull tiny-random HuggingFace models at runtime and # take ~tens of seconds per model on a cold cache — opt-in via -# `mix test --only conformance`. (Soak tests deliberately stay in the +# `mix test --only conformance`. `:qwen3_full` is an even heavier +# conformance variant that downloads ~1.5 GB of weights, so it is +# excluded even from `--only conformance`; run explicitly via +# `mix test --only qwen3_full`. (Soak tests deliberately stay in the # default suite; see `test/soak/memory_test.exs` for the rationale.) -ExUnit.start(max_cases: 1, exclude: [:conformance]) +ExUnit.start(max_cases: 1, exclude: [:conformance, :qwen3_full])