diff --git a/sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/utils.py b/sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/utils.py index 5ad2e366a3..9c9394bdb1 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_server/djl_serving/utils.py +++ b/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( diff --git a/sagemaker-serve/tests/unit/model_server/test_djl_utils.py b/sagemaker-serve/tests/unit/model_server/test_djl_utils.py index a66f04cbf4..fca3cc16cf 100644 --- a/sagemaker-serve/tests/unit/model_server/test_djl_utils.py +++ b/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()