Fine-tune curated hub models

Make sure that your user's IAM role has the SageMaker AI TrainHubModel permission attached. For more information, see Adding and removing IAM identity permissions in the AWS IAM User Guide.

You should attach a policy like the following example to your user's IAM role:

Have a private curated hub with a model reference to a JumpStart model that you want to fine-tune. For more information about creating a private hub, see Create a private model hub. To learn how to add publicly available JumpStart models to your private hub, see Add models to a private hub.

Have a training dataset that you want to use for fine-tuning the model. The dataset should be in the appropriate training format for the model that you want to fine-tune.

Make sure that you have the latest version (at least 3.0.0) of the SageMaker Python SDK installed. For more information, see Use Version 3.x of the SageMaker Python SDK on the Read the Docs website.

!pip install --upgrade sagemaker

Import the AWS SDK for Python (Boto3) and the modules you need from the SageMaker Python SDK.

import boto3
from sagemaker.train import ModelTrainer
from sagemaker.train.configs import Compute, InputData
from sagemaker.core.jumpstart.configs import JumpStartConfig
from sagemaker.core.helper.session_helper import Session

Initialize a Boto3 session, a SageMaker AI client, and a SageMaker Python SDK session.

sagemaker_client = boto3.Session(region_name=<AWS-region>).client("sagemaker")
sm_session = Session(sagemaker_client=sagemaker_client)

Create a ModelTrainer using from_jumpstart_config and provide the JumpStart model ID, the name of your hub that contains the model reference, and your SageMaker Python SDK session. To list the available model IDs programmatically, use the hub.list_sagemaker_public_hub_models() method in the SageMaker Python SDK.

Optionally, you can specify the instance_type and instance_count fields when creating the ModelTrainer. If you don't, the training job uses the default instance type and count for the model you're using.

You can also optionally specify the output_path to the Amazon S3 location where you want to store the fine-tuned model weights. If you don't specify the output_path, then uses a default SageMaker AI Amazon S3 bucket for the region in your account, named with the following format: sagemaker-<region>-<account-id>.

jumpstart_config = JumpStartConfig(
    model_id="meta-textgeneration-llama-3-2-1b",
    hub_name=<your-hub-name>,
    # For gated models, set accept_eula=True to accept the end-user license agreement.
    # This example uses False. Change it to True to accept the EULA for gated models.
    accept_eula=False,
)
model_trainer = ModelTrainer.from_jumpstart_config(
    jumpstart_config=jumpstart_config,
    # Optional: specify your desired instance type and count for the training job.
    # compute=Compute(instance_type="ml.g5.2xlarge", instance_count=1),
    # Optional: specify a custom S3 location to store the fine-tuned model artifacts.
    # output_path="s3://<output-path-for-model-artifacts>",
)

Create an InputData object. Set channel_name to train and data_source to the location of your fine-tuning dataset. In the following example, replace <your-fine-tuning-dataset> with the Amazon S3 URI of your dataset. If you have additional considerations, such as using local mode or multiple training data channels, see SageMaker Train in the SageMaker Python SDK documentation on the Read the Docs website.

training_input = InputData(
    channel_name="train",
    data_source="s3://<your-fine-tuning-dataset>",
)

Call the model trainer's train() method and pass in your training data. Declare EULA acceptance on the JumpStartConfig (shown in the previous step), not in the train() call.

model_trainer.train(input_data_config=[training_input])