Enable quantization and compilation in the same optimization job via ModelBuilder and add validations to block compilation jobs using TRTLLM an Llama-3.1.

Joseph Zhang · Joseph Zhang · commit 7ec16e6df038 · 2024-11-07T11:37:10.000-08:00
diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -718,24 +718,36 @@ 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, override_env = _extract_optimization_config_and_env(
-            quantization_config, compilation_config
+        # optimization_config can contain configs for both quantization and compilation
+        optimization_config, quantization_override_env, compilation_override_env = (
+            _extract_optimization_config_and_env(quantization_config, compilation_config)
         )
-        if not optimization_config and is_compilation:
-            override_env = override_env or pysdk_model_env_vars
-            optimization_config = {
-                "ModelCompilationConfig": {
-                    "OverrideEnvironment": override_env,
-                }
-            }
+        if (
+            not optimization_config or not optimization_config.get("ModelCompilationConfig")
+        ) and is_compilation:
+            # Ensure optimization_config exists
+            if not optimization_config:
+                optimization_config = {}
+
+            # Fallback to default if override_env is None or empty
+            if not compilation_override_env:
+                compilation_override_env = pysdk_model_env_vars
+
+            # Update optimization_config with ModelCompilationConfig
+            override_compilation_config = (
+                {"OverrideEnvironment": compilation_override_env}
+                if compilation_override_env
+                else {}
+            )
+            optimization_config["ModelCompilationConfig"] = override_compilation_config
 
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
@@ -766,7 +778,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,
         }
@@ -789,7 +801,13 @@ def _optimize_for_jumpstart(
                     "AcceptEula": True
                 }
 
-        optimization_env_vars = _update_environment_variables(optimization_env_vars, override_env)
+        optimization_env_vars = _update_environment_variables(
+            optimization_env_vars,
+            {
+                **(quantization_override_env or {}),
+                **(compilation_override_env or {}),
+            },
+        )
         if optimization_env_vars:
             self.pysdk_model.env.update(optimization_env_vars)
         if quantization_config or is_compilation:
diff --git a/src/sagemaker/serve/builder/model_builder.py b/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
@@ -1279,6 +1276,28 @@ def _model_builder_optimize_wrapper(
             )
 
         if input_args:
+            optimization_instance_type = input_args["DeploymentInstanceType"]
+
+            # Compilation using TRTLLM and Llama-3.1 is currently not supported.
+            # TRTLLM is used by Neo if the following are provided:
+            #  1) a GPU instance type
+            #  2) compilation config
+            gpu_instance_families = ["g4", "g5", "p4d"]
+            is_gpu_instance = optimization_instance_type and any(
+                gpu_instance_family in optimization_instance_type
+                for gpu_instance_family in gpu_instance_families
+            )
+
+            # HF Model ID format = "meta-llama/Meta-Llama-3.1-8B"
+            # JS Model ID format = "meta-textgeneration-llama-3-1-8b"
+            llama_3_1_keywords = ["llama-3.1", "llama-3-1"]
+            is_llama_3_1 = self.model and any(
+                keyword in self.model.lower() for keyword in llama_3_1_keywords
+            )
+
+            if is_gpu_instance and self.model and is_llama_3_1 and self.is_compiled:
+                raise ValueError("Compilation is not supported for Llama-3.1 with a GPU instance.")
+
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)
             return _generate_optimized_model(self.pysdk_model, job_status)
@@ -1342,11 +1361,18 @@ def _optimize_for_hf(
             model_source = _generate_model_source(self.pysdk_model.model_data, False)
             create_optimization_job_args["ModelSource"] = model_source
 
-            optimization_config, override_env = _extract_optimization_config_and_env(
-                quantization_config, compilation_config
+            optimization_config, quantization_override_env, compilation_override_env = (
+                _extract_optimization_config_and_env(quantization_config, compilation_config)
+            )
+            create_optimization_job_args["OptimizationConfigs"] = [
+                {k: v} for k, v in optimization_config.items()
+            ]
+            self.pysdk_model.env.update(
+                {
+                    **(quantization_override_env or {}),
+                    **(compilation_override_env or {}),
+                }
             )
-            create_optimization_job_args["OptimizationConfigs"] = [optimization_config]
-            self.pysdk_model.env.update(override_env)
 
             output_config = {"S3OutputLocation": output_path}
             if kms_key:
diff --git a/src/sagemaker/serve/utils/optimize_utils.py b/src/sagemaker/serve/utils/optimize_utils.py
@@ -260,7 +260,7 @@ def _is_s3_uri(s3_uri: Optional[str]) -> bool:
 
 def _extract_optimization_config_and_env(
     quantization_config: Optional[Dict] = None, compilation_config: Optional[Dict] = None
-) -> Optional[Tuple[Optional[Dict], Optional[Dict]]]:
+) -> Optional[Tuple[Optional[Dict], Optional[Dict], Optional[Dict]]]:
     """Extracts optimization config and environment variables.
 
     Args:
@@ -271,15 +271,25 @@ def _extract_optimization_config_and_env(
         Optional[Tuple[Optional[Dict], Optional[Dict]]]:
             The optimization config and environment variables.
     """
-    if quantization_config:
-        return {"ModelQuantizationConfig": quantization_config}, quantization_config.get(
-            "OverrideEnvironment"
-        )
-    if compilation_config:
-        return {"ModelCompilationConfig": compilation_config}, compilation_config.get(
-            "OverrideEnvironment"
-        )
-    return None, None
+    optimization_config = {}
+    quantization_override_env = (
+        quantization_config.get("OverrideEnvironment", {}) if quantization_config else None
+    )
+    compilation_override_env = (
+        compilation_config.get("OverrideEnvironment", {}) if compilation_config else None
+    )
+
+    if quantization_config is not None:
+        optimization_config["ModelQuantizationConfig"] = quantization_config
+
+    if compilation_config is not None:
+        optimization_config["ModelCompilationConfig"] = compilation_config
+
+    # Return both dicts and environment variable if either is present
+    if optimization_config:
+        return optimization_config, quantization_override_env, compilation_override_env
+
+    return None, None, None
 
 
 def _custom_speculative_decoding(
diff --git a/tests/unit/sagemaker/serve/builder/test_js_builder.py b/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(
@@ -1383,3 +1442,103 @@ def test_optimize_compile_for_jumpstart_with_neuron_env(
         self.assertEqual(optimized_model.env["OPTION_ROLLING_BATCH"], "auto")
         self.assertEqual(optimized_model.env["OPTION_MAX_ROLLING_BATCH_SIZE"], "4")
         self.assertEqual(optimized_model.env["OPTION_NEURON_OPTIMIZE_LEVEL"], "2")
+
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_gated_model",
+        return_value=True,
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder.JumpStartModel")
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch("sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model")
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.prepare_tgi_js_resources",
+        return_value=({"model_type": "t5", "n_head": 71}, True),
+    )
+    def test_optimize_compile_for_jumpstart_without_compilation_config(
+        self,
+        mock_prepare_for_tgi,
+        mock_pre_trained_model,
+        mock_is_jumpstart_model,
+        mock_js_model,
+        mock_is_gated_model,
+        mock_serve_settings,
+        mock_telemetry,
+    ):
+        mock_sagemaker_session = Mock()
+        mock_metadata_config = Mock()
+        mock_sagemaker_session.wait_for_optimization_job.side_effect = (
+            lambda *args: mock_optimization_job_response
+        )
+
+        mock_metadata_config.resolved_config = {
+            "supported_inference_instance_types": ["ml.inf2.48xlarge"],
+            "hosting_neuron_model_id": "huggingface-llmneuron-mistral-7b",
+        }
+
+        mock_js_model.return_value = MagicMock()
+        mock_js_model.return_value.env = {
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "ENDPOINT_SERVER_TIMEOUT": "3600",
+            "MODEL_CACHE_ROOT": "/opt/ml/model",
+            "SAGEMAKER_ENV": "1",
+            "HF_MODEL_ID": "/opt/ml/model",
+            "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
+        }
+
+        mock_pre_trained_model.return_value = MagicMock()
+        mock_pre_trained_model.return_value.env = dict()
+        mock_pre_trained_model.return_value.config_name = "config_name"
+        mock_pre_trained_model.return_value.model_data = mock_model_data
+        mock_pre_trained_model.return_value.image_uri = mock_tgi_image_uri
+        mock_pre_trained_model.return_value.list_deployment_configs.return_value = (
+            DEPLOYMENT_CONFIGS
+        )
+        mock_pre_trained_model.return_value.deployment_config = DEPLOYMENT_CONFIGS[0]
+        mock_pre_trained_model.return_value._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,
+        )
+
+        optimized_model = model_builder.optimize(
+            accept_eula=True,
+            instance_type="ml.inf2.24xlarge",
+            output_path="s3://bucket/code/",
+        )
+
+        self.assertEqual(
+            optimized_model.image_uri,
+            mock_optimization_job_response["OptimizationOutput"]["RecommendedInferenceImage"],
+        )
+        self.assertEqual(
+            optimized_model.model_data["S3DataSource"]["S3Uri"],
+            mock_optimization_job_response["OutputConfig"]["S3OutputLocation"],
+        )
+        self.assertEqual(optimized_model.env["SAGEMAKER_PROGRAM"], "inference.py")
+        self.assertEqual(optimized_model.env["ENDPOINT_SERVER_TIMEOUT"], "3600")
+        self.assertEqual(optimized_model.env["MODEL_CACHE_ROOT"], "/opt/ml/model")
+        self.assertEqual(optimized_model.env["SAGEMAKER_ENV"], "1")
+        self.assertEqual(optimized_model.env["HF_MODEL_ID"], "/opt/ml/model")
+        self.assertEqual(optimized_model.env["SAGEMAKER_MODEL_SERVER_WORKERS"], "1")
diff --git a/tests/unit/sagemaker/serve/builder/test_model_builder.py b/tests/unit/sagemaker/serve/builder/test_model_builder.py
diff --git a/tests/unit/sagemaker/serve/utils/test_optimize_utils.py b/tests/unit/sagemaker/serve/utils/test_optimize_utils.py
