Refactoring

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -43,7 +43,6 @@
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
 from sagemaker.serve.utils.tuning import (
     _pretty_print_benchmark_results,
-    sharded_supported,
     _serial_benchmark,
     _concurrent_benchmark,
     _more_performant,
@@ -278,15 +277,6 @@ def _tune_for_js(self, max_tuning_duration: int = 1800):
             self.js_model_config
         )
 
-        if len(admissible_tensor_parallel_degrees) > 1 and not sharded_supported(
-            self.model, self.js_model_config
-        ):
-            admissible_tensor_parallel_degrees = [1]
-            logger.warning(
-                "Sharded across multiple GPUs is not supported for this model. "
-                "Model can only be sharded across [1] GPU"
-            )
-
         benchmark_results = {}
         best_tuned_combination = None
         timeout = datetime.now() + timedelta(seconds=max_tuning_duration)

src/sagemaker/serve/utils/tuning.py

@@ -263,24 +263,3 @@ def _more_performant(best_tuned_configuration: list, tuned_configuration: list)
             return True
         return False
     return tuned_avg_latency <= best_avg_latency
-
-
-def sharded_supported(model_id: str, config_dict: dict) -> bool:
-    """Check if sharded is supported for this ``Model``"""
-    model_type = config_dict.get("model_type", None)
-
-    if model_type is None:
-        return False
-
-    if model_id.startswith("facebook/galactica"):
-        return True
-
-    if model_type in ["bloom", "mpt", "ssm", "gpt_neox", "phi", "phi-msft", "opt", "t5"]:
-        return True
-
-    if model_type in ["RefinedWeb", "RefinedWebModel", "falcon"] and not config_dict.get(
-        "alibi", False
-    ):
-        return True
-
-    return False

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

@@ -470,64 +470,3 @@ def test_tune_for_djl_js_local_container_invoke_ex(
 
         tuned_model = model.tune()
         assert tuned_model.env == mock_djl_model_serving_properties
-
-    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
-    @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",
-        return_value=MagicMock(),
-    )
-    @patch(
-        "sagemaker.serve.builder.jumpstart_builder.prepare_djl_js_resources",
-        return_value=(
-            mock_set_serving_properties,
-            {"model_type": "sharded_not_enabled", "n_head": 71},
-            True,
-        ),
-    )
-    @patch("sagemaker.serve.builder.jumpstart_builder._get_ram_usage_mb", return_value=1024)
-    @patch(
-        "sagemaker.serve.builder.jumpstart_builder._get_nb_instance", return_value="ml.g5.24xlarge"
-    )
-    @patch(
-        "sagemaker.serve.builder.jumpstart_builder._get_admissible_tensor_parallel_degrees",
-        return_value=[4, 2, 1],
-    )
-    @patch(
-        "sagemaker.serve.utils.tuning._serial_benchmark",
-        side_effect=[(5, 5, 25), (5.4, 5.4, 20), (5.2, 5.2, 15)],
-    )
-    @patch(
-        "sagemaker.serve.utils.tuning._concurrent_benchmark",
-        side_effect=[(0.9, 1), (0.10, 4), (0.13, 2)],
-    )
-    def test_tune_for_djl_js_local_container_sharded_not_enabled(
-        self,
-        mock_concurrent_benchmarks,
-        mock_serial_benchmarks,
-        mock_admissible_tensor_parallel_degrees,
-        mock_get_nb_instance,
-        mock_get_ram_usage_mb,
-        mock_prepare_for_tgi,
-        mock_pre_trained_model,
-        mock_is_jumpstart_model,
-        mock_telemetry,
-    ):
-        builder = ModelBuilder(
-            model=mock_model_id,
-            schema_builder=mock_schema_builder,
-            mode=Mode.LOCAL_CONTAINER,
-        )
-
-        mock_pre_trained_model.return_value.image_uri = mock_djl_image_uri
-
-        model = builder.build()
-        builder.serve_settings.telemetry_opt_out = True
-
-        mock_pre_trained_model.return_value.env = mock_djl_model_serving_properties
-
-        tuned_model = model.tune()
-        assert tuned_model.env == mock_djl_most_performant_model_serving_properties