9311df2b29* [WIP] balanced device map for studio * gpus as a request parameter * API for multi GPU stuff * return multi gpu util in new API * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use balanced_low0 instead of balanced * Use balanced_low0 instead of balanced * Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests - balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string) - get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES gracefully instead of crashing on int() parse - _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs - test_gpu_selection.py: add missing get_visible_gpu_utilization import and add required job_id arg to start_training() calls * Smart GPU determinism using estimates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * disallow gpu selection for gguf for now * cleanup * Slightly larger baseline * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Treat empty list as auto * Verbose logging/debug * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup and revert unnecessary deletions * Cleanup excessive logs and guard against disk/cpu offload * auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * support for non cuda gpus * Fix multi-GPU auto-selection memory accounting The multi_gpu_factor was applied uniformly to all GPUs including the first one, which unfairly penalizes single-GPU capacity when transitioning to multi-GPU. This created a discontinuity where a model that barely fits 1 GPU would suddenly require 2 GPUs because the first GPU's free memory was discounted by 20%. Now the first GPU keeps its full free memory, and only additional GPUs have an overhead factor (0.85) applied to account for inter-GPU communication and sharding overhead. This gives more accurate auto-selection and avoids unnecessary multi-GPU for models that comfortably fit on one device. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add sandbox tests for multi-GPU selection logic 24 tests covering model size estimation, memory requirements, automatic GPU selection, device map generation, GPU ID validation, and multi-GPU overhead accounting. All tests use mocks so they run without GPUs on Linux, macOS, and Windows. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry 1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer) instead of the FP16 1.3x multiplier. This prevents over-sharding quantized models across unnecessary GPUs. 2. When model size estimation fails, auto_select_gpu_ids now falls back to all visible GPUs instead of returning None (which could default to single-GPU loading for an unknown-size model). 3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as None) instead of rejecting it as an unsupported explicit request. 4. Training retry path for "could not get source code" now preserves the gpu_ids parameter so the retry lands on the same GPUs. 5. Updated sandbox tests to cover the new 4-bit inference estimate branch. * Remove accidentally added unsloth-zoo submodule * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix UUID/MIG visibility and update test expectations 1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the visibility APIs now return "unresolved" with empty device lists instead of exposing all physical GPUs. This prevents the UI from showing GPUs that the backend process cannot actually use. 2. test_gpu_selection.py: Updated test expectations to match the new multi-GPU overhead accounting (first GPU at full capacity, 0.85x for additional GPUs) and 4-bit inference memory estimation formula. All 60 tests now pass. * Add CPU/disk offload guard to audio inference path The audio model loading branch returned before the common get_offloaded_device_map_entries() check, so audio models loaded with a multi-GPU device_map that spilled layers to CPU/disk would be accepted instead of rejected. Now audio loads also verify no modules are offloaded. * Improve VRAM requirement estimates * Replace balanced_low_0 with balanced * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refine calculations for slightly easier nums * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * adjust estimates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use nums instead of obj to avoid seralisation error * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden nvidia-smi parsing and fix fallback GPU list 1. nvidia.py: Wrap int() casts for GPU index and memory in try/except so MIG slices, N/A values, or unexpected nvidia-smi output skip the unparseable row instead of aborting the entire GPU list. 2. nvidia.py: Handle GPU names containing commas by using the last field as memory instead of a fixed positional index. 3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified VRAM data) instead of raw devices list, which could include GPUs with null VRAM that were excluded from the ranking. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * consolidate raise_if_offload * Improve MoE support. Guard against nvidia-smi failures * Improve MoE support. Guard against nvidia-smi failures * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case 1. vram_estimation.py: compute_lora_params now includes shared experts (n_shared_experts) alongside routed experts when computing MoE LoRA adapter parameters. Previously only n_experts were counted, causing the estimator to undercount adapter, optimizer, and gradient memory for DeepSeek/GLM-style models with shared experts. 2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which reports system-wide VRAM usage) instead of memory_allocated (which only reports this process's PyTorch allocations). This prevents auto-selection from treating a GPU as mostly free when another process is consuming VRAM. Falls back to memory_allocated when mem_get_info is unavailable. 3. hardware.py: apply_gpu_ids([]) now returns early instead of setting CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty list inherits the parent visibility, same as None. 4. hardware.py: Upgraded fallback_all GPU selection log from debug to warning so operators are notified when the model likely will not fit in available VRAM. * Guard nvidia-smi subprocess calls against OSError and TimeoutExpired get_visible_gpu_utilization and get_backend_visible_gpu_info now catch OSError (nvidia-smi not found) and TimeoutExpired internally instead of relying on callers to wrap every invocation. Returns the standard available=False sentinel on failure so the torch-based fallback in hardware.py can take over. * Guard get_primary_gpu_utilization and reset GPU caches between tests 1. nvidia.py: get_primary_gpu_utilization now catches OSError and TimeoutExpired internally, matching the pattern already used in get_visible_gpu_utilization and get_backend_visible_gpu_info. All three nvidia-smi callers are now self-contained. 2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the module-level _physical_gpu_count and _visible_gpu_count caches in tearDown. Applied to all test classes that exercise GPU selection, device map, or visibility functions. This prevents stale cache values from leaking between tests and causing flaky results on machines with real GPUs. * Fix nvidia-smi fallback regression and physical GPU count validation 1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and get_backend_visible_gpu_info now check result.get("available") before returning the nvidia-smi result. When nvidia-smi is unavailable or returns no data (e.g., containers without nvidia-smi, UUID/MIG masks), the functions fall through to the torch-based fallback instead of returning an empty result. This fixes a regression where the internal exception handling in nvidia.py prevented the caller's except block from triggering the fallback. 2. hardware.py: resolve_requested_gpu_ids now separates negative-ID validation from physical upper-bound validation. The physical count check is only enforced when it is plausibly a true physical count (i.e., higher than the largest parent-visible ID), since torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the visible count, not the physical total. The parent-visible-set check remains authoritative in all cases. This prevents valid physical IDs like [2, 3] from being rejected as "out of range" when nvidia-smi is unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only 2 devices. * Fix UUID/MIG torch fallback to enumerate devices by ordinal When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers, get_parent_visible_gpu_ids() returns [] because the tokens are non-numeric. The torch fallback in get_visible_gpu_utilization() and get_backend_visible_gpu_info() previously passed that empty list to _torch_get_per_device_info(), getting nothing back. Now both functions detect the empty-list case and fall back to enumerating torch-visible ordinals (0..device_count-1) with index_kind="relative". This means the UI and auto-selection still see real device data in Kubernetes, MIG, and Slurm-style UUID environments where nvidia-smi output cannot be mapped to physical indices. Updated test_uuid_parent_visibility to verify the new torch fallback path returns available=True with relative ordinals. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <danielhanchen@gmail.com>
1125 lines
42 KiB
Python
1125 lines
42 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
import os
|
|
import re
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from core.training.training import TrainingBackend
|
|
from models.inference import LoadRequest
|
|
from models.training import TrainingStartRequest
|
|
from utils.hardware import (
|
|
apply_gpu_ids,
|
|
DeviceType,
|
|
auto_select_gpu_ids,
|
|
estimate_required_model_memory_gb,
|
|
get_backend_visible_gpu_info,
|
|
get_device_map,
|
|
get_offloaded_device_map_entries,
|
|
get_parent_visible_gpu_ids,
|
|
get_visible_gpu_utilization,
|
|
prepare_gpu_selection,
|
|
resolve_requested_gpu_ids,
|
|
)
|
|
import utils.hardware.hardware as _hw_module
|
|
|
|
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
|
|
def _load_route_module(name: str, relative_path: str):
|
|
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class _GpuCacheResetMixin:
|
|
"""Reset module-level GPU caches between tests to prevent state leaks."""
|
|
|
|
def tearDown(self):
|
|
_hw_module._physical_gpu_count = None
|
|
_hw_module._visible_gpu_count = None
|
|
|
|
|
|
class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_parent_visibility_defaults_to_physical_enumeration(self):
|
|
with (
|
|
patch.dict(os.environ, {}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4),
|
|
):
|
|
self.assertEqual(get_parent_visible_gpu_ids(), [0, 1, 2, 3])
|
|
self.assertEqual(resolve_requested_gpu_ids(None), [0, 1, 2, 3])
|
|
|
|
def test_parent_visibility_uses_cuda_visible_devices(self):
|
|
with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True):
|
|
self.assertEqual(get_parent_visible_gpu_ids(), [1, 3])
|
|
self.assertEqual(resolve_requested_gpu_ids(None), [1, 3])
|
|
|
|
def test_parent_visibility_uses_empty_numeric_ids_for_uuid_masks(self):
|
|
with (
|
|
patch.dict(
|
|
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True
|
|
),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
self.assertEqual(get_parent_visible_gpu_ids(), [])
|
|
|
|
def test_invalid_requests_raise_clear_value_errors(self):
|
|
cases = [
|
|
([1, 1], "duplicate GPU IDs"),
|
|
([-1], "Rejected IDs: [-1]"),
|
|
([99], "Rejected IDs: [99]"),
|
|
([0], "outside the parent-visible set [1, 3]"),
|
|
]
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
for gpu_ids, message in cases:
|
|
with self.subTest(gpu_ids = gpu_ids):
|
|
with self.assertRaisesRegex(ValueError, re.escape(message)):
|
|
resolve_requested_gpu_ids(gpu_ids)
|
|
|
|
def test_explicit_ids_must_be_physical_not_relative(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
self.assertEqual(resolve_requested_gpu_ids([1, 3]), [1, 3])
|
|
|
|
def test_explicit_ids_are_rejected_for_uuid_parent_visibility(self):
|
|
with (
|
|
patch.dict(
|
|
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True
|
|
),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
with self.assertRaisesRegex(
|
|
ValueError, "unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG"
|
|
):
|
|
resolve_requested_gpu_ids([1])
|
|
|
|
def test_empty_list_is_treated_as_auto(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
self.assertEqual(resolve_requested_gpu_ids([]), [1, 3])
|
|
|
|
def test_apply_gpu_ids_only_updates_cuda_visible_devices(self):
|
|
with patch.dict(
|
|
os.environ,
|
|
{"CUDA_VISIBLE_DEVICES": "1,3", "TEST_PARENT_ENV": "keep-me"},
|
|
clear = True,
|
|
):
|
|
apply_gpu_ids([5, 6])
|
|
|
|
self.assertEqual(os.environ["CUDA_VISIBLE_DEVICES"], "5,6")
|
|
self.assertEqual(os.environ["TEST_PARENT_ENV"], "keep-me")
|
|
|
|
|
|
class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_visible_gpu_utilization_filters_to_parent_visible_ids(self):
|
|
smi_output = "\n".join(
|
|
[
|
|
"0, 10, 30, 1000, 10000, 50, 100",
|
|
"1, 20, 40, 2000, 10000, 60, 120",
|
|
"3, 30, 50, 3000, 10000, 70, 140",
|
|
]
|
|
)
|
|
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.nvidia.subprocess.run") as mock_run,
|
|
):
|
|
mock_run.return_value = SimpleNamespace(
|
|
returncode = 0,
|
|
stdout = smi_output,
|
|
)
|
|
result = get_visible_gpu_utilization()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [1, 3])
|
|
self.assertEqual(result["index_kind"], "physical")
|
|
self.assertEqual([device["index"] for device in result["devices"]], [1, 3])
|
|
self.assertEqual(result["devices"][0]["visible_ordinal"], 0)
|
|
self.assertEqual(result["devices"][1]["visible_ordinal"], 1)
|
|
self.assertEqual(result["devices"][0]["gpu_utilization_pct"], 20.0)
|
|
self.assertEqual(result["devices"][1]["power_utilization_pct"], 50.0)
|
|
|
|
def test_backend_visible_gpu_info_preserves_physical_indices(self):
|
|
smi_output = "\n".join(
|
|
[
|
|
"0, GPU Zero, 10000",
|
|
"1, GPU One, 20000",
|
|
"3, GPU Three, 30000",
|
|
]
|
|
)
|
|
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.nvidia.subprocess.run") as mock_run,
|
|
):
|
|
mock_run.return_value = SimpleNamespace(
|
|
returncode = 0,
|
|
stdout = smi_output,
|
|
)
|
|
result = get_backend_visible_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [1, 3])
|
|
self.assertEqual(result["index_kind"], "physical")
|
|
self.assertEqual([device["index"] for device in result["devices"]], [1, 3])
|
|
self.assertEqual(result["devices"][0]["visible_ordinal"], 0)
|
|
self.assertEqual(result["devices"][1]["visible_ordinal"], 1)
|
|
self.assertEqual(result["devices"][0]["name"], "GPU One")
|
|
self.assertAlmostEqual(result["devices"][1]["memory_total_gb"], 29.3, places = 1)
|
|
|
|
def test_uuid_parent_visibility_falls_back_to_torch(self):
|
|
"""UUID/MIG masks should fall through nvidia to torch fallback and
|
|
still report visible devices using relative ordinals."""
|
|
fake_torch_devices = [
|
|
{
|
|
"index": 0,
|
|
"visible_ordinal": 0,
|
|
"name": "GPU-A",
|
|
"total_gb": 24.0,
|
|
"used_gb": 2.0,
|
|
},
|
|
{
|
|
"index": 1,
|
|
"visible_ordinal": 1,
|
|
"name": "GPU-B",
|
|
"total_gb": 24.0,
|
|
"used_gb": 3.0,
|
|
},
|
|
]
|
|
with (
|
|
patch.dict(
|
|
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True
|
|
),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._torch_get_per_device_info",
|
|
return_value = fake_torch_devices,
|
|
),
|
|
):
|
|
result = get_backend_visible_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [])
|
|
self.assertEqual(len(result["devices"]), 2)
|
|
self.assertEqual(result["index_kind"], "relative")
|
|
|
|
def test_mlx_visible_gpu_info_is_best_effort_relative(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX),
|
|
patch(
|
|
"utils.hardware.hardware.get_gpu_memory_info",
|
|
return_value = {
|
|
"available": True,
|
|
"device_name": "Apple Silicon",
|
|
"total_gb": 64.0,
|
|
"allocated_gb": 8.0,
|
|
"utilization_pct": 12.5,
|
|
},
|
|
),
|
|
):
|
|
result = get_backend_visible_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["index_kind"], "relative")
|
|
self.assertEqual(result["devices"][0]["index"], 0)
|
|
self.assertEqual(result["devices"][0]["visible_ordinal"], 0)
|
|
|
|
|
|
class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_get_device_map_uses_explicit_gpu_selection(self):
|
|
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA):
|
|
self.assertEqual(get_device_map(None), "sequential")
|
|
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(
|
|
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True
|
|
),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
):
|
|
self.assertEqual(get_device_map(None), "balanced")
|
|
|
|
def test_get_offloaded_device_map_entries_returns_only_cpu_and_disk(self):
|
|
model = SimpleNamespace(
|
|
hf_device_map = {
|
|
"model.embed_tokens": 0,
|
|
"model.layers.0": 1,
|
|
"model.layers.1": "cpu",
|
|
"lm_head": "disk",
|
|
}
|
|
)
|
|
|
|
self.assertEqual(
|
|
get_offloaded_device_map_entries(model),
|
|
{
|
|
"model.layers.1": "cpu",
|
|
"lm_head": "disk",
|
|
},
|
|
)
|
|
|
|
def test_get_offloaded_device_map_entries_handles_models_without_device_map(self):
|
|
self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {})
|
|
|
|
def test_estimate_required_memory_formulas(self):
|
|
eight_gb = 8 * (1024**3)
|
|
|
|
with patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (eight_gb, "config"),
|
|
):
|
|
# FP16 inference: 8GB * 1.3 = 10.4GB
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
load_in_4bit = False,
|
|
)
|
|
self.assertAlmostEqual(required_gb, 10.4, places = 3)
|
|
self.assertEqual(metadata["model_size_source"], "config")
|
|
|
|
# 4bit inference: base_4bit = 8/3.2 = 2.5GB
|
|
# required = 2.5 + max(2.5*0.3, 2.0) = 2.5 + 2.0 = 4.5GB
|
|
required_gb, _ = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
load_in_4bit = True,
|
|
)
|
|
self.assertAlmostEqual(required_gb, 4.5, places = 2)
|
|
|
|
# Full FT fallback: model_size * 3.5 + overhead
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test", training_type = "Full Finetuning"
|
|
)
|
|
self.assertEqual(metadata.get("estimation_mode"), "fallback")
|
|
self.assertGreater(required_gb, 25.0)
|
|
self.assertLess(required_gb, 40.0)
|
|
|
|
# LoRA fp16 fallback: model_size + lora_overhead + activations + overhead
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = False,
|
|
)
|
|
self.assertEqual(metadata.get("estimation_mode"), "fallback")
|
|
self.assertGreater(required_gb, 8.0)
|
|
self.assertLess(required_gb, 15.0)
|
|
|
|
# QLoRA 4-bit fallback: compressed weights + lora overhead + activations + overhead
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
self.assertEqual(metadata.get("estimation_mode"), "fallback")
|
|
self.assertGreater(required_gb, 3.0)
|
|
self.assertLess(required_gb, 8.0)
|
|
|
|
# Larger model: 16GB fp16
|
|
sixteen_gb = 16 * (1024**3)
|
|
with patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (sixteen_gb, "config"),
|
|
):
|
|
required_gb, _ = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
# QLoRA for 16GB model should be < 12 GB
|
|
self.assertGreater(required_gb, 5.0)
|
|
self.assertLess(required_gb, 12.0)
|
|
|
|
def test_estimate_fp16_model_size_bytes_uses_vllm_fallback_last(self):
|
|
config = object()
|
|
with (
|
|
patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "unsloth/test",
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_hf_safetensors_total_params",
|
|
return_value = None,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
return_value = config,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._estimate_fp16_model_size_bytes_from_config",
|
|
return_value = None,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_local_weight_size_bytes",
|
|
return_value = None,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._estimate_fp16_model_size_bytes_from_vllm_utils",
|
|
return_value = 1234,
|
|
),
|
|
):
|
|
model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes(
|
|
"unsloth/test"
|
|
)
|
|
|
|
self.assertEqual(model_size_bytes, 1234)
|
|
self.assertEqual(source, "vllm_utils")
|
|
|
|
def test_auto_select_gpu_ids_chooses_smallest_fitting_subset(self):
|
|
fake_devices = {
|
|
"devices": [
|
|
{"index": 0, "vram_total_gb": 16.0, "vram_used_gb": 4.0},
|
|
{"index": 1, "vram_total_gb": 16.0, "vram_used_gb": 6.0},
|
|
{"index": 2, "vram_total_gb": 16.0, "vram_used_gb": 7.0},
|
|
]
|
|
}
|
|
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
14.0,
|
|
{"required_gb": 14.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_visible_gpu_utilization",
|
|
return_value = fake_devices,
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "auto")
|
|
# First GPU full (12GB) + second GPU with overhead (10*0.85=8.5) = 20.5GB
|
|
self.assertAlmostEqual(metadata["usable_gb"], 20.5, places = 3)
|
|
|
|
def test_auto_select_gpu_ids_falls_back_to_all_visible(self):
|
|
fake_devices = {
|
|
"devices": [
|
|
{"index": 0, "vram_total_gb": 12.0, "vram_used_gb": 2.0},
|
|
{"index": 1, "vram_total_gb": 12.0, "vram_used_gb": 2.0},
|
|
]
|
|
}
|
|
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
30.0,
|
|
{"required_gb": 30.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_visible_gpu_utilization",
|
|
return_value = fake_devices,
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "fallback_all")
|
|
# First GPU full (10GB) + second GPU with overhead (10*0.85=8.5) = 18.5GB
|
|
self.assertAlmostEqual(metadata["usable_gb"], 18.5, places = 3)
|
|
|
|
def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self):
|
|
with (
|
|
patch(
|
|
"utils.hardware.hardware.resolve_requested_gpu_ids",
|
|
return_value = [2, 3],
|
|
),
|
|
patch("utils.hardware.hardware.auto_select_gpu_ids") as mock_auto_select,
|
|
):
|
|
selected, metadata = prepare_gpu_selection(
|
|
[2, 3],
|
|
model_name = "unsloth/test",
|
|
)
|
|
|
|
self.assertEqual(selected, [2, 3])
|
|
self.assertEqual(metadata["selection_mode"], "explicit")
|
|
mock_auto_select.assert_not_called()
|
|
|
|
def test_prepare_gpu_selection_treats_empty_list_as_auto(self):
|
|
with patch(
|
|
"utils.hardware.hardware.auto_select_gpu_ids",
|
|
return_value = ([0, 1], {"selection_mode": "auto"}),
|
|
) as mock_auto_select:
|
|
selected, metadata = prepare_gpu_selection(
|
|
[],
|
|
model_name = "unsloth/test",
|
|
)
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "auto")
|
|
mock_auto_select.assert_called_once()
|
|
|
|
def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self):
|
|
with (
|
|
patch.dict(
|
|
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
14.0,
|
|
{"required_gb": 14.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
):
|
|
selected, metadata = prepare_gpu_selection(
|
|
None,
|
|
model_name = "unsloth/test",
|
|
)
|
|
|
|
self.assertIsNone(selected)
|
|
self.assertEqual(metadata["selection_mode"], "inherit_parent_visible")
|
|
self.assertIsNone(metadata["selected_gpu_ids"])
|
|
|
|
|
|
class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_training_backend_resolves_explicit_gpu_ids_before_spawn(self):
|
|
backend = TrainingBackend()
|
|
|
|
class DummyProcess:
|
|
pid = 12345
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
class DummyThread:
|
|
def start(self):
|
|
return None
|
|
|
|
dummy_queue = object()
|
|
|
|
with (
|
|
patch(
|
|
"core.training.training.prepare_gpu_selection",
|
|
return_value = ([1, 2], {"selection_mode": "explicit"}),
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Queue",
|
|
side_effect = [dummy_queue, dummy_queue],
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Process", return_value = DummyProcess()
|
|
) as mock_process,
|
|
patch(
|
|
"core.training.training.threading.Thread", return_value = DummyThread()
|
|
),
|
|
):
|
|
backend.start_training(
|
|
job_id = "test-job-1",
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
gpu_ids = [1, 2],
|
|
)
|
|
|
|
config = mock_process.call_args.kwargs["kwargs"]["config"]
|
|
self.assertEqual(config["gpu_ids"], [1, 2])
|
|
self.assertEqual(config["resolved_gpu_ids"], [1, 2])
|
|
self.assertEqual(config["gpu_selection"]["selection_mode"], "explicit")
|
|
|
|
def test_training_backend_auto_selects_gpu_ids_when_omitted(self):
|
|
backend = TrainingBackend()
|
|
|
|
class DummyProcess:
|
|
pid = 12345
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
class DummyThread:
|
|
def start(self):
|
|
return None
|
|
|
|
dummy_queue = object()
|
|
|
|
with (
|
|
patch(
|
|
"core.training.training.prepare_gpu_selection",
|
|
return_value = ([0, 1], {"selection_mode": "auto"}),
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Queue",
|
|
side_effect = [dummy_queue, dummy_queue],
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Process", return_value = DummyProcess()
|
|
) as mock_process,
|
|
patch(
|
|
"core.training.training.threading.Thread", return_value = DummyThread()
|
|
),
|
|
):
|
|
backend.start_training(
|
|
job_id = "test-job-2",
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
gpu_ids = None,
|
|
)
|
|
|
|
config = mock_process.call_args.kwargs["kwargs"]["config"]
|
|
self.assertIsNone(config["gpu_ids"])
|
|
self.assertEqual(config["resolved_gpu_ids"], [0, 1])
|
|
self.assertEqual(config["gpu_selection"]["selection_mode"], "auto")
|
|
|
|
def test_training_backend_preserves_uuid_parent_visibility_in_auto_mode(self):
|
|
backend = TrainingBackend()
|
|
|
|
class DummyProcess:
|
|
pid = 12345
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
class DummyThread:
|
|
def start(self):
|
|
return None
|
|
|
|
dummy_queue = object()
|
|
|
|
with (
|
|
patch.dict(
|
|
os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Queue",
|
|
side_effect = [dummy_queue, dummy_queue],
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Process", return_value = DummyProcess()
|
|
) as mock_process,
|
|
patch(
|
|
"core.training.training.threading.Thread", return_value = DummyThread()
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
14.0,
|
|
{"required_gb": 14.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
):
|
|
backend.start_training(
|
|
job_id = "test-job-uuid-auto",
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
gpu_ids = None,
|
|
)
|
|
|
|
config = mock_process.call_args.kwargs["kwargs"]["config"]
|
|
self.assertIsNone(config["resolved_gpu_ids"])
|
|
self.assertEqual(
|
|
config["gpu_selection"]["selection_mode"], "inherit_parent_visible"
|
|
)
|
|
|
|
def test_inference_orchestrator_resolves_explicit_gpu_ids_before_spawn(self):
|
|
class DummyThread:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
with patch("core.inference.orchestrator.threading.Thread", DummyThread):
|
|
from core.inference.orchestrator import InferenceOrchestrator
|
|
|
|
orchestrator = InferenceOrchestrator()
|
|
|
|
config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None)
|
|
|
|
with (
|
|
patch(
|
|
"core.inference.orchestrator.prepare_gpu_selection",
|
|
return_value = ([1], {"selection_mode": "explicit"}),
|
|
),
|
|
patch.object(orchestrator, "_ensure_subprocess_alive", return_value = False),
|
|
patch.object(orchestrator, "_spawn_subprocess") as mock_spawn,
|
|
patch.object(
|
|
orchestrator,
|
|
"_wait_response",
|
|
return_value = {"success": True, "model_info": {}},
|
|
),
|
|
patch(
|
|
"utils.transformers_version.needs_transformers_5", return_value = False
|
|
),
|
|
):
|
|
self.assertTrue(orchestrator.load_model(config = config, gpu_ids = [1]))
|
|
|
|
sub_config = mock_spawn.call_args.args[0]
|
|
self.assertEqual(sub_config["gpu_ids"], [1])
|
|
self.assertEqual(sub_config["resolved_gpu_ids"], [1])
|
|
self.assertEqual(sub_config["gpu_selection"]["selection_mode"], "explicit")
|
|
|
|
def test_inference_orchestrator_auto_selects_gpu_ids_when_omitted(self):
|
|
class DummyThread:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
with patch("core.inference.orchestrator.threading.Thread", DummyThread):
|
|
from core.inference.orchestrator import InferenceOrchestrator
|
|
|
|
orchestrator = InferenceOrchestrator()
|
|
|
|
config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None)
|
|
|
|
with (
|
|
patch(
|
|
"core.inference.orchestrator.prepare_gpu_selection",
|
|
return_value = ([0], {"selection_mode": "auto"}),
|
|
),
|
|
patch.object(orchestrator, "_ensure_subprocess_alive", return_value = False),
|
|
patch.object(orchestrator, "_spawn_subprocess") as mock_spawn,
|
|
patch.object(
|
|
orchestrator,
|
|
"_wait_response",
|
|
return_value = {"success": True, "model_info": {}},
|
|
),
|
|
patch(
|
|
"utils.transformers_version.needs_transformers_5", return_value = False
|
|
),
|
|
):
|
|
self.assertTrue(orchestrator.load_model(config = config, gpu_ids = None))
|
|
|
|
sub_config = mock_spawn.call_args.args[0]
|
|
self.assertIsNone(sub_config["gpu_ids"])
|
|
self.assertEqual(sub_config["resolved_gpu_ids"], [0])
|
|
self.assertEqual(sub_config["gpu_selection"]["selection_mode"], "auto")
|
|
|
|
|
|
class TestRouteErrors(unittest.TestCase):
|
|
def test_prepare_gpu_selection_rejects_gpu_ids_on_non_cuda_backend(self):
|
|
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
|
|
with self.assertRaises(ValueError) as exc_info:
|
|
prepare_gpu_selection([0], model_name = "unsloth/test")
|
|
|
|
self.assertIn("only supported on CUDA devices", str(exc_info.exception))
|
|
|
|
def test_inference_route_rejects_gpu_ids_for_gguf(self):
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_gguf_gpu_ids_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = None,
|
|
gguf_file = "/tmp/test.gguf",
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = None,
|
|
identifier = "unsloth/test.gguf",
|
|
display_name = "unsloth/test.gguf",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
with patch.object(
|
|
inference_route.ModelConfig,
|
|
"from_identifier",
|
|
return_value = model_config,
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route.load_model(request, current_subject = "test-user")
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("GGUF", exc_info.exception.detail)
|
|
|
|
def test_training_route_returns_400_for_invalid_gpu_ids(self):
|
|
training_route = _load_route_module(
|
|
"training_route_module_for_test",
|
|
"routes/training.py",
|
|
)
|
|
request = TrainingStartRequest(
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
format_type = "alpaca",
|
|
gpu_ids = [99],
|
|
)
|
|
|
|
class DummyBackend:
|
|
current_job_id = None
|
|
|
|
def is_training_active(self):
|
|
return False
|
|
|
|
def start_training(self, **kwargs):
|
|
raise ValueError("Invalid gpu_ids [99]")
|
|
|
|
with (
|
|
patch.object(
|
|
training_route, "get_training_backend", return_value = DummyBackend()
|
|
),
|
|
patch(
|
|
"core.inference.get_inference_backend",
|
|
return_value = SimpleNamespace(active_model_name = None),
|
|
),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
training_route.start_training(request, current_subject = "test-user")
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("gpu_ids [99]", exc_info.exception.detail)
|
|
|
|
def test_training_route_returns_400_for_uuid_parent_visibility_gpu_ids(self):
|
|
training_route = _load_route_module(
|
|
"training_route_module_for_uuid_parent_visibility_test",
|
|
"routes/training.py",
|
|
)
|
|
request = TrainingStartRequest(
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
format_type = "alpaca",
|
|
gpu_ids = [1],
|
|
)
|
|
|
|
class DummyBackend:
|
|
current_job_id = None
|
|
|
|
def is_training_active(self):
|
|
return False
|
|
|
|
def start_training(self, **kwargs):
|
|
raise ValueError(
|
|
"Invalid gpu_ids [1]: explicit physical GPU IDs are unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries"
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
training_route, "get_training_backend", return_value = DummyBackend()
|
|
),
|
|
patch(
|
|
"core.inference.get_inference_backend",
|
|
return_value = SimpleNamespace(active_model_name = None),
|
|
),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
training_route.start_training(request, current_subject = "test-user")
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("UUID/MIG", exc_info.exception.detail)
|
|
|
|
def test_inference_route_returns_400_for_invalid_gpu_ids(self):
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test", gpu_ids = [99])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = False,
|
|
is_lora = False,
|
|
path = None,
|
|
identifier = "unsloth/test",
|
|
display_name = "unsloth/test",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
class DummyInferenceBackend:
|
|
active_model_name = None
|
|
models = {}
|
|
|
|
def load_model(self, **kwargs):
|
|
raise ValueError("Invalid gpu_ids [99]")
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route.ModelConfig,
|
|
"from_identifier",
|
|
return_value = model_config,
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_inference_backend",
|
|
return_value = DummyInferenceBackend(),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_llama_cpp_backend",
|
|
return_value = SimpleNamespace(is_loaded = False),
|
|
),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route.load_model(request, current_subject = "test-user")
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("gpu_ids [99]", exc_info.exception.detail)
|
|
|
|
def test_inference_route_returns_400_for_uuid_parent_visibility_gpu_ids(self):
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_uuid_parent_visibility_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test", gpu_ids = [1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = False,
|
|
is_lora = False,
|
|
path = None,
|
|
identifier = "unsloth/test",
|
|
display_name = "unsloth/test",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
class DummyInferenceBackend:
|
|
active_model_name = None
|
|
models = {}
|
|
|
|
def load_model(self, **kwargs):
|
|
raise ValueError(
|
|
"Invalid gpu_ids [1]: explicit physical GPU IDs are unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries"
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route.ModelConfig,
|
|
"from_identifier",
|
|
return_value = model_config,
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_inference_backend",
|
|
return_value = DummyInferenceBackend(),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_llama_cpp_backend",
|
|
return_value = SimpleNamespace(is_loaded = False),
|
|
),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route.load_model(request, current_subject = "test-user")
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("UUID/MIG", exc_info.exception.detail)
|
|
|
|
|
|
class TestRaiseIfOffloaded(unittest.TestCase):
|
|
def test_no_offload_is_noop(self):
|
|
from utils.hardware import raise_if_offloaded
|
|
|
|
model = SimpleNamespace(hf_device_map = {"model.embed_tokens": 0, "lm_head": 1})
|
|
raise_if_offloaded(model, "balanced", "Test")
|
|
|
|
def test_cpu_offload_raises(self):
|
|
from utils.hardware import raise_if_offloaded
|
|
|
|
model = SimpleNamespace(
|
|
hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"}
|
|
)
|
|
with self.assertRaisesRegex(ValueError, "offloaded"):
|
|
raise_if_offloaded(model, "balanced", "Test")
|
|
|
|
def test_no_device_map_attr_is_noop(self):
|
|
from utils.hardware import raise_if_offloaded
|
|
|
|
raise_if_offloaded(SimpleNamespace(), "sequential", "Test")
|
|
|
|
|
|
class TestMinGpuVram(unittest.TestCase):
|
|
def test_min_gpu_vram_decreases_with_more_gpus(self):
|
|
from utils.hardware.vram_estimation import (
|
|
ModelArchConfig,
|
|
TrainingVramConfig,
|
|
estimate_training_vram,
|
|
)
|
|
|
|
arch = ModelArchConfig(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
)
|
|
config = TrainingVramConfig(
|
|
training_method = "qlora",
|
|
load_in_4bit = True,
|
|
)
|
|
breakdown = estimate_training_vram(arch, config)
|
|
v1 = breakdown.min_gpu_vram(1)
|
|
v2 = breakdown.min_gpu_vram(2)
|
|
v4 = breakdown.min_gpu_vram(4)
|
|
self.assertGreater(v1, v2)
|
|
self.assertGreater(v2, v4)
|
|
self.assertGreater(v4, 0)
|
|
|
|
def test_total_equals_min_gpu_vram_1(self):
|
|
from utils.hardware.vram_estimation import (
|
|
ModelArchConfig,
|
|
TrainingVramConfig,
|
|
estimate_training_vram,
|
|
)
|
|
|
|
arch = ModelArchConfig(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
)
|
|
config = TrainingVramConfig(
|
|
training_method = "qlora",
|
|
load_in_4bit = True,
|
|
)
|
|
breakdown = estimate_training_vram(arch, config)
|
|
self.assertEqual(breakdown.total, breakdown.min_gpu_vram(1))
|
|
|
|
|
|
class TestPerGpuFitGuardAllCounts(unittest.TestCase):
|
|
def test_min_per_gpu_generated_for_all_visible_counts(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (8 * (1024**3), "config"),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "unsloth/test",
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
return_value = SimpleNamespace(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
),
|
|
),
|
|
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 6),
|
|
):
|
|
_, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
|
|
self.assertEqual(metadata.get("estimation_mode"), "detailed")
|
|
breakdown = metadata["vram_breakdown"]
|
|
for n in range(1, 7):
|
|
self.assertIn(f"min_per_gpu_{n}", breakdown)
|
|
|
|
|
|
class TestAutoSelectWithNoneRequired(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_auto_select_falls_back_when_estimate_unavailable(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (None, {"model_size_source": "unavailable"}),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_parent_visible_gpu_spec",
|
|
return_value = {
|
|
"raw": "0,1",
|
|
"numeric_ids": [0, 1],
|
|
"supports_explicit_gpu_ids": True,
|
|
},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_parent_visible_gpu_ids",
|
|
return_value = [0, 1],
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "fallback_all")
|
|
|
|
|
|
class TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_auto_select_returns_non_cuda_for_xpu(self):
|
|
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertIsNone(selected)
|
|
self.assertEqual(metadata["selection_mode"], "non_cuda")
|
|
|
|
def test_prepare_gpu_selection_rejects_explicit_ids_on_xpu(self):
|
|
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")
|