Deploying Granite Code models in Amazon SageMaker

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.

The IBM Granite Code models are a family of high-performance, foundational language models pre-trained on over 3 trillion tokens of code and natural language data across 116 programming languages. These models range from 3 billion to 34 billion parameters and come in base and instruction-following variants.

What sets the Granite Code models apart is their strong performance on a wide range of code intelligence tasks like code generation, translation, analysis, and refactoring - often outperforming larger open-source models.

IBM has released the Granite Code models to open source under the permissive Apache 2.0 license, enabling their use for both research and commercial purposes with no restrictions. The models are available on Hugging Face.

Hugging Face is a popular open source hub for machine learning (ML) models. AWS and Hugging Face have a partnership that allows a seamless integration through SageMaker with a set of AWS Deep Learning Containers (DLCs) for training and inference in PyTorch or TensorFlow, and Hugging Face estimators and predictors for the SageMaker Python SDK. SageMaker features and capabilities help developers and data scientists get started with natural language processing (NLP) on AWS with ease.

In this notebook, we will deploy the Granite models on Amazon SageMaker for accelerating legacy code conversion and modernisation use cases.

Deploying Granite Code models in Amazon SageMaker

Prepare the environment

Create SageMaker Estimator

Link to the Granite Code models in HuggingFace: https://huggingface.co/ibm-granite

Next we configure the model object by specifying a unique name, the image_uri for the managed TGI container, and the execution role for the endpoint. Additionally, we specify a number of environment variables including the HF_MODEL_ID which corresponds to the model from the HuggingFace Hub that will be deployed, and the HF_TASK which configures the inference task to be performed by the model.

You should also define SM_NUM_GPUS, which specifies the tensor parallelism degree of the model. Tensor parallelism can be used to split the model across multiple GPUs, which is necessary when working with LLMs that are too big for a single GPU. Here, you should set SM_NUM_GPUS to the number of available GPUs on your selected instance type. For example, in this tutorial, we set SM_NUM_GPUS to 4 because our selected instance type ml.g5.12xlarge has 4 available GPUs.

The HuggingFaceModel handles downloading the Granite model and dependencies, packaging them into a Docker container, and deploying to a SageMaker inference endpoint.

After we have created the HuggingFaceModel we deploy it to Amazon SageMaker using the deploy method. We deploy the model with the ml.g5.12xlarge instance type, which has 4 NVIDIA A10G GPUs and 96GB of GPU memory.

SageMaker will now create our endpoint and deploy the model to it. This can takes a 10-15 minutes.

Run inference with the model

Now that we have the Granite Code model loaded and deployed to a SageMaker endpoint, we can start generating or converting code. We use the predict method from the predictor to run inference on our endpoint. We can inference with different parameters to impact the generation. Parameters can be defined as in the parameters attribute of the payload.

Example 1: Code Generation

In this example, we want to write a function in the Python programming language that reverses a string.

The output contains Python code similar to the following snippet:

def reverse_string(my_string):
    return my_string[::-1]

Be sure to test the generated code to verify that it works as you expect.

For example, if you run reversed("good morning"), the result is gninrom doog.

Example 2: Code Conversion

In this example, we want to convert code from one programming language to another. The prompt below converts a code snippet from C++ to Python.

You can send the prompt to the Granite Code model loaded and deployed to the SageMaker endpoint, and adjust the following hyperparameters.

The output contains Python code similar to the following snippet:

Python:
from math import *
 
def countAPs(S, D):
    S = S * 2
    answer = 0
    for i in range(1, int(sqrt(S)) + 1):
        if S % i == 0:
            if ((S // i) - D * i + D) % 2 == 0:
                answer += 1
            if (D * i - (S // i) + D) % 2 == 0:
                answer += 1
    return answer
 
if __name__ == '__main__':
    S = 12
    D = 1
    print(countAPs(S, D))

Be sure to test the generated code to verify that it works as you expect.

Example 3: Code Conversion (C to Java)

In this example, you want to convert code from one programming language to another. The prompt below converts a code snippet from C to Java.

Specifically, we cover common programming constructs like linked lists and file I/O operations. The C code is converted to Java while preserving the functionality and logic. In the Java code, we utilize classes, objects, and Java-specific APIs like FileWriter and BufferedReader to achieve similar results as the C code.

• In the first example, the C code implements a singly linked list data structure. It defines a Node struct containing an integer data value and a pointer to the next node. The code provides functions to create a new node, add a node to the end of the list, and print the list.
• In the second example, the C code demonstrates how to write data to a file and then read the data back from the file. It opens a file named "example.txt" in write mode, writes a string to the file, and closes the file. Then, it opens the same file in read mode, reads the contents into a buffer, and prints the buffer to the console.

The output contains Java code similar to the following snippet:

import java.io.*;

public class FileExample {
    public static void main(String[] args) {
        try (FileWriter writer = new FileWriter("example.txt");
             BufferedReader reader = new BufferedReader(new FileReader("example.txt"))) {

            writer.write("This is an example of writing to a file.");

            String line;
            while ((line = reader.readLine())!= null) {
                System.out.println(line);
            }
        } catch (IOException e) {
            System.out.println("An error occurred: " + e.getMessage());
        }
    }
}

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.