elixir-nx / bumblebee

Pre-trained Neural Network models in Axon (+ 🤗 Models integration)
Apache License 2.0
1.26k stars 90 forks source link

CUDA 12.2 support #351

Closed lambdaofgod closed 4 months ago

lambdaofgod commented 4 months ago

I can't make bumblebee work with CUDA 12.2.

It would be good to include information on whether it is supported in README because I can create tensors, but running models fails which is confusing.

I tried this setup in livebook:

Mix.install(
  [
    {:nx, "~> 0.7.0"},
    {:kino, "~> 0.11"},
    {:explorer, "~> 0.8.0"},
    {:exla, "~> 0.7"},
    {:kino_bumblebee, "~> 0.5.0"}
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA]
    ],
    exla: [
      default_client: :cuda,
      clients: [
        host: [platform: :host],
        cuda: [platform: :cuda]
      ]
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda120"
  ]
)

This cell works

t = Nx.tensor([[1, 2], [3, 4]])

and the tensor is on GPU as expected:

#Nx.Tensor<
  s64[2][2]
  EXLA.Backend<cuda:0, 0.250714851.616693825.138807>
  [
    [1, 2],
    [3, 4]
  ]
>

But when I try to run BART for zero-shot classification

{:ok, model} = Bumblebee.load_model({:hf, "facebook/bart-large-mnli"})
{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "facebook/bart-large-mnli"})

{:ok, generation_config} = Bumblebee.load_generation_config({:hf, "facebook/bart-large-mnli"})
serving = Bumblebee.Text.ZeroShotClassification.zero_shot_classification(model, tokenizer, ["machine learning", "data visualization"])

text_input = Kino.Input.text("Text", default: "This cat is so cute.")

Nx.Serving.run(serving, "object detection")

The livebook terminates, and the cell output is


14:58:02.697 [warning] libdevice is required by this HLO module but was not found at ./libdevice.10.bc

14:58:02.708 [warning] libdevice is required by this HLO module but was not found at ./libdevice.10.bc

system info

Elixir and packages

Erlang/OTP 25 [erts-13.0] [source] [64-bit] [smp:32:32] [ds:32:32:10] [async-threads:1] [jit:ns]

Elixir 1.15.2 (compiled with Erlang/OTP 24)

Livebook 0.12.1

nvidia-smi

Mon Feb 26 14:55:08 2024
+---------------------------------------------------------------------------------------+ | NVIDIA-SMI 535.113.01 Driver Version: 535.113.01 CUDA Version: 12.2 | |-----------------------------------------+----------------------+----------------------+ | 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 GeForce RTX 3090 On | 00000000:08:00.0 On | N/A | | 0% 44C P8 43W / 420W | 960MiB / 24576MiB | 11% Default | | | | N/A | +-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=======================================================================================| | 0 N/A N/A 2654 G /usr/lib/xorg/Xorg 531MiB | | 0 N/A N/A 128882 G nvidia-settings 48MiB | | 0 N/A N/A 147766 G ...re/Steam/ubuntu12_64/steamwebhelper 5MiB | | 0 N/A N/A 194239 G ...,262144 --variations-seed-version=1 305MiB | | 0 N/A N/A 194464 G ...ForSitePerProcess --enable-crashpad 51MiB | +---------------------------------------------------------------------------------------+

seanmor5 commented 4 months ago

Can you search for libdevice on /usr/lib/cuda or elsewhere? If you do a find on your system for libdevice and if you find it, then add that directory to:

XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/lib/cuda

as a system environment variable before running the livebook and see if it fixes your issue

lambdaofgod commented 4 months ago

Yes, it works. Isn't there another way to set this though? I tried passing it to config (adding a line XLA_FLAGS: "--xla_gpu_cuda_data_dir=/usr/lib/cuda") and it doesn't see it.

jonatanklosko commented 4 months ago

@lambdaofgod you can use System.put_env/2 at the very beginning, or Mix.install([...], system_env: %{"XLA_FLAGS" => "--xla_gpu_cuda_data_dir=/usr/lib/cuda"}), or better yet set it as global env var in Livebook settings, or set it in your .bashrc config.