From 3b5a49776b8f054e95c099631593603b008d51b2 Mon Sep 17 00:00:00 2001
From: Datta Nimmaturi <venkatadattasainimmaturi@gmail.com>
Date: Tue, 31 Mar 2026 13:54:41 +0530
Subject: [PATCH] [studio] multi gpu: revert to balanced for inference. (#4698)

* Revert to balanced for inference

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused for_inference parameter from get_device_map

Since inference and training both use "balanced" now, the for_inference
flag is dead code. Remove it from the function signature, the call site
in inference.py, and simplify the tests accordingly.

* Remove redundant TestDeviceMapForInference test class

TestGpuAutoSelection already covers the same multi-gpu and single-gpu
device_map assertions. The TestDeviceMapForInference class was left
over from when for_inference had distinct behavior.

* Remove redundant test_get_device_map_multi_gpu_uses_balanced

Its assertions ([0,1] -> balanced, [0] -> sequential) are already
covered by test_get_device_map_uses_explicit_gpu_selection.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
---
 studio/backend/core/inference/inference.py |  2 +-
 studio/backend/tests/test_gpu_selection.py | 22 ----------------------
 studio/backend/utils/hardware/hardware.py  |  4 +---
 3 files changed, 2 insertions(+), 26 deletions(-)

diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py
index ddd485525d..867bdefc62 100644
--- a/studio/backend/core/inference/inference.py
+++ b/studio/backend/core/inference/inference.py
@@ -268,7 +268,7 @@ class InferenceBackend:
                 return False
 
             self.loading_models.add(model_name)
-            device_map = get_device_map(gpu_ids, for_inference = True)
+            device_map = get_device_map(gpu_ids)
             logger.info(
                 f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
             )
diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py
index 275b6f33d6..c6f26037af 100644
--- a/studio/backend/tests/test_gpu_selection.py
+++ b/studio/backend/tests/test_gpu_selection.py
@@ -252,11 +252,6 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase):
             self.assertEqual(get_device_map([0]), "sequential")
             self.assertEqual(get_device_map([0, 1]), "balanced")
 
-    def test_get_device_map_multi_gpu_uses_balanced(self):
-        with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA):
-            self.assertEqual(get_device_map([0, 1]), "balanced")
-            self.assertEqual(get_device_map([0]), "sequential")
-
     def test_get_device_map_uses_all_inherited_visible_gpus_for_uuid_masks(self):
         with (
             patch.dict(
@@ -1106,20 +1101,3 @@ class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase):
         with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU):
             with self.assertRaisesRegex(ValueError, "only supported on CUDA"):
                 prepare_gpu_selection([0], model_name = "unsloth/test")
-
-
-class TestDeviceMapForInference(_GpuCacheResetMixin, unittest.TestCase):
-    def test_inference_uses_balanced_low_0(self):
-        with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA):
-            self.assertEqual(
-                get_device_map([0, 1], for_inference = True), "balanced_low_0"
-            )
-
-    def test_training_uses_balanced(self):
-        with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA):
-            self.assertEqual(get_device_map([0, 1], for_inference = False), "balanced")
-
-    def test_single_gpu_always_sequential(self):
-        with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA):
-            self.assertEqual(get_device_map([0], for_inference = True), "sequential")
-            self.assertEqual(get_device_map([0], for_inference = False), "sequential")
diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py
index 742e8f6b7e..b6d3faf6d7 100644
--- a/studio/backend/utils/hardware/hardware.py
+++ b/studio/backend/utils/hardware/hardware.py
@@ -1294,8 +1294,6 @@ def apply_gpu_ids(gpu_ids) -> None:
 
 def get_device_map(
     gpu_ids: Optional[list[int]] = None,
-    *,
-    for_inference: bool = False,
 ) -> str:
     """Return the Hugging Face ``device_map`` string for model loading.
 
@@ -1327,7 +1325,7 @@ def get_device_map(
                 multi_gpu = True
 
         if multi_gpu:
-            return "balanced_low_0" if for_inference else "balanced"
+            return "balanced"
 
     return "sequential"