Reduce StableDiffusion memory usage

[josevalim]

A list of ideas to explore:

• Lazy transfers (so we don't load data into the GPU at once)
• FP16 on load
• FP16 policies on Axon
• Attention slicing (no longer applicable Remove attention slicing from docs huggingface/diffusers#4487)
• Flash attention (JAX version) (see notes in Refactor attention implementation #300)
• DPM-Solver++ (more schedulers here, here, and in the comments below) (another PyTorch implementation)
• TokenMerging
• LCM+LoRA
• DeepCache (not applicable Reduce StableDiffusion memory usage #147 (comment))

[josevalim]

More on attention: https://pytorch.org/blog/flash-decoding/

[bfolkens]

I'd also suggest FlashAttention-2 and Medusa

[josevalim]

Alternative to DPM Solver: https://arxiv.org/abs/2311.05556

[josevalim]

More notes on optimizations here:

• https://huggingface.co/docs/diffusers/main/en/stable_diffusion
• https://huggingface.co/docs/diffusers/main/en/optimization/fp16
• https://huggingface.co/docs/diffusers/main/en/optimization/memory

[jonatanklosko]

I tested SD v1-4 on a GPU using the new lower precision options params_variant: "fp16", type: :bf16. Here are a couple runs:

Type	Steps	Batch, Images	Time	Memory	Lazy transfers
bf16	20	1, 1	0.7s	4669MiB	No
bf16	20	1, 4	2.2s	8769MiB	No
f32	20	1, 1	1.3s	8759MiB	No
f32	20	1, 4	4.3s	16951MiB	No
bf16	20	1, 1	3.7s	6957MiB	Yes
f32	20	1, 1	8.2s	13379MiB	Yes

Note that the reported memory is just the final memory after using preallocate: false, so it's not ideally reliable. XLA even does memory reservations at compilation time, my guess is that it runs some example operations to pick preferable algorithm or fine tune algorithm parameters. That said, it seems clear that bf16 reduces both memory and time roughly by a factor of 2. Weirdly, lazy transfers seem to bump the memory usage (but it doesn't mean that much memory is required in practice, it's just XLA bumping the reservation, see below).

Source (first row)

# Stable Diffusion testing

```elixir
Mix.install([
  {:nx, github: "elixir-nx/nx", sparse: "nx", override: true},
  {:exla, github: "elixir-nx/nx", sparse: "exla", override: true},
  {:axon, github: "elixir-nx/axon", override: true},
  {:kino, "~> 0.11.3"},
  {:bumblebee, github: "elixir-nx/bumblebee"}
])

Application.put_env(:exla, :clients,
  host: [platform: :host],
  cuda: [platform: :cuda, preallocate: false]
  # cuda: [platform: :cuda, memory_fraction: 0.3]
  # cuda: [platform: :cuda]
)

Application.put_env(:exla, :preferred_clients, [:cuda, :host])

Nx.global_default_backend({EXLA.Backend, client: :host})
```

## init

```elixir
with {output, 0} <- System.shell("nvidia-smi --query-gpu=memory.total,memory.used --format=csv") do
  IO.puts(output)
end
```

<!-- livebook:{"branch_parent_index":0} -->

## Stable Diffusion fp16

```elixir
repository_id = "CompVis/stable-diffusion-v1-4"

{:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "openai/clip-vit-large-patch14"})

{:ok, clip} =
  Bumblebee.load_model({:hf, repository_id, subdir: "text_encoder"},
    params_variant: "fp16",
    type: :bf16
  )

{:ok, unet} =
  Bumblebee.load_model({:hf, repository_id, subdir: "unet"},
    params_variant: "fp16",
    type: :bf16
  )

{:ok, vae} =
  Bumblebee.load_model({:hf, repository_id, subdir: "vae"},
    architecture: :decoder,
    params_variant: "fp16",
    type: :bf16
  )

{:ok, scheduler} = Bumblebee.load_scheduler({:hf, repository_id, subdir: "scheduler"})

clip = update_in(clip.params, &Nx.backend_copy(&1, {EXLA.Backend, client: :cuda}))
unet = update_in(unet.params, &Nx.backend_copy(&1, {EXLA.Backend, client: :cuda}))
vae = update_in(vae.params, &Nx.backend_copy(&1, {EXLA.Backend, client: :cuda}))

serving =
  Bumblebee.Diffusion.StableDiffusion.text_to_image(clip, unet, vae, tokenizer, scheduler,
    num_steps: 20,
    num_images_per_prompt: 1,
    compile: [batch_size: 1, sequence_length: 60],
    defn_options: [compiler: EXLA]
  )

Kino.start_child({Nx.Serving, name: SD, serving: serving})
```

```elixir
prompt = "numbat, forest, high quality, detailed, digital art"

output = Nx.Serving.batched_run(SD, prompt)

for result <- output.results do
  Kino.Image.new(result.image)
end
|> Kino.Layout.grid(columns: 2)
```

[jonatanklosko]

I experimented with different values of memory_fraction as an upper limit. For the first entry in the table above:

• lazy_transfers: :always - 3GiB (3.4s)
• manual backend_copy - 4.6GiB (0.7s)
• preallocate_params: true it's 6.2GiB (0.7s)

So lazy transfers do help a bit, but imply a significant slowdown.

What's interesting though is that preallocate_params requires more memory than manual backend_copy. It's even more surprising given that the OOM happens at serving runtime, not during the params preallocation.

[josevalim]

preallocate/jit will transfer the data twice, one as arguments, one as return type. So we probably need a new callback/abstraction to make this easier :D

[jonatanklosko]

FTR fixed in #317, now preallocate_params: true effectively does backend_copy :)

[josevalim]

I have added an entry for LCM+Lora, @wtedw may have input here (and we may need to update/release a Axon before). /cc @seanmor5

[seanmor5]

I think we should update Axon to better support LoRA, I have a draft in place right now but I have to revisit it to make it work as I intend :)