Skip to content
Merged
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
8 changes: 3 additions & 5 deletions sagemaker-train/src/sagemaker/train/base_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@
from sagemaker.train.defaults import TrainDefaults
from sagemaker.train.utils import _get_unique_name

logger = logging.getLogger(__name__)


class BaseTrainer(ABC):
"""Abstract base class for all SageMaker training workflows.

Expand Down Expand Up @@ -156,7 +159,6 @@ def _fetch_full_recipe_template(self) -> Optional[Dict[str, Any]]:

Returns None if the template can't be fetched (fallback to synthetic template).
"""
logger = logging.getLogger(__name__)
frt = getattr(self.hyperparameters, '_full_recipe_template', None) if hasattr(self, 'hyperparameters') else None
if isinstance(frt, dict):
return frt
Expand Down Expand Up @@ -392,8 +394,6 @@ def _train_serverful_smtj(self, training_dataset=None, validation_dataset=None,
)
from sagemaker.train.defaults import TrainDefaults

logger = logging.getLogger(__name__)

sagemaker_session = TrainDefaults.get_sagemaker_session(
sagemaker_session=self.sagemaker_session
)
Expand Down Expand Up @@ -848,8 +848,6 @@ def _train_hyperpod(self, training_dataset=None, validation_dataset=None,
(SFT, DPO, RLVR).
"""

logger = logging.getLogger(__name__)

sagemaker_session = TrainDefaults.get_sagemaker_session(
sagemaker_session=self.sagemaker_session
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -611,6 +611,9 @@ def test_non_spec_keys_flow_into_train_hyperparameters(
overrides={"training_config": {"sequence_length": 8192}},
)

# For serverless (ie compute=None): Non-spec keys that aren't overridden don't flow to hyperparams
trainer.compute = MagicMock()

with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
patch("sagemaker.train.sft_trainer._create_input_data_config") as mock_input, \
Expand Down Expand Up @@ -640,6 +643,54 @@ def test_non_spec_keys_flow_into_train_hyperparameters(
assert "warmup_ratio" in final_hp
assert final_hp["warmup_ratio"] == "0.1"

@patch("sagemaker.train.sft_trainer._validate_eula_for_gated_model", return_value=False)
@patch("sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn")
@patch("sagemaker.train.sft_trainer._validate_and_resolve_model_package_group", return_value="my-group")
@patch("sagemaker.train.sft_trainer._resolve_model_and_name", return_value=("model_obj", "nova-lite-v2"))
def test_serverless_non_spec_keys_dont_flow_into_train_hyperparameters(
self, mock_resolve, mock_validate_group, mock_get_options, mock_eula,
mock_hyperparams_with_full_template
):
"""Non-spec keys from full template are included in final training hyperparameters."""
mock_get_options.return_value = (mock_hyperparams_with_full_template, "model-arn", False)

from sagemaker.train.sft_trainer import SFTTrainer

trainer = SFTTrainer(
model="nova-lite-v2",
model_package_group="my-group",
training_dataset="s3://bucket/train.jsonl",
overrides={"training_config": {"sequence_length": 8191}},
)

with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
patch("sagemaker.train.sft_trainer._create_input_data_config") as mock_input, \
patch("sagemaker.train.sft_trainer._convert_input_data_to_channels", return_value=[]), \
patch("sagemaker.train.sft_trainer._create_output_config", return_value=MagicMock()), \
patch("sagemaker.train.sft_trainer._create_serverless_config", return_value=MagicMock()), \
patch("sagemaker.train.sft_trainer._create_mlflow_config", return_value=None), \
patch("sagemaker.train.sft_trainer._create_model_package_config", return_value=None), \
patch("sagemaker.train.sft_trainer._validate_hyperparameter_values"), \
patch("sagemaker.train.sft_trainer._get_jumpstart_tags", return_value=[]):

mock_session = MagicMock()
mock_session.boto_session.region_name = "us-west-2"
mock_defaults.get_sagemaker_session.return_value = mock_session
mock_defaults.get_role.return_value = "arn:aws:iam::123:role/Role"
mock_tj.create.return_value = MagicMock()

trainer.train(wait=False)

call_kwargs = mock_tj.create.call_args[1]
final_hp = call_kwargs["hyper_parameters"]

# Non-spec key is now in final hyperparameters
assert "sequence_length" in final_hp
assert final_hp["sequence_length"] == "8191"
# Other full template keys also present
assert "warmup_ratio" not in final_hp

@patch("sagemaker.train.sft_trainer._validate_eula_for_gated_model", return_value=False)
@patch("sagemaker.train.sft_trainer._get_fine_tuning_options_and_model_arn")
@patch("sagemaker.train.sft_trainer._validate_and_resolve_model_package_group", return_value="my-group")
Expand All @@ -660,6 +711,9 @@ def test_nested_keys_flow_into_train_hyperparameters(
overrides={"training_config": {"lr_scheduler": {"warmup_steps": 30}}},
)

# For serverless (ie compute=None): Non-spec keys that aren't overridden don't flow to hyperparams
trainer.compute = MagicMock()

with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
patch("sagemaker.train.sft_trainer._create_input_data_config") as mock_input, \
Expand Down Expand Up @@ -725,6 +779,9 @@ def test_deeply_nested_peft_keys_flow_into_hyperparameters(
overrides={"training_config": {"peft": {"lora_tuning": {"alpha": 128}}}},
)

# For serverless (ie compute=None): Non-spec keys that aren't overridden don't flow to hyperparams
trainer.compute = MagicMock()

with patch("sagemaker.train.sft_trainer.TrainingJob") as mock_tj, \
patch("sagemaker.train.sft_trainer.TrainDefaults") as mock_defaults, \
patch("sagemaker.train.sft_trainer._create_input_data_config"), \
Expand Down
Loading