apache / apache/beam

[Feature Request][Tracking]: Use accelerate from Hugging Face to optimize loading Pytorch models

Open
#24,340 0 comments 0 reactions 0 assignees View on GitHub
ml new feature P3 python run-inference
Dominant language
Java
Stars
8.7k
Forks
4.7k
Avg merge
1d 20h
Merged PRs (30d)
196

Description

### What would you like to happen?

[Accelerate](https://huggingface.co/blog/accelerate-large-models) module from HuggingFace is used to optimize model loading for PyTorch, Tensorflow.

The model loading pipeline in torch happens as below
* Create the model
* Load in memory its weights (in an object usually called state_dict)
* Load those weights in the created model
* Move the model on the device for inference

This is not a smart way to load the model and HuggingFace accelerate helps mitigate this issue.
As per [documentation](https://huggingface.co/docs/accelerate/usage_guides/big_modeling#limits-and-further-development), this is still in development and needs at least one GPU to run this API but as per the docs this could be fixed in the future.

### Issue Priority

Priority: 3

### Issue Component

Component: run-inference

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.