Serve GPT4ALL-J-6b on SageMaker with transformer-neuronx using DJL container.

This notebook was run on ml.m5.xlarge but we will deploy using the ml.inf2.8xlarge instance

In this notebook, we explore how to host a large language model on SageMaker using the Large Model Inference container that packages DJL model server and transformers-neuronx, an open source library that serves the purpose of running transformer inference workflows using AWS Neuron SDK. Neuron SDK is required to to run deep learning workloads on AWS Inferentia and AWS Trainium based instances . We use DJLServing as the model serving solution in this example.

In this notebook, we deploy the open source GPT4ALL-J-6B(https://huggingface.co/nomic-ai/gpt4all-j) model across NeuronCores on an ml.inf2.24xlarge instance. The model is downloaded, converted to Inf2 format,partitioned and loaded to the Inferentia Chips.

This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook.

Licence agreement

• View license information https://huggingface.co/nomic-ai/gpt4all-j before using the model.
• This notebook is a sample notebook and not intended for production use. Please refer to the licence at https://github.com/aws/mit-0.

Import the relevant libraries and configure several global variables using boto3

Create SageMaker compatible Model artifact, upload model to S3 .

SageMaker Large Model Inference containers can be used to host models without providing your own inference code. You can also provide a 'model.py' file if you like to customize in the model loading and inference process. SageMaker needs the model artifacts to be in a Tarball format. In this example, we provide the serving.properties.

The tarball is in the following format

code
├──── 
│   └── serving.properties

• serving.properties is the configuration file that can be used to configure the model server.

Start preparing model artifacts

In LMI contianer, we expect some artifacts to help setting up the model

serving.properties (required): Defines the model server settings model.py (optional): A python file to define the core inference logic requirements.txt (optional): Any additional pip wheel need to install

Create serving.properties

This is a configuration file to indicate to DJL Serving which model and configurations you want to use

Here is a list of settings that we use in this configuration file -

• engine: The engine for DJL to use. In this case, we intend to use transformers-neuronx and hence set it to Python.
• option.entryPoint: The entrypoint python file or module. Since here we are using the default handler, we will configure it to use transformers-neuronx
• option.model_id : The model id of a pretrained model hosted inside a model repository on huggingface.co (https://huggingface.co/models). The container uses this model id to download the corresponding model repository on huggingface.co.
• option.n_positions: Number of tokens model can accept.

If you want to download teh model from S3, you can instead set the parameteroption.model_idto the Amazon S3 bucket that contains the model. When this is set, the container leverages s5cmd to download the model from s3. This is extremely fast and useful when downloading large models.

• option.tensor_parallel_degree: This property value sets the number of devices over which the tensor parallel modules are distributed. For example, if we have 12 Neruon Core and you do tensor parallel degree of 4. LMI will distribute to 3 copies of models and each of them takes 4 cores

For more details on the configuration options and an exhaustive list, you can refer the documentation - https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints-large-model-configuration.html.

Create the model tarball and upload it to a S3 bucket

Create SageMaker endpoint

1. Create the Model using the Image container and the Model Tarball uploaded earlier

2. Create the endpoint config using the following key parameters

   a) Instance Type is ml.inf2.8xlarge

   b) ContainerStartupHealthCheckTimeoutInSeconds is 3600 to ensure health check starts after the model is ready

3. The container downloads the model into the /tmp space on the container because SageMaker maps the /tmp to the Amazon Elastic Block Store (Amazon EBS). We need to add a volume_size parameter to ensure /tmp directory has enough space to download and compile the model.

4. Create the end point using the endpoint config created

Note that the model will be compiled to INF2 format during the deployment time. You can confirm this in the endpoint CloudWatch logs. You can ignore the below warning when deploying the endpoint.

Test the endpoint

This is a generative model so we pass in a Text as a prompt.

Conclusion

In this post, we demonstrated how to use SageMaker large model inference containers to host GPT4ALL-J-6B.

Clean Up

Notebook CI Test Results

This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.