fix(serve): size tensor parallelism from the target instance's GPUs

sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/utils.py

@@ -15,7 +15,9 @@
 
 def _get_default_tensor_parallel_degree(hf_model_config: dict, gpu_count: int = None) -> int:
     """Placeholder docstring"""
-    available_gpus = _get_available_gpus()
+    # gpu_count is the GPU count of the target instance; only fall back to the GPUs on the
+    # machine running the build when it is not known (e.g. local container mode).
+    available_gpus = None if gpu_count else _get_available_gpus()
     if not available_gpus and not gpu_count:
         return None
 
@@ -28,7 +30,7 @@ def _get_default_tensor_parallel_degree(hf_model_config: dict, gpu_count: int =
     if not attention_heads:
         return 1
 
-    tot_gpus = len(available_gpus) if available_gpus else gpu_count
+    tot_gpus = gpu_count if gpu_count else len(available_gpus)
     for i in (n + 1 for n in reversed(range(tot_gpus))):
         if attention_heads % i == 0:
             logger.info(

sagemaker-serve/tests/unit/model_server/test_djl_utils.py

@@ -39,6 +39,22 @@ def test_get_default_tensor_parallel_degree(self, mock_gpus):
         result = _get_default_tensor_parallel_degree(hf_config)
         self.assertIsNotNone(result)
 
+    @patch("sagemaker.serve.model_server.djl_serving.utils._get_available_gpus")
+    def test_get_default_tensor_parallel_degree_prefers_target_gpu_count(self, mock_gpus):
+        """The target instance's GPU count wins over the GPUs on the build machine."""
+        hf_config = {"num_attention_heads": 32}
+
+        mock_gpus.return_value = [0]  # 1-GPU notebook building for a 4-GPU instance
+        self.assertEqual(_get_default_tensor_parallel_degree(hf_config, gpu_count=4), 4)
+
+        mock_gpus.return_value = list(range(8))  # 8-GPU box building for a 1-GPU instance
+        self.assertEqual(_get_default_tensor_parallel_degree(hf_config, gpu_count=1), 1)
+
+    @patch("sagemaker.serve.model_server.djl_serving.utils._get_available_gpus")
+    def test_get_default_tensor_parallel_degree_uses_local_gpus_without_gpu_count(self, mock_gpus):
+        mock_gpus.return_value = [0, 1]
+        self.assertEqual(_get_default_tensor_parallel_degree({"num_attention_heads": 32}), 2)
+
 
 if __name__ == "__main__":
     unittest.main()