aws / aws/sagemaker-huggingface-inference-toolkit

Support passing model_kwargs to pipeline

Open
#85 1 comment 2 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.