From 84ee26b932235f8b92156177eb6cda43202d2560 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Fri, 10 Jul 2026 15:14:32 -0700 Subject: [PATCH] fix(train): send bounded HyperParameters override map for serverless customization jobs Serverless (SMTJ) model-customization jobs began failing at CreateTrainingJob with a ValidationException ("hyperParameters ... Member must have length less than or equal to 100") after the Nova release added recipe/overrides support. _apply_recipe_to_hyperparameters flattened the entire resolved Hub recipe (200+ leaf keys: KL config, vLLM params, LoRA settings, etc.) into the hyper_parameters map sent to the API. The serverless HyperParameters map is an override map capped at 100 entries, so the full flatten blew past the limit. Fix: for the serverless path (new serverless=True flag), emit a bounded override map containing only Hub spec-allowlisted keys plus leaves the user explicitly changed (via overrides=, a recipe file, or .hyperparameters.*). Unchanged recipe defaults are dropped since the base recipe is applied server-side. Explicit user overrides of non-spec keys are still forwarded so they are never silently dropped. Applied to RLVR/SFT/DPO/MTRL trainers; the serverful path keeps full-recipe flattening (it renders into the recipe YAML). Refs P467902218. --- .../src/sagemaker/train/base_trainer.py | 84 +++++++++++++++-- .../src/sagemaker/train/dpo_trainer.py | 8 +- .../sagemaker/train/multi_turn_rl_trainer.py | 10 +- .../src/sagemaker/train/rlvr_trainer.py | 8 +- .../src/sagemaker/train/sft_trainer.py | 8 +- .../train/test_serverful_recipe_validation.py | 94 +++++++++++++++++++ .../train/test_trainer_recipe_integration.py | 46 +++++---- 7 files changed, 225 insertions(+), 33 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index a453720f53..ff813f8482 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -236,24 +236,45 @@ def _patch_resolved_recipe(self, resolved: Dict[str, Any]) -> None: if dotpath: _set_nested_value(resolved, dotpath, value) - def _apply_recipe_to_hyperparameters(self, final_hyperparameters: Dict[str, Any]) -> Dict[str, Any]: + def _apply_recipe_to_hyperparameters( + self, + final_hyperparameters: Dict[str, Any], + serverless: bool = False, + ) -> Dict[str, Any]: """Apply resolved recipe values to final_hyperparameters dict. If recipe/overrides were provided, or if the user set hyperparameters directly via ``.hyperparameters.*``, merges resolved recipe values into - the hyperparameters dict. All leaf values from the resolved recipe are - applied — including keys not in the Hub spec subset — enabling - power users to override any parameter in the full recipe. - Values are converted to strings (matching the SageMaker API - expectation for hyperparameter values). + the hyperparameters dict. Values are converted to strings (matching the + SageMaker API expectation for hyperparameter values). + + By default (serverful SMTJ) every leaf value from the resolved recipe is + applied — including keys not in the Hub spec subset — because those + flattened values are rendered back into the recipe YAML that is handed to + ``ModelTrainer.from_recipe``. + + When ``serverless`` is True the result becomes an *override map*: the Hub + spec-allowlisted keys plus any leaf the user explicitly changed (via the + overrides dict, a user recipe file, or direct ``.hyperparameters.*`` + assignments). The recipe's unchanged non-spec defaults are dropped + because the base recipe is applied server-side. This is required because + ``CreateTrainingJob`` caps ``HyperParameters`` at 100 entries; flattening + the entire resolved recipe (200+ leaf keys such as KL/vLLM/LoRA settings) + blew past that limit and failed with a ``ValidationException`` — see + P467902218. Explicit user overrides are still forwarded so they are never + silently dropped. Args: final_hyperparameters: The hyperparameters dict from to_dict(). + serverless: When True, emit a bounded override map (spec keys + user + deltas) instead of the fully flattened recipe, to stay within the + API's 100-entry HyperParameters limit. Returns: The updated hyperparameters dict with recipe values applied. """ - if not hasattr(self, 'hyperparameters') or not isinstance(getattr(self.hyperparameters, '_specs', None), dict): + specs = getattr(self.hyperparameters, '_specs', None) if hasattr(self, 'hyperparameters') else None + if not isinstance(specs, dict): return final_hyperparameters try: @@ -261,13 +282,58 @@ def _apply_recipe_to_hyperparameters(self, final_hyperparameters: Dict[str, Any] except NoRecipeError: return final_hyperparameters + # For serverless, only spec-allowlisted keys and leaves the user + # explicitly changed are forwarded; everything else is applied + # server-side from the base recipe. + allowed = None + if serverless: + allowed = set(specs.keys()) | self._user_changed_recipe_keys() + flat = flatten_resolved_recipe(resolved) for k, v in flat.items(): - if v is not None: - final_hyperparameters[k] = str(v) if not isinstance(v, str) else v + if v is None: + continue + if allowed is not None and k not in allowed: + continue + final_hyperparameters[k] = str(v) if not isinstance(v, str) else v return final_hyperparameters + def _user_changed_recipe_keys(self) -> set: + """Return the set of leaf key names the user explicitly changed. + + Combines keys from the programmatic overrides dict, the user recipe YAML, + and direct ``.hyperparameters.*`` assignments. Used by the serverless + path to decide which non-spec recipe leaves to forward (so explicit + overrides are never silently dropped) while excluding the recipe's + unchanged defaults. + """ + changed = set() + + # Direct hyperparameter assignments (.hyperparameters.x = val) + user_set = getattr(self.hyperparameters, '_user_set', None) if hasattr(self, 'hyperparameters') else None + if isinstance(user_set, set): + changed |= user_set + + # Programmatic overrides dict (nested); collect its leaf key names. + overrides = getattr(self, '_overrides', None) + if isinstance(overrides, dict): + changed |= set(flatten_resolved_recipe(overrides).keys()) + + # User recipe YAML (local or S3); collect its leaf key names. + recipe_path = getattr(self, '_recipe_path', None) + if recipe_path: + try: + from sagemaker.train.recipe_resolver import _load_user_recipe + user_recipe = _load_user_recipe(recipe_path) + changed |= set(flatten_resolved_recipe(user_recipe).keys()) + except Exception as e: # pragma: no cover - defensive + logging.getLogger(__name__).debug( + "Could not read user recipe to determine changed keys: %s", e + ) + + return changed + def _validate_instance_count(self, instance_count, sagemaker_session): """Validate instance/node count against allowed values from SMHP recipe.""" smhp_replicas_enum = _get_smhp_replicas_enum( diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index 387e7b226a..3231b75587 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/sagemaker-train/src/sagemaker/train/dpo_trainer.py @@ -301,8 +301,12 @@ def train(self, final_hyperparameters = self.hyperparameters.to_dict() - # Apply recipe/overrides if provided (overrides > recipe > Hub defaults) - final_hyperparameters = self._apply_recipe_to_hyperparameters(final_hyperparameters) + # Apply recipe/overrides if provided (overrides > recipe > Hub defaults). + # Serverless HyperParameters is an override map capped at 100 entries by + # the API, so send only spec keys + user-changed leaves (P467902218). + final_hyperparameters = self._apply_recipe_to_hyperparameters( + final_hyperparameters, serverless=True + ) # Resolve is_multimodal: auto-detect from training dataset if not explicitly set if self.is_multimodal is None: effective_training_dataset = training_dataset or self.training_dataset diff --git a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py index 5c85e2f6ae..305ae6c9b3 100644 --- a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py +++ b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py @@ -279,8 +279,14 @@ def train( self._final_hyperparameters = self.hyperparameters.to_dict() - # Apply recipe/overrides if provided (overrides > recipe > Hub defaults) - self._final_hyperparameters = self._apply_recipe_to_hyperparameters(self._final_hyperparameters) + # Apply recipe/overrides if provided (overrides > recipe > Hub defaults). + # Restrict to Hub-allowlisted overridable keys so the serverless + # HyperParameters override map stays within the API's 100-entry limit + # and the _build_training_config delta filter compares like-for-like + # against the spec defaults snapshot (P467902218). + self._final_hyperparameters = self._apply_recipe_to_hyperparameters( + self._final_hyperparameters, serverless=True + ) _validate_hyperparameter_values(self._final_hyperparameters) diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py index acb7277eea..c14429a1f4 100644 --- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py @@ -461,8 +461,12 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, final_hyperparameters = self.hyperparameters.to_dict() - # Apply recipe/overrides if provided (overrides > recipe > Hub defaults) - final_hyperparameters = self._apply_recipe_to_hyperparameters(final_hyperparameters) + # Apply recipe/overrides if provided (overrides > recipe > Hub defaults). + # Serverless HyperParameters is an override map capped at 100 entries by + # the API, so send only spec keys + user-changed leaves (P467902218). + final_hyperparameters = self._apply_recipe_to_hyperparameters( + final_hyperparameters, serverless=True + ) # Resolve is_multimodal: auto-detect from training dataset if not explicitly set if self.is_multimodal is None: effective_training_dataset = training_dataset or self.training_dataset diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py index f16e097d8e..720082de54 100644 --- a/sagemaker-train/src/sagemaker/train/sft_trainer.py +++ b/sagemaker-train/src/sagemaker/train/sft_trainer.py @@ -354,8 +354,12 @@ def train(self, training_dataset: Optional[Union[str, DataSet]] = None, validati final_hyperparameters = self.hyperparameters.to_dict() - # Apply recipe/overrides if provided (overrides > recipe > Hub defaults) - final_hyperparameters = self._apply_recipe_to_hyperparameters(final_hyperparameters) + # Apply recipe/overrides if provided (overrides > recipe > Hub defaults). + # Serverless HyperParameters is an override map capped at 100 entries by + # the API, so send only spec keys + user-changed leaves (P467902218). + final_hyperparameters = self._apply_recipe_to_hyperparameters( + final_hyperparameters, serverless=True + ) # Resolve is_multimodal: auto-detect from training dataset if not explicitly set if self.is_multimodal is None: effective_training_dataset = training_dataset or self.training_dataset diff --git a/sagemaker-train/tests/unit/train/test_serverful_recipe_validation.py b/sagemaker-train/tests/unit/train/test_serverful_recipe_validation.py index f49e7c8462..981b9f0b85 100644 --- a/sagemaker-train/tests/unit/train/test_serverful_recipe_validation.py +++ b/sagemaker-train/tests/unit/train/test_serverful_recipe_validation.py @@ -372,6 +372,100 @@ def test_returns_unchanged_when_nothing_to_resolve(self): assert result == {"existing": "val"} + def test_serverless_drops_unchanged_non_spec_recipe_keys(self): + """Serverless path drops the recipe's unchanged non-spec defaults. + + Regression test for P467902218 — flattening the full resolved recipe + (200+ leaf keys) into the serverless HyperParameters override map blew + past the API's 100-entry limit and failed with a ValidationException. + With serverless=True the unchanged non-spec recipe defaults are applied + server-side and must not be sent, keeping the map bounded. + """ + trainer = _ConcreteTrainer() + trainer._recipe_path = None + trainer._overrides = None + trainer._resolved_recipe_cache = None + + hp_mock = MagicMock() + hp_mock._specs = {"max_steps": {"type": "integer"}, "lr": {"type": "float"}} + # No direct hyperparameter assignments. + hp_mock._user_set = set() + trainer.hyperparameters = hp_mock + + # A full recipe: two overridable spec keys plus 200 non-spec defaults + # (KL/vLLM/LoRA settings) that the user never changed. + training_config = {"max_steps": 50, "lr": 0.001} + training_config.update({f"internal_recipe_key_{i}": i for i in range(200)}) + resolved = {"training_config": training_config} + + with patch.object(trainer, "get_resolved_recipe", return_value=resolved): + result = trainer._apply_recipe_to_hyperparameters( + {"existing_key": "val"}, serverless=True + ) + + # Only the pre-existing key + the two spec keys survive; the 200 unchanged + # non-spec recipe leaves are dropped, keeping the map under the 100 limit. + assert result == {"existing_key": "val", "max_steps": "50", "lr": "0.001"} + assert len(result) < 100 + assert not any(k.startswith("internal_recipe_key_") for k in result) + + def test_serverless_forwards_explicit_user_overrides(self): + """Serverless path forwards non-spec keys the user explicitly overrode. + + A user override of a non-spec recipe key (e.g. peft.lora_tuning.alpha) + must never be silently dropped — otherwise the job would train with the + wrong value. Only the recipe's *unchanged* defaults are pruned. + """ + trainer = _ConcreteTrainer() + trainer._recipe_path = None + # User explicitly overrode a non-spec, deeply nested key. + trainer._overrides = {"training_config": {"peft": {"lora_tuning": {"alpha": 128}}}} + trainer._resolved_recipe_cache = None + + hp_mock = MagicMock() + hp_mock._specs = {"max_steps": {"type": "integer"}} + hp_mock._user_set = set() + trainer.hyperparameters = hp_mock + + resolved = { + "training_config": { + "max_steps": 50, + "peft": {"lora_tuning": {"alpha": 128, "rank": 16}}, + "unchanged_default": "keepserverside", + } + } + + with patch.object(trainer, "get_resolved_recipe", return_value=resolved): + result = trainer._apply_recipe_to_hyperparameters({}, serverless=True) + + # Spec key + explicitly overridden non-spec key are forwarded... + assert result["max_steps"] == "50" + assert result["alpha"] == "128" + # ...but unchanged non-spec defaults (rank, unchanged_default) are dropped. + assert "rank" not in result + assert "unchanged_default" not in result + + def test_serverful_applies_all_recipe_keys(self): + """Serverful path (serverless=False) still flattens the full recipe.""" + trainer = _ConcreteTrainer() + trainer._recipe_path = None + trainer._overrides = None + trainer._resolved_recipe_cache = None + + hp_mock = MagicMock() + hp_mock._specs = {"max_steps": {"type": "integer"}} + hp_mock._user_set = set() + trainer.hyperparameters = hp_mock + + resolved = {"training_config": {"max_steps": 50, "internal_recipe_key": "megatron"}} + + with patch.object(trainer, "get_resolved_recipe", return_value=resolved): + result = trainer._apply_recipe_to_hyperparameters({}) + + # Non-spec recipe keys are preserved for the serverful path, because + # they are rendered back into the recipe YAML. + assert result == {"max_steps": "50", "internal_recipe_key": "megatron"} + # --------------------------------------------------------------------------- # Tests: disable_output_compression diff --git a/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py b/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py index bb873a6b8b..b9cbbc6a41 100644 --- a/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py +++ b/sagemaker-train/tests/unit/train/test_trainer_recipe_integration.py @@ -599,7 +599,13 @@ def test_non_spec_keys_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.""" + """Explicitly overridden non-spec keys flow into serverless hyperparameters. + + Serverless HyperParameters is a bounded override map (P467902218): a + non-spec key the user *overrode* (sequence_length) must be forwarded, + but the recipe's *unchanged* non-spec defaults (warmup_ratio) must be + dropped since they are applied server-side. + """ mock_get_options.return_value = (mock_hyperparams_with_full_template, "model-arn", False) from sagemaker.train.sft_trainer import SFTTrainer @@ -633,12 +639,13 @@ def test_non_spec_keys_flow_into_train_hyperparameters( call_kwargs = mock_tj.create.call_args[1] final_hp = call_kwargs["hyper_parameters"] - # Non-spec key is now in final hyperparameters + # Explicitly overridden non-spec key is forwarded assert "sequence_length" in final_hp assert final_hp["sequence_length"] == "8192" - # Other full template keys also present - assert "warmup_ratio" in final_hp - assert final_hp["warmup_ratio"] == "0.1" + # Unchanged non-spec template defaults are NOT sent (applied + # server-side) — this is what keeps the map under the 100-entry limit. + assert "warmup_ratio" not in final_hp + assert "gradient_accumulation_steps" 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") @@ -648,7 +655,12 @@ def test_nested_keys_flow_into_train_hyperparameters( self, mock_resolve, mock_validate_group, mock_get_options, mock_eula, mock_hyperparams_with_full_template ): - """Nested recipe keys (lr_scheduler.warmup_steps) are flattened into final hyperparameters.""" + """Overridden nested recipe keys (lr_scheduler.warmup_steps) flatten into serverless hyperparameters. + + The overridden leaf (warmup_steps) is forwarded; the sibling unchanged + default (min_lr) is dropped since the base recipe applies it server-side + (P467902218). + """ mock_get_options.return_value = (mock_hyperparams_with_full_template, "model-arn", False) from sagemaker.train.sft_trainer import SFTTrainer @@ -682,12 +694,11 @@ def test_nested_keys_flow_into_train_hyperparameters( call_kwargs = mock_tj.create.call_args[1] final_hp = call_kwargs["hyper_parameters"] - # Nested key overridden and flattened to leaf name + # Overridden nested key flattened to leaf name and forwarded assert "warmup_steps" in final_hp assert final_hp["warmup_steps"] == "30" - # Other nested leaf retains default from full template - assert "min_lr" in final_hp - assert final_hp["min_lr"] == "1e-06" + # Sibling unchanged nested default is dropped (applied server-side) + assert "min_lr" 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") @@ -696,7 +707,11 @@ def test_nested_keys_flow_into_train_hyperparameters( def test_deeply_nested_peft_keys_flow_into_hyperparameters( self, mock_resolve, mock_validate_group, mock_get_options, mock_eula, ): - """Deeply nested keys (peft.lora_tuning.alpha) flatten into hyperparameters.""" + """Overridden deeply nested keys (peft.lora_tuning.alpha) flatten into serverless hyperparameters. + + Only the explicitly overridden leaf (alpha) is forwarded; unchanged + deeply nested defaults (rank, peft_scheme) are dropped (P467902218). + """ mock_hp = MagicMock() mock_hp._specs = { "learning_rate": {"default": 1e-4, "type": "float", "min": 1e-7, "max": 1.0}, @@ -747,12 +762,11 @@ def test_deeply_nested_peft_keys_flow_into_hyperparameters( call_kwargs = mock_tj.create.call_args[1] final_hp = call_kwargs["hyper_parameters"] - # Deeply nested override flattened + # Deeply nested override flattened and forwarded assert final_hp["alpha"] == "128" - # Unchanged deeply nested default preserved - assert final_hp["rank"] == "16" - # String-only leaf that's deeply nested - assert final_hp["peft_scheme"] == "lora" + # Unchanged deeply nested defaults are dropped (applied server-side) + assert "rank" not in final_hp + assert "peft_scheme" 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")