Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions sagemaker-train/src/sagemaker/train/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,4 +119,7 @@ def __getattr__(name):
elif name == "HyperPodCompute":
from sagemaker.core.training.configs import HyperPodCompute
return HyperPodCompute
elif name == "list_hyperparameters":
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
return list_hyperparameters
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
54 changes: 54 additions & 0 deletions sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1734,3 +1734,57 @@ def extract_image_from_hyperpod_template(template_content: str) -> Optional[str]
if image_match:
return image_match.group(1).strip()
return None


def list_hyperparameters(
model: str,
technique: Union[str, CustomizationTechnique] = "SFT",
training_type: Union[str, TrainingType] = "LORA",
hub_name: Optional[str] = None,
sagemaker_session: Optional[Session] = None,
) -> FineTuningOptions:
"""List available hyperparameters for a model and fine-tuning technique.

Returns a FineTuningOptions object containing all tunable parameters with
their defaults, types, and valid ranges, without requiring a fully
constructed trainer.

Args:
model: SageMakerHub model name (e.g. "huggingface-llm-qwen2-5-7b-instruct").
technique: Customization technique. One of "SFT", "DPO", "RLVR",
"RLAIF", "CPT", or a CustomizationTechnique enum value.
training_type: Training type. One of "LORA", "FULL", or a
TrainingType enum value.
hub_name: Hub to query. Defaults to "SageMakerPublicHub".
sagemaker_session: Optional SageMaker session. If not provided,
a default session is created.

Returns:
FineTuningOptions: Object with .get_info() for display and attribute
access for programmatic use.

Example:
>>> from sagemaker.train import list_hyperparameters
>>> hp = list_hyperparameters("huggingface-llm-qwen2-5-7b-instruct", "SFT", "LORA")
>>> hp.get_info() # Display all parameters with defaults and ranges
>>> hp.get_info("learning_rate") # Display info for a single parameter
"""
technique_val = (
technique.value if isinstance(technique, CustomizationTechnique) else technique
)
training_type_val = (
training_type if isinstance(training_type, str) else training_type.value
)

session = sagemaker_session or TrainDefaults.get_sagemaker_session(
sagemaker_session=None
)

options, _, _ = _get_fine_tuning_options_and_model_arn(
model_name=model,
customization_technique=technique_val,
training_type=TrainingType(training_type_val),
sagemaker_session=session,
hub_name=hub_name,
)
return options
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Integration tests for list_hyperparameters utility."""
from __future__ import absolute_import

import pytest
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
from sagemaker.train.common import FineTuningOptions, TrainingType, CustomizationTechnique


class TestListHyperparametersInteg:
"""Integration tests for list_hyperparameters against live SageMakerPublicHub."""

def test_sft_lora_returns_expected_params(self):
"""SFT LORA returns a FineTuningOptions with known hyperparameters."""
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "SFT", "LORA")

assert isinstance(hp, FineTuningOptions)
assert "learning_rate" in hp._specs
assert "global_batch_size" in hp._specs
assert "lora_rank" in hp._specs
assert hp._specs["learning_rate"]["type"] == "float"

def test_dpo_lora_has_additional_params(self):
"""DPO LORA returns params including adam_beta (not present in SFT)."""
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "DPO", "LORA")

assert isinstance(hp, FineTuningOptions)
assert "adam_beta" in hp._specs
assert "learning_rate" in hp._specs

def test_rlvr_lora_returns_params(self):
"""RLVR LORA returns FineTuningOptions with RL-specific params."""
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "RLVR", "LORA")

assert isinstance(hp, FineTuningOptions)
assert "learning_rate" in hp._specs
assert len(hp._specs) > 10

def test_accepts_enum_values(self):
"""Accepts CustomizationTechnique and TrainingType enums."""
hp = list_hyperparameters(
"meta-textgeneration-llama-3-2-1b-instruct",
CustomizationTechnique.SFT,
TrainingType.LORA,
)

assert isinstance(hp, FineTuningOptions)
assert "learning_rate" in hp._specs

def test_get_info_does_not_raise(self):
"""get_info() runs without error on returned object."""
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "SFT", "LORA")
# Should print without raising
hp.get_info("learning_rate")

def test_invalid_model_raises(self):
"""Non-existent model raises an error."""
with pytest.raises(Exception):
list_hyperparameters("nonexistent-model-xyz-123", "SFT", "LORA")

def test_invalid_technique_raises(self):
"""Technique not available for model raises an error."""
# PPO is not available on Llama 3.2
with pytest.raises(Exception):
list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "PPO", "LORA")
Original file line number Diff line number Diff line change
Expand Up @@ -1362,3 +1362,95 @@ def test_uses_shared_regex_from_reward_verifier(self):
# Both call sites must share the same compiled pattern, not copies.
from sagemaker.train.common_utils import rlvr_reward_verifier
assert fu.LAMBDA_ARN_REGEX is rlvr_reward_verifier.LAMBDA_ARN_REGEX


class TestListHyperparameters:
"""Tests for the list_hyperparameters public API."""

@patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
@patch('boto3.client')
def test_list_hyperparameters_returns_finetuning_options(self, mock_boto_client, mock_get_hub_content):
"""list_hyperparameters returns a FineTuningOptions object with correct params."""
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
from sagemaker.train.common import FineTuningOptions

mock_get_hub_content.return_value = {
'hub_content_arn': "arn:aws:sagemaker:us-west-2:123456789012:model/test-model",
'hub_content_document': {
"GatedBucket": False,
"RecipeCollection": [
{
"CustomizationTechnique": "SFT",
"SmtjRecipeTemplateS3Uri": "s3://bucket/template.json",
"SmtjOverrideParamsS3Uri": "s3://bucket/params.json",
"Peft": "LORA"
}
]
}
}

mock_s3_client = Mock()
mock_boto_client.return_value = mock_s3_client
mock_s3_client.get_object.return_value = {
"Body": Mock(read=Mock(return_value=json.dumps({
"learning_rate": {"type": "float", "default": 0.0001, "min": 5e-7, "max": 0.001, "required": True},
"global_batch_size": {"type": "integer", "default": 8, "required": True},
"max_epochs": {"type": "integer", "default": 5, "min": 1, "max": 100, "required": True},
}).encode()))
}

mock_session = Mock()
mock_session.boto_session.region_name = "us-west-2"
mock_session.boto_session.client.return_value = mock_s3_client

with patch('sagemaker.train.common_utils.finetune_utils.TrainDefaults.get_sagemaker_session', return_value=mock_session):
result = list_hyperparameters("test-model", "SFT", "LORA", sagemaker_session=mock_session)

assert isinstance(result, FineTuningOptions)
assert result.learning_rate == 0.0001
assert result.global_batch_size == 8
assert result.max_epochs == 5

@patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
@patch('boto3.client')
def test_list_hyperparameters_accepts_enum_values(self, mock_boto_client, mock_get_hub_content):
"""list_hyperparameters accepts both string and enum values for technique/training_type."""
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
from sagemaker.train.common import CustomizationTechnique, TrainingType

mock_get_hub_content.return_value = {
'hub_content_arn': "arn:aws:sagemaker:us-west-2:123456789012:model/test-model",
'hub_content_document': {
"GatedBucket": False,
"RecipeCollection": [
{
"CustomizationTechnique": "DPO",
"SmtjRecipeTemplateS3Uri": "s3://bucket/template.json",
"SmtjOverrideParamsS3Uri": "s3://bucket/params.json",
"Peft": "LORA"
}
]
}
}

mock_s3_client = Mock()
mock_boto_client.return_value = mock_s3_client
mock_s3_client.get_object.return_value = {
"Body": Mock(read=Mock(return_value=json.dumps({
"learning_rate": {"type": "float", "default": 0.0001, "required": True},
}).encode()))
}

mock_session = Mock()
mock_session.boto_session.region_name = "us-west-2"
mock_session.boto_session.client.return_value = mock_s3_client

with patch('sagemaker.train.common_utils.finetune_utils.TrainDefaults.get_sagemaker_session', return_value=mock_session):
result = list_hyperparameters(
"test-model",
CustomizationTechnique.DPO,
TrainingType.LORA,
sagemaker_session=mock_session,
)

assert result.learning_rate == 0.0001
Loading