Qwen3-0.6B quantized generation on Emily

Mix.install(
  [
    {:emily, "~> 1.3"},
    {:bumblebee, "~> 0.7"},
    {:tokenizers, "~> 0.5"},
    {:nx, "~> 0.12"},
    {:axon, "~> 0.8"},
    {:kino, "~> 0.14"}
  ],
  config: [
    nx: [default_backend: Emily.Backend]
  ]
)

Overview

This notebook loads Qwen/Qwen3-0.6B through Bumblebee, quantizes every dense layer's kernel to int4 via MLX affine group-wise quantization, and greedy-decodes a completion. It also demonstrates Emily.Stream for concurrent serving on a shared model.

The checkpoint is ~1.5 GB on first fetch. Budget several minutes for the cold run.

Dense baseline

{:ok, %{model: model, params: params, spec: spec}} =
  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"})

model_info = %{model: model, params: params, spec: spec}

config =
  Bumblebee.configure(generation_config,
    max_new_tokens: 32,
    strategy: %{type: :greedy_search}
  )

serving =
  Bumblebee.Text.generation(model_info, tokenizer, config,
    defn_options: [compiler: Emily.Compiler, native: true, native_fallback: :raise]
  )

%{results: [%{text: dense_text}]} =
  Nx.Serving.run(serving, "The quick brown fox jumps over the lazy dog.")

dense_text

native: true, native_fallback: :raise on :defn_options lowers the per-token decode through Emily's native Expr compiler and asserts it stays fully native — the quantized serving below passes the same options, so the int4 quantized matmul lowers native too.

Quantization transform

Emily.Quantization.Transform.quantize/3 does the rewrite in two steps:

1. Graph rewrite — Axon.rewrite_nodes/2 replaces each :dense node with a :quantized_dense layer whose forward pass dispatches through Emily.Quantization.Layers.quantized_dense/4.
2. Parameter quantization — each dense kernel in the Axon.ModelState is swapped for a %Emily.QuantizedWeight{}, and the quantized weights are evaluated so they no longer hold on to the bf16 kernels they came from.

The transform needs Axon, which is an optional dependency of Emily; Bumblebee already pulls it in.

Notes. The transform rewrites top-level :dense nodes; a model with dense layers nested inside other Axon ops needs a recursive rewriter. The transpose: true default stores weights as [out, in] (MLX / AWQ convention, groups along the reduction axis); set false if you're feeding a checkpoint that's already laid out the other way. Bits must be one of [1, 2, 3, 4, 6, 8] for MLX affine quantization. Pass except: ["lm_head"] (or any other layer-name substrings) to keep some layers in full precision.

Quantized inference

{qmodel, qparams} =
  Emily.Quantization.Transform.quantize(model, params,
    bits: 4,
    group_size: 128,
    transpose: true
  )

qmodel_info = %{model: qmodel, params: qparams, spec: spec}

qserving =
  Bumblebee.Text.generation(qmodel_info, tokenizer, config,
    defn_options: [compiler: Emily.Compiler, native: true, native_fallback: :raise]
  )

%{results: [%{text: quant_text}]} =
  Nx.Serving.run(qserving, "The quick brown fox jumps over the lazy dog.")

quant_text

The quantized output will drift from the dense baseline — int4 noise across every linear is expected. The test at test/emily/conformance/qwen3_quant_full_test.exs pins a deterministic reference string for regression-testing the quantization stack.

Concurrent serving via Emily.Stream

For concurrent inference on a shared model, each serving worker should own its own MLX command queue. Emily.Stream.with_stream/2 does that per-process:

stream = Emily.Stream.new(:gpu)

task1 =
  Task.async(fn ->
    Emily.Stream.with_stream(stream, fn ->
      Nx.Serving.run(qserving, "Question 1?")
    end)
  end)

task2 =
  Task.async(fn ->
    Emily.Stream.with_stream(Emily.Stream.new(:gpu), fn ->
      Nx.Serving.run(qserving, "Question 2?")
    end)
  end)

{Task.await(task1, :infinity), Task.await(task2, :infinity)}

Each Emily.Stream maps to its own Metal command queue. Weights are shared across streams — no duplication — so the memory cost of adding a stream is the Metal command buffer, not the model.

Create streams once at worker init, not per-request.