Deploy large models for inference with TorchServe

This tutorial demonstrates how to deploy large models and serve inference in Amazon SageMaker AI with TorchServe on GPUs. This example deploys the OPT-30b model to an ml.g5 instance. You can modify this to work with other models and instance types. Replace the italicized placeholder text in the examples with your own information.

TorchServe is a powerful open platform for large distributed model inference. By supporting popular libraries like PyTorch, native PiPPy, DeepSpeed, and HuggingFace Accelerate, it offers uniform handler APIs that remain consistent across distributed large model and non-distributed model inference scenarios. For more information, see TorchServe’s large model inference documentation.

Deep learning containers with TorchServe

To deploy a large model with TorchServe on SageMaker AI, you can use one of the SageMaker AI deep learning containers (DLCs). By default, TorchServe is installed in all AWS PyTorch DLCs. During model loading, TorchServe can install specialized libraries tailored for large models such as PiPPy, Deepspeed, and Accelerate.

The following table lists all of the SageMaker AI DLCs with TorchServe.

Getting started

Before deploying your model, complete the prerequisites. You can also configure your model parameters and customize the handler code.

Prerequisites

To get started, ensure that you have the following prerequisites:

Configure model settings and parameters

TorchServe uses torchrun to set up the distributed environment for model parallel processing. TorchServe has the capability to support multiple workers for a large model. By default, TorchServe uses a round-robin algorithm to assign GPUs to a worker on a host. In the case of large model inference, the number of GPUs assigned to each worker is automatically calculated based on the number of GPUs specified in the model_config.yaml file. The environment variable CUDA_VISIBLE_DEVICES, which specifies the GPU device IDs that are visible at a given time, is set based this number.

For example, suppose there are 8 GPUs on a node and one worker needs 4 GPUs on a node (nproc_per_node=4). In this case, TorchServe assigns four GPUs to the first worker (CUDA_VISIBLE_DEVICES="0,1,2,3") and four GPUs to the second worker (CUDA_VISIBLE_DEVICES="4,5,6,7”).

In addition to this default behavior, TorchServe provides the flexibility for users to specify GPUs for a worker. For instance, if you set the variable deviceIds: [2,3,4,5] in the model config YAML file, and set nproc_per_node=2, then TorchServe assigns CUDA_VISIBLE_DEVICES=”2,3” to the first worker and CUDA_VISIBLE_DEVICES="4,5” to the second worker.

In the following model_config.yaml example, we configure both front-end and back-end parameters for the OPT-30b model. The configured front-end parameters are parallelType, deviceType, deviceIds and torchrun. For more detailed information about the front-end parameters you can configure, see the PyTorch GitHub documentation. The back-end configuration is based on a YAML map that allows for free-style customization. For the back-end parameters, we define the DeepSpeed configuration and additional parameters used by custom handler code.

# TorchServe front-end parameters
minWorkers: 1
maxWorkers: 1
maxBatchDelay: 100
responseTimeout: 1200
parallelType: "tp"
deviceType: "gpu"
# example of user specified GPU deviceIds
deviceIds: [0,1,2,3] # sets CUDA_VISIBLE_DEVICES

torchrun:
    nproc-per-node: 4

# TorchServe back-end parameters
deepspeed:
    config: ds-config.json
    checkpoint: checkpoints.json

handler: # parameters for custom handler code
    model_name: "facebook/opt-30b"
    model_path: "model/models--facebook--opt-30b/snapshots/ceea0a90ac0f6fae7c2c34bcb40477438c152546"
    max_length: 50
    max_new_tokens: 10
    manual_seed: 40

Customize handlers

TorchServe offers base handlers and handler utilities for large model inference built with popular libraries. The following example demonstrates how the custom handler class TransformersSeqClassifierHandler extends BaseDeepSpeedHandler and uses the handler utilities. For a full code example, see the custom_handler.py code on the PyTorch GitHub documentation.

class TransformersSeqClassifierHandler(BaseDeepSpeedHandler, ABC):
    """
    Transformers handler class for sequence, token classification and question answering.
    """

    def __init__(self):
        super(TransformersSeqClassifierHandler, self).__init__()
        self.max_length = None
        self.max_new_tokens = None
        self.tokenizer = None
        self.initialized = False

    def initialize(self, ctx: Context):
        """In this initialize function, the HF large model is loaded and
        partitioned using DeepSpeed.
        Args:
            ctx (context): It is a JSON Object containing information
            pertaining to the model artifacts parameters.
        """
        super().initialize(ctx)
        model_dir = ctx.system_properties.get("model_dir")
        self.max_length = int(ctx.model_yaml_config["handler"]["max_length"])
        self.max_new_tokens = int(ctx.model_yaml_config["handler"]["max_new_tokens"])
        model_name = ctx.model_yaml_config["handler"]["model_name"]
        model_path = ctx.model_yaml_config["handler"]["model_path"]
        seed = int(ctx.model_yaml_config["handler"]["manual_seed"])
        torch.manual_seed(seed)

        logger.info("Model %s loading tokenizer", ctx.model_name)

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        config = AutoConfig.from_pretrained(model_name)
        with torch.device("meta"):
            self.model = AutoModelForCausalLM.from_config(
                config, torch_dtype=torch.float16
            )
        self.model = self.model.eval()

        ds_engine = get_ds_engine(self.model, ctx)
        self.model = ds_engine.module
        logger.info("Model %s loaded successfully", ctx.model_name)
        self.initialized = True

    def preprocess(self, requests):
        """
        Basic text preprocessing, based on the user's choice of application mode.
        Args:
            requests (list): A list of dictionaries with a "data" or "body" field, each
                            containing the input text to be processed.
        Returns:
            tuple: A tuple with two tensors: the batch of input ids and the batch of
                attention masks.
        """

    def inference(self, input_batch):
        """
        Predicts the class (or classes) of the received text using the serialized transformers
        checkpoint.
        Args:
            input_batch (tuple): A tuple with two tensors: the batch of input ids and the batch
                                of attention masks, as returned by the preprocess function.
        Returns:
            list: A list of strings with the predicted values for each input text in the batch.
        """
        
    def postprocess(self, inference_output):
        """Post Process Function converts the predicted response into Torchserve readable format.
        Args:
            inference_output (list): It contains the predicted response of the input text.
        Returns:
            (list): Returns a list of the Predictions and Explanations.
        """

Prepare your model artifacts

Before deploying your model on SageMaker AI, you must package your model artifacts. For large models, we recommend that you use the PyTorch torch-model-archiver tool with the argument --archive-format no-archive, which skips compressing model artifacts. The following example saves all of the model artifacts to a new folder named opt/.

torch-model-archiver --model-name opt --version 1.0 --handler custom_handler.py --extra-files ds-config.json -r requirements.txt --config-file opt/model-config.yaml --archive-format no-archive

Once the opt/ folder is created, download the OPT-30b model to the folder using the PyTorch Download_model tool.

cd opt
python path_to/Download_model.py --model_path model --model_name facebook/opt-30b --revision main

Lastly, upload the model artifacts to an Amazon S3 bucket.

aws s3 cp opt {your_s3_bucket}/opt --recursive

You should now have model artifacts stored in Amazon S3 that are ready to deploy to a SageMaker AI endpoint.

Deploy the model using the SageMaker Python SDK

After preparing your model artifacts, you can deploy your model to a SageMaker AI Hosting endpoint. This section describes how to deploy a single large model to an endpoint and make streaming response predictions. For more information about streaming responses from endpoints, see Invoke real-time endpoints.

To deploy your model, complete the following steps:

You have now deployed your model to a SageMaker AI endpoint and should be able to invoke it for responses. For more information about SageMaker AI real-time endpoints, see Single-model endpoints.