[CPU offload hooks] hooks with overlapped transfers and computations

Open
  • Is something wrong?
#3,267 1 comment 1 reaction 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
25/100
Issue type
Feature
Clarity
Needs clarification
Activity status
Stale
Tech stack
python
Domain
machine-learning, performance

Research direction

Start by reading the linked gist and the referenced diffusers pipeline_utils.py implementation of enable_sequential_cpu_offload(). Compare the proposed CUDA-stream hook with the current behavior and identify how it could fit in accelerate; the issue names no target file or test, so completion criteria still need to be established.

Written by the indexing model from the issue text.

Description

feature request

diffusers relies on cpu_offload() a lot for implementing enable_sequential_cpu_offload(). It offloads the modules of a model to CPU when they are not being used and only pops them on to the GPU when it's needed for computation.

As one can notice, the cost of these frequent transfers blocks the underlying computation and hence it leads to quite a bit of increased latency. But it also tremendously helps in running very big models on consumer hardware (very important as a good diffusion model is actually composed of multiple big models).

So, the question is can we overlap communication with computation? https://gist.github.com/gau-nernst/9408e13c32d3c6e7025d92cce6cba140 implements a hook that leverages CUDA streams to implement the same functionality as enable_sequential_cpu_offload() but is significantly faster. See the results:

image

As @SunMarc and I were discussing, it'd be extremely cool to have a similar hook supported in accelerate so that we can make diffusion models, in particular, more accessible without having to completely give away speed.

Cc: @DN6 @a-r-r-o-w as well.

Dominant language
Python
Stars
9.9k
Forks
1.6k
Avg merge
20h 23m
Merged PRs (30d)
12

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from huggingface/accelerate

All issues in huggingface/accelerate

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.