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

GPT-J 6B has been sitting around as the open-source answer to GPT-3 for half a year now, and plenty of researchers and developers have wanted to use it. The problem was never the model itself. The problem was getting it running somewhere that didn't fall over the moment you tried to serve real traffic.

Here's the snag: GPT-J's weights alone eat up roughly 24GB, and loading it the standard way through Transformers' from_pretrained method took Philipp Schmid, who wrote this up for Hugging Face, 3 minutes and 32 seconds on a P3.2xlarge EC2 instance. Even with the model already sitting on disk, that dropped only to 1 minute 23 seconds. That's a problem because Amazon SageMaker enforces a 60-second window for a request to load the model and return a prediction. Miss that window and your endpoint isn't scalable, full stop.

The fix turns out to be almost embarrassingly simple: skip from_pretrained for loading and instead use PyTorch's native torch.save() and torch.load(). Schmid tested this on BERT first and saw load times fall from 1.97 seconds to 0.166 seconds, about 12x faster. Applied to GPT-J, that same trick cut load time from 83 seconds to 7.7 seconds, a 10.5x improvement that finally puts the model comfortably inside SageMaker's time budget. The catch, and it's a real one, is that pickled PyTorch models tied to torch.save are brittle across library versions. Save with transformers 4.13.2 and try loading with 4.15.0, and you might just break everything.

Once that hurdle's cleared, the actual SageMaker deployment is almost anticlimactic. Package the model into a model.tar.gz, either using Hugging Face's pre-uploaded artifact or a provided conversion script, point the HuggingFaceModel class at that S3 URI, and deploy to something as modest as a g4dn.xlarge instance running an NVIDIA T4, roughly $500 a month. Inference after the initial request lands around 3 seconds for greedy decoding, climbing to somewhere between 15 and 45 seconds for few-shot prompts depending on generation length and parameters like beam search or temperature.

What's notable here isn't the novelty of the underlying idea, engineers have used torch.save shortcuts before, but the fact that a 6-billion-parameter language model can now run in production on unremarkable, relatively cheap hardware without a research team babysitting it. That's the gap between "cool open-source release" and "thing you can actually ship," and it's been narrower for smaller models for a while. GPT-J just caught up.