Bumblebee: Slow load_model in GenServer, slow Nx.Serving.run in exs file

mgwidmann

I’ve run into an issue that I don’t understand. Running on an M1 mac with 32 GB of ram trying to execute a very small LLM and I’m seeing huge times to do certain function calls.

In my GenServer, I have the following:

  # @model "meta-llama/Llama-3.1-8B-Instruct"
  # @model "mistralai/Mistral-7B-Instruct-v0.3"
  # @model "HuggingFaceTB/SmolLM2-1.7B-Instruct"
  @model "meta-llama/Llama-3.2-1B-Instruct"

  def handle_cast(:load_model, _nothing) do
    Logger.info("Async model load beginning...", [service: "RagLLM"])

    repo = {:hf, @model, auth_token: System.fetch_env!("HF_TOKEN")}

    {:ok, model_info} = Bumblebee.load_model(repo, type: :bf16)
    Logger.info("Model loaded", [service: "RagLLM"])
    {:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
    Logger.info("Tokenizer loaded", [service: "RagLLM"])
    {:ok, generation_config} = Bumblebee.load_generation_config(repo)
    Logger.info("Generation Config loaded", [service: "RagLLM"])

    generation_config = Bumblebee.configure(generation_config, max_new_tokens: 100)

    Logger.info("Model, tokenizer defined, bumblebee configured", [service: "RagLLM"])

    serving =
      Bumblebee.Text.generation(model_info, tokenizer, generation_config,
        compile: [batch_size: 1, sequence_length: 6000]
      )

    Logger.info("Serving created", [service: "RagLLM"])

    {:ok, _server} = Nx.Serving.start_link(serving: serving, name: LLMServing, batch_timeout: 100)

    {:noreply, serving}
  end

The load_model call can take 5 minutes to finish the load_model call. I’m using EMLX for my mac by setting the following in my application start function:

    if :os.type() == {:unix, :darwin} do
      IO.puts("Loading EMLX for MacOS")
      Nx.default_backend({EMLX.Backend, device: :gpu})
      Nx.Defn.default_options(compiler: EMLX)
    else
      Nx.default_backend({EXLA.Backend, device: :gpu})
      Nx.Defn.default_options(compiler: EXLA)
    end

But in an invoke_model.exs file, it gets through this line instantly but hangs on the NX.Serving.run/2 call.

if :os.type() == {:unix, :darwin} do
  Application.ensure_all_started(:emlx)
  Nx.default_backend({EMLX.Backend, device: :gpu})
  Nx.Defn.default_options(compiler: EMLX)
else
  Application.ensure_all_started(:exla)
end
Application.ensure_all_started(:bumblebee)

:observer.start()

repo = {:hf, "meta-llama/Llama-3.2-1B-Instruct", auth_token: System.fetch_env!("HF_TOKEN")}

{:ok, model_info} = Bumblebee.load_model(repo, type: :bf16)
IO.puts("Model loaded")
{:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
IO.puts("Tokenizer loaded")
{:ok, generation_config} = Bumblebee.load_generation_config(repo)
IO.puts("Generation Config loaded")

generation_config = Bumblebee.configure(generation_config, max_new_tokens: 100)

IO.puts("Model, tokenizer defined, bumblebee configured")

serving =
  Bumblebee.Text.generation(model_info, tokenizer, generation_config,
    compile: [batch_size: 1, sequence_length: 6000]
  )
IO.puts("Starting serving")

{:ok, _server} = Nx.Serving.start_link(serving: serving, name: LLMServing, batch_timeout: 100)

:timer.sleep(10_000)

IO.puts("Simple run call") # Last output that shows up

IO.inspect(Nx.Serving.run(serving, "What is the meaning of life?"))

IO.puts("Running model...")

IO.inspect(Nx.Serving.batched_run(LLMServing, "What is the meaning of life?"))

IO.puts("Done!")

I am not sure what I am doing wrong. Any ideas?

13 comments

/nx #bumblebee #llm

0 659 13

2025-08-13 20:31:11 UTC

Showing Posts 1 to 10

Show Best Posts
Show All (oldest first)
Show All (newest first)

jonatanklosko

Creator of Livebook

You want this:

-Nx.default_backend({EMLX.Backend, device: :gpu})
+Nx.global_default_backend({EMLX.Backend, device: :gpu})

default_backend applies only to the calling process, global_default_backend applies to all.

but hangs on the NX.Serving.run/2 call

What about the batch_run? For how long does it hang? Does it work in the same setup if you use EXLA instead?

Post #1

mgwidmann

Yes, I can’t get this function to complete either. I’ve waited over 30 minutes and nothing.

Post #2

mgwidmann

I’m trying even on an EC2 machine with 23 GB of GPU memory and can’t load the "meta-llama/Llama-3.2-1B-Instruct" model.

20:19:19.812 [error] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR

20:19:19.812 [error] Memory usage: 2097610752 bytes free, 23696375808 bytes total.

20:19:19.813 [error] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR

20:19:19.813 [error] Memory usage: 2097610752 bytes free, 23696375808 bytes total.
** (RuntimeError) DNN library initialization failed. Look at the errors above for more details.
    (exla 0.9.2) lib/exla/mlir/module.ex:147: EXLA.MLIR.Module.unwrap!/1
    (exla 0.9.2) lib/exla/mlir/module.ex:124: EXLA.MLIR.Module.compile/5
    (stdlib 7.0) timer.erl:599: :timer.tc/2
    (exla 0.9.2) lib/exla/defn.ex:432: anonymous fn/14 in EXLA.Defn.compile/8
    (exla 0.9.2) lib/exla/mlir/context_pool.ex:10: anonymous fn/3 in EXLA.MLIR.ContextPool.checkout/1
    (nimble_pool 1.1.0) lib/nimble_pool.ex:462: NimblePool.checkout!/4
    (exla 0.9.2) lib/exla/defn/locked_cache.ex:36: EXLA.Defn.LockedCache.run/2
    (stdlib 7.0) timer.erl:599: :timer.tc/2

Post #3

jonatanklosko

Creator of Livebook

Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR

Do you have cuDNN installed on the machine? Which version?

Note that you can also try EXLA with CPU, if you have enough RAM. It will take a while, but if it finishes, we will known that the issue is likely in EMLX.

Post #4

mgwidmann

It appears yes it is already installed. I used the deep learning image, so most things should already be there I expect.

sudo apt-get install zlib1g
Reading package lists... Done
Building dependency tree... Done
Reading state information... Done
zlib1g is already the newest version (1:1.2.11.dfsg-2ubuntu9.2).
zlib1g set to manually installed.
0 upgraded, 0 newly installed, 0 to remove and 5 not upgraded.

Post #5

jonatanklosko

Creator of Livebook

What is the cuDNN version? You can check using something like apt list --installed | grep libcudnn.

Post #6

mgwidmann

I was able to get it working but changing my PATH and LD_LIBRARY_PATH to be the latest 12.8 but it uses 20 G of GPU memory when the huggingface calculator says it should only require 1-4 G (18 G to train) and I’m only doing inference so its using way more memory than needed it seems to me.

Every 2.0s: nvidia-smi                                                                                                         ip-172-31-34-180: Wed Jul 30 20:43:31 2025

Wed Jul 30 20:43:31 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.172.08             Driver Version: 570.172.08     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA A10G                    On  |   00000000:00:1E.0 Off |                    0 |
|  0%   39C    P0             63W /  300W |   20603MiB /  23028MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A           54025      C   ...g/28.0/erts-16.0/bin/beam.smp      20594MiB |
+-----------------------------------------------------------------------------------------+

It does seem like EMXL is having issues for me. On this machine it takes 2 minutes to load the model and execute. I also tried on an arm64 machine but couldn’t get it to work, perhaps it will if I make similar fixes like I did here.

Post #7

mgwidmann

Made the same changes to the arm64 machine and no, getting warnings and a crash. Used only 13 G of GPU memory though, out of the 15 G it has.

Every 2.0s: nvidia-smi                                                                                                          ip-172-31-42-52: Wed Jul 30 20:53:08 2025

Wed Jul 30 20:53:08 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.133.20             Driver Version: 570.133.20     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA T4G                     On  |   00000000:00:1F.0 Off |                    0 |
| N/A   65C    P0             39W /   70W |   13543MiB /  15360MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A          173457      C   ...g/28.0/erts-16.0/bin/beam.smp      13540MiB |
+-----------------------------------------------------------------------------------------+

20:56:05.810 [info] InUse at ec968ea00000 of size 2097152 next 141
** (RuntimeError) Out of memory while trying to allocate 11793234816 bytes.
    (exla 0.9.2) lib/exla/executable.ex:130: EXLA.Executable.unwrap!/1
    (exla 0.9.2) lib/exla/executable.ex:31: EXLA.Executable.run/3
    (exla 0.9.2) lib/exla/defn.ex:312: EXLA.Defn.maybe_outfeed/7
    (stdlib 7.0) timer.erl:599: :timer.tc/2
    (exla 0.9.2) lib/exla/defn.ex:244: anonymous fn/7 in EXLA.Defn.__compile__/4
    (nx 0.9.2) lib/nx/defn.ex:332: anonymous fn/4 in Nx.Defn.compile/3
    (bumblebee 0.6.3) lib/bumblebee/text/text_generation.ex:73: anonymous fn/4 in Bumblebee.Text.TextGeneration.generation/4
    (nx 0.9.2) lib/nx/serving.ex:1833: anonymous fn/2 in Nx.Serving.Default.handle_batch/3

So looks like it needs the full 20 G it is using on the other machine. I don’t understand why it needs so much nor do I understand why it doesn’t work with EMLX

Post #8

mgwidmann

I came across this issue, I wonder if it is the root cause since its still open. It could be crashed and not reporting anywhere.

https://github.com/elixir-nx/emlx/issues/58

Post #9

mgwidmann

Forcing {:nx, "~> 0.10.0", override: true}, got the examples working in that issue, so perhaps there have been updates that EMLX should update to. Just not my issue for some reason…