aws / aws/sagemaker-huggingface-inference-toolkit

Sagemaker endpoint doesn't use GPU (instance ml.g4dn.xlarge)

Open
#117 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
270
Forks
60
PR merge metrics
No merged PRs in 30d

Description

I have spent the whole day trying to deploy a custom HF model to Sagemaker endpoint and making sure it uses the GPU, and I had no luck, hoping to get some insight here.

Here's my script for the model deployment
```
img_url_old_lib='763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04'

huggingface_model = HuggingFaceModel(
model_data=model_uri,
role=role,
source_dir='code',
entry_point='inference.py',
name='hf-inference-1-13-gpu',
image_uri=img_url_old_lib
)

predictor = huggingface_model.deploy(
initial_instance_count=1, # number of instances
instance_type='ml.g4dn.xlarge', #has 1 GPU
endpoint_name='hf-inference-1-13-gpu',
)
```
And here's my `code/inference.py` script
```
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def model_fn(model_dir, context=None):
"""
Load the model for inference
"""

model_path = os.path.join(model_dir, 'model/')

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
print("Loaded Processor")

model = VisionEncoderDecoderModel.from_pretrained(model_path)
print("Loaded Model")

model_dict = {'model': model.to(device), 'processor': processor}

return model_dict

def predict_fn(images, model, context=None):
"""
Apply model to the incoming request
"""

images = [Image.open(io.BytesIO(content)) for content in images]
print("Opened Image")

processor = model['processor']
model = model['model']

pixel_values = processor(images, return_tensors="pt").pixel_values
pixel_values = pixel_values.to(device)
generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)
print("Generated Text: " + str(generated_text))

return generated_text
```

I've read the threads [here](https://discuss.huggingface.co/t/sagemaker-endpoint-not-using-gpu-for-pygmalionai/37361) and [here](https://discuss.huggingface.co/t/how-do-i-deploy-a-hub-model-to-sagemaker-and-give-it-a-gpu-not-elastic-inference/14727/3), and followed the suggestions made by @philschmid , I tried changing the version of the `transformers_version` arg variable but it still doesn't use the GPU (see pic below). I tested the model in the SM notebook using the same GPU instance(ml.g4dn.xlarge) and I can confirm the inference code does use the GPU as expected. So I'm not sure why when it's deployed to the endpoint with the docker image it doesn't use the GPU.

I'd appreciate any help on this, thanks!
![GPU monitor](https://github.com/aws/sagemaker-huggingface-inference-toolkit/assets/92494420/bf462ae5-c8ee-494f-979a-f5cfc92115f0)

Contributor guide

Open the contributing guide

Research direction

Start with the deployment parameters and code/inference.py, then check whether torch.cuda.is_available() and GPU monitoring differ inside the deployed endpoint image. Compare the deployed inference behavior with the notebook run on ml.g4dn.xlarge. Done means identifying why the endpoint does not use its GPU and confirming GPU-backed inference after deployment.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, docker, huggingface, python, pytorch
Domain
backend-api-design, cloud, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.