ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[studio] multi gpu: revert to balanced for inference. (#4698)

Browse source 
* 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>
This commit is contained in:
Datta Nimmaturi 2026-03-31 13:54:41 +05:30 • committed by GitHub
commit 3b5a49776b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 2 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/core/inference/inference.py
View file

@ -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)"
)

22
studio/backend/tests/test_gpu_selection.py
View file

@ -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")

4
studio/backend/utils/hardware/hardware.py
View file

@ -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"