Livebook cuda 12.2 XLA out of memory at 11006MiB

polypush135

How to restrict memory usage on EXLA?

This is my memory usage before loading EXLA.

nvidia-smi
Sun Sep  3 13:13:11 2023       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   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 3060        On  | 00000000:01:00.0  On |                  N/A |
| 45%   53C    P5              30W / 170W |    398MiB / 12288MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1914      G   /usr/lib/xorg/Xorg                          131MiB |
|    0   N/A  N/A      2128      G   /usr/bin/gnome-shell                         63MiB |
|    0   N/A  N/A      7334      G   ...irefox/3068/usr/lib/firefox/firefox      193MiB |
+---------------------------------------------------------------------------------------+

And after.

❯ nvidia-smi
Sun Sep  3 13:18:13 2023       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   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 3060        On  | 00000000:01:00.0  On |                  N/A |
|  0%   43C    P3              30W / 170W |  11446MiB / 12288MiB |      2%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1914      G   /usr/lib/xorg/Xorg                          131MiB |
|    0   N/A  N/A      2128      G   /usr/bin/gnome-shell                         65MiB |
|    0   N/A  N/A      7334      G   ...irefox/3068/usr/lib/firefox/firefox      229MiB |
|    0   N/A  N/A     15749      C   ...ang/25.2.1/erts-13.1.4/bin/beam.smp    11006MiB |
+---------------------------------------------------------------------------------------+

As you can see Firefox and gnome bumped a little bit and the beam is trying to use the rest.

The issue is I get it to run once, but if I re evaluate to say change the seed or steps I then run out of memory.

13:20:48.865 [info] Total bytes in pool: 11350867968 memory_limit_: 11350867968 available bytes: 0 curr_region_allocation_bytes_: 22701735936

13:20:48.865 [info] Stats: 
Limit:                     11350867968
InUse:                     11229623808
MaxInUse:                  11261113856
NumAllocs:                        6164
MaxAllocSize:               3514198016
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0


13:20:48.865 [warning] ****************************************************************************************************

13:20:48.865 [error] Execution of replica 0 failed: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 6553600 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:    6.25MiB
              constant allocation:         0B
        maybe_live_out allocation:    6.25MiB
     preallocated temp allocation:         0B
                 total allocation:   12.50MiB
              total fragmentation:         0B (0.00%)
Peak buffers:
	Buffer 1:
		Size: 6.25MiB
		Entry Parameter Subshape: f32[1638400]
		==========================

	Buffer 2:
		Size: 6.25MiB
		XLA Label: copy
		Shape: f32[1280,1280]
		==========================

	Buffer 3:
		Size: 8B
		XLA Label: tuple
		Shape: (f32[1280,1280])
		==========================

So is there a way to limit the MaxInUse for xla from the start? 10 should be enough right? Its trying to grab 11 now.

20 comments

/livebook #troubleshooting #cuda

1 2297 20

2023-10-23 11:03:14 UTC

Showing Posts 1 to 10

Show Best Posts
Show All (oldest first)
Show All (newest first)

jonatanklosko

Creator of Livebook

XLA reserves memory upfront and then allocates within that reservation as needed. This behaviour can be customized with client options preallocate: false or other :memory_fraction. However, I don’t think this will help with the OOM error.

We are still yet to do more optimisations for stable diffusion, but two things you can try this:

Load the parameters into the CPU with Bumblebee.load_model(..., backend: {EXLA.Backend, client: :host})
Enable lazy transfers in serving defn options: defn_options: [compiler: EXLA, lazy_transfers: :always]

This way, instead of placing all parameters on the GPU, they will be transferred as needed.

Also make sure to try with batch size 1.

Post #1

nutheory

@jonatanklosko what you wrote in your response makes total sense but it seems to either conflict or is just a different way of saying what is in the Bumblebee docs here … bumblebee/examples/phoenix at main · elixir-nx/bumblebee · GitHub

My question is … Is there a reason or benefit to doing it one way or another? Because I tried it the way its explained in the bumblebee doc and I’m still getting a OOM error.

Post #2

jonatanklosko

Creator of Livebook

@nutheory they aren’t conflicting, but aim to do the same thing in a slightly different way. :preallocate_params is a recent addition, I updated the docs to mention it instead. Recently I’ve also added this docs page that goes into more details on how to minimise memory usage by trading off inference time : )

Post #3

nutheory

@jonatanklosko thanks for clarifying that. So I got everything working in a docker container using EXLA on a A100(ubuntu/cuda) at Lambda Labs (im working on a proof of concept). Im noticing exla is about 3x the inference time of oobabooga run on the same server. I feel like maybe this could be a pytorch/libtorch/torchx difference? I really wanna test that theory though, but i cant figure out how to build libtorch with cuda 11.6, it results in a working docker build that breaks (on app startup) from having a mismatch/old build of CuDNN (8.3.2) that is older than the oldest image I could find on docker hub (8.4). I dont know if you can help or point me in the right direction … but would love help from anyone with more experience in docker/cuda/libtorch.

Post #4

jonatanklosko

Creator of Livebook

If you paste how you build the serving (which options you use) perhaps there’s something that can be improved.

One thing that can make a difference is quantization, if the Python version you test uses it. We don’t support quantization in Bumblebee yet, but there’s ongoing work in EXLA that should allow it in the future.

Post #5

nutheory

@jonatanklosko ah **** you’re right i was using quantization, good catch. So the LLM im using is NousResearch/Llama-2-7b-hf … I tried two scenarios … 1) I loaded up the .bin files (25gb) to the cpu… 2) I loaded up the .safetensors (12gb) with preallocate: true… both gave the same inference time.

def load_text_serving do
model = System.fetch_env!(“TEXT_GENERATION_MODEL”)
auth_token = System.fetch_env!(“HF_TOKEN”)

{:ok, txt_model} =
  Bumblebee.load_model({:hf, model, auth_token: auth_token},
    backend: {EXLA.Backend, client: :host},
    params_filename: "model.safetensors"
  )

{:ok, tokenizer} =
  Bumblebee.load_tokenizer({:hf, model, auth_token: auth_token})

{:ok, generation_config} =
  Bumblebee.load_generation_config({:hf, model, auth_token: auth_token})

config =
  Bumblebee.configure(generation_config,
    min_new_tokens: 120,
    max_new_tokens: 600,
    strategy: %{type: :contrastive_search, top_k: 4, alpha: 0.6}
  )

Bumblebee.Text.generation(txt_model, tokenizer, config,
  compile: [batch_size: 1, sequence_length: 1000],
  defn_options: [compiler: EXLA, preallocate_params: true]
)

end

Post #6

jonatanklosko

Creator of Livebook

With sequence_length: 1000 we always pad the input text to 1000 tokens. You can create a couple variants of the computation for different input lengths by doing sequence_length: [100, 300, 1000]. Note that this does not require more memory.

As for the f16 .safetensors, you can try changing axon policy to force f16 at every step, though I haven’t tested if this makes a difference:

policy = Axon.MixedPrecision.create_policy(params: {:f, 16}, compute: {:f, 16}, output: {:f, 16})
model_info = update_in(model_info.model, &Axon.MixedPrecision.apply_policy(&1, policy))

Post #7

jonatanklosko

Creator of Livebook

Ah wait, you have preallocate_params: true inside :defn_options, it should be outside : )

Post #8

nutheory

this is why i got confused …

defn_options: [compiler: EXLA], preallocate_params: true - move and keep all parameters to the GPU upfront. This requires the most memory, but should provide the fastest inference time.
defn_options: [compiler: EXLA] - copy all parameters to the GPU before each computation and discard afterwards. This requires less memory, but the copying increases the inference time.
defn_options: [compiler: EXLA, lazy_transfers: :always] - lazily copy parameters to the GPU during the computation as needed. This requires the least memory, at the cost of inference time.

https://hexdocs.pm/bumblebee/llama.html

Post #9