Unable to configure XLA_TARGET=cuda118 to use GPU

[elixir-nx]

Unable to configure XLA_TARGET=cuda118 to use GPU

[zolrath]

When I encountered this I had run the Livebook before I had set the XLA_TARGET, so XLA was already compiled with only CPU support. Even after adding the XLA_TARGET and ensuring I had CUDA and cuDNN installed properly I would get 14:30:19.336 [info] TfrtCpuClient created..

I solved this by deleting the xla and mix folders in ~/.cache, as I couldn't see a clean way to run mix deps.clean xla --build in Livebook.

[josevalim]

@jonatanklosko there is an env var called MIX_INSTALL_FORCE, we can set it to "true" before we evaluate.

[seanmor5]

What if you set cuda explicitly in your config?

  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA, client: :cuda]
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda"
  ]

[josevalim]

You can also run: EXLA.Client.get_supported_platforms() and see if CUDA is listed. If it isn't, for some reason XLA is not identifying your GPU device.

[josevalim]

And additionally EXLA.Client.default_name() to get the name of the default one.

[josevalim]

@jonatanklosko we should consider the XLA_TARGET in the cache key, right?

[xrd]

@zolrath Yes, that is what must have happened to me as well. I did get CUDA at first, and then could not switch back to it no matter what I tried. Thanks everyone for your help, and thanks @seanmor5 and @josevalim for an amazing project!

[jonatanklosko]

@josevalim ohh I've just realised that setting the env var and restarting runtime is not enough (unless the env var is set in Mix.install options). That's because we changed XLA to look at the initial XLA_TARGET, rather than every time it compiles. So in a mix project that would require mix deps.clean xla --build and for mix install force: true.

[josevalim]

Do you remember why we did such change?

[jonatanklosko]

Actually what I was remembering is XLA_BUILD (elixir-nx/nx#845 (comment)). The issue with XLA_TARGET is that we compile dependencies only once, specifically exla.

Inside the exla project if you call mix compile with a different XLA_BUILD it will pick up the new version, but if you call mix compile in a project that depends on exla, it wouldn't pick up the change, since exla doesn't recompile.

[josevalim]

But it works if we pass it to Mix.install/2 or force build, right?

maybe we should have an option to “force build” from Livebook UI? Maybe a drop down on the “Reconnect and setup”?

[xrd]

I'm able to use cuda now, shall I close it?

[josevalim]

We will close it @xrd as soon as we figure out how to improve the user experience :)

[josevalim]

Actually, one way to improve this would be to add:

system_env: %{"XLA_TARGET" => System.get_env("XLA_TARGET")}

this will enforce system differences to be picked up.

[jonatanklosko]

@josevalim yeah this works, my concern is that once someone sets the env var, they won't remember or bother to always add it to the notebooks. Then someone else using those notebooks may run into similar issue and it is not obvious that's how they could fix it.

[josevalim]

So I guess the best option is some UI to force installation the next install?

[jonatanklosko]

That could be useful in other cases too. Would we modify the code on the fly to include force: true?

[josevalim]

Moved to Livebook backlog: livebook-dev/livebook#1579