Deploy GPT-J 6B for inference using Hugging Face Transformers and Amazon SageMaker

Summary

Deploying EleutherAI's 6 billion parameter GPT-J model for production inference presents latency hurdles due to large memory footprints and lengthy startup times. Loading the model via standard methods takes up to several minutes, conflicting with strict real-time response limits such as Amazon SageMaker's 60-second threshold. To overcome this limitation, the model is serialized using PyTorch's native save mechanisms, packaged into a compressed archive with supporting assets, and stored on Amazon S3. This alternative loading workflow reduces GPT-J load times down to 7.7 seconds. An Amazon SageMaker real-time endpoint is then deployed on an NVIDIA T4 GPU instance using the Hugging Face Inference Toolkit.

Context

EleutherAI's 6 billion parameter GPT-J model has a ~24GB memory footprint in float32 and requires significant RAM to load. Default loading methods take 1 minute and 23 seconds from disk or over 3 minutes over network, exceeding Amazon SageMaker's 60-second response limit for real-time inference.

Approach / What changed

Convert the float16 model using torch.save to serialize the entire module, packaging it into a model.tar.gz archive on S3 alongside tokenizers. Deploy this artifact to an Amazon SageMaker real-time endpoint on an ml.g4dn.xlarge instance using HuggingFaceModel, and send parameterized inference payloads.

Takeaways

• Using torch.save and torch.load drops GPT-J loading time from 83 seconds to 7.7 seconds, representing an approximately 10.5x speedup.
• Serializing modules via PyTorch binds data to specific classes, requiring matched PyTorch and Transformers library versions across saving and loading environments.
• SageMaker real-time endpoints require responses within a 60-second window, making fast initial model loading critical for endpoint reliability.

Related reading

• huggingface.co · Mar 16, 2022

  Accelerate BERT inference with Hugging Face Transformers and AWS Inferentia

  Production deployments of BERT and Transformer architectures often face cost and latency challenges because these models are significantly larger and more computationally intensive than traditional algorithms. To optimize text classification workloads, developers can compile vanilla PyTorch models for AWS Inferentia using the AWS Neuron SDK and its tracing utilities. Because the Neuron SDK requires static tensor dimensions, the model is traced with fixed input lengths and packaged alongside a custom inference script configuring one Neuron Core per worker. Deploying the resulting artifacts to an Amazon SageMaker ml.inf1.xlarge endpoint yields an average latency of 5 to 6 milliseconds for a sequence length of 128 across 10,000 synchronous evaluation requests.

  Philipp Schmid

• huggingface.co · Jul 8, 2021

  Deploy Hugging Face models easily with Amazon SageMaker

  Hugging Face and Amazon introduced an inference solution integrating Hugging Face Transformers with Amazon SageMaker. The SageMaker Hugging Face Inference Toolkit and specialized Deep Learning Containers enable deployment of trained models or publicly available Hub models to managed production endpoints. Users configure endpoints with minimal code using the SageMaker Python SDK, standard pipelines, or custom inference scripts that override default behaviors. The architecture supports models stored in Amazon S3 as well as direct references via HF_MODEL_ID and HF_TASK environment variables. Deployments benefit from native AWS infrastructure features, including built-in monitoring, Identity and Access Management permission controls, and Virtual Private Cloud connectivity.