Create an AutoML job to fine-tune text generation models using the API

Large language models (LLMs) excel in multiple generative tasks, including text generation, summarization, completion, question answering, and more. Their performance can be attributed to their significant size and extensive training on diverse datasets and various tasks. However, specific domains, such as healthcare and financial services, may require customized fine-tuning to adapt to unique data and use cases. By tailoring their training to their particular domain, LLMs can improve their performance and provide more accurate outputs for targeted applications.

Autopilot offers the capability to fine-tune a selection of pre-trained generative text models. In particular, Autopilot supports the instruction-based fine tuning of a selection of general-purpose large language models (LLMs) powered by JumpStart.

Fine-tuning a pre-trained model requires a specific dataset of clear instructions that guide the model on how to generate output or behave for that task. The model learns from the dataset, adjusting its parameters to conform to the provided instructions. Instruction-based fine-tuning involves using labeled examples formatted as prompt-response pairs and phrased as instructions. For more information about fine-tuning, see Fine-tune a foundation model.

The following guidelines outline the process of creating an Amazon SageMaker Autopilot job as a pilot experiment to fine-tune text generation LLMs using the SageMaker AI API Reference.

To create an Autopilot experiment programmatically for fine-tuning an LLM, you can call the CreateAutoMLJobV2 API in any language supported by Amazon SageMaker Autopilot or the AWS CLI.

For information about how this API action translates into a function in the language of your choice, see the See Also section of CreateAutoMLJobV2 and choose an SDK. As an example, for Python users, see the full request syntax of create_auto_ml_job_v2 in AWS SDK for Python (Boto3).

Once your LLM is fine-tuned, you can evaluate its performance by accessing various ROUGE scores through the BestCandidate when making a DescribeAutoMLJobV2 API call. The model also provides information about its training and validation loss as well as perplexity. For a comprehensive list of metrics for evaluating the quality of the text generated by the fine-tuned models, see Metrics for fine-tuning large language models in Autopilot.

Prerequisites

Before using Autopilot to create a fine-tuning experiment in SageMaker AI, make sure to take the following steps:

Once this is done, you can use the ARN of this execution role in Autopilot API requests.

Required parameters

When calling CreateAutoMLJobV2 to create an Autopilot experiment for LLM fine-tuning, you must provide the following values:

The following is an example of the full request format used when making an API call to CreateAutoMLJobV2 for fine-tuning a (Falcon7BInstruct) model.

{
   "AutoMLJobName": "<job_name>",
   "AutoMLJobInputDataConfig": [ 
      { 
         "ChannelType": "training",
         "CompressionType": "None",
         "ContentType": "text/csv", 
         "DataSource": { 
            "S3DataSource": { 
               "S3DataType": "S3Prefix",
               "S3Uri": "s3://<bucket_name>/<input_data>.csv"
            }
         }
      }
   ],
  "OutputDataConfig": {
      "S3OutputPath": "s3://<bucket_name>/output",
      "KmsKeyId": "arn:aws:kms:<region>:<account_id>:key/<key_value>"
   },
   "RoleArn":"arn:aws:iam::<account_id>:role/<sagemaker_execution_role_name>",
   "AutoMLProblemTypeConfig": {
        "TextGenerationJobConfig": {
            "BaseModelName": "Falcon7BInstruct"
       }
   }
}

All other parameters are optional.

Optional parameters

The following sections provide details of some optional parameters that you can pass to your fine-tuning AutoML job.

"AutoMLProblemTypeConfig": {
  "TextGenerationJobConfig": {
    "BaseModelName": "Falcon7B",
    "TextGenerationHyperParameters": {"epochCount":"5", "learningRate":"0.000001", "batchSize": "32", "learningRateWarmupSteps": "10"}
  }
}