fix(train): send bounded HyperParameters override map for serverless customization jobs

sagemaker-train/src/sagemaker/train/base_trainer.py

@@ -236,38 +236,104 @@ 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:
             resolved = self.get_resolved_recipe()
         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(

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

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)
 

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

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

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

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")