Skip to content

Commit b2f88c5

Browse files
joshuatownerJoshua Towner
andauthored
feat(train): add list_hyperparameters() for pre-trainer HP discovery (#6149)
* feat(train): add list_hyperparameters() for pre-trainer HP discovery Add a public utility function that returns available hyperparameters for a model/technique/training_type combination without requiring a fully constructed trainer object. This enables tools and scripts to discover valid hyperparameter names, defaults, and ranges before setting up training infrastructure (model package groups, datasets, roles, etc.). Motivation: COE 398545 identified that hardcoded HP names in downstream consumers break when recipe templates rename parameters. Dynamic discovery at code-generation time prevents this class of failure. Usage: from sagemaker.train import list_hyperparameters hp = list_hyperparameters('model-name', 'SFT', 'LORA') hp.get_info() # display all params hp.get_info('learning_rate') # display one param * test: use meta-textgeneration-llama-3-2-1b-instruct in integ tests Switch to the same model used by ~80% of existing integ tests to avoid deprecation risk. Llama 3.2 1B is the most battle-tested model in the repo's test infrastructure. --------- Co-authored-by: Joshua Towner <josh@tonwer.xyz>
1 parent d3801bb commit b2f88c5

4 files changed

Lines changed: 225 additions & 0 deletions

File tree

‎sagemaker-train/src/sagemaker/train/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,4 +119,7 @@ def __getattr__(name):
119119
elif name == "HyperPodCompute":
120120
from sagemaker.core.training.configs import HyperPodCompute
121121
return HyperPodCompute
122+
elif name == "list_hyperparameters":
123+
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
124+
return list_hyperparameters
122125
raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

‎sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1880,3 +1880,57 @@ def extract_image_from_hyperpod_template(template_content: str) -> Optional[str]
18801880
if image_match:
18811881
return image_match.group(1).strip()
18821882
return None
1883+
1884+
1885+
def list_hyperparameters(
1886+
model: str,
1887+
technique: Union[str, CustomizationTechnique] = "SFT",
1888+
training_type: Union[str, TrainingType] = "LORA",
1889+
hub_name: Optional[str] = None,
1890+
sagemaker_session: Optional[Session] = None,
1891+
) -> FineTuningOptions:
1892+
"""List available hyperparameters for a model and fine-tuning technique.
1893+
1894+
Returns a FineTuningOptions object containing all tunable parameters with
1895+
their defaults, types, and valid ranges, without requiring a fully
1896+
constructed trainer.
1897+
1898+
Args:
1899+
model: SageMakerHub model name (e.g. "huggingface-llm-qwen2-5-7b-instruct").
1900+
technique: Customization technique. One of "SFT", "DPO", "RLVR",
1901+
"RLAIF", "CPT", or a CustomizationTechnique enum value.
1902+
training_type: Training type. One of "LORA", "FULL", or a
1903+
TrainingType enum value.
1904+
hub_name: Hub to query. Defaults to "SageMakerPublicHub".
1905+
sagemaker_session: Optional SageMaker session. If not provided,
1906+
a default session is created.
1907+
1908+
Returns:
1909+
FineTuningOptions: Object with .get_info() for display and attribute
1910+
access for programmatic use.
1911+
1912+
Example:
1913+
>>> from sagemaker.train import list_hyperparameters
1914+
>>> hp = list_hyperparameters("huggingface-llm-qwen2-5-7b-instruct", "SFT", "LORA")
1915+
>>> hp.get_info() # Display all parameters with defaults and ranges
1916+
>>> hp.get_info("learning_rate") # Display info for a single parameter
1917+
"""
1918+
technique_val = (
1919+
technique.value if isinstance(technique, CustomizationTechnique) else technique
1920+
)
1921+
training_type_val = (
1922+
training_type if isinstance(training_type, str) else training_type.value
1923+
)
1924+
1925+
session = sagemaker_session or TrainDefaults.get_sagemaker_session(
1926+
sagemaker_session=None
1927+
)
1928+
1929+
options, _, _ = _get_fine_tuning_options_and_model_arn(
1930+
model_name=model,
1931+
customization_technique=technique_val,
1932+
training_type=TrainingType(training_type_val),
1933+
sagemaker_session=session,
1934+
hub_name=hub_name,
1935+
)
1936+
return options
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Integration tests for list_hyperparameters utility."""
14+
from __future__ import absolute_import
15+
16+
import pytest
17+
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
18+
from sagemaker.train.common import FineTuningOptions, TrainingType, CustomizationTechnique
19+
20+
21+
class TestListHyperparametersInteg:
22+
"""Integration tests for list_hyperparameters against live SageMakerPublicHub."""
23+
24+
def test_sft_lora_returns_expected_params(self):
25+
"""SFT LORA returns a FineTuningOptions with known hyperparameters."""
26+
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "SFT", "LORA")
27+
28+
assert isinstance(hp, FineTuningOptions)
29+
assert "learning_rate" in hp._specs
30+
assert "global_batch_size" in hp._specs
31+
assert "lora_rank" in hp._specs
32+
assert hp._specs["learning_rate"]["type"] == "float"
33+
34+
def test_dpo_lora_has_additional_params(self):
35+
"""DPO LORA returns params including adam_beta (not present in SFT)."""
36+
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "DPO", "LORA")
37+
38+
assert isinstance(hp, FineTuningOptions)
39+
assert "adam_beta" in hp._specs
40+
assert "learning_rate" in hp._specs
41+
42+
def test_rlvr_lora_returns_params(self):
43+
"""RLVR LORA returns FineTuningOptions with RL-specific params."""
44+
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "RLVR", "LORA")
45+
46+
assert isinstance(hp, FineTuningOptions)
47+
assert "learning_rate" in hp._specs
48+
assert len(hp._specs) > 10
49+
50+
def test_accepts_enum_values(self):
51+
"""Accepts CustomizationTechnique and TrainingType enums."""
52+
hp = list_hyperparameters(
53+
"meta-textgeneration-llama-3-2-1b-instruct",
54+
CustomizationTechnique.SFT,
55+
TrainingType.LORA,
56+
)
57+
58+
assert isinstance(hp, FineTuningOptions)
59+
assert "learning_rate" in hp._specs
60+
61+
def test_get_info_does_not_raise(self):
62+
"""get_info() runs without error on returned object."""
63+
hp = list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "SFT", "LORA")
64+
# Should print without raising
65+
hp.get_info("learning_rate")
66+
67+
def test_invalid_model_raises(self):
68+
"""Non-existent model raises an error."""
69+
with pytest.raises(Exception):
70+
list_hyperparameters("nonexistent-model-xyz-123", "SFT", "LORA")
71+
72+
def test_invalid_technique_raises(self):
73+
"""Technique not available for model raises an error."""
74+
# PPO is not available on Llama 3.2
75+
with pytest.raises(Exception):
76+
list_hyperparameters("meta-textgeneration-llama-3-2-1b-instruct", "PPO", "LORA")

‎sagemaker-train/tests/unit/train/common_utils/test_finetune_utils.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1604,3 +1604,95 @@ def test_returns_none_when_enum_empty_list(self, mock_spec):
16041604
def test_returns_none_on_exception(self, mock_spec):
16051605
mock_spec.side_effect = RuntimeError("hub content unavailable")
16061606
assert self._call() is None
1607+
1608+
1609+
class TestListHyperparameters:
1610+
"""Tests for the list_hyperparameters public API."""
1611+
1612+
@patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
1613+
@patch('boto3.client')
1614+
def test_list_hyperparameters_returns_finetuning_options(self, mock_boto_client, mock_get_hub_content):
1615+
"""list_hyperparameters returns a FineTuningOptions object with correct params."""
1616+
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
1617+
from sagemaker.train.common import FineTuningOptions
1618+
1619+
mock_get_hub_content.return_value = {
1620+
'hub_content_arn': "arn:aws:sagemaker:us-west-2:123456789012:model/test-model",
1621+
'hub_content_document': {
1622+
"GatedBucket": False,
1623+
"RecipeCollection": [
1624+
{
1625+
"CustomizationTechnique": "SFT",
1626+
"SmtjRecipeTemplateS3Uri": "s3://bucket/template.json",
1627+
"SmtjOverrideParamsS3Uri": "s3://bucket/params.json",
1628+
"Peft": "LORA"
1629+
}
1630+
]
1631+
}
1632+
}
1633+
1634+
mock_s3_client = Mock()
1635+
mock_boto_client.return_value = mock_s3_client
1636+
mock_s3_client.get_object.return_value = {
1637+
"Body": Mock(read=Mock(return_value=json.dumps({
1638+
"learning_rate": {"type": "float", "default": 0.0001, "min": 5e-7, "max": 0.001, "required": True},
1639+
"global_batch_size": {"type": "integer", "default": 8, "required": True},
1640+
"max_epochs": {"type": "integer", "default": 5, "min": 1, "max": 100, "required": True},
1641+
}).encode()))
1642+
}
1643+
1644+
mock_session = Mock()
1645+
mock_session.boto_session.region_name = "us-west-2"
1646+
mock_session.boto_session.client.return_value = mock_s3_client
1647+
1648+
with patch('sagemaker.train.common_utils.finetune_utils.TrainDefaults.get_sagemaker_session', return_value=mock_session):
1649+
result = list_hyperparameters("test-model", "SFT", "LORA", sagemaker_session=mock_session)
1650+
1651+
assert isinstance(result, FineTuningOptions)
1652+
assert result.learning_rate == 0.0001
1653+
assert result.global_batch_size == 8
1654+
assert result.max_epochs == 5
1655+
1656+
@patch('sagemaker.train.common_utils.finetune_utils._get_hub_content_metadata')
1657+
@patch('boto3.client')
1658+
def test_list_hyperparameters_accepts_enum_values(self, mock_boto_client, mock_get_hub_content):
1659+
"""list_hyperparameters accepts both string and enum values for technique/training_type."""
1660+
from sagemaker.train.common_utils.finetune_utils import list_hyperparameters
1661+
from sagemaker.train.common import CustomizationTechnique, TrainingType
1662+
1663+
mock_get_hub_content.return_value = {
1664+
'hub_content_arn': "arn:aws:sagemaker:us-west-2:123456789012:model/test-model",
1665+
'hub_content_document': {
1666+
"GatedBucket": False,
1667+
"RecipeCollection": [
1668+
{
1669+
"CustomizationTechnique": "DPO",
1670+
"SmtjRecipeTemplateS3Uri": "s3://bucket/template.json",
1671+
"SmtjOverrideParamsS3Uri": "s3://bucket/params.json",
1672+
"Peft": "LORA"
1673+
}
1674+
]
1675+
}
1676+
}
1677+
1678+
mock_s3_client = Mock()
1679+
mock_boto_client.return_value = mock_s3_client
1680+
mock_s3_client.get_object.return_value = {
1681+
"Body": Mock(read=Mock(return_value=json.dumps({
1682+
"learning_rate": {"type": "float", "default": 0.0001, "required": True},
1683+
}).encode()))
1684+
}
1685+
1686+
mock_session = Mock()
1687+
mock_session.boto_session.region_name = "us-west-2"
1688+
mock_session.boto_session.client.return_value = mock_s3_client
1689+
1690+
with patch('sagemaker.train.common_utils.finetune_utils.TrainDefaults.get_sagemaker_session', return_value=mock_session):
1691+
result = list_hyperparameters(
1692+
"test-model",
1693+
CustomizationTechnique.DPO,
1694+
TrainingType.LORA,
1695+
sagemaker_session=mock_session,
1696+
)
1697+
1698+
assert result.learning_rate == 0.0001

0 commit comments

Comments
 (0)