Memory explosion when accessing a large map while using Nx+EXLA

Hey there,

I’m currently benchmarking some sentence transformers (see Nx vs. Python performance for sentence-transformer encoding) and stumbled upon something weird:

Basically I’m encoding a sentence into a vector using Bumblebee (Nx+EXLA) and then calculate a cosine similarity with a list of pre-calculated vectors (~150k). I’m benchmarking the whole process to see what performance I can achieve:

[%{embedding: query}] = Nx.Serving.batched_run(SentenceTransformer, ["this is a test input"])

sim =
  for chunk <- vectors do
    Bumblebee.Utils.Nx.cosine_similarity(query, chunk)
  end
  |> Nx.concatenate()

{similarity, labels} = Nx.top_k(sim, k: 10)

indexes = Nx.to_flat_list(labels)
scores = Nx.to_flat_list(similarity)

for {idx, score} <- Enum.zip(indexes, scores) do
  # accessing the sentence_map here leads to the issue
  %{sentence: sentence_map[idx], score: score}
end

This works fine until I try to map the resulting top k indexes back to the sentences using a 150k key map of index → sentence entries. When doing this, EXLA suddenly starts to consume a LOT of memory:

image1704×1118 84.9 KB

I think that the memory usage is in EXLA as is not visible in the observer:

image1924×1424 157 KB

Running on CUDA seems to confirm this, it even runs out of memory completely:

This is my script to reproduce:

# System.put_env("XLA_TARGET", "cuda118")

Mix.install([
  {:bumblebee, github: "elixir-nx/bumblebee", ref: "23de64b1b88ed3aad266025c207f255312b80ba6"},
  {:nx, github: "elixir-nx/nx", sparse: "nx", override: true},
  {:exla, github: "elixir-nx/nx", sparse: "exla", override: true},
  {:axon, "~> 0.5.1"},
  {:kino, "~> 0.9"}
])

Nx.global_default_backend(EXLA.Backend)
# Nx.Defn.global_default_options(compiler: EXLA, client: :cuda)
Nx.Defn.global_default_options(compiler: EXLA, client: :host)

model_name = "sentence-transformers/all-MiniLM-L6-v2"
{:ok, model_info} = Bumblebee.load_model({:hf, model_name})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, model_name})

serving =
  Bumblebee.Text.TextEmbedding.text_embedding(model_info, tokenizer,
    compile: [batch_size: 64, sequence_length: 128],
    defn_options: [compiler: EXLA],
    output_attribute: :hidden_state,
    output_pool: :mean_pooling
  )

Kino.start_child({Nx.Serving, serving: serving, name: SentenceTransformer, batch_size: 64, batch_timeout: 50})

defmodule ConcurrentBench do
  def run(fun, concurrency \\ System.schedulers_online(), timeout \\ 10_000) do
    # use an erlang counter to count the number of function invocations
    counter = :counters.new(1, [:write_concurrency])

    # returns time in microseconds
    {taken, _} =
      :timer.tc(fn ->
        tasks =
          for _i <- 1..concurrency do
            Task.async(fn ->
              Stream.repeatedly(fn ->
                fun.()
                # only count after the function ran successfully
                :counters.add(counter, 1, 1)
              end)
              |> Stream.run()
            end)
          end

        results = Task.yield_many(tasks, timeout)

        # kill all processes
        Enum.map(results, fn {task, res} ->
          res || Task.shutdown(task, :brutal_kill)
        end)
      end)

    runs = :counters.get(counter, 1)
    ips = runs / (taken / 1_000_000)

    %{runs: runs, ips: ips}
  end
end

n = 150000

sentence_map = Map.new(1..n, fn i -> {i, i} end)
[%{embedding: vector}] = Nx.Serving.batched_run(SentenceTransformer, ["This is a test sentence"])

IO.puts("encoded sample input")

 
vectors =
  Stream.duplicate(vector, n)
  |> Stream.chunk_every(10000)
  |> Stream.map(fn chunk -> Nx.stack(chunk) end)
  |> Enum.to_list()

IO.puts("created dummy comparison vectors")
IO.puts("Running concurrent bench now")

ConcurrentBench.run(
  fn ->
    [%{embedding: query}] = Nx.Serving.batched_run(SentenceTransformer, ["this is a test input"])

    sim =
      for chunk <- vectors do
        Bumblebee.Utils.Nx.cosine_similarity(query, chunk)
      end
      |> Nx.concatenate()

    {similarity, labels} = Nx.top_k(sim, k: 10)

    indexes = Nx.to_flat_list(labels)
    scores = Nx.to_flat_list(similarity)

    # uncomment this
    # for {idx, score} <- Enum.zip(indexes, scores) do
    #   %{sentence: sentence_map[idx], score: score}
    # end

    nil
  end,
  16, 60_000
) |> IO.inspect()

See the “uncomment this” section.

I guess this is a bug? At least accessing the map should not lead to EXLA allocating memory, right?

Thanks for any suggestions!