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?