Create an inference optimization job

To begin creating an optimization job

1. In SageMaker AI Studio, create an optimization job through any of the following paths:

   • To create a job for a JumpStart model, do the following:

     1. In the navigation menu, choose JumpStart.

     2. On the All public models page, choose a model provider, and then choose one of the models that supports optimization.

     3. On the model details page, choose Optimize. This button is enabled only for models that support optimization.

     4. On the Create inference optimization job page, some JumpStart models require you to sign an end-user license agreement (EULA) before you can proceed. If requested, review the license terms in the License agreement section. If the terms are acceptable for your use case, select the checkbox for I accept the EULA, and read the terms and conditions.

   • To create a job for a fine-tuned JumpStart model, do the following:

     1. In the navigation menu, under Jobs, choose Training.

     2. On the Training Jobs page, choose the name of a job that you used to fine tune a JumpStart model. These jobs have the type JumpStart training in the Job type column.

     3. On the details page for the training job, choose Optimize.

   • To create a job for a custom model, do the following:

     1. In the navigation menu, under Jobs, choose Inference optimization.

     2. Choose Create new job.

     3. On the Create inference optimization job page, choose Add model.

     4. In the Add model window, choose Custom Model.

     5. For Custom model name, enter a name.

     6. For S3 URI, enter the URI for the location in Amazon S3 where you've stored your model artifacts.

2. On the Create inference optimization job page, for Job name, you can accept the default name that SageMaker AI assigns. Or, to enter a custom job name, choose the Job name field, and choose Enter job name.

To set the optimization configurations

1. For Deployment instance type, choose the instance type that you want to optimize the model for.

   The instance type affects what optimization techniques you can choose. For most types that use GPU hardware, the supported techniques are Quantization and Speculative decoding. If you choose an instance that uses custom silicon, like the AWS Inferentia instance ml.inf2.8xlarge, the supported technique is Compilation, which you can use to compile the model for that specific hardware type.

2. Select one or more of the optimization techniques that Studio provides:

   • If you select Quantization, choose a data type for Precision data type.

   • If you select Speculative decoding, choose one of the following options:

     • Use SageMaker AI draft model – Choose to use the draft model that SageMaker AI provides.

       Note

       If you choose to use the SageMaker AI draft model, you must also enable network isolation. Studio provides this option under Security.

     • Choose JumpStart draft model – Choose to select a model from the JumpStart catalog to use as your draft model.

     • Choose your own draft model – Choose to use your own draft model, and provide the S3 URI that locates it.

   • If you choose Fast model loading, Studio shows the OPTION_TENSOR_PARALLEL_DEGREE environment variable. Use the Value field to set the degree of tensor parallelism. The value must evenly divide the number of GPUs in the instance you chose for Deployment instance type. For example, to shard your model while using an instance with 8 GPUs, use the values 2, 4, or 8.

   • If you set Deployment instance type to an AWS Inferentia or AWS Trainium instance, Studio might show that Compilation is the one supported option. In that case, Studio selects this option for you.

3. For Output, enter the URI of a location in Amazon S3. There, SageMaker AI stores the artifacts of the optimized model that your job creates.

4. (Optional) Expand Advanced options for more fine-grained control over settings such as the IAM role, VPC, and environment variables. For more information, see Advanced options below.

5. When you're finished configuring the job, choose Create job.

   Studio shows the job details page, which shows the job status and all of its settings.

To set up your project

1. In your application code, import the necessary libraries. The following example imports the SDK for Python (Boto3). It also imports the classes from the SageMaker AI Python SDK that you use to define and work with models:

   import boto3
   from sagemaker.serve.builder.model_builder import ModelBuilder
   from sagemaker.serve.builder.schema_builder import SchemaBuilder
   from sagemaker.session import Session
   from pathlib import Path

2. Initialize a SageMaker AI session. The following example uses the Session() class:

   sagemaker_session = Session()

To define your model

1. Create a SchemaBuilder instance, and provide input and output samples. You supply this instance to the ModelBuilder class when you define a model. With it, SageMaker AI automatically generates the marshalling functions for serializing and deserializing the input and output.

   For more information about using the SchemaBuilder and ModelBuilder classes, see Create a model in Amazon SageMaker AI with ModelBuilder.

   The following example provides sample input and output strings to the SchemaBuilder class:

   response = "Jupiter is the largest planet in the solar system. It is the fifth planet from the sun."
   sample_input = {
       "inputs": "What is the largest planet in the solar system?",
       "parameters": {"max_new_tokens": 128, "top_p": 0.9, "temperature": 0.6},
   }
   sample_output = [{"generated_text": response}]
   schema_builder = SchemaBuilder(sample_input, sample_output)

2. Define your model to SageMaker AI. The following example sets the parameters to initialize a ModelBuilder instance:

   model_builder = ModelBuilder(
       model="jumpstart-model-id",
       schema_builder=schema_builder,
       sagemaker_session=sagemaker_session,
       role_arn=sagemaker_session.get_caller_identity_arn(),
   )

   This example uses a JumpStart model. Replace jumpstart-model-id with the ID of a JumpStart model, such as meta-textgeneration-llama-3-70b.

   Note

   If you want to optimize with speculative decoding, and you want to use the SageMaker AI draft, you must enable network isolation. To enable it, include the following argument when you initialize a ModelBuilder instance:

   enable_network_isolation=True,

   For more information about network isolation, see Network Isolation.

To optimize with quantization

1. To run a quantization job, use the optimize() method, and set the quantization_config argument. The following example sets OPTION_QUANTIZE as an environment variable in the optimization container:

   optimized_model = model_builder.optimize(
       instance_type="instance-type",
       accept_eula=True,
       quantization_config={
           "OverrideEnvironment": {
               "OPTION_QUANTIZE": "awq",
           },
       },
       output_path="s3://output-path",
   )

   In this example, replace instance-type with an ML instance, such as ml.p4d.24xlarge. Replace s3://output-path with the path to the S3 location where you store the optimized model that the job creates.

   The optimize() method returns a Model object, which you can use to deploy your model to an endpoint.

2. When the job completes, deploy the model. The following example uses the deploy() method:

   predictor = optimized_model.deploy(
       instance_type="instance-type", 
       accept_eula=True,
   )

   In this example, replace instance-type with an ML instance, such as ml.p4d.24xlarge.

   The deploy() method returns a predictor object, which you can use to send inference requests to the endpoint that hosts the model.

To optimize with speculative decoding using the SageMaker AI draft model

When you optimize your model with speculative decoding, you can choose to use a draft model that SageMaker AI provides, or you can use your own. The following examples use the SageMaker AI draft model.

1. To run a speculative decoding job, use the optimize() method, and set the speculative_decoding_config argument. The following example sets the ModelProvider key to SAGEMAKER to use the draft model that SageMaker AI provides.

   optimized_model = model_builder.optimize(
       instance_type="instance-type",
       accept_eula=True,
       speculative_decoding_config={
           "ModelProvider": "SAGEMAKER",
       },
   )

   In this example, replace instance-type with an ML instance, such as ml.p4d.24xlarge.

   The optimize() method returns a Model object, which you can use to deploy your model to an endpoint.

2. When the job completes, deploy the model. The following example uses the deploy() method:

   predictor = optimized_model.deploy(accept_eula=True)

   The deploy() method returns a predictor object, which you can use to send inference requests to the endpoint that hosts the model.

To optimize with speculative decoding using a custom draft model

Before you can provide your custom draft model to SageMaker AI, you must first upload the model artifacts to Amazon S3.

The following examples demonstrate one possible way to provide a custom draft model. The examples download the draft model from the Hugging Face Hub, upload it to Amazon S3, and provide the S3 URI to the speculative_decoding_config argument.

1. If you want to download a model from the Hugging Face Hub, add the huggingface_hub library to your project, and download a model with the snapshot_download() method. The following example downloads a model to a local directory:

   import huggingface_hub
   
   huggingface_hub.snapshot_download(
       repo_id="model-id",
       revision="main",
       local_dir=download-dir,
       token=hf-access-token,
   )

   In this example, replace model-id with the ID of a model the Hugging Face Hub, such as meta-llama/Meta-Llama-3-8B. Replace download-dir with a local directory. Replace hf-access-token with your user access token. To learn how to get your access token, see User access tokens in the Hugging Face documentation.

   For more information about the huggingface_hub library, see Hub client library in the Hugging Face documentation.

2. To make your downloaded model available to SageMaker AI, upload it to Amazon S3. The following example uploads the model with the sagemaker_session object:

   custom_draft_model_uri = sagemaker_session.upload_data(
       path=hf_local_download_dir.as_posix(),
       bucket=sagemaker_session.default_bucket(),
       key_prefix="prefix",
   )

   In this example, replace prefix with a qualifier that helps you distinguish the draft model in S3, such as spec-dec-custom-draft-model.

   The upload_data() method returns the S3 URI for the model artifacts.

3. To run a speculative decoding job, use the optimize() method, and set the speculative_decoding_config argument. The following example sets the ModelSource key to the S3 URI of the custom draft model:

   optimized_model = model_builder.optimize(
       instance_type="instance-type",
       accept_eula=True,
       speculative_decoding_config={
           "ModelSource": custom_draft_model_uri + "/",
       },
   )

   In this example, replace instance-type with an ML instance, such as ml.p4d.24xlarge.

   The optimize() method returns a Model object, which you can use to deploy your model to an endpoint.

4. When the job completes, deploy the model. The following example uses the deploy() method:

   predictor = optimized_model.deploy(accept_eula=True)

   The deploy() method returns a predictor object, which you can use to send inference requests to the endpoint that hosts the model.

To optimize with compilation

1. To run a compilation job, use the optimize() method, and set the compilation_config argument. The following example uses the OverrideEnvironment key to set the necessary environment variables in the optimization container:

   optimized_model = model_builder.optimize(
       instance_type="instance-type",
       accept_eula=True,
       compilation_config={
           "OverrideEnvironment": {
               "OPTION_TENSOR_PARALLEL_DEGREE": "24",
               "OPTION_N_POSITIONS": "8192",
               "OPTION_DTYPE": "fp16",
               "OPTION_ROLLING_BATCH": "auto",
               "OPTION_MAX_ROLLING_BATCH_SIZE": "4",
               "OPTION_NEURON_OPTIMIZE_LEVEL": "2",
           }
       },
       output_path="s3://output-path",
   )

   In this example, set instance-type to an ML instance type with accelerated hardware. For example, for accelerated inference with AWS Inferentia, you could set the type to an Inf2 instance, such as ml.inf2.48xlarge. Replace s3://output-path with the path to the S3 location where you store the optimized model that the job creates.

2. When the job completes, deploy the model. The following example uses the deploy() method:

   predictor = optimized_model.deploy(accept_eula=True)

   The deploy() method returns a predictor object, which you can use to send inference requests to the endpoint that hosts the model.

To test your model with an inference request

• To send a test inference request to your deployed model, use the predict() method of a predictor object. The following example passes the sample_input variable that was also passed to the SchemaBuilder class in the examples to define your model:

  predictor.predict(sample_input)

  The sample input has the prompt, "What is the largest planet in the solar system?". The predict() method returns the response that the model generated, as shown by the following example:

  {'generated_text': ' Jupiter is the largest planet in the solar system. It is the fifth planet from the sun. It is a gas giant with . . .'}