Categories: FAANG

Amazon SageMaker JumpStart adds fine-tuning support for models in a private model hub

ML 18118 image001

Amazon SageMaker JumpStart is a machine learning (ML) hub that provides pre-trained models, solution templates, and algorithms to help developers quickly get started with machine learning. Within SageMaker JumpStart, the private model hub feature allows organizations to create their own internal repository of ML models, enabling teams to share and manage models securely within their organization.

Today, we are announcing an enhanced private hub feature with several new capabilities that give organizations greater control over their ML assets. These enhancements include the ability to fine-tune SageMaker JumpStart models directly within the private hub, support for adding and managing custom-trained models, deep linking capabilities for associated notebooks, and improved model version management. These new features streamline the ML workflow by combining the convenience of pre-built solutions with the flexibility of custom development, while maintaining enterprise-grade security and governance.

For enterprise customers, the ability to curate and fine-tune both pre-built and custom models is crucial for successful AI implementation. Model curation provides quality control, compliance, and security while preventing duplicate efforts across teams. When enterprises fine-tune curated models, they can specialize general-purpose solutions for their specific industry needs and gain competitive advantages through improved performance on their proprietary data. Similarly, the ability to fine-tune custom models enables organizations to continuously improve their AI solutions, adapt to changing business conditions, and preserve institutional knowledge, while maintaining cost-efficiency.

A common enterprise scenario involves centralized data science teams developing foundation models (FMs), evaluating the performance against open source FMs, and iterating on performance. After they develop their custom FM, it can serve as a baseline for the entire organization, and individual departments—such as legal, finance, or customer service—can fine-tune these models using their department-specific data that might be subject to different privacy requirements or access controls. This hub-and-spoke approach to model development maximizes resource efficiency while allowing for specialized optimization at the department level. This comprehensive approach to model management, now supported by the enhanced private hub features in SageMaker JumpStart, enables enterprises to balance standardization with customization while maintaining proper governance and control over their ML assets.

Solution overview

SageMaker JumpStart has introduced several new enhancements to its private model hub feature, allowing administrators greater control and flexibility in managing their organization’s ML models. These enhancements include:

Fine-tuning of models referenced in the private hub – Administrators can now add models from the SageMaker JumpStart catalog to their private hub and fine-tune them using Amazon SageMaker training jobs, without having to create the models from scratch.
Support for custom models – In addition to the pre-trained SageMaker JumpStart models, administrators can now add their own custom-trained models to the private hub and fine-tune them as needed.
Deep linking of notebooks – Administrators can now deep link to specific notebooks associated with the models in the private hub, making it straightforward for users to access and work with the models.
Updating models in the private hub – The private hub now supports updating models over time as new versions or iterations become available, allowing organizations to stay current with the latest model improvements.

These new capabilities give AWS customers more control over their ML infrastructure and enable faster model deployment and experimentation, while still maintaining the appropriate access controls and permissions within their organization.

In the following sections, we provide guidance on how to use these new private model hub features using the Amazon SageMaker SDK and Amazon SageMaker Studio console.

To learn more about how to manage models using private hubs, see Manage Amazon SageMaker JumpStart foundation model access with private hubs.

Prerequisites

To use the SageMaker Python SDK and run the code associated with this post, you need the following prerequisites:

An AWS account that contains your AWS resources
An AWS Identity and Access Management (IAM) role with access to SageMaker Studio notebooks
SageMaker JumpStart enabled in a SageMaker Studio domain

Create a private hub, curate models, and configure access control

This section provides a step-by-step guide for administrators to create a private hub, curate models, and configure access control for your organization’s users.

Because the feature has been integrated in the latest SageMaker Python SDK, to use the model granular access control feature with a private hub, let’s first update the SageMaker Python SDK:
```
!pip3 install sagemaker —force-reinstall —quiet
```

Next, import the SageMaker and Boto3 libraries:

import boto3 from sagemaker
import Session from sagemaker.session
import Hub

Configure your private hub:

HUB_NAME="CompanyHub"
HUB_DISPLAY_NAME="Allowlisted Models"
HUB_DESCRIPTION="These are allowlisted models taken from the SageMaker Public Hub"
REGION="<your_region_name>" # for example, "us-west-2"

In the preceding code, HUB_NAME specifies the name of your hub. HUB_DISPLAY_NAME is the display name for your hub that will be shown to users in UI experiences. HUB_DESCRIPTION is the description for your hub that will be shown to users.

Use an AWS Region where SageMaker JumpStart is available, as of March 2025: us-west-2, us-east-1, us-east-2, eu-west-1, eu-central-1, eu-central-2, eu-north-1, eu-south-2, me-south-1, me-central-1, ap-south-1, ap-south-2, eu-west-3, af-south-1, sa-east-1, ap-east-1, ap-northeast-2, ap-northeast-3, ap-southeast-3, ap-southeast-4, ap-southeast-5, ap-southeast-7, eu-west-2, eu-south-1, ap-northeast-1, us-west-1, ap-southeast-1, ap-southeast-2, ca-central-1, ca-west-1, cn-north-1, cn-northwest-1, il-central-1, mx-central-1, us-gov-east-1, us-gov-west-1.

Set up a Boto3 client for SageMaker:

sm_client = boto3.client('sagemaker')
session = Session(sagemaker_client=sm_client)
session.get_caller_identity_arn()

Check if the following policies have been already added to your admin IAM role; if not, you can add them as inline policies (use the Region configured in Step 3):

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "s3:ListBucket",
                "s3:GetObject",
                "s3:GetObjectTagging"
            ],
            "Resource": [
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>",
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>/*"
            ],
            "Effect": "Allow"
        }
    ]
}

In addition to setting up IAM permissions to the admin role, you need to scope down permissions for your users so they can’t access public contents.

Use the following policy to deny access to the public hub for your users. These can be added as inline policies in the user’s IAM role (use the Region configured in Step 3):

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": "s3:*",
            "Effect": "Deny",
            "Resource": [
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>",
                "arn:aws:s3:::jumpstart-cache-prod-<REGION>/*"
            ],
            "Condition": {
                "StringNotLike": {"s3:prefix": ["*.ipynb", "*/eula.txt"]}
            }
        },
        {
            "Action": "sagemaker:*",
            "Effect": "Deny",
            "Resource": [
                "arn:aws:sagemaker:<REGION>:aws:hub/SageMakerPublicHub",
                "arn:aws:sagemaker:<REGION>:aws:hub-content/SageMakerPublicHub/*/*"
            ]
        }
    ]
}

After you have set up the private hub configuration and permissions, you’re ready to create the private hub.

Use the following code to create the private hub within your AWS account in the Region you specified earlier:

hub = Hub(hub_name=HUB_NAME, sagemaker_session=session)

try:
  hub.create(
      description=HUB_DESCRIPTION,
      display_name=HUB_DISPLAY_NAME
  )
  print(f"Successfully created Hub with name {HUB_NAME} in {REGION}")
except Exception as e:
  if "ResourceInUse" in str(e):
    print(f"A hub with the name {HUB_NAME} already exists in your account.")
  else:
    raise e

Use describe() to verify the configuration of your hub. After your private hub is set up, you can add a reference to models from the SageMaker JumpStart public hub to your private hub. No model artifacts need to be managed by the customer. The SageMaker team will manage version or security updates. For a list of available models, refer to Built-in Algorithms with pre-trained Model Table.

To search programmatically, run the following command:

from sagemaker.jumpstart.filters import Or

filter_value = Or(
"framework == meta",
"framework == deepseek"
)
models = []
next_token = None

while True:
    response = hub.list_sagemaker_public_hub_models(
        filter=filter_value,
        next_token=next_token
    )
    models.extend(response["hub_content_summaries"])
    next_token = response.get("next_token")
    
    if not next_token:
        break
print(models)

The filter argument is optional. For a list of filters you can apply, refer to the following GitHub repo.

Use the retrieved models from the preceding command to create model references for your private hub:

for model in models:
    print(f"Adding {model.get('hub_content_name')} to Hub")
    hub.create_model_reference(model_arn=model.get("hub_content_arn"), 
                               model_name=model.get("hub_content_name"))

The SageMaker JumpStart private hub offers other useful features for managing and interacting with the curated models. Administrators can check the metadata of a specific model using the hub.describe_model(model_name=<model_name>) command. To list the available models in the private hub, you can use a simple loop:

response = hub.list_models()
models = response["hub_content_summaries"]
while response["next_token"]:
    response = hub.list_models(next_token=response["next_token"])
    models.extend(response["hub_content_summaries"])

for model in models:
    print(model.get('HubContentArn'))

If you need to remove a specific model reference from the private hub, use the following command:

hub.delete_model_reference("<model_name>")

If you want to delete the private hub from your account and Region, you will need to delete all the HubContents first, then delete the private hub. Use the following code:

for model in models:
    hub.delete_model_reference(model_name=model.get('HubContentName'))
    
hub.delete()

Fine-tune models referenced in the private hub

This section walks through how to interact with allowlisted models in SageMaker JumpStart. We demonstrate how to list available models, identify a model from the public hub, and fine-tune the model using the SageMaker Python SDK as well as the SageMaker Studio UI.

User experience using the SageMaker Python SDK

To interact with your models using the SageMaker Python SDK, complete the following steps:

Just like the admin process, the first step is to force reinstall the SageMaker Python SDK:
```
!pip3 install sagemaker —force-reinstall —quiet
```

When interacting with the SageMaker SDK functions, add references to the hub_arn:

model_id="meta-vlm-llama-3-2-11b-vision"
model_version="2.1.8"
hub_arn="<YourHubARN>"

from sagemaker import hyperparameters

my_hyperparameters = hyperparameters.retrieve_default(
    model_id=model_id, model_version=model_version, hub_arn=hub_arn
)
print(my_hyperparameters)
hyperparameters.validate(
    model_id=model_id, model_version=model_version, hyperparameters=my_hyperparameters, hub_arn=hub_arn
)

You can then start a training job by specifying the model ID, version, and hub name:

from sagemaker.jumpstart.estimator import JumpStartEstimator

estimator = JumpStartEstimator(
    model_id=model_id,
    hub_name=hub_arn,
    model_version=model_version,
    environment={"accept_eula": "false"},  # Please change {"accept_eula": "true"}
    disable_output_compression=True,
    instance_type="ml.p4d.24xlarge",
    hyperparameters=my_hyperparameters,
)
estimator.fit({"training": train_data_location})

For a custom model, see the example notebooks in GitHub.

User experience in SageMaker Studio

Complete the following steps to interact with allowlisted models using SageMaker Studio:

On the SageMaker Studio console, choose JumpStart in the navigation pane or in the Prebuilt and automated solutions section.
Choose one of model hubs you have access to.

If the user has access to multiple hubs, you will see a list of hubs, as shown in the following screenshot.

If the user has access to only one hub, you will be redirected to the model list.

To fine-tune a model, choose Train (this option will be enabled if it’s supported).
Modify your training job configurations like training data, instance type, and hyperparameters, and choose Submit.

Deep link notebooks in the private hub

You can now also access the notebook associated with the model in your curated hub.

Choose your model, then choose Preview notebooks.
Choose Open in JupyterLab to start the deep link workflow.
Select a running JupyterLab space and choose Open notebook.

You will need to upgrade your space to use a SageMaker distribution of at least 2.4.1. For more information on how to upgrade your SageMaker distribution, see Update the SageMaker Distribution Image.

This will automatically open the selected notebook in your JupyterLab instance, with your private HubName inputted into the necessary classes.

Update models in the private hub

Modify your existing private HubContent by calling the new sagemaker:UpdateHubContent API. You can now update an existing HubContent version in-place without needing to delete and re-add it. We don’t support updating the HubContentDocument at this time because there can be backward-incompatible changes that are introduced that fundamentally alter the performance and usage of the model itself. Refer to the public API documentation for more details.

client.update_hub_content(
    hub_content_name="my-model",
    hub_content_version="1.0.0",
    hub_content_type="Model",
    hub_name="my-hub",
    support_status="DEPRECATED"
)

Additionally, you can modify your ModelReferences by calling the new sagemaker:UpdateHubContentReference API. Refer to the public API documentation for more usage details.

client.update_hub_content_reference(
    hub_content_name="your-model",
    hub_content_type="ModelReference",
    hub_name="my-hub",
    min_version="1.2.0"
)

Conclusion

This post demonstrated the new enhancements to the SageMaker JumpStart private model hub feature, which gives enterprise customers greater control and flexibility in managing their ML assets. The key capabilities introduced include the ability to fine-tune pre-built SageMaker JumpStart models directly within the private hub, support for importing and fine-tuning custom-trained models, deep linking to associated notebooks for streamlined access and collaboration, and improved model version management through APIs. These features enable enterprises to curate a centralized repository of trusted, specialized ML models, while still providing the flexibility for individual teams and departments to fine-tune and adapt these models to their specific needs. The seamless integration with SageMaker Studio further streamlines the model development and deployment workflow, empowering enterprises to accelerate their ML initiatives while maintaining the appropriate security and control over their ML assets.

Now that you’ve seen how the enhanced private model hub features in Amazon SageMaker JumpStart can give your organization greater control and flexibility over managing your machine learning assets, start leveraging these capabilities to curate a centralized repository of trusted models and accelerate your AI initiatives.

About the Authors

Marc Karp is an ML Architect with the Amazon SageMaker Service team. He focuses on helping customers design, deploy, and manage ML workloads at scale. In his spare time, he enjoys traveling and exploring new places.

Niris Okram is a senior academic research specialist solutions architect at AWS. He has extensive experience working with public, private and research customers on various fields related to cloud. He is passionate about designing and building systems to accelerate the customer’s mission on AWS cloud.

Benjamin Crabtree is a software engineer with the Amazon SageMaker and Bedrock teams. He is passionate about democratizing the new and frequent breakthroughs in AI. Ben received his undergraduate degree from the University of Michigan and now lives in Brooklyn, NY.

Banu Nagasundaram leads product, engineering, and strategic partnerships for SageMaker JumpStart, SageMaker’s machine learning and GenAI hub. She is passionate about building solutions that help customers accelerate their AI journey and unlock business value.

AI Generated Robotic Content