Hi, I’m trying to get my head around Nx and it’s configuration to run on CUDA.

# Configures Nx default backend and default defn options
config :nx, :default_backend, EXLA.Backend
config :nx, :default_defn_options, compiler: EXLA

I also set the XLA_TARGET environment variable to cuda118 both when building and running my app.

From the following output, everything looks fine (at least to me ).

iex> System.get_env "XLA_TARGET"
"cuda118"
iex> Application.get_all_env :nx
[default_backend: EXLA.Backend, default_defn_options: [compiler: EXLA]]
iex> EXLA.Client.get_supported_platforms
%{host: 4, cuda: 1, interpreter: 1}
iex> EXLA.Client.default_name
:cuda

From the EXLA.Client outputs, it seems the default client should be :cuda, it’s also listed in the supported platforms.

When running a serving with this config, I would assume that it picks my default client. But it’s not the case:

{:ok, model_info} =
  Bumblebee.load_model({:hf, "patrickjohncyh/fashion-clip"},
    module: Bumblebee.Text.ClipText,
    architecture: :for_embedding
  )

{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "patrickjohncyh/fashion-clip"})

serving =  Bumblebee.Text.text_embedding(model_info, tokenizer, output_attribute: :embedding)

%Nx.Serving{
  module: Nx.Serving.Default,
  arg: #Function<1.26226188/2 in Bumblebee.Text.TextEmbedding.text_embedding/3>,
  client_preprocessing: #Function<2.26226188/1 in Bumblebee.Text.TextEmbedding.text_embedding/3>,
  client_postprocessing: #Function<3.26226188/2 in Bumblebee.Text.TextEmbedding.text_embedding/3>,
  streaming: nil,
  batch_size: nil,
  distributed_postprocessing: &Function.identity/1,
  process_options: [batch_keys: [:default]],
  defn_options: [compiler: EXLA]
}

The following code does run the serving on the GPU but on the CPU.

Nx.run serving, "a pair of white Nike Air Max shoes in white"

It seems I have to explicitly set client: :cuda for my serving to run on the GPU. Either by configuring the :defn_options for my serving or setting it globally.

%{serving | defn_options: [compiler: EXLA, client: :cuda]}
# or
Nx.Defn.global_default_options [compiler: EXLA, client: :cuda]

From the EXLA.Backend docs I was assuming that servings would use the default XLA client.