Enable quantization and compilation in the same optimization job via ModelBuilder.

Author: Joseph Zhang

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -718,20 +718,23 @@ def _optimize_for_jumpstart(
                 f"Model '{self.model}' requires accepting end-user license agreement (EULA)."
             )
 
-        is_compilation = (not quantization_config) and (
-            (compilation_config is not None) or _is_inferentia_or_trainium(instance_type)
+        is_compilation = (compilation_config is not None) or _is_inferentia_or_trainium(
+            instance_type
         )
 
         pysdk_model_env_vars = dict()
         if is_compilation:
             pysdk_model_env_vars = self._get_neuron_model_env_vars(instance_type)
 
+        # optimization_config can contain configs for both quantization and compilation
         optimization_config, override_env = _extract_optimization_config_and_env(
             quantization_config, compilation_config
         )
-        if not optimization_config and is_compilation:
+        if (
+            not optimization_config or not optimization_config.get("ModelCompilationConfig")
+        ) and is_compilation:
             override_env = override_env or pysdk_model_env_vars
-            optimization_config = {
+            optimization_config["ModelCompilationConfig"] = {
                 "ModelCompilationConfig": {
                     "OverrideEnvironment": override_env,
                 }
@@ -766,7 +769,7 @@ def _optimize_for_jumpstart(
             "OptimizationJobName": job_name,
             "ModelSource": model_source,
             "DeploymentInstanceType": self.instance_type,
-            "OptimizationConfigs": [optimization_config],
+            "OptimizationConfigs": [{k: v} for k, v in optimization_config.items()],
             "OutputConfig": output_config,
             "RoleArn": self.role_arn,
         }

src/sagemaker/serve/builder/model_builder.py

@@ -1235,9 +1235,6 @@ def _model_builder_optimize_wrapper(
         if self.mode != Mode.SAGEMAKER_ENDPOINT:
             raise ValueError("Model optimization is only supported in Sagemaker Endpoint Mode.")
 
-        if quantization_config and compilation_config:
-            raise ValueError("Quantization config and compilation config are mutually exclusive.")
-
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
         self.instance_type = instance_type or self.instance_type
         self.role_arn = role_arn or self.role_arn
@@ -1345,7 +1342,9 @@ def _optimize_for_hf(
             optimization_config, override_env = _extract_optimization_config_and_env(
                 quantization_config, compilation_config
             )
-            create_optimization_job_args["OptimizationConfigs"] = [optimization_config]
+            create_optimization_job_args["OptimizationConfigs"] = [
+                {k: v} for k, v in optimization_config.items()
+            ]
             self.pysdk_model.env.update(override_env)
 
             output_config = {"S3OutputLocation": output_path}

src/sagemaker/serve/utils/optimize_utils.py

@@ -271,14 +271,25 @@ def _extract_optimization_config_and_env(
         Optional[Tuple[Optional[Dict], Optional[Dict]]]:
             The optimization config and environment variables.
     """
+    optimization_config = {}
+    quantization_override_env = {}
+    compilation_override_env = {}
+
     if quantization_config:
-        return {"ModelQuantizationConfig": quantization_config}, quantization_config.get(
-            "OverrideEnvironment"
-        )
+        optimization_config["ModelQuantizationConfig"] = quantization_config
+        quantization_override_env = quantization_config.get("OverrideEnvironment")
+
     if compilation_config:
-        return {"ModelCompilationConfig": compilation_config}, compilation_config.get(
-            "OverrideEnvironment"
-        )
+        optimization_config["ModelCompilationConfig"] = compilation_config
+        compilation_override_env = compilation_config.get("OverrideEnvironment")
+
+    # Return both dicts and environment variable if either is present
+    if optimization_config:
+        return optimization_config, {
+            **(quantization_override_env or {}),
+            **(compilation_override_env or {}),
+        }
+
     return None, None
 
 

tests/unit/sagemaker/serve/builder/test_js_builder.py

@@ -1198,6 +1198,65 @@ def test_optimize_quantize_for_jumpstart(
 
         self.assertIsNotNone(out_put)
 
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_optimize_quantize_and_compile_for_jumpstart(
+        self,
+        mock_serve_settings,
+        mock_telemetry,
+    ):
+        mock_sagemaker_session = Mock()
+        mock_metadata_config = Mock()
+        mock_metadata_config.resolved_config = {
+            "supported_inference_instance_types": ["ml.inf2.48xlarge"],
+            "hosting_neuron_model_id": "huggingface-llmneuron-mistral-7b",
+        }
+
+        mock_pysdk_model = Mock()
+        mock_pysdk_model.env = {"SAGEMAKER_ENV": "1"}
+        mock_pysdk_model.model_data = mock_model_data
+        mock_pysdk_model.image_uri = mock_tgi_image_uri
+        mock_pysdk_model.list_deployment_configs.return_value = DEPLOYMENT_CONFIGS
+        mock_pysdk_model.deployment_config = DEPLOYMENT_CONFIGS[0]
+        mock_pysdk_model.config_name = "config_name"
+        mock_pysdk_model._metadata_configs = {"config_name": mock_metadata_config}
+
+        sample_input = {
+            "inputs": "The diamondback terrapin or simply terrapin is a species "
+            "of turtle native to the brackish coastal tidal marshes of the",
+            "parameters": {"max_new_tokens": 1024},
+        }
+        sample_output = [
+            {
+                "generated_text": "The diamondback terrapin or simply terrapin is a "
+                "species of turtle native to the brackish coastal "
+                "tidal marshes of the east coast."
+            }
+        ]
+
+        model_builder = ModelBuilder(
+            model="meta-textgeneration-llama-3-70b",
+            schema_builder=SchemaBuilder(sample_input, sample_output),
+            sagemaker_session=mock_sagemaker_session,
+        )
+
+        model_builder.pysdk_model = mock_pysdk_model
+
+        out_put = model_builder._optimize_for_jumpstart(
+            accept_eula=True,
+            quantization_config={
+                "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
+            },
+            compilation_config={"OverrideEnvironment": {"OPTION_TENSOR_PARALLEL_DEGREE": "2"}},
+            env_vars={
+                "OPTION_TENSOR_PARALLEL_DEGREE": "1",
+                "OPTION_MAX_ROLLING_BATCH_SIZE": "2",
+            },
+            output_path="s3://bucket/code/",
+        )
+
+        self.assertIsNotNone(out_put)
+
     @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
     @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
     @patch(

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -2650,21 +2650,75 @@ def test_optimize_local_mode(self, mock_get_serve_setting):
             ),
         )
 
+    @patch.object(ModelBuilder, "_prepare_for_mode")
     @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
-    def test_optimize_exclusive_args(self, mock_get_serve_setting):
-        mock_sagemaker_session = Mock()
+    def test_optimize_for_hf_with_both_quantization_and_compilation(
+        self,
+        mock_get_serve_setting,
+        mock_prepare_for_mode,
+    ):
+        mock_prepare_for_mode.side_effect = lambda *args, **kwargs: (
+            {
+                "S3DataSource": {
+                    "CompressionType": "None",
+                    "S3DataType": "S3Prefix",
+                    "S3Uri": "s3://bucket/code/code/",
+                }
+            },
+            {"DTYPE": "bfloat16"},
+        )
+
+        mock_pysdk_model = Mock()
+        mock_pysdk_model.model_data = None
+        mock_pysdk_model.env = {"HF_MODEL_ID": "meta-llama/Meta-Llama-3-8B-Instruc"}
+
         model_builder = ModelBuilder(
-            model="meta-textgeneration-llama-3-70b",
-            sagemaker_session=mock_sagemaker_session,
+            model="meta-llama/Meta-Llama-3-8B-Instruct",
+            env_vars={"HF_TOKEN": "token"},
+            model_metadata={
+                "CUSTOM_MODEL_PATH": "s3://bucket/path/",
+            },
+            role_arn="role-arn",
+            instance_type="ml.g5.2xlarge",
         )
 
-        self.assertRaisesRegex(
-            ValueError,
-            "Quantization config and compilation config are mutually exclusive.",
-            lambda: model_builder.optimize(
-                quantization_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
-                compilation_config={"OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}},
-            ),
+        model_builder.pysdk_model = mock_pysdk_model
+
+        out_put = model_builder._optimize_for_hf(
+            job_name="job_name-123",
+            quantization_config={
+                "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
+            },
+            compilation_config={"OverrideEnvironment": {"OPTION_TENSOR_PARALLEL_DEGREE": "2"}},
+            output_path="s3://bucket/code/",
+        )
+
+        self.assertEqual(model_builder.env_vars["HF_TOKEN"], "token")
+        self.assertEqual(model_builder.role_arn, "role-arn")
+        self.assertEqual(model_builder.instance_type, "ml.g5.2xlarge")
+        self.assertEqual(model_builder.pysdk_model.env["OPTION_QUANTIZE"], "awq")
+        self.assertEqual(model_builder.pysdk_model.env["OPTION_TENSOR_PARALLEL_DEGREE"], "2")
+        self.assertEqual(
+            out_put,
+            {
+                "OptimizationJobName": "job_name-123",
+                "DeploymentInstanceType": "ml.g5.2xlarge",
+                "RoleArn": "role-arn",
+                "ModelSource": {"S3": {"S3Uri": "s3://bucket/code/code/"}},
+                "OptimizationConfigs": [
+                    {
+                        "ModelQuantizationConfig": {
+                            "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"}
+                        }
+                    },
+                    {
+                        "ModelCompilationConfig": {
+                            "OverrideEnvironment": {"OPTION_TENSOR_PARALLEL_DEGREE": "2"}
+                        }
+                    },
+                ],
+                "OutputConfig": {"S3OutputLocation": "s3://bucket/code/"},
+            },
         )
 
     @patch.object(ModelBuilder, "_prepare_for_mode")