aws / aws/sagemaker-huggingface-inference-toolkit
Support passing model_kwargs to pipeline
- Dominant language
- Python
- Stars
- 270
- Forks
- 60
- PR merge metrics
- No merged PRs in 30d
Description
I'm trying to deploy BLIP-2 (specifically `Salesforce/blip2-opt-2.7b`) to a Sagemaker (SM) endpoint, but coming up against some problems.
We can deploy this model by tar'ing the model artifacts as `model.tar.gz` and hosting on S3, but creating a ~9GB tar file is time-consuming and leads to slow deployment feedback loops.
Alternatively, the toolkit has experimental support for downloading models from 🤗Hub on start, which is a more time/space efficient.
However, this functionality only supports passing `HF_TASK` and `HF_MODEL_ID` as env vars. In order to run inference on this model using GPU's available on SM (T4/A10) we need to pass additional `model_kwargs` as:
```
pipe = pipeline(model="Salesforce/blip2-opt-2.7b", model_kwargs={"load_in_8bit": True})
```
A potential solution to this would be:
On [line 104 of handler_service.py](https://github.com/aws/sagemaker-huggingface-inference-toolkit/blob/27275f40a2bbff85bb507646e6a3ef866d0599af/src/sagemaker_huggingface_inference_toolkit/handler_service.py#L104) the ability to pass `kwargs` has not been implemented, but the function `get_pipeline` allows for `kwargs`.
Contributor guide
Research direction
Start with src/sagemaker_huggingface_inference_toolkit/handler_service.py around line 104 and inspect how get_pipeline accepts kwargs. Trace how HF_TASK and HF_MODEL_ID are read, then determine how model_kwargs should reach the pipeline for the SageMaker startup path. Done means the deployment path can pass model_kwargs such as load_in_8bit to the pipeline.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, huggingface, python
- Domain
- cloud, machine-learning
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100