huggingface / huggingface/accelerate
[CPU offload hooks] hooks with overlapped transfers and computations
- Dominant language
- Python
- Stars
- 9.9k
- Forks
- 1.5k
- Avg merge
- 5d 2h
- Merged PRs (30d)
- 27
Description
`diffusers` relies on `cpu_offload()` a lot for implementing [`enable_sequential_cpu_offload()`](https://github.com/huggingface/diffusers/blob/fdec8bd6754f8ae5428fb542f08707e0a5aba24e/src/diffusers/pipelines/pipeline_utils.py#L1092). 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:
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.
Contributor guide
Assessment
This issue has not been assessed yet.