Deploy Stable LM Zephyr 3B Model on SageMaker
How to Deploy a Light-weight LLM on SageMaker
Published Sep 30, 2024
In the rapidly evolving landscape of Generative AI, choosing the right language model for your project can be a game-changer. Recently, I had the opportunity to deploy the Stable LM Zephyr 3B Model on AWS SageMaker for a customer, and the experience was both challenging and rewarding. This blog post will walk you through our journey, highlighting why we chose this particular model and how we leveraged AWS SageMaker for deployment.
With its compact size, low resource requirements, and strong performance, this model can be an ideal choice for many LLM use cases where efficiency is crucial. While its lightweight design makes it accessible, deploying it on SageMaker presented unique challenges, from managing the model’s configuration optimizing inference. In this post, we’ll share our experiences and provide insights into effectively deploying LLMs on the cloud.
The Stable LM Zephyr 3B, developed by Stability AI, offers a unique combination of efficiency, performance and ethical considerations that sets it apart from the crowded field of language models :
- Efficiency: With just 3 billion parameters, this lightweight LLM provides an excellent balance between computational requirements and performance, making it suitable for a variety of applications, even in resource-constrained environments.
- Competitive Performance: Despite its relatively small size, the Zephyr 3B delivers strong results across various natural language processing tasks, as demonstrated by its competitive scores on benchmarks like MT-Bench and Alpaca Benchmark.
- Enhanced Instruction Following: The model is fine-tuned for instruction following and Q&A-type tasks using Direct Preference Optimization (DPO). This extension of the Stable LM 3B-4e1t model ensures improved responsiveness and accuracy in user interactions.
- Commitment to Ethical AI: Released with a focus on safety, reliability, and appropriateness, the model is designed to support responsible AI use.
Choosing Stable LM Zephyr 3B, allowed us to leverage these advantages while utilizing AWS SageMaker’s scalability and robustness. In the following sections, we'll dive into the deployment process, challenges we faced, and the insights we gained along the way. Refer to [1] for the details of this model including lineage, training insights, performance benchmark, license, etc.
The model size is pretty big (~5.3 GB). You need to choose a big-sized SageMaker instance to run the following Jupyter notebook. In my experiment, I chose ml.m5.2xlarge instance and it took around 5 minutes to complete the deployment.
Use the following code to test the model on the Inference Endpoint.
The model output is like this:
The model is pretty new - released on 7th Dec, 2023. It is not available yet in SageMaker Jumpstart. It provides the instruction code to deploy the model on SageMaker by using the class HuggingFaceModel. Unfortunately, this method doesn’t work due to the incompatible version of the key library - transformers. The latest HuggingFace Inference Containers only contains transformers version 4.37.0 (refer to [3]). With this version, you can successfully deploy the model, but when you perform the inference, you will get the following error:
The checkpoint you are trying to load has model type `stablelm` but Transformers does not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date.
We spent a lot of time figuring out which Container image has a compatible transformers version. We also tried to use latest SageMaker Framework Containers for PyTorch 2.3.0 as the base image, and build a custom Docker image by installing the latest transformers. However, for some unknown reason, we cannot successfully invoke the SageMaker Inference Endpoint with this custom image. It went into an unknown state, the health check cannot be passed, and in the end, it failed due to timeout.
Finally, we figured out the latest HuggingFace Text Generation Inference (TGI) Container v2.0-hf-tgi-2.2.0-pt-2.3.0-inf-gpu-py310 meets all of the dependency requirements:
- PyTorch 2.3.0
- Python 3.10
- Transformers 4.43.1
With this Container image, we are able to successfully deploy the Inference Endpoint and perform prediction against the deployed model.
With the persistent effort and close collaboration, we successfully identified the correct approach and the base Docker image, enabling us to deploy the Stable LM Zephyr 3B model on AWS SageMaker and perform inference effectively. This breakthrough not only solved our immediate challenge but also opened up new possibilities. For those looking to use this model and expand on this foundation, you can use this base image to build your custom model, adding any additional dependencies your specific use case may require. Our journey demonstrates that even when dealing with a new model, lacking prior deployment knowledge on the Internet, and facing challenges with the mistakes in the official documentation, complex AI deployment tasks can still be accomplished successfully in cloud environments like AWS SageMaker through perseverance and innovative problem-solving.
- AWS DLC Releases for TGI, https://github.com/aws/deep-learning-containers/releases?q=tgi+AND+gpu&expanded=true
Pengfei Zhang - Senior Solutions Architect, AWS
Pengfei Zhang is a Senior Solutions Architect in Singapore Startups, providing architectural consultations to diverse early and scale-up AWS customers. He has recently guided priority customers through their Generative AI journey on AWS. He is specialized in Serverless and AIML technology domains.
Rajat Kumar Sinha - Senior AI Engineer, Intellect.co
Rajat Kumar Sinha is a Senior AI Engineer with over 7 years of experience in computer vision and natural language processing (NLP). Over the past two years, he has focused on Generative AI, leveraging cutting-edge techniques to develop innovative solutions across various domains. Currently, he is utilizing his expertise at Intellect to revolutionize mental health support, making it more accessible and personalized for users. His work aims to bridge the gap between scientific research and real-world applications, bringing the latest advancements in AI closer to people’s lives and well-being.