From 9311df2b2956893387ef926d684dd1e1cea73cc5 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Mon, 30 Mar 2026 15:03:15 +0530 Subject: [PATCH] [Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602) * [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 --- studio/backend/core/inference/inference.py | 45 +- studio/backend/core/inference/llama_cpp.py | 23 +- studio/backend/core/inference/orchestrator.py | 11 + studio/backend/core/inference/worker.py | 4 + studio/backend/core/training/trainer.py | 26 +- studio/backend/core/training/training.py | 18 + studio/backend/core/training/worker.py | 4 + studio/backend/main.py | 81 +- studio/backend/models/inference.py | 4 + studio/backend/models/training.py | 6 + studio/backend/routes/inference.py | 13 + studio/backend/routes/training.py | 16 +- studio/backend/tests/test_gpu_selection.py | 1125 ++++++++++++++++ .../tests/test_gpu_selection_sandbox.py | 544 ++++++++ studio/backend/tests/test_utils.py | 32 +- studio/backend/tests/test_vram_estimation.py | 695 ++++++++++ .../backend/utils/hardware/VRAM_ESTIMATION.md | 161 +++ studio/backend/utils/hardware/__init__.py | 37 + studio/backend/utils/hardware/hardware.py | 1143 +++++++++++++++-- studio/backend/utils/hardware/nvidia.py | 279 ++++ .../backend/utils/hardware/vram_estimation.py | 501 ++++++++ 21 files changed, 4541 insertions(+), 227 deletions(-) create mode 100644 studio/backend/tests/test_gpu_selection.py create mode 100644 studio/backend/tests/test_gpu_selection_sandbox.py create mode 100644 studio/backend/tests/test_vram_estimation.py create mode 100644 studio/backend/utils/hardware/VRAM_ESTIMATION.md create mode 100644 studio/backend/utils/hardware/nvidia.py create mode 100644 studio/backend/utils/hardware/vram_estimation.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 1a265690ff..ddd485525d 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -18,7 +18,14 @@ from typing import Optional, Union, Generator, Tuple from utils.models import ModelConfig, get_base_model_from_lora from utils.paths import is_model_cached from utils.utils import format_error_message -from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory +from utils.hardware import ( + get_device, + clear_gpu_cache, + log_gpu_memory, + get_device_map, + raise_if_offloaded, + get_visible_gpu_count, +) from core.inference.audio_codecs import AudioCodecManager from io import StringIO import structlog @@ -241,6 +248,7 @@ class InferenceBackend: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """ Load any model: base, LoRA adapter, text, or vision. @@ -260,6 +268,10 @@ class InferenceBackend: return False self.loading_models.add(model_name) + device_map = get_device_map(gpu_ids, for_inference = True) + logger.info( + f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" + ) self.models[model_name] = { "is_vision": config.is_vision, @@ -290,6 +302,7 @@ class InferenceBackend: config.path, auto_model = CsmForConditionalGeneration, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -325,6 +338,7 @@ class InferenceBackend: config.path, dtype = torch.float32, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -345,6 +359,7 @@ class InferenceBackend: llm_path, dtype = torch.float32, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -361,6 +376,7 @@ class InferenceBackend: config.path, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -378,6 +394,7 @@ class InferenceBackend: whisper_language = "English", whisper_task = "transcribe", load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -405,6 +422,7 @@ class InferenceBackend: model_name = config.path, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -420,6 +438,11 @@ class InferenceBackend: audio_type, self.device, model_repo_path = model_repo_path ) + # Reject CPU/disk offload for audio models too + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) + self.active_model_name = model_name self.loading_models.discard(model_name) logger.info(f"Successfully loaded audio model: {model_name}") @@ -441,6 +464,7 @@ class InferenceBackend: max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -497,6 +521,7 @@ class InferenceBackend: max_seq_length = max_seq_length, dtype = dtype, load_in_4bit = load_in_4bit, + device_map = device_map, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, ) @@ -507,6 +532,10 @@ class InferenceBackend: self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = tokenizer + raise_if_offloaded( + self.models[model_name]["model"], device_map, "Inference" + ) + # Load chat template info self._load_chat_template_info(model_name) @@ -615,6 +644,7 @@ class InferenceBackend: dtype = None, load_in_4bit: bool = True, hf_token: Optional[str] = None, + gpu_ids: Optional[list[int]] = None, ) -> Tuple[bool, Optional[str], Optional[str]]: """ Final Corrected Version: @@ -639,7 +669,12 @@ class InferenceBackend: base_model_name, None, is_lora = False ) if not self.load_model( - base_config, max_seq_length, dtype, load_in_4bit, hf_token + base_config, + max_seq_length, + dtype, + load_in_4bit, + hf_token, + gpu_ids = gpu_ids, ): return False, None, None @@ -1037,12 +1072,12 @@ class InferenceBackend: input_text, add_special_tokens = False, return_tensors = "pt", - ).to(self.device) + ).to(model.device) else: # Text-only for vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to( - self.device + model.device ) # Stream with TextIteratorStreamer + background thread @@ -1182,7 +1217,7 @@ class InferenceBackend: return_dict = True, return_tensors = "pt", truncation = False, - ).to(self.device) + ).to(model.device) try: from transformers import TextIteratorStreamer diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index ca39054ec0..f5361a6e8c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -293,7 +293,8 @@ class LlamaCppBackend: continue gpus.append((idx, free_mib)) return gpus - except Exception: + except Exception as e: + logger.debug(f"Failed to query GPU free memory via nvidia-smi: {e}") return [] @staticmethod @@ -334,6 +335,11 @@ class LlamaCppBackend: return sorted(selected), False # Model is too large even for all GPUs, let --fit handle it + logger.debug( + "Model does not fit in available GPU memory, falling back to --fit", + model_size_mib = round(model_size_mib, 2), + ranked_gpus = ranked, + ) return None, True # ── KV cache VRAM estimation ───────────────────────────────────── @@ -392,6 +398,11 @@ class LlamaCppBackend: If the model weights alone don't fit, returns min_ctx unchanged. """ if not self._can_estimate_kv(): + logger.debug( + "Skipping context fit because KV cache metadata is unavailable", + requested_ctx = requested_ctx, + available_mib = available_mib, + ) return requested_ctx budget_bytes = available_mib * 1024 * 1024 * 0.70 @@ -405,6 +416,12 @@ class LlamaCppBackend: # Model weights alone exceed budget -- can't help by reducing ctx. # Return requested_ctx unchanged; --fit will handle VRAM management. if model_footprint >= budget_bytes: + logger.debug( + "Model footprint exceeds GPU budget before KV cache", + requested_ctx = requested_ctx, + available_mib = available_mib, + model_size_gb = round(model_footprint / (1024**3), 2), + ) return requested_ctx # Binary search for max context that fits @@ -1082,6 +1099,10 @@ class LlamaCppBackend: # Can't estimate KV -- fall back to file-size-only check. # Without KV estimation we cannot prove a hardware cap, so # keep the ceiling at the native context (already the default). + logger.debug( + "Falling back to file-size-only GPU selection", + model_size_gb = round(model_size / (1024**3), 2), + ) gpu_indices, use_fit = self._select_gpus(model_size, gpus) if effective_ctx < original_ctx: diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 6ff7fd2cbf..78cc60e1c6 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -27,6 +27,7 @@ import uuid from io import BytesIO from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union +from utils.hardware import prepare_gpu_selection logger = get_logger(__name__) @@ -571,6 +572,7 @@ class InferenceOrchestrator: load_in_4bit: bool = True, hf_token: Optional[str] = None, trust_remote_code: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """Load a model for inference. @@ -594,7 +596,16 @@ class InferenceOrchestrator: "hf_token": hf_token or "", "gguf_variant": getattr(config, "gguf_variant", None), "trust_remote_code": trust_remote_code, + "gpu_ids": gpu_ids, } + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + gpu_ids, + model_name = model_name, + hf_token = hf_token, + load_in_4bit = load_in_4bit, + ) + sub_config["resolved_gpu_ids"] = resolved_gpu_ids + sub_config["gpu_selection"] = gpu_selection # Always kill existing subprocess and spawn fresh. # Reusing a subprocess after unsloth patches torch internals diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index afe0ecc458..b3ce43795c 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -29,6 +29,7 @@ from pathlib import Path from typing import Any logger = get_logger(__name__) +from utils.hardware import apply_gpu_ids def _activate_transformers_version(model_name: str) -> None: @@ -178,6 +179,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: load_in_4bit = load_in_4bit, hf_token = hf_token, trust_remote_code = trust_remote_code, + gpu_ids = config.get("resolved_gpu_ids"), ) if success: @@ -501,6 +503,8 @@ def run_inference_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) + apply_gpu_ids(config.get("resolved_gpu_ids")) + model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 2324916236..9b4a14f09f 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -33,7 +33,14 @@ if sys.platform in ("win32", "darwin"): sys.path.insert(0, _compile_cache) import torch -from utils.hardware import clear_gpu_cache, safe_num_proc, dataset_map_num_proc +from utils.hardware import ( + clear_gpu_cache, + safe_num_proc, + dataset_map_num_proc, + get_device_map, + raise_if_offloaded, + get_visible_gpu_count, +) torch._dynamo.config.recompile_limit = 64 from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported @@ -487,6 +494,7 @@ class UnslothTrainer: is_dataset_audio: bool = False, trust_remote_code: bool = False, full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, ) -> bool: """Load model for training (supports both text and vision models)""" self.load_in_4bit = load_in_4bit # Store for training_meta.json @@ -624,6 +632,11 @@ class UnslothTrainer: self._update_progress(error = friendly, is_training = False) return False + device_map = get_device_map(gpu_ids) + logger.info( + f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)" + ) + # Branch based on model type if self._audio_type == "csm": # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False @@ -636,6 +649,7 @@ class UnslothTrainer: dtype = None, auto_model = CsmForConditionalGeneration, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -651,6 +665,7 @@ class UnslothTrainer: model_name = model_name, dtype = None, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, auto_model = WhisperForConditionalGeneration, whisper_language = "English", @@ -672,6 +687,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -711,6 +727,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = torch.float32, # Spark-TTS requires float32 load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -725,6 +742,7 @@ class UnslothTrainer: model_name, max_seq_length = max_seq_length, load_in_4bit = False, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -741,6 +759,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -754,6 +773,7 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, # Auto-detect load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, @@ -786,12 +806,15 @@ class UnslothTrainer: max_seq_length = max_seq_length, dtype = None, # Auto-detect load_in_4bit = load_in_4bit, + device_map = device_map, full_finetuning = full_finetuning, token = hf_token, trust_remote_code = trust_remote_code, ) logger.info("Loaded text model") + raise_if_offloaded(self.model, device_map, "Studio training") + if self.should_stop: return False @@ -824,6 +847,7 @@ class UnslothTrainer: is_dataset_audio = is_dataset_audio, trust_remote_code = trust_remote_code, full_finetuning = full_finetuning, + gpu_ids = gpu_ids, ) error_msg = str(e) error_lower = error_msg.lower() diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 4439e4e173..2c8f9a21db 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -28,6 +28,7 @@ from pathlib import Path from typing import Optional, Tuple, Any import matplotlib.pyplot as plt +from utils.hardware import prepare_gpu_selection logger = get_logger(__name__) @@ -185,6 +186,7 @@ class TrainingBackend: "enable_tensorboard": kwargs.get("enable_tensorboard", False), "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), "trust_remote_code": kwargs.get("trust_remote_code", False), + "gpu_ids": kwargs.get("gpu_ids"), } # Derive load_in_4bit from training_type @@ -192,6 +194,22 @@ class TrainingBackend: config["load_in_4bit"] = False # Spawn subprocess — use locals so state is untouched on failure + resolved_gpu_ids, gpu_selection = prepare_gpu_selection( + kwargs.get("gpu_ids"), + model_name = config["model_name"], + hf_token = config["hf_token"] or None, + training_type = config["training_type"], + load_in_4bit = config["load_in_4bit"], + batch_size = config.get("batch_size", 4), + max_seq_length = config.get("max_seq_length", 2048), + lora_rank = config.get("lora_r", 16), + target_modules = config.get("target_modules"), + gradient_checkpointing = config.get("gradient_checkpointing", "unsloth"), + optimizer = config.get("optim", "adamw_8bit"), + ) + config["resolved_gpu_ids"] = resolved_gpu_ids + config["gpu_selection"] = gpu_selection + from .worker import run_training_process event_queue = _CTX.Queue() diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 891dfca8f7..e68a6c7aee 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -29,6 +29,7 @@ import urllib.error import urllib.request logger = get_logger(__name__) +from utils.hardware import apply_gpu_ids _CAUSAL_CONV1D_RELEASE_TAG = "v1.6.1.post4" @@ -367,6 +368,8 @@ def run_training_process( env = os.getenv("ENVIRONMENT_TYPE", "production"), ) + apply_gpu_ids(config.get("resolved_gpu_ids")) + model_name = config["model_name"] # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -682,6 +685,7 @@ def run_training_process( is_dataset_image = config.get("is_dataset_image", False), is_dataset_audio = config.get("is_dataset_audio", False), trust_remote_code = config.get("trust_remote_code", False), + gpu_ids = config.get("resolved_gpu_ids"), ) if not success or trainer.should_stop: if trainer.should_stop: diff --git a/studio/backend/main.py b/studio/backend/main.py index 67908d8617..c18f18a743 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -67,7 +67,12 @@ from routes import ( ) from auth import storage from auth.authentication import get_current_subject -from utils.hardware import detect_hardware, get_device, DeviceType +from utils.hardware import ( + detect_hardware, + get_device, + DeviceType, + get_backend_visible_gpu_info, +) import utils.hardware.hardware as _hw_module from utils.cache_cleanup import clear_unsloth_compiled_cache @@ -230,69 +235,14 @@ async def shutdown_server( async def get_system_info(): """Get system information""" import platform - import subprocess import psutil - from utils.hardware import get_device, get_gpu_memory_info, DeviceType + from utils.hardware import get_device - # GPU Info — query nvidia-smi for physical GPUs, filtered by - # CUDA_VISIBLE_DEVICES when set (the frontend uses this for GGUF - # fit estimation and llama-server respects CVD too). - import os - - gpu_info: dict = {"available": False, "devices": []} - - device = get_device() - if device == DeviceType.CUDA: - # Parse CUDA_VISIBLE_DEVICES allowlist - allowed_indices = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None and cvd.strip(): - try: - allowed_indices = set(int(x.strip()) for x in cvd.split(",")) - except ValueError: - pass # Non-numeric (e.g. GPU-uuid), show all - - try: - result = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=index,name,memory.total", - "--format=csv,noheader,nounits", - ], - capture_output = True, - text = True, - timeout = 10, - ) - if result.returncode == 0: - for line in result.stdout.strip().splitlines(): - parts = [p.strip() for p in line.split(",")] - if len(parts) == 3: - idx = int(parts[0]) - if allowed_indices is not None and idx not in allowed_indices: - continue - gpu_info["devices"].append( - { - "index": idx, - "name": parts[1], - "memory_total_gb": round(int(parts[2]) / 1024, 2), - } - ) - gpu_info["available"] = len(gpu_info["devices"]) > 0 - except Exception: - pass - - # Fallback to torch-based single-GPU detection - if not gpu_info["available"]: - mem_info = get_gpu_memory_info() - if mem_info.get("available"): - gpu_info["available"] = True - gpu_info["devices"].append( - { - "index": mem_info.get("device", 0), - "name": mem_info.get("device_name", "Unknown"), - "memory_total_gb": round(mem_info.get("total_gb", 0), 2), - } - ) + visibility_info = get_backend_visible_gpu_info() + gpu_info = { + "available": visibility_info["available"], + "devices": visibility_info["devices"], + } # CPU & Memory memory = psutil.virtual_memory() @@ -311,6 +261,13 @@ async def get_system_info(): } +@app.get("/api/system/gpu-visibility") +async def get_gpu_visibility( + current_subject: str = Depends(get_current_subject), +): + return get_backend_visible_gpu_info() + + @app.get("/api/system/hardware") async def get_hardware_info(): """Return GPU name, total VRAM, and key ML package versions.""" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index a20c2052aa..aabfba9b3a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -44,6 +44,10 @@ class LoadRequest(BaseModel): None, description = "KV cache data type for both K and V (e.g. 'f16', 'bf16', 'q8_0', 'q4_1', 'q5_1')", ) + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.", + ) class UnloadRequest(BaseModel): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 68791aa7a8..eeb98c872e 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -128,6 +128,12 @@ class TrainingStartRequest(BaseModel): enable_tensorboard: bool = Field(False, description = "Enable TensorBoard logging") tensorboard_dir: Optional[str] = Field(None, description = "TensorBoard directory") + # GPU selection + gpu_ids: Optional[List[int]] = Field( + None, + description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries.", + ) + class TrainingJobResponse(BaseModel): """Immediate response when training is initiated""" diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7d48198d42..1a94256059 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -206,8 +206,17 @@ async def load_model( detail = f"Invalid model identifier: {request.model_path}", ) + # Normalize gpu_ids: empty list means auto-selection, same as None + effective_gpu_ids = request.gpu_ids if request.gpu_ids else None + # ── GGUF path: load via llama-server ────────────────────── if config.is_gguf: + if effective_gpu_ids is not None: + raise HTTPException( + status_code = 400, + detail = "gpu_ids is not supported for GGUF models yet.", + ) + llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() @@ -369,6 +378,7 @@ async def load_model( load_in_4bit = load_in_4bit, hf_token = request.hf_token, trust_remote_code = request.trust_remote_code, + gpu_ids = effective_gpu_ids, ) if not success: @@ -420,6 +430,9 @@ async def load_model( except HTTPException: raise + except ValueError as e: + logger.warning("Rejected inference GPU selection: %s", e) + raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: logger.error(f"Error loading model: {e}", exc_info = True) msg = str(e) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 4cfb060dee..e625408bad 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -88,14 +88,22 @@ async def get_hardware_utilization( Get a live snapshot of GPU hardware utilization. Designed to be polled by the frontend during training. - Returns GPU utilization %, temperature, VRAM usage, and power draw - via nvidia-smi for maximum accuracy. + Returns live GPU memory usage information for the active backend. """ from utils.hardware import get_gpu_utilization return get_gpu_utilization() +@router.get("/hardware/visible") +async def get_visible_hardware_utilization( + current_subject: str = Depends(get_current_subject), +): + from utils.hardware import get_visible_gpu_utilization + + return get_visible_gpu_utilization() + + @router.post("/start") async def start_training( request: TrainingStartRequest, @@ -202,6 +210,7 @@ async def start_training( "enable_tensorboard": request.enable_tensorboard, "tensorboard_dir": request.tensorboard_dir or "", "trust_remote_code": request.trust_remote_code, + "gpu_ids": request.gpu_ids, } # Training page has no trust_remote_code toggle — the value comes from @@ -269,6 +278,9 @@ async def start_training( error = None, ) + except ValueError as e: + logger.warning("Rejected training GPU selection: %s", e) + raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: logger.error(f"Error starting training: {e}", exc_info = True) raise HTTPException( diff --git a/studio/backend/tests/test_gpu_selection.py b/studio/backend/tests/test_gpu_selection.py new file mode 100644 index 0000000000..275b6f33d6 --- /dev/null +++ b/studio/backend/tests/test_gpu_selection.py @@ -0,0 +1,1125 @@ +# 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") diff --git a/studio/backend/tests/test_gpu_selection_sandbox.py b/studio/backend/tests/test_gpu_selection_sandbox.py new file mode 100644 index 0000000000..830a98a2fb --- /dev/null +++ b/studio/backend/tests/test_gpu_selection_sandbox.py @@ -0,0 +1,544 @@ +#!/usr/bin/env python3 +""" +Sandbox test for multi-GPU selection logic. + +Tests the core GPU selection, memory estimation, and device_map logic +in an isolated environment. Can be run on Linux, macOS, and Windows +without requiring actual GPUs -- all hardware calls are mocked. + +Usage: + python -m pytest studio/backend/tests/test_gpu_selection_sandbox.py -v + # or directly: + python studio/backend/tests/test_gpu_selection_sandbox.py +""" + +import os +import sys +import unittest +from pathlib import Path +from unittest.mock import patch, MagicMock + +# Ensure backend is on sys.path +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +def _make_fake_config( + vocab_size = 32000, + hidden_size = 4096, + intermediate_size = 11008, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + tie_word_embeddings = False, +): + """Create a fake HF config-like object for estimation tests.""" + from types import SimpleNamespace + + return SimpleNamespace( + vocab_size = vocab_size, + hidden_size = hidden_size, + intermediate_size = intermediate_size, + num_hidden_layers = num_hidden_layers, + num_attention_heads = num_attention_heads, + num_key_value_heads = num_key_value_heads, + tie_word_embeddings = tie_word_embeddings, + ) + + +class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase): + """Test the config-based model size estimation.""" + + def test_llama_8b_size_reasonable(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + + config = _make_fake_config( + vocab_size = 128256, + hidden_size = 4096, + intermediate_size = 14336, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + tie_word_embeddings = False, + ) + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNotNone(size) + size_gb = size / (1024**3) + # Llama 3.1 8B should be ~15GB in fp16 + self.assertGreater(size_gb, 12) + self.assertLess(size_gb, 20) + + def test_small_model(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + + config = _make_fake_config( + vocab_size = 32000, + hidden_size = 2048, + intermediate_size = 5504, + num_hidden_layers = 22, + num_attention_heads = 32, + num_key_value_heads = 4, + ) + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNotNone(size) + size_gb = size / (1024**3) + # ~1B model should be ~2GB in fp16 + self.assertGreater(size_gb, 1) + self.assertLess(size_gb, 5) + + def test_returns_none_for_incomplete_config(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + from types import SimpleNamespace + + config = SimpleNamespace(vocab_size = 32000) # Missing most fields + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNone(size) + + def test_moe_model(self): + from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config + from types import SimpleNamespace + + config = SimpleNamespace( + vocab_size = 152064, + hidden_size = 3584, + intermediate_size = 18944, + num_hidden_layers = 28, + num_attention_heads = 28, + num_key_value_heads = 4, + tie_word_embeddings = False, + num_local_experts = 64, + moe_intermediate_size = 2560, + ) + size = _estimate_fp16_model_size_bytes_from_config(config) + self.assertIsNotNone(size) + size_gb = size / (1024**3) + # MoE model with 64 experts should be large + self.assertGreater(size_gb, 50) + + +class TestEstimateRequiredModelMemory(unittest.TestCase): + """Test memory requirement estimation.""" + + def test_inference_fp16_uses_1_3x(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (10 * (1024**3), "config"), # 10GB model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = None, # inference + load_in_4bit = False, + ) + self.assertIsNotNone(required) + self.assertAlmostEqual(required, 13.0, places = 0) + self.assertEqual(meta["mode"], "inference") + + def test_inference_4bit_uses_reduced_estimate(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (30 * (1024**3), "config"), # 30GB fp16 model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = None, # inference + load_in_4bit = True, + ) + self.assertIsNotNone(required) + # 4bit base = 30/3.2 = 9.375GB, required = 9.375 + max(9.375*0.3, 2) = 12.19GB + self.assertAlmostEqual(required, 12.2, places = 0) + + def test_4bit_training_reduces_base(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (30 * (1024**3), "config"), # 30GB fp16 model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = "LoRA/QLoRA", + load_in_4bit = True, + ) + self.assertIsNotNone(required) + # fallback: base=30/3.2=9.375, lora=30*0.04=1.2, act=30*0.15=4.5, cuda=1.4 + self.assertAlmostEqual(required, 16.5, places = 0) + + def test_full_finetune_uses_3_5x(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (10 * (1024**3), "config"), # 10GB model + ): + required, meta = estimate_required_model_memory_gb( + "test/model", + training_type = "Full Finetuning", + ) + self.assertIsNotNone(required) + # fallback: 10 * 3.5 + 1.4 cuda overhead = 36.4 + self.assertAlmostEqual(required, 36.4, places = 0) + + def test_returns_none_when_unavailable(self): + from utils.hardware.hardware import estimate_required_model_memory_gb + + with patch( + "utils.hardware.hardware.estimate_fp16_model_size_bytes", + return_value = (None, "unavailable"), + ): + required, meta = estimate_required_model_memory_gb("test/model") + self.assertIsNone(required) + + +class TestAutoSelectGpuIds(unittest.TestCase): + """Test automatic GPU selection based on model size and free memory.""" + + def _make_utilization(self, devices): + """Create a fake utilization response.""" + return { + "available": True, + "devices": [ + { + "index": idx, + "vram_total_gb": total, + "vram_used_gb": total - free, + } + for idx, total, free in devices + ], + } + + def test_single_gpu_sufficient(self): + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 10.0, + { + "mode": "inference", + "required_gb": 10.0, + "model_size_source": "config", + "model_size_gb": 7.7, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1,2,3", + "numeric_ids": [0, 1, 2, 3], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1, 2, 3]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 80.0, 75.0), + (1, 80.0, 78.0), + (2, 80.0, 70.0), + (3, 80.0, 72.0), + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # Should pick GPU 1 (most free memory: 78GB) -- enough for 10GB + self.assertEqual(len(selected), 1) + self.assertEqual(selected[0], 1) + + def test_two_gpus_needed(self): + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 50.0, + { + "mode": "inference", + "required_gb": 50.0, + "model_size_source": "config", + "model_size_gb": 38.0, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 40.0, 30.0), # 30GB free + (1, 40.0, 35.0), # 35GB free + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB + self.assertEqual(len(selected), 2) + + def test_non_cuda_returns_none(self): + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): + selected, meta = auto_select_gpu_ids("test/model") + self.assertIsNone(selected) + self.assertEqual(meta["selection_mode"], "non_cuda") + + +class TestGetDeviceMap(unittest.TestCase): + """Test device_map string generation.""" + + def test_single_gpu_returns_sequential(self): + from utils.hardware.hardware import get_device_map + import utils.hardware.hardware as hw + + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0", + "numeric_ids": [0], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_visible_gpu_count", return_value = 1), + ): + dm = get_device_map(gpu_ids = [0]) + self.assertEqual(dm, "sequential") + + def test_multi_gpu_returns_balanced(self): + from utils.hardware.hardware import get_device_map + import utils.hardware.hardware as hw + + with patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA): + dm = get_device_map(gpu_ids = [0, 1]) + self.assertEqual(dm, "balanced") + + def test_cpu_returns_sequential(self): + from utils.hardware.hardware import get_device_map + import utils.hardware.hardware as hw + + with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU): + dm = get_device_map(gpu_ids = None) + self.assertEqual(dm, "sequential") + + +class TestResolveRequestedGpuIds(unittest.TestCase): + """Test GPU ID validation.""" + + def test_none_returns_parent_visible(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + result = resolve_requested_gpu_ids(None) + self.assertEqual(result, [2, 3]) + + def test_empty_list_returns_parent_visible(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + result = resolve_requested_gpu_ids([]) + self.assertEqual(result, [2, 3]) + + def test_duplicates_rejected(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([1, 1]) + + def test_out_of_range_rejected(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear = False), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4), + ): + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([5]) + + def test_uuid_env_var_rejects_explicit_ids(self): + from utils.hardware.hardware import resolve_requested_gpu_ids + + with ( + patch.dict( + os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False + ), + patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8), + ): + with self.assertRaises(ValueError): + resolve_requested_gpu_ids([0]) + + +class TestApplyGpuIds(unittest.TestCase): + """Test CUDA_VISIBLE_DEVICES environment variable setting.""" + + def test_apply_list(self): + from utils.hardware.hardware import apply_gpu_ids + + with patch.dict(os.environ, {}, clear = False): + apply_gpu_ids([3, 5]) + self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), "3,5") + + def test_apply_none_does_nothing(self): + from utils.hardware.hardware import apply_gpu_ids + + original = os.environ.get("CUDA_VISIBLE_DEVICES") + apply_gpu_ids(None) + self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), original) + + +class TestMultiGpuOverheadAccounting(unittest.TestCase): + """Test that multi-GPU overhead is applied correctly. + + The first GPU should keep its full free memory, and only + additional GPUs should have the overhead factor applied. + """ + + def _make_utilization(self, devices): + return { + "available": True, + "devices": [ + { + "index": idx, + "vram_total_gb": total, + "vram_used_gb": total - free, + } + for idx, total, free in devices + ], + } + + def test_first_gpu_not_penalized(self): + """A model that just fits on 1 GPU should not require 2 GPUs.""" + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + # Model requires 79GB, GPU has 80GB free + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 79.0, + { + "mode": "inference", + "required_gb": 79.0, + "model_size_source": "config", + "model_size_gb": 60.0, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 80.0, 80.0), + (1, 80.0, 80.0), + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # Should fit on 1 GPU (80GB >= 79GB) + self.assertEqual(len(selected), 1) + + def test_second_gpu_has_overhead(self): + """When 2 GPUs are needed, the second one's contribution is reduced.""" + from utils.hardware.hardware import auto_select_gpu_ids + import utils.hardware.hardware as hw + + # Model requires 110GB. First GPU has 80GB, second has 40GB. + # With overhead: 80 + 40*0.85 = 114GB -- just enough + with ( + patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA), + patch.object( + hw, + "estimate_required_model_memory_gb", + return_value = ( + 110.0, + { + "mode": "inference", + "required_gb": 110.0, + "model_size_source": "config", + "model_size_gb": 85.0, + }, + ), + ), + patch.object( + hw, + "_get_parent_visible_gpu_spec", + return_value = { + "raw": "0,1", + "numeric_ids": [0, 1], + "supports_explicit_gpu_ids": True, + }, + ), + patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]), + patch.object( + hw, + "get_visible_gpu_utilization", + return_value = self._make_utilization( + [ + (0, 80.0, 80.0), + (1, 80.0, 40.0), + ] + ), + ), + ): + selected, meta = auto_select_gpu_ids("test/model") + # Should use both GPUs + self.assertEqual(len(selected), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 3c33b33cb3..50557c6718 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -285,7 +285,7 @@ class TestLogGpuMemory: def test_does_not_raise(self): log_gpu_memory("test") - def test_logs_gpu_info_when_available(self, caplog): + def test_logs_gpu_info_when_available(self, capfd): fake_info = { "available": True, "backend": "cuda", @@ -295,35 +295,27 @@ class TestLogGpuMemory: "utilization_pct": 12.5, "free_gb": 14.0, } - import structlog - from loggers import get_logger - with ( - patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ), - caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"), + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info ): log_gpu_memory("unit-test") - assert "unit-test" in caplog.text - assert "CUDA" in caplog.text - assert "FakeGPU" in caplog.text + captured = capfd.readouterr() + assert "unit-test" in captured.out + assert "CUDA" in captured.out + assert "FakeGPU" in captured.out - def test_logs_cpu_fallback_when_no_gpu(self, caplog): + def test_logs_cpu_fallback_when_no_gpu(self, capfd): fake_info = {"available": False, "backend": "cpu"} - import structlog - from loggers import get_logger - with ( - patch( - "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info - ), - caplog.at_level(logging.INFO, logger = "utils.hardware.hardware"), + with patch( + "utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info ): log_gpu_memory("cpu-test") - assert "No GPU available" in caplog.text + captured = capfd.readouterr() + assert "No GPU available" in captured.out # ========== format_error_message() ========== diff --git a/studio/backend/tests/test_vram_estimation.py b/studio/backend/tests/test_vram_estimation.py new file mode 100644 index 0000000000..0be067310d --- /dev/null +++ b/studio/backend/tests/test_vram_estimation.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import unittest +from types import SimpleNamespace + +from utils.hardware.vram_estimation import ( + ModelArchConfig, + TrainingVramConfig, + extract_arch_config, + compute_model_weights_bytes, + compute_total_params, + compute_lora_params, + compute_lora_adapter_bytes, + compute_optimizer_bytes, + compute_gradient_bytes, + compute_activation_bytes, + estimate_training_vram, + DEFAULT_TARGET_MODULES, +) + + +def _gb(b: int) -> float: + return b / (1024**3) + + +LLAMA_8B = 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, +) + +QWEN_05B = ModelArchConfig( + hidden_size = 896, + num_hidden_layers = 24, + num_attention_heads = 14, + num_key_value_heads = 2, + intermediate_size = 4864, + vocab_size = 151936, + tie_word_embeddings = True, +) + +MOE_CONFIG = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, +) + +DEEPSEEK_V3 = ModelArchConfig( + hidden_size = 7168, + num_hidden_layers = 61, + num_attention_heads = 128, + num_key_value_heads = 128, + intermediate_size = 18432, + vocab_size = 129280, + tie_word_embeddings = False, + num_experts = 256, + moe_intermediate_size = 2048, + n_shared_experts = 1, + num_dense_layers = 3, + q_lora_rank = 1536, + kv_lora_rank = 512, + qk_nope_head_dim = 128, + qk_rope_head_dim = 64, + v_head_dim = 128, +) + +QWEN3_MOE_30B = ModelArchConfig( + hidden_size = 2048, + num_hidden_layers = 48, + num_attention_heads = 32, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 151936, + tie_word_embeddings = True, + num_experts = 128, + moe_intermediate_size = 768, + n_shared_experts = 0, + num_dense_layers = 0, +) + +GLM4_MOE = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + num_dense_layers = 1, +) + +GPT_OSS = ModelArchConfig( + hidden_size = 6144, + num_hidden_layers = 64, + num_attention_heads = 64, + num_key_value_heads = 8, + intermediate_size = 2880, + vocab_size = 200064, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = None, + n_shared_experts = 0, + num_dense_layers = 0, +) + + +class TestExtractArchConfig(unittest.TestCase): + def test_basic_config(self): + hf_config = 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, + ) + arch = extract_arch_config(hf_config) + self.assertIsNotNone(arch) + self.assertEqual(arch.hidden_size, 4096) + self.assertEqual(arch.num_hidden_layers, 32) + self.assertEqual(arch.num_key_value_heads, 8) + self.assertIsNone(arch.num_experts) + + def test_vlm_text_config(self): + text_cfg = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 24, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 32000, + tie_word_embeddings = True, + ) + hf_config = SimpleNamespace(text_config = text_cfg) + arch = extract_arch_config(hf_config) + self.assertIsNotNone(arch) + self.assertEqual(arch.hidden_size, 2048) + + def test_moe_detection(self): + hf_config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_local_experts = 8, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 8) + + def test_missing_fields_returns_none(self): + hf_config = SimpleNamespace(hidden_size = 4096) + arch = extract_arch_config(hf_config) + self.assertIsNone(arch) + + def test_intermediate_size_list(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 24, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = [8192, 8192], + vocab_size = 32000, + tie_word_embeddings = True, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.intermediate_size, 8192) + + +class TestModelWeightsBytes(unittest.TestCase): + def test_llama_8b_fp16(self): + weight_bytes = compute_model_weights_bytes(LLAMA_8B, "full", False) + weight_gb = _gb(weight_bytes) + self.assertGreater(weight_gb, 14.0) + self.assertLess(weight_gb, 18.0) + + def test_llama_8b_qlora_4bit(self): + weight_bytes = compute_model_weights_bytes(LLAMA_8B, "qlora", True) + weight_gb = _gb(weight_bytes) + self.assertGreater(weight_gb, 4.0) + self.assertLess(weight_gb, 7.0) + + def test_4bit_smaller_than_fp16(self): + fp16 = compute_model_weights_bytes(LLAMA_8B, "full", False) + q4 = compute_model_weights_bytes(LLAMA_8B, "qlora", True) + self.assertLess(q4, fp16) + ratio = fp16 / q4 + self.assertGreater(ratio, 2.0) + self.assertLess(ratio, 4.0) + + def test_moe_larger_than_dense(self): + dense = compute_model_weights_bytes(LLAMA_8B, "full", False) + moe = compute_model_weights_bytes(MOE_CONFIG, "full", False) + self.assertGreater(moe, dense * 3) + + +class TestLoraParams(unittest.TestCase): + def test_llama_8b_default_modules_rank16(self): + lora_p = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + total_p = compute_total_params(LLAMA_8B) + ratio = lora_p / total_p + self.assertGreater(ratio, 0.005) + self.assertLess(ratio, 0.05) + + def test_higher_rank_more_params(self): + r16 = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + r64 = compute_lora_params(LLAMA_8B, 64, DEFAULT_TARGET_MODULES) + self.assertAlmostEqual(r64 / r16, 4.0, places = 1) + + def test_fewer_modules_fewer_params(self): + all_mods = compute_lora_params(LLAMA_8B, 16, DEFAULT_TARGET_MODULES) + qv_only = compute_lora_params(LLAMA_8B, 16, ["q_proj", "v_proj"]) + self.assertLess(qv_only, all_mods) + + def test_moe_mlp_modules_scale_with_experts(self): + dense_lora = compute_lora_params( + LLAMA_8B, 16, ["gate_proj", "up_proj", "down_proj"] + ) + moe_lora = compute_lora_params( + MOE_CONFIG, 16, ["gate_proj", "up_proj", "down_proj"] + ) + ratio = moe_lora / dense_lora + self.assertAlmostEqual(ratio, 8.0, delta = 0.5) + + def test_attention_modules_same_for_moe(self): + dense_attn = compute_lora_params( + LLAMA_8B, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) + moe_attn = compute_lora_params( + MOE_CONFIG, 16, ["q_proj", "k_proj", "v_proj", "o_proj"] + ) + self.assertEqual(dense_attn, moe_attn) + + +class TestOptimizerBytes(unittest.TestCase): + def test_adamw_8bit(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_8bit"), 4_000_000) + + def test_adamw_torch(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "adamw_torch"), 6_000_000) + + def test_sgd(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "sgd"), 4_000_000) + + def test_unknown_defaults_to_4(self): + self.assertEqual(compute_optimizer_bytes(1_000_000, "some_new_opt"), 4_000_000) + + +class TestGradientBytes(unittest.TestCase): + def test_fp16_gradients(self): + self.assertEqual(compute_gradient_bytes(1_000_000), 2_000_000) + + +class TestActivationBytes(unittest.TestCase): + def test_no_gc_scales_with_layers(self): + act_none = compute_activation_bytes(LLAMA_8B, 2, 2048, "none") + act_gc = compute_activation_bytes(LLAMA_8B, 2, 2048, "true") + self.assertGreater(act_none, act_gc * 10) + + def test_unsloth_gc_smaller_than_standard(self): + act_true = compute_activation_bytes(LLAMA_8B, 2, 2048, "true") + act_unsloth = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") + self.assertLess(act_unsloth, act_true) + + def test_lora_activations_smaller_than_full_ft(self): + full_ft = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = False) + lora = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth", is_lora = True) + self.assertLess(lora, full_ft) + + def test_scales_with_batch_size(self): + act_bsz2 = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") + act_bsz4 = compute_activation_bytes(LLAMA_8B, 4, 2048, "unsloth") + self.assertAlmostEqual(act_bsz4 / act_bsz2, 2.0, delta = 0.1) + + def test_scales_with_seq_len(self): + act_2k = compute_activation_bytes(LLAMA_8B, 2, 2048, "unsloth") + act_4k = compute_activation_bytes(LLAMA_8B, 2, 4096, "unsloth") + self.assertAlmostEqual(act_4k / act_2k, 2.0, delta = 0.1) + + +class TestEstimateTrainingVram(unittest.TestCase): + def test_llama_8b_qlora_reasonable_total(self): + config = TrainingVramConfig( + training_method = "qlora", + batch_size = 2, + max_seq_length = 2048, + lora_rank = 16, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(LLAMA_8B, config) + total_gb = _gb(breakdown.total) + self.assertGreater(total_gb, 5.0) + self.assertLess(total_gb, 12.0) + + def test_llama_8b_full_ft_reasonable_total(self): + config = TrainingVramConfig( + training_method = "full", + batch_size = 2, + max_seq_length = 2048, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = False, + ) + breakdown = estimate_training_vram(LLAMA_8B, config) + total_gb = _gb(breakdown.total) + self.assertGreater(total_gb, 50.0) + self.assertLess(total_gb, 75.0) + + def test_qlora_much_less_than_full_ft(self): + qlora_config = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + batch_size = 2, + max_seq_length = 2048, + ) + full_config = TrainingVramConfig( + training_method = "full", + load_in_4bit = False, + batch_size = 2, + max_seq_length = 2048, + ) + qlora = estimate_training_vram(LLAMA_8B, qlora_config) + full = estimate_training_vram(LLAMA_8B, full_config) + self.assertLess(qlora.total, full.total / 3) + + def test_qwen_05b_qlora_fits_in_4gb(self): + config = TrainingVramConfig( + training_method = "qlora", + batch_size = 2, + max_seq_length = 2048, + lora_rank = 16, + gradient_checkpointing = "unsloth", + optimizer = "adamw_8bit", + load_in_4bit = True, + ) + breakdown = estimate_training_vram(QWEN_05B, config) + total_gb = _gb(breakdown.total) + self.assertLess(total_gb, 5.0) + + def test_breakdown_components_positive(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + self.assertGreater(breakdown.model_weights, 0) + self.assertGreater(breakdown.lora_adapters, 0) + self.assertGreater(breakdown.optimizer_states, 0) + self.assertGreater(breakdown.gradients, 0) + self.assertGreater(breakdown.activations, 0) + self.assertGreater(breakdown.cuda_overhead, 0) + + def test_full_ft_no_lora_adapters(self): + config = TrainingVramConfig(training_method = "full", load_in_4bit = False) + breakdown = estimate_training_vram(LLAMA_8B, config) + self.assertEqual(breakdown.lora_adapters, 0) + + def test_to_gb_dict_keys(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + gb_dict = breakdown.to_gb_dict() + expected_keys = { + "model_weights_gb", + "lora_adapters_gb", + "optimizer_states_gb", + "gradients_gb", + "activations_gb", + "cuda_overhead_gb", + "total_gb", + } + self.assertEqual(set(gb_dict.keys()), expected_keys) + + def test_total_equals_sum_of_parts(self): + config = TrainingVramConfig(training_method = "qlora", load_in_4bit = True) + breakdown = estimate_training_vram(LLAMA_8B, config) + parts_sum = ( + breakdown.model_weights + + breakdown.lora_adapters + + breakdown.optimizer_states + + breakdown.gradients + + breakdown.activations + + breakdown.cuda_overhead + ) + self.assertEqual(breakdown.total, parts_sum) + + def test_larger_batch_increases_total(self): + small = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + batch_size = 1, + ) + large = TrainingVramConfig( + training_method = "qlora", + load_in_4bit = True, + batch_size = 8, + ) + small_v = estimate_training_vram(LLAMA_8B, small) + large_v = estimate_training_vram(LLAMA_8B, large) + self.assertGreater(large_v.total, small_v.total) + + def test_adamw_fp32_uses_more_optimizer_memory(self): + opt8 = TrainingVramConfig( + training_method = "full", + load_in_4bit = False, + optimizer = "adamw_8bit", + ) + opt32 = TrainingVramConfig( + training_method = "full", + load_in_4bit = False, + optimizer = "adamw_torch", + ) + v8 = estimate_training_vram(LLAMA_8B, opt8) + v32 = estimate_training_vram(LLAMA_8B, opt32) + self.assertAlmostEqual( + v32.optimizer_states / v8.optimizer_states, 1.5, delta = 0.1 + ) + + +class TestExtractArchConfigMoE(unittest.TestCase): + def test_deepseek_v3_shared_experts(self): + hf_config = SimpleNamespace( + hidden_size = 7168, + num_hidden_layers = 61, + num_attention_heads = 128, + num_key_value_heads = 128, + intermediate_size = 18432, + vocab_size = 129280, + tie_word_embeddings = False, + n_routed_experts = 256, + moe_intermediate_size = 2048, + n_shared_experts = 1, + first_k_dense_replace = 3, + q_lora_rank = 1536, + kv_lora_rank = 512, + qk_nope_head_dim = 128, + qk_rope_head_dim = 64, + v_head_dim = 128, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 256) + self.assertEqual(arch.n_shared_experts, 1) + self.assertEqual(arch.num_dense_layers, 3) + self.assertEqual(arch.q_lora_rank, 1536) + self.assertEqual(arch.kv_lora_rank, 512) + + def test_qwen3_moe_decoder_sparse_step(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 48, + num_attention_heads = 32, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 151936, + tie_word_embeddings = True, + num_local_experts = 128, + moe_intermediate_size = 768, + decoder_sparse_step = 1, + mlp_only_layers = [], + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 128) + self.assertEqual(arch.num_dense_layers, 0) + self.assertIsNone(arch.q_lora_rank) + + def test_qwen3_moe_with_mlp_only_layers(self): + hf_config = SimpleNamespace( + hidden_size = 2048, + num_hidden_layers = 24, + num_attention_heads = 16, + num_key_value_heads = 4, + intermediate_size = 8192, + vocab_size = 151936, + tie_word_embeddings = True, + num_local_experts = 60, + moe_intermediate_size = 1408, + decoder_sparse_step = 1, + mlp_only_layers = [0, 1, 2, 3], + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_dense_layers, 4) + + def test_glm4_moe_first_k_dense(self): + hf_config = SimpleNamespace( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + n_routed_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + first_k_dense_replace = 1, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_dense_layers, 1) + self.assertEqual(arch.n_shared_experts, 1) + + def test_gpt_oss_no_moe_intermediate(self): + hf_config = SimpleNamespace( + hidden_size = 6144, + num_hidden_layers = 64, + num_attention_heads = 64, + num_key_value_heads = 8, + intermediate_size = 2880, + vocab_size = 200064, + tie_word_embeddings = False, + num_local_experts = 128, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.num_experts, 128) + self.assertIsNone(arch.moe_intermediate_size) + self.assertEqual(arch.num_dense_layers, 0) + + def test_backward_compat_no_new_fields(self): + hf_config = 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, + ) + arch = extract_arch_config(hf_config) + self.assertEqual(arch.n_shared_experts, 0) + self.assertEqual(arch.num_dense_layers, 0) + self.assertIsNone(arch.q_lora_rank) + + +class TestSharedExperts(unittest.TestCase): + def test_shared_experts_increase_weight_bytes(self): + no_shared = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 64, + moe_intermediate_size = 1407, + n_shared_experts = 0, + ) + with_shared = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 32, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 64, + moe_intermediate_size = 1407, + n_shared_experts = 2, + ) + w_no = compute_model_weights_bytes(no_shared, "full", False) + w_yes = compute_model_weights_bytes(with_shared, "full", False) + self.assertGreater(w_yes, w_no) + delta_per_layer = 4096 * 1407 * 3 * 2 + expected_delta = delta_per_layer * 32 * 2 + actual_delta = w_yes - w_no + self.assertAlmostEqual( + actual_delta, expected_delta, delta = expected_delta * 0.01 + ) + + def test_deepseek_v3_params_in_range(self): + total = compute_total_params(DEEPSEEK_V3) + total_b = total / 1e9 + self.assertGreater(total_b, 600) + self.assertLess(total_b, 750) + + +class TestMLA(unittest.TestCase): + def test_mla_different_from_standard(self): + from utils.hardware.vram_estimation import _compute_attn_elements + + mla_arch = DEEPSEEK_V3 + std_arch = ModelArchConfig( + hidden_size = 7168, + num_hidden_layers = 61, + num_attention_heads = 128, + num_key_value_heads = 128, + intermediate_size = 18432, + vocab_size = 129280, + ) + mla_attn = _compute_attn_elements(mla_arch) + std_attn = _compute_attn_elements(std_arch) + self.assertNotEqual(mla_attn, std_attn) + + def test_mla_lora_produces_values(self): + lora_p = compute_lora_params(DEEPSEEK_V3, 16, ["q_proj", "v_proj", "o_proj"]) + self.assertGreater(lora_p, 0) + + +class TestDenseMoEMix(unittest.TestCase): + def test_dense_layers_change_total(self): + all_moe = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + num_dense_layers = 0, + ) + mixed = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 46, + num_attention_heads = 96, + num_key_value_heads = 8, + intermediate_size = 10944, + vocab_size = 151552, + tie_word_embeddings = False, + num_experts = 128, + moe_intermediate_size = 1408, + n_shared_experts = 1, + num_dense_layers = 1, + ) + w_all = compute_model_weights_bytes(all_moe, "full", False) + w_mixed = compute_model_weights_bytes(mixed, "full", False) + self.assertNotEqual(w_all, w_mixed) + + def test_glm4_moe_params_reasonable(self): + total = compute_total_params(GLM4_MOE) + total_b = total / 1e9 + self.assertGreater(total_b, 80) + self.assertLess(total_b, 120) + + def test_qwen3_moe_30b_params_reasonable(self): + total = compute_total_params(QWEN3_MOE_30B) + total_b = total / 1e9 + self.assertGreater(total_b, 20) + self.assertLess(total_b, 50) + + def test_gpt_oss_uses_intermediate_size(self): + total = compute_total_params(GPT_OSS) + total_b = total / 1e9 + self.assertGreater(total_b, 350) + self.assertLess(total_b, 500) + + def test_lora_dense_vs_moe_layers_differ(self): + all_moe = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 10, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, + moe_intermediate_size = 1024, + num_dense_layers = 0, + ) + mixed = ModelArchConfig( + hidden_size = 4096, + num_hidden_layers = 10, + num_attention_heads = 32, + num_key_value_heads = 8, + intermediate_size = 14336, + vocab_size = 32000, + tie_word_embeddings = False, + num_experts = 8, + moe_intermediate_size = 1024, + num_dense_layers = 5, + ) + lora_all = compute_lora_params( + all_moe, 16, ["gate_proj", "up_proj", "down_proj"] + ) + lora_mix = compute_lora_params(mixed, 16, ["gate_proj", "up_proj", "down_proj"]) + self.assertNotEqual(lora_all, lora_mix) + + +if __name__ == "__main__": + unittest.main() diff --git a/studio/backend/utils/hardware/VRAM_ESTIMATION.md b/studio/backend/utils/hardware/VRAM_ESTIMATION.md new file mode 100644 index 0000000000..26072b208f --- /dev/null +++ b/studio/backend/utils/hardware/VRAM_ESTIMATION.md @@ -0,0 +1,161 @@ +# VRAM Estimation for Training + +``` +Total VRAM = Weights + LoRA Adapters + Optimizer + Gradients + Activations + CUDA Overhead +``` + +| Symbol | Meaning | +|--------|---------| +| `H` | `hidden_size` | +| `L` | `num_hidden_layers` | +| `V` | `vocab_size` | +| `K` | `(H / num_attention_heads) * num_key_value_heads` | +| `M` | `intermediate_size` (or `moe_intermediate_size`) | +| `E` | `num_experts` (1 for dense) | +| `r` | LoRA rank | +| `B` | `per_device_train_batch_size` | +| `S` | `max_seq_length` | + +--- + +## 1. Model Weights + +``` +QKVO = (H + K + K + H) * H +MLP = H * M * 3 * E + (E * H if E > 1 else 0) + +Quantizable = (QKVO + MLP) * L +Non-quantizable = 2*H*L + V*H + (V*H if not tie_embeddings else 0) +``` + +| Mode | Bytes | +|------|-------| +| QLoRA 4-bit | `Quantizable * 2 / 3.2 + Non-quantizable * 2` | +| LoRA / Full fp16 | `(Quantizable + Non-quantizable) * 2` | + +The 3.2 factor (`16/5`) accounts for BNB NF4 blockwise scales. + +## 2. LoRA Adapters + +| Module | A | B | +|--------|---|---| +| q_proj | `H×r` | `r×H` | +| k_proj | `H×r` | `r×K` | +| v_proj | `H×r` | `r×K` | +| o_proj | `H×r` | `r×H` | +| gate_proj | `H×r` | `r×M` | +| up_proj | `H×r` | `r×M` | +| down_proj | `M×r` | `r×H` | + +MLP modules multiply by `E` for MoE. + +``` +LoRA_bytes = sum(A + B per selected module) * L * 2 +``` + +## 3. Optimizer States (calibrated) + +| Optimizer | Bytes/param | Notes | +|-----------|------------|-------| +| `adamw_8bit` | 4 | BNB upcasts to fp32 during step | +| `adamw_torch` | 6 | Fused, no master copy | +| `paged_adamw_32bit` | 8 | Full fp32 states | +| `sgd` | 4 | | + +Trainable params = all params (Full FT) or LoRA params only. + +## 4. Gradients + +``` +Gradient_bytes = trainable_params * 2 (fp16, accumulated in-place) +``` + +## 5. Activations + +Per-layer (from `unsloth_zoo/vllm_utils.py`): +``` +Per_layer = (S*B*(H+K+K) + S*B*2 + S*B*(M+M)) * 2 * 1.25 +``` + +| GC Mode | Full FT | LoRA/QLoRA | +|---------|---------|------------| +| none | `L` layers | `L` layers | +| true (HF) | 2.0 | 1.0 | +| unsloth | 1.5 | 1.0 | + +## 6. Floors + +Gradients and activations have minimum floors at **15% of model weight memory** to account for autograd overhead, attention score matrices, NCCL buffers, mixed-precision scaling, and PyTorch fragmentation. + +``` +gradient_bytes = max(computed, weights * 0.15) +activation_bytes = max(computed, weights * 0.15 * B/2) +``` + +## 7. CUDA Overhead + +**1.4 GB** fixed — CUDA driver + PyTorch runtime, calibrated on RTX 5070 Ti. + +## 8. Multi-GPU Overhead + +When sharding across multiple GPUs, each additional GPU (beyond the first) contributes only **85%** of its free VRAM to the usable pool. The 15% discount accounts for NCCL all-reduce buffers, PCIe/NVLink transfer overhead, synchronization barriers, and memory fragmentation from non-uniform shard sizes. Calibrated empirically on 2-8 GPU setups with NVLink and PCIe topologies. + +``` +usable_gb = free[gpu_0] + sum(free[gpu_i] * 0.85 for i in 1..N) +``` + +--- + +## Reference Table (bsz=2, seq=2048, rank=16, GC=unsloth, adamw_8bit) + +| Model | Weights | LoRA | Optim | Grad | Act | CUDA | Total | +|-------|---------|------|-------|------|-----|------|-------| +| 0.5B QLoRA | 0.5 | 0.0 | 0.0 | 0.1 | 0.1 | 1.4 | **2.1** | +| 1B QLoRA | 1.1 | 0.0 | 0.0 | 0.2 | 0.2 | 1.4 | **2.9** | +| 3B QLoRA | 2.4 | 0.0 | 0.1 | 0.5 | 0.5 | 1.4 | **4.9** | +| 8B QLoRA | 6.0 | 0.1 | 0.2 | 1.2 | 1.2 | 1.4 | **10.1** | +| 8B LoRA fp16 | 15.0 | 0.1 | 0.2 | 3.0 | 3.0 | 1.4 | **22.6** | +| 8B Full FT | 15.0 | — | 29.9 | 15.0 | 3.0 | 1.4 | **64.2** | +| 32B LoRA fp16 | 61.0 | 0.2 | 0.5 | 12.2 | 12.2 | 1.4 | **87.6** | +| 72B QLoRA | 45.5 | 0.4 | 0.8 | 9.1 | 9.1 | 1.4 | **66.3** | + +## E2E Validation (Llama-3.2-1B, B200 emulating 24GB) + +| Config | Estimated | Actual (nvsmi) | Error | +|--------|----------|----------------|-------| +| QLoRA bsz=2 seq=512 | 2.55 GB | 2.65 GB | -3.7% | +| QLoRA bsz=2 seq=2048 | 2.60 GB | 2.65 GB | -1.8% | +| QLoRA bsz=4 seq=2048 | 2.65 GB | 2.65 GB | +0.0% | +| LoRA fp16 bsz=2 | 3.84 GB | 3.88 GB | -1.0% | +| Full FT adamw_8bit | 10.89 GB | 10.80 GB | +0.8% | +| Full FT adamw_torch | 13.19 GB | 12.93 GB | +2.0% | + +*Note: e2e numbers predate the 15% floors, which add safety margin on top.* + +--- + +## Parameter Flow + +``` +Frontend -> routes/{training,inference}.py + -> prepare_gpu_selection(gpu_ids, model_name, ...) + | + +-- gpu_ids is explicit (e.g. [5,6,7]) + | -> resolve_requested_gpu_ids: validate against parent-visible set + | -> return all requested GPUs (model sharded across all of them) + | + +-- gpu_ids is None or [] + -> auto_select_gpu_ids: estimate VRAM, pick minimum GPUs needed + -> estimate_required_model_memory_gb -> estimate_training_vram + -> greedy selection: rank GPUs by free VRAM, add until model fits + + -> get_device_map(resolved_gpu_ids) + -> "balanced" if >1 GPU, "sequential" otherwise + + -> worker subprocess: apply_gpu_ids(resolved_gpu_ids) + -> sets CUDA_VISIBLE_DEVICES before torch/CUDA init +``` + +Threaded params: `batch_size`, `max_seq_length`, `lora_r`, `target_modules`, `gradient_checkpointing`, `optim`. + +Source: `studio/backend/utils/hardware/vram_estimation.py` diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index f86a56d186..aaa0452406 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -18,11 +18,31 @@ from .hardware import ( get_gpu_summary, get_package_versions, get_gpu_utilization, + get_visible_gpu_utilization, + get_backend_visible_gpu_info, get_physical_gpu_count, get_visible_gpu_count, + get_parent_visible_gpu_ids, + resolve_requested_gpu_ids, + estimate_fp16_model_size_bytes, + estimate_required_model_memory_gb, + auto_select_gpu_ids, + prepare_gpu_selection, safe_num_proc, safe_thread_num_proc, dataset_map_num_proc, + get_device_map, + get_offloaded_device_map_entries, + raise_if_offloaded, + apply_gpu_ids, +) + +from .vram_estimation import ( + ModelArchConfig, + TrainingVramConfig, + VramBreakdown, + extract_arch_config, + estimate_training_vram, ) __all__ = [ @@ -38,9 +58,26 @@ __all__ = [ "get_gpu_summary", "get_package_versions", "get_gpu_utilization", + "get_visible_gpu_utilization", + "get_backend_visible_gpu_info", "get_physical_gpu_count", "get_visible_gpu_count", + "get_parent_visible_gpu_ids", + "resolve_requested_gpu_ids", + "estimate_fp16_model_size_bytes", + "estimate_required_model_memory_gb", + "auto_select_gpu_ids", + "prepare_gpu_selection", "safe_num_proc", "safe_thread_num_proc", "dataset_map_num_proc", + "get_device_map", + "get_offloaded_device_map_entries", + "raise_if_offloaded", + "apply_gpu_ids", + "ModelArchConfig", + "TrainingVramConfig", + "VramBreakdown", + "extract_arch_config", + "estimate_training_vram", ] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 61ee8a0967..742e8f6b7e 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -16,10 +16,12 @@ Usage: ... """ +import os import platform import structlog from loggers import get_logger from enum import Enum +from pathlib import Path from typing import Optional, Dict, Any logger = get_logger(__name__) @@ -32,6 +34,7 @@ class DeviceType(str, Enum): """Supported compute backends. Inherits from str so it serializes cleanly in JSON.""" CUDA = "cuda" + XPU = "xpu" MLX = "mlx" CPU = "cpu" @@ -96,6 +99,17 @@ def detect_hardware() -> DeviceType: print(f"Hardware detected: CUDA — {device_name}") return DEVICE + # --- XPU: Intel GPU --- + if _has_torch(): + import torch + + if hasattr(torch, "xpu") and torch.xpu.is_available(): + DEVICE = DeviceType.XPU + CHAT_ONLY = False + device_name = torch.xpu.get_device_name(0) + print(f"Hardware detected: XPU — {device_name}") + return DEVICE + # --- MLX: Apple Silicon --- if is_apple_silicon() and _has_mlx(): DEVICE = DeviceType.MLX @@ -140,6 +154,11 @@ def clear_gpu_cache(): torch.cuda.synchronize() torch.cuda.empty_cache() torch.cuda.ipc_collect() + elif device == DeviceType.XPU: + import torch + + torch.xpu.synchronize() + torch.xpu.empty_cache() elif device == DeviceType.MLX: # MLX manages memory automatically; no explicit cache clear needed. # mlx.core has no empty_cache equivalent — gc.collect() above is enough. @@ -180,6 +199,33 @@ def get_gpu_memory_info() -> Dict[str, Any]: logger.error(f"Error getting CUDA GPU info: {e}") return {"available": False, "backend": device.value, "error": str(e)} + # ---- XPU path (Intel GPU) ---- + if device == DeviceType.XPU: + try: + import torch + + idx = torch.xpu.current_device() + props = torch.xpu.get_device_properties(idx) + + total = props.total_memory + allocated = torch.xpu.memory_allocated(idx) + reserved = torch.xpu.memory_reserved(idx) + + return { + "available": True, + "backend": device.value, + "device": idx, + "device_name": props.name, + "total_gb": total / (1024**3), + "allocated_gb": allocated / (1024**3), + "reserved_gb": reserved / (1024**3), + "free_gb": (total - allocated) / (1024**3), + "utilization_pct": (allocated / total) * 100, + } + except Exception as e: + logger.error("Error getting XPU GPU info: %s", e) + return {"available": False, "backend": device.value, "error": str(e)} + # ---- MLX path (Apple Silicon) ---- if device == DeviceType.MLX: try: @@ -280,134 +326,199 @@ def get_package_versions() -> Dict[str, Optional[str]]: return versions -# ========== Live GPU Utilization (nvidia-smi) ========== +# ========== Torch-based GPU fallbacks (AMD ROCm, Intel XPU, nvidia-smi missing) ========== + + +def _torch_get_device_module(): + """Return the appropriate torch device module (cuda or xpu) and its name.""" + device = get_device() + import torch + + if device == DeviceType.CUDA: + return torch.cuda, "cuda" + if device == DeviceType.XPU and hasattr(torch, "xpu"): + return torch.xpu, "xpu" + return None, None + + +def _torch_get_physical_gpu_count() -> Optional[int]: + mod, _ = _torch_get_device_module() + if mod is None: + return None + try: + return mod.device_count() + except Exception: + return None + + +def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]: + """Query torch for per-GPU name, total VRAM, and used VRAM.""" + mod, _ = _torch_get_device_module() + if mod is None: + return [] + + devices = [] + for ordinal, phys_idx in enumerate(device_indices): + try: + # torch uses 0-based ordinals relative to CUDA_VISIBLE_DEVICES + props = mod.get_device_properties(ordinal) + total_bytes = props.total_memory + # Prefer mem_get_info (reports system-wide usage, not just this + # process) so auto-selection accounts for other GPU consumers. + if hasattr(mod, "mem_get_info"): + free_bytes, total_bytes = mod.mem_get_info(ordinal) + used_bytes = total_bytes - free_bytes + else: + used_bytes = mod.memory_allocated(ordinal) + devices.append( + { + "index": phys_idx, + "visible_ordinal": ordinal, + "name": props.name, + "total_gb": round(total_bytes / (1024**3), 2), + "used_gb": round(used_bytes / (1024**3), 2), + } + ) + except Exception as e: + logger.debug("torch device query failed for ordinal %d: %s", ordinal, e) + return devices + + +# ========== Live GPU Utilization ========== def get_gpu_utilization() -> Dict[str, Any]: - """ - Return a live snapshot of GPU utilization via ``nvidia-smi``. - - Designed to be polled by the frontend during training (not streaming). - Uses ``nvidia-smi --query-gpu`` which is the most accurate source for - utilization %, temperature, and power draw – stats that PyTorch does - not expose. - - Returns dict with keys: - available – bool, whether stats could be retrieved - gpu_utilization_pct – GPU core utilization % - temperature_c – GPU temperature in °C - vram_used_gb – VRAM currently used (GiB) - vram_total_gb – VRAM total (GiB) - vram_utilization_pct – VRAM used / total * 100 - power_draw_w – current power draw (W) - power_limit_w – power limit (W) - power_utilization_pct – power draw / limit * 100 - """ + """Return a live snapshot of device utilization information.""" device = get_device() - if device != DeviceType.CUDA: - return {"available": False, "backend": device.value} - - def _parse_smi_value(raw: str): - """Parse a single nvidia-smi CSV value. Returns float or None for [N/A].""" - raw = raw.strip() - if not raw or raw == "[N/A]": - return None + if device == DeviceType.CUDA: try: - return float(raw) - except (ValueError, TypeError): - return None + from . import nvidia - # ── nvidia-smi (most complete source) ─────────────────────── - smi_data = {} - try: - import subprocess - - result = subprocess.run( - [ - "nvidia-smi", - "--query-gpu=utilization.gpu,temperature.gpu," - "memory.used,memory.total,power.draw,power.limit", - "--format=csv,noheader,nounits", - ], - capture_output = True, - text = True, - timeout = 5, - ) - - if result.returncode == 0 and result.stdout.strip(): - # nvidia-smi outputs one line per GPU; take GPU 0 - first_line = result.stdout.strip().splitlines()[0] - parts = [p.strip() for p in first_line.split(",")] - if len(parts) >= 6: - smi_data = { - "gpu_util": _parse_smi_value(parts[0]), - "temp": _parse_smi_value(parts[1]), - "vram_used_mb": _parse_smi_value(parts[2]), - "vram_total_mb": _parse_smi_value(parts[3]), - "power_draw": _parse_smi_value(parts[4]), - "power_limit": _parse_smi_value(parts[5]), - } - - except FileNotFoundError: - logger.debug("nvidia-smi not found, falling back to torch.cuda") - except Exception as e: - logger.warning(f"nvidia-smi query failed: {e}") - - # ── Backfill VRAM from torch.cuda if nvidia-smi returned [N/A] ── - vram_used_mb = smi_data.get("vram_used_mb") - vram_total_mb = smi_data.get("vram_total_mb") - - if vram_used_mb is None or vram_total_mb is None: - try: - import torch - - idx = torch.cuda.current_device() - props = torch.cuda.get_device_properties(idx) - if vram_total_mb is None: - vram_total_mb = props.total_memory / (1024**2) # bytes → MiB - if vram_used_mb is None: - vram_used_mb = torch.cuda.memory_allocated(idx) / (1024**2) + result = nvidia.get_primary_gpu_utilization() + if result.get("available"): + result["backend"] = device.value + return result except Exception as e: - logger.debug(f"torch.cuda VRAM backfill failed: {e}") + logger.warning("nvidia-smi utilization query failed: %s", e) - # ── Build response ────────────────────────────────────────── - gpu_util = smi_data.get("gpu_util") - temp = smi_data.get("temp") - power_draw = smi_data.get("power_draw") - power_limit = smi_data.get("power_limit") + mem = get_gpu_memory_info() + if device != DeviceType.CPU and mem.get("available"): + return { + "available": True, + "backend": device.value, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } - vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None - vram_total_gb = ( - round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None - ) - vram_pct = ( - round((vram_used_mb / vram_total_mb) * 100, 1) - if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 - else None - ) - power_pct = ( - round((power_draw / power_limit) * 100, 1) - if power_draw is not None and power_limit and power_limit > 0 - else None - ) + return {"available": False, "backend": device.value} - # If we got at least something useful, report available - has_any = any(v is not None for v in [gpu_util, temp, vram_used_gb, power_draw]) - if not has_any: - return {"available": False, "backend": device.value} + +def get_visible_gpu_utilization() -> Dict[str, Any]: + device = get_device() + + if device == DeviceType.CUDA: + parent_visible_spec = _get_parent_visible_gpu_spec() + try: + from . import nvidia + + result = nvidia.get_visible_gpu_utilization( + parent_visible_spec["numeric_ids"], + parent_cuda_visible_devices = parent_visible_spec["raw"], + ) + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("nvidia-smi visible GPU utilization query failed: %s", e) + + # Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel) + if device in (DeviceType.CUDA, DeviceType.XPU): + parent_ids = get_parent_visible_gpu_ids() + # When parent_visible_ids is empty (UUID/MIG mask or no CVD set), + # enumerate torch-visible ordinals so the UI still shows devices. + if parent_ids: + torch_indices = parent_ids + index_kind = "physical" + else: + visible_count = _torch_get_physical_gpu_count() or 0 + torch_indices = list(range(visible_count)) + index_kind = "relative" + torch_devices = _torch_get_per_device_info(torch_indices) + if torch_devices: + devices = [] + for td in torch_devices: + total = td["total_gb"] + used = td["used_gb"] + devices.append( + { + "index": td["index"], + "index_kind": index_kind, + "visible_ordinal": td["visible_ordinal"], + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": used, + "vram_total_gb": total, + "vram_utilization_pct": round((used / total) * 100, 1) + if total > 0 + else None, + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ) + return { + "available": True, + "backend": device.value, + "parent_visible_gpu_ids": parent_ids, + "devices": devices, + "index_kind": index_kind, + } + + if device == DeviceType.MLX: + mem = get_gpu_memory_info() + if not mem.get("available"): + return { + "available": False, + "backend": device.value, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", + } + return { + "available": True, + "backend": device.value, + "parent_visible_gpu_ids": [0], + "devices": [ + { + "index": 0, + "index_kind": "relative", + "visible_ordinal": 0, + "gpu_utilization_pct": None, + "temperature_c": None, + "vram_used_gb": round(mem.get("allocated_gb", 0), 2), + "vram_total_gb": round(mem.get("total_gb", 0), 2), + "vram_utilization_pct": round(mem.get("utilization_pct", 0), 1), + "power_draw_w": None, + "power_limit_w": None, + "power_utilization_pct": None, + } + ], + "index_kind": "relative", + } return { - "available": True, + "available": False, "backend": device.value, - "gpu_utilization_pct": gpu_util, - "temperature_c": temp, - "vram_used_gb": vram_used_gb, - "vram_total_gb": vram_total_gb, - "vram_utilization_pct": vram_pct, - "power_draw_w": power_draw, - "power_limit_w": power_limit, - "power_utilization_pct": power_pct, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", } @@ -417,37 +528,712 @@ _physical_gpu_count: Optional[int] = None _visible_gpu_count: Optional[int] = None +def _get_parent_visible_gpu_spec() -> Dict[str, Any]: + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") + + if cuda_visible is None: + return { + "raw": None, + "numeric_ids": list(range(get_physical_gpu_count())), + "supports_explicit_gpu_ids": True, + } + + cuda_visible = cuda_visible.strip() + if cuda_visible == "" or cuda_visible == "-1": + return { + "raw": cuda_visible, + "numeric_ids": [], + "supports_explicit_gpu_ids": True, + } + + tokens = [value.strip() for value in cuda_visible.split(",") if value.strip()] + try: + numeric_ids = [int(value) for value in tokens] + except ValueError: + return { + "raw": cuda_visible, + "numeric_ids": None, + "supports_explicit_gpu_ids": False, + } + + return { + "raw": cuda_visible, + "numeric_ids": numeric_ids, + "supports_explicit_gpu_ids": True, + } + + +def get_parent_visible_gpu_ids() -> list[int]: + parent_visible_ids = _get_parent_visible_gpu_spec()["numeric_ids"] + return list(parent_visible_ids) if parent_visible_ids is not None else [] + + +def resolve_requested_gpu_ids(gpu_ids: Optional[list[int]]) -> list[int]: + parent_visible_spec = _get_parent_visible_gpu_spec() + parent_visible_ids = get_parent_visible_gpu_ids() + physical_gpu_count = get_physical_gpu_count() + + if gpu_ids is None: + return parent_visible_ids + + requested_ids = list(gpu_ids) + if len(requested_ids) == 0: + return parent_visible_ids + + if not parent_visible_spec["supports_explicit_gpu_ids"]: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: explicit physical GPU IDs are " + f"unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries " + f"({parent_visible_spec['raw']!r}). Omit gpu_ids to use the " + "parent-visible devices." + ) + + if len(set(requested_ids)) != len(requested_ids): + raise ValueError( + f"Invalid gpu_ids {requested_ids}: duplicate GPU IDs are not allowed. " + f"Parent-visible GPUs: {parent_visible_ids}" + ) + + # Reject negative IDs unconditionally. + negative_ids = [gpu_id for gpu_id in requested_ids if gpu_id < 0] + if negative_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: GPU IDs must be non-negative. " + f"Rejected IDs: {negative_ids}. Parent-visible GPUs: {parent_visible_ids}" + ) + + # Only enforce the physical upper bound when we have a reliable count + # from nvidia-smi. When the count comes from torch, it reflects visible + # devices (filtered by CUDA_VISIBLE_DEVICES), not the physical total, + # so high physical indices like 3 would be falsely rejected on a + # CUDA_VISIBLE_DEVICES="2,3" machine that reports device_count()=2. + # The parent-visible check below is authoritative in all cases. + if physical_gpu_count > 0 and parent_visible_ids: + max_parent_id = max(parent_visible_ids) + if physical_gpu_count > max_parent_id: + # Count is plausibly physical (not just visible), so enforce it + out_of_range = [ + gpu_id for gpu_id in requested_ids if gpu_id >= physical_gpu_count + ] + if out_of_range: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: IDs must be physical GPU IDs " + f"between 0 and {physical_gpu_count - 1}. " + f"Rejected IDs: {out_of_range}. Parent-visible GPUs: {parent_visible_ids}" + ) + + disallowed_ids = [ + gpu_id for gpu_id in requested_ids if gpu_id not in parent_visible_ids + ] + if disallowed_ids: + raise ValueError( + f"Invalid gpu_ids {requested_ids}: requested GPUs {disallowed_ids} are " + f"outside the parent-visible set {parent_visible_ids}" + ) + + return requested_ids + + +def _resolve_model_identifier_for_gpu_estimate( + model_name: str, hf_token: Optional[str] = None +) -> str: + try: + from utils.models.model_config import ModelConfig + + config = ModelConfig.from_identifier(model_name, hf_token = hf_token) + if config and config.is_lora and config.base_model: + return config.base_model + return config.identifier if config else model_name + except Exception as e: + logger.debug( + "Could not resolve base model for GPU estimate '%s': %s", model_name, e + ) + return model_name + + +def _get_local_weight_size_bytes(model_name: str) -> Optional[int]: + model_path = Path(model_name) + if not model_path.exists(): + return None + + weight_exts = (".safetensors", ".bin", ".pt", ".pth") + total = 0 + for file in model_path.rglob("*"): + if file.is_file() and file.suffix in weight_exts: + total += file.stat().st_size + return total if total > 0 else None + + +def _get_hf_safetensors_total_params( + model_name: str, hf_token: Optional[str] = None +) -> Optional[int]: + try: + from huggingface_hub import model_info as hf_model_info + + info = hf_model_info(model_name, token = hf_token) + safetensors = getattr(info, "safetensors", None) + if isinstance(safetensors, dict): + total = safetensors.get("total") + if total: + return int(total) + except Exception as e: + logger.warning("Could not get safetensors metadata for '%s': %s", model_name, e) + return None + + +def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = None): + try: + from transformers import AutoConfig + + trust_remote_code = model_name.lower().startswith("unsloth/") + return AutoConfig.from_pretrained( + model_name, + token = hf_token, + trust_remote_code = trust_remote_code, + ) + except Exception as e: + logger.warning("Could not load config for '%s': %s", model_name, e) + return None + + +def _estimate_fp16_model_size_bytes_from_config(config) -> Optional[int]: + from .vram_estimation import extract_arch_config, compute_total_params + + arch = extract_arch_config(config) + if arch is None: + return None + return compute_total_params(arch) * 2 + + +def _estimate_fp16_model_size_bytes_from_vllm_utils(config) -> Optional[int]: + if config is None: + return None + + previous_unsloth_present = os.environ.get("UNSLOTH_IS_PRESENT") + os.environ["UNSLOTH_IS_PRESENT"] = "1" + try: + from unsloth_zoo import vllm_utils as _vllm_utils + + synthetic_total_bytes = 1024 * (1024**3) + original_get_mem_info = _vllm_utils.get_mem_info + try: + _vllm_utils.get_mem_info = lambda: ( + synthetic_total_bytes, + synthetic_total_bytes, + ) + _, _, _, memory_left_for_kv_cache_gb = ( + _vllm_utils.approximate_vllm_memory_usage( + config, + load_in_4bit = False, + load_in_8bit = False, + max_seq_length = 1, + gpu_memory_utilization = 1.0, + enable_lora = False, + account_for_gradients = False, + cuda_graph_overhead = False, + ) + ) + finally: + _vllm_utils.get_mem_info = original_get_mem_info + except Exception as e: + logger.debug("Could not estimate model size via vllm_utils: %s", e) + return None + finally: + if previous_unsloth_present is None: + os.environ.pop("UNSLOTH_IS_PRESENT", None) + else: + os.environ["UNSLOTH_IS_PRESENT"] = previous_unsloth_present + + model_size_gb = 1024.0 - memory_left_for_kv_cache_gb + if model_size_gb <= 0: + return None + return int(round(model_size_gb * (1024**3))) + + +def estimate_fp16_model_size_bytes( + model_name: str, hf_token: Optional[str] = None +) -> tuple[Optional[int], str]: + estimate_model = _resolve_model_identifier_for_gpu_estimate( + model_name, hf_token = hf_token + ) + + total_params = None + if "/" in estimate_model and not Path(estimate_model).exists(): + total_params = _get_hf_safetensors_total_params( + estimate_model, hf_token = hf_token + ) + if total_params: + return int(total_params * 2), "safetensors" + + config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token) + if config is not None: + config_bytes = _estimate_fp16_model_size_bytes_from_config(config) + if config_bytes is not None: + return config_bytes, "config" + + local_bytes = _get_local_weight_size_bytes(estimate_model) + if local_bytes is not None: + return local_bytes, "weight_bytes" + + vllm_bytes = _estimate_fp16_model_size_bytes_from_vllm_utils(config) + if vllm_bytes is not None: + return vllm_bytes, "vllm_utils" + + return None, "unavailable" + + +def estimate_required_model_memory_gb( + model_name: str, + *, + hf_token: Optional[str] = None, + training_type: Optional[str] = None, + load_in_4bit: bool = True, + batch_size: int = 4, + max_seq_length: int = 2048, + lora_rank: int = 16, + target_modules: Optional[list] = None, + gradient_checkpointing: str = "unsloth", + optimizer: str = "adamw_8bit", +) -> tuple[Optional[float], Dict[str, Any]]: + from .vram_estimation import ( + TrainingVramConfig, + extract_arch_config, + estimate_training_vram, + CUDA_OVERHEAD_BYTES, + QUANT_4BIT_FACTOR, + DEFAULT_TARGET_MODULES, + ) + + model_size_bytes, source = estimate_fp16_model_size_bytes( + model_name, hf_token = hf_token + ) + metadata: Dict[str, Any] = { + "mode": "inference" if training_type is None else "training", + "model_size_source": source, + } + if model_size_bytes is None: + metadata["required_gb"] = None + return None, metadata + + model_size_gb = model_size_bytes / (1024**3) + metadata["model_size_gb"] = round(model_size_gb, 3) + min_buffer_gb = 2.0 + + if training_type is None: + if load_in_4bit: + base_4bit_gb = model_size_gb / QUANT_4BIT_FACTOR + required_gb = base_4bit_gb + max(base_4bit_gb * 0.3, min_buffer_gb) + else: + required_gb = model_size_gb * 1.3 + metadata["required_gb"] = round(required_gb, 3) + return required_gb, metadata + + training_method = ( + "full" + if training_type == "Full Finetuning" + else ("qlora" if load_in_4bit else "lora") + ) + vram_config = TrainingVramConfig( + training_method = training_method, + batch_size = batch_size, + max_seq_length = max_seq_length, + lora_rank = lora_rank, + target_modules = target_modules or list(DEFAULT_TARGET_MODULES), + gradient_checkpointing = gradient_checkpointing, + optimizer = optimizer, + load_in_4bit = load_in_4bit, + ) + + estimate_model = _resolve_model_identifier_for_gpu_estimate( + model_name, hf_token = hf_token + ) + config = _load_config_for_gpu_estimate(estimate_model, hf_token = hf_token) + arch = extract_arch_config(config) if config is not None else None + + if arch is not None: + breakdown = estimate_training_vram(arch, vram_config) + required_gb = breakdown.total / (1024**3) + metadata["required_gb"] = round(required_gb, 3) + metadata["estimation_mode"] = "detailed" + metadata["vram_breakdown"] = breakdown.to_gb_dict() + max_gpus = max(1, get_visible_gpu_count()) + for n_gpus in range(1, max_gpus + 1): + metadata["vram_breakdown"][f"min_per_gpu_{n_gpus}"] = round( + breakdown.min_gpu_vram(n_gpus) / (1024**3), 3 + ) + return required_gb, metadata + + # Fallback when model config is unavailable + overhead_gb = CUDA_OVERHEAD_BYTES / (1024**3) + if training_method == "full": + required_gb = model_size_gb * 3.5 + overhead_gb + elif training_method == "qlora": + base_4bit_gb = model_size_gb / QUANT_4BIT_FACTOR + lora_overhead_gb = model_size_gb * 0.04 + act_gb = model_size_gb * 0.15 * (batch_size / 4) * (max_seq_length / 2048) + required_gb = base_4bit_gb + lora_overhead_gb + act_gb + overhead_gb + else: + lora_overhead_gb = model_size_gb * 0.04 + act_gb = model_size_gb * 0.15 * (batch_size / 4) * (max_seq_length / 2048) + required_gb = model_size_gb + lora_overhead_gb + act_gb + overhead_gb + + metadata["required_gb"] = round(required_gb, 3) + metadata["estimation_mode"] = "fallback" + return required_gb, metadata + + +def auto_select_gpu_ids( + model_name: str, + *, + hf_token: Optional[str] = None, + training_type: Optional[str] = None, + load_in_4bit: bool = True, + batch_size: int = 4, + max_seq_length: int = 2048, + lora_rank: int = 16, + target_modules: Optional[list] = None, + gradient_checkpointing: str = "unsloth", + optimizer: str = "adamw_8bit", +) -> tuple[Optional[list[int]], Dict[str, Any]]: + metadata: Dict[str, Any] = {"selection_mode": "auto"} + + if get_device() != DeviceType.CUDA: + metadata["selection_mode"] = "non_cuda" + return None, metadata + + required_gb, estimate_metadata = estimate_required_model_memory_gb( + model_name, + hf_token = hf_token, + training_type = training_type, + load_in_4bit = load_in_4bit, + batch_size = batch_size, + max_seq_length = max_seq_length, + lora_rank = lora_rank, + target_modules = target_modules, + gradient_checkpointing = gradient_checkpointing, + optimizer = optimizer, + ) + metadata.update(estimate_metadata) + parent_visible_spec = _get_parent_visible_gpu_spec() + metadata["parent_cuda_visible_devices"] = parent_visible_spec["raw"] + + if not parent_visible_spec["supports_explicit_gpu_ids"]: + metadata["selection_mode"] = "inherit_parent_visible" + metadata["selected_gpu_ids"] = None + return None, metadata + + if required_gb is None: + # Cannot estimate model size -- fall back to all visible GPUs + # rather than risk loading on a single GPU that may not have + # enough memory. + parent_ids = get_parent_visible_gpu_ids() + metadata["selection_mode"] = "fallback_all" + metadata["selected_gpu_ids"] = parent_ids + return parent_ids, metadata + + utilization = get_visible_gpu_utilization() + devices = utilization.get("devices", []) + parent_ids = get_parent_visible_gpu_ids() + + if not devices: + metadata["selection_mode"] = "fallback_all" + metadata["selected_gpu_ids"] = parent_ids + return parent_ids, metadata + + gpu_candidates = [] + for device in devices: + total_gb = device.get("vram_total_gb") + used_gb = device.get("vram_used_gb") + if total_gb is None or used_gb is None: + continue + free_gb = max(total_gb - used_gb, 0.0) + gpu_candidates.append( + { + "index": device["index"], + "free_gb": free_gb, + } + ) + + if not gpu_candidates: + metadata["selection_mode"] = "fallback_all" + metadata["selected_gpu_ids"] = parent_ids + return parent_ids, metadata + + ranked = sorted(gpu_candidates, key = lambda item: (-item["free_gb"], item["index"])) + free_by_index = {item["index"]: item["free_gb"] for item in ranked} + selected: list[int] = [] + usable_gb = 0.0 + # Multi-GPU sharding has overhead from inter-GPU communication (NCCL + # all-reduce, PCIe/NVLink transfers, synchronization barriers), so each + # additional GPU contributes less than its raw free memory. The first GPU + # keeps its full capacity (no cross-device overhead). 0.85 was calibrated + # empirically on 2-8 GPU setups with NVLink and PCIe topologies -- the + # 15% discount accounts for NCCL buffers (~2-5% of VRAM), pipeline bubble + # overhead, and memory fragmentation from non-uniform shard sizes. + multi_gpu_overhead = 0.85 + + # Per-GPU check: activations don't shard, so each GPU needs its weight + # shard + full activation cost. Use precomputed min_per_gpu_N values. + vram_breakdown = estimate_metadata.get("vram_breakdown", {}) + + for candidate in ranked: + selected.append(candidate["index"]) + if len(selected) == 1: + usable_gb = candidate["free_gb"] + else: + first_gpu_id = selected[0] + usable_gb = free_by_index[first_gpu_id] + sum( + free_by_index[gpu_id] * multi_gpu_overhead for gpu_id in selected[1:] + ) + + total_fits = usable_gb >= required_gb + + per_gpu_fits = True + if total_fits and len(selected) > 1: + min_key = f"min_per_gpu_{len(selected)}" + min_per_gpu_gb = vram_breakdown.get(min_key) + if min_per_gpu_gb is not None: + smallest_free = min(free_by_index[gpu_id] for gpu_id in selected) + per_gpu_fits = smallest_free >= min_per_gpu_gb + + if total_fits and per_gpu_fits: + metadata["usable_gb"] = round(usable_gb, 3) + metadata["selection_mode"] = "auto" + metadata["selected_gpu_ids"] = selected + logger.debug( + "Selected GPUs automatically", + model_name = model_name, + selected_gpu_ids = selected, + usable_gb = metadata["usable_gb"], + required_gb = metadata.get("required_gb"), + multi_gpu_overhead = multi_gpu_overhead, + ) + return selected, metadata + + # Use only GPUs with verified VRAM data (from gpu_candidates, not raw devices) + fallback_all = ( + [c["index"] for c in gpu_candidates] if gpu_candidates else parent_ids + ) + metadata["selection_mode"] = "fallback_all" + if ranked: + fallback_usable = ranked[0]["free_gb"] + sum( + c["free_gb"] * multi_gpu_overhead for c in ranked[1:] + ) + else: + fallback_usable = 0.0 + metadata["usable_gb"] = round(fallback_usable, 3) + metadata["selected_gpu_ids"] = fallback_all + logger.warning( + "Falling back to all visible GPUs -- model may not fit", + model_name = model_name, + selected_gpu_ids = fallback_all, + usable_gb = metadata["usable_gb"], + required_gb = metadata.get("required_gb"), + multi_gpu_overhead = multi_gpu_overhead, + ) + return fallback_all, metadata + + +def prepare_gpu_selection( + gpu_ids: Optional[list[int]], + *, + model_name: str, + hf_token: Optional[str] = None, + training_type: Optional[str] = None, + load_in_4bit: bool = True, + batch_size: int = 4, + max_seq_length: int = 2048, + lora_rank: int = 16, + target_modules: Optional[list] = None, + gradient_checkpointing: str = "unsloth", + optimizer: str = "adamw_8bit", +) -> tuple[Optional[list[int]], Dict[str, Any]]: + """Resolve which physical GPUs to use for a model load. + + GPU selection modes: + - **Explicit** (``gpu_ids=[5, 6, 7]``): the caller chooses exact GPUs. + All listed GPUs are used and the model is sharded across them via + ``device_map="balanced"``, regardless of whether the model would fit + on fewer GPUs. IDs are validated against the parent-visible set. + - **Auto** (``gpu_ids=None`` or ``[]``): ``auto_select_gpu_ids`` estimates + VRAM requirements and picks the *minimum* number of GPUs needed, + preferring GPUs with the most free memory. + + The returned ``gpu_ids`` list is later passed to ``get_device_map()`` which + maps it to a Hugging Face ``device_map`` string, and to ``apply_gpu_ids()`` + in the worker subprocess which narrows ``CUDA_VISIBLE_DEVICES`` before any + torch/CUDA initialisation. + """ + if gpu_ids and get_device() != DeviceType.CUDA: + raise ValueError( + f"gpu_ids {list(gpu_ids)} is only supported on CUDA devices, " + f"but the current backend is '{get_device().value}'." + ) + + if gpu_ids: + resolved = resolve_requested_gpu_ids(gpu_ids) + metadata = { + "selection_mode": "explicit", + "selected_gpu_ids": resolved, + } + return resolved, metadata + + selected_gpu_ids, metadata = auto_select_gpu_ids( + model_name, + hf_token = hf_token, + training_type = training_type, + load_in_4bit = load_in_4bit, + batch_size = batch_size, + max_seq_length = max_seq_length, + lora_rank = lora_rank, + target_modules = target_modules, + gradient_checkpointing = gradient_checkpointing, + optimizer = optimizer, + ) + return selected_gpu_ids, metadata + + def get_physical_gpu_count() -> int: """ - Return the number of physical NVIDIA GPUs on the machine. + Return the number of physical GPUs on the machine. - Uses ``nvidia-smi -L`` which is NOT affected by CUDA_VISIBLE_DEVICES, - so it always reflects the true hardware count. + Uses ``nvidia-smi -L`` on NVIDIA (unaffected by CUDA_VISIBLE_DEVICES), + with a torch-based fallback for AMD ROCm and Intel XPU. Result is cached after the first call. """ global _physical_gpu_count if _physical_gpu_count is not None: return _physical_gpu_count - try: - import subprocess + device = get_device() - result = subprocess.run( - ["nvidia-smi", "-L"], - capture_output = True, - text = True, - timeout = 5, - ) - if result.returncode == 0 and result.stdout.strip(): - _physical_gpu_count = len(result.stdout.strip().splitlines()) - else: - _physical_gpu_count = 1 - except Exception: + if device == DeviceType.CUDA: + try: + from . import nvidia + + count = nvidia.get_physical_gpu_count() + if count is not None: + _physical_gpu_count = count + return _physical_gpu_count + except Exception: + pass + # nvidia-smi unavailable or failed — fall back to torch + count = _torch_get_physical_gpu_count() + _physical_gpu_count = count if count is not None else 1 + return _physical_gpu_count + + if device == DeviceType.XPU: + count = _torch_get_physical_gpu_count() + _physical_gpu_count = count if count is not None else 1 + return _physical_gpu_count + + if device == DeviceType.MLX: _physical_gpu_count = 1 + return _physical_gpu_count + + _physical_gpu_count = 0 return _physical_gpu_count +def get_backend_visible_gpu_info() -> Dict[str, Any]: + device = get_device() + if device in (DeviceType.CUDA, DeviceType.XPU): + parent_visible_ids = get_parent_visible_gpu_ids() + # Try nvidia-smi first (NVIDIA only) + if device == DeviceType.CUDA: + try: + from . import nvidia + + parent_visible_spec = _get_parent_visible_gpu_spec() + result = nvidia.get_backend_visible_gpu_info( + parent_visible_spec["numeric_ids"], + parent_visible_spec["raw"], + ) + if result.get("available"): + result["backend"] = device.value + return result + except Exception as e: + logger.warning("Backend GPU visibility query failed: %s", e) + + # Torch fallback (AMD ROCm, Intel XPU, nvidia-smi missing/failed) + # When parent_visible_ids is empty (UUID/MIG mask), enumerate by + # torch ordinal so the UI still shows devices. + if parent_visible_ids: + torch_indices = parent_visible_ids + index_kind = "physical" + else: + visible_count = _torch_get_physical_gpu_count() or 0 + torch_indices = list(range(visible_count)) + index_kind = "relative" + torch_devices = _torch_get_per_device_info(torch_indices) + if torch_devices: + devices = [ + { + "index": td["index"], + "index_kind": index_kind, + "visible_ordinal": td["visible_ordinal"], + "name": td["name"], + "memory_total_gb": td["total_gb"], + } + for td in torch_devices + ] + return { + "available": True, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": parent_visible_ids, + "devices": devices, + "index_kind": index_kind, + } + + return { + "available": False, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": parent_visible_ids, + "devices": [], + "index_kind": "physical", + } + + if device == DeviceType.MLX: + mem = get_gpu_memory_info() + if not mem.get("available"): + return { + "available": False, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", + } + return { + "available": True, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": [0], + "devices": [ + { + "index": 0, + "index_kind": "relative", + "visible_ordinal": 0, + "name": mem.get("device_name", "MLX"), + "memory_total_gb": round(mem.get("total_gb", 0), 2), + } + ], + "index_kind": "relative", + } + + return { + "available": False, + "backend": device.value, + "backend_cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "relative", + } + + def get_visible_gpu_count() -> int: """ Return the number of GPUs visible to this process. @@ -460,8 +1246,6 @@ def get_visible_gpu_count() -> int: if _visible_gpu_count is not None: return _visible_gpu_count - import os - cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES") if cuda_visible is not None: # "" means zero GPUs, "0" means 1, "0,1,2" means 3 @@ -476,13 +1260,103 @@ def get_visible_gpu_count() -> int: try: import torch - _visible_gpu_count = torch.cuda.device_count() + if get_device() == DeviceType.XPU and hasattr(torch, "xpu"): + _visible_gpu_count = torch.xpu.device_count() + else: + _visible_gpu_count = torch.cuda.device_count() except Exception: _visible_gpu_count = get_physical_gpu_count() return _visible_gpu_count +def apply_gpu_ids(gpu_ids) -> None: + if gpu_ids is None: + return + + # Empty list means "no GPUs visible" -- treat the same as None + # (inherit parent) to avoid setting CUDA_VISIBLE_DEVICES="" which + # disables CUDA entirely and crashes downstream torch calls. + if isinstance(gpu_ids, (list, tuple)) and len(gpu_ids) == 0: + return + + global _visible_gpu_count + + if isinstance(gpu_ids, (list, tuple)): + value = ",".join(str(g) for g in gpu_ids) + else: + value = str(gpu_ids) + + os.environ["CUDA_VISIBLE_DEVICES"] = value + _visible_gpu_count = None + logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value) + + +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. + + Returns ``"balanced"`` (shard evenly across GPUs) when: + - ``gpu_ids`` explicitly lists >1 GPU, **or** + - ``CUDA_VISIBLE_DEVICES`` uses UUID/MIG identifiers (non-numeric) and + more than one GPU is visible (fallback: we cannot resolve numeric IDs, + so we assume the caller intends multi-GPU). + + Returns ``"sequential"`` (single device) in all other cases, including + non-CUDA backends (CPU, MLX). + + Callers should use ``prepare_gpu_selection()`` upstream to determine the + ``gpu_ids`` list -- that function handles the smart auto-selection of the + minimum number of GPUs needed for a given model. + """ + device = get_device() + if device == DeviceType.CUDA: + multi_gpu = gpu_ids is not None and len(gpu_ids) > 1 + + if not multi_gpu: + # UUID/MIG masks cannot be split into numeric IDs, so if multiple + # GPUs are visible we assume multi-GPU sharding is intended. + parent_visible_spec = _get_parent_visible_gpu_spec() + if ( + parent_visible_spec["numeric_ids"] is None + and get_visible_gpu_count() > 1 + ): + multi_gpu = True + + if multi_gpu: + return "balanced_low_0" if for_inference else "balanced" + + return "sequential" + + +def get_offloaded_device_map_entries(model) -> dict[str, str]: + hf_device_map = getattr(model, "hf_device_map", None) + if not isinstance(hf_device_map, dict): + return {} + return { + module_name: placement + for module_name, placement in hf_device_map.items() + if placement in ("cpu", "disk") + } + + +def raise_if_offloaded(model, device_map: str, context: str = "Loading") -> None: + """Raise ``ValueError`` if *model* has modules offloaded to CPU or disk.""" + offloaded = get_offloaded_device_map_entries(model) + if not offloaded: + return + example = ", ".join( + f"{name}={placement}" for name, placement in list(offloaded.items())[:5] + ) + raise ValueError( + f"{context} does not support models loaded with CPU or disk offload. " + f"device_map='{device_map}' produced offloaded modules: {example}" + ) + + def safe_num_proc(desired: Optional[int] = None) -> int: """ Return a safe ``num_proc`` for ``dataset.map()`` calls. @@ -507,7 +1381,6 @@ def safe_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer ≥ 1. """ - import os import sys # Windows and macOS use 'spawn' for multiprocessing -- the overhead of @@ -546,8 +1419,6 @@ def safe_thread_num_proc(desired: Optional[int] = None) -> int: Returns: A safe integer >= 1. """ - import os - if desired is None or not isinstance(desired, int): desired = max(1, (os.cpu_count() or 1) // 3) diff --git a/studio/backend/utils/hardware/nvidia.py b/studio/backend/utils/hardware/nvidia.py new file mode 100644 index 0000000000..dc5295c302 --- /dev/null +++ b/studio/backend/utils/hardware/nvidia.py @@ -0,0 +1,279 @@ +# 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 subprocess +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + + +def _parse_smi_value(raw: str): + raw = raw.strip() + if not raw or raw == "[N/A]": + return None + try: + return float(raw) + except (ValueError, TypeError): + return None + + +def _build_gpu_metrics( + vram_used_mb, + vram_total_mb, + power_draw, + power_limit, + **extra, +) -> dict[str, Any]: + return { + **extra, + "vram_used_gb": round(vram_used_mb / 1024, 2) + if vram_used_mb is not None + else None, + "vram_total_gb": round(vram_total_mb / 1024, 2) + if vram_total_mb is not None + else None, + "vram_utilization_pct": round((vram_used_mb / vram_total_mb) * 100, 1) + if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0 + else None, + "power_draw_w": power_draw, + "power_limit_w": power_limit, + "power_utilization_pct": round((power_draw / power_limit) * 100, 1) + if power_draw is not None and power_limit and power_limit > 0 + else None, + } + + +def _visible_ordinal_map( + parent_visible_ids: Optional[list[int]], +) -> Optional[dict[int, int]]: + if parent_visible_ids is None: + return None + return {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)} + + +def get_physical_gpu_count() -> Optional[int]: + """Return physical GPU count via nvidia-smi, or None on failure.""" + try: + result = subprocess.run( + ["nvidia-smi", "-L"], + capture_output = True, + text = True, + timeout = 5, + ) + if result.returncode == 0 and result.stdout.strip(): + return len(result.stdout.strip().splitlines()) + logger.warning( + "nvidia-smi -L returned code %d; caller should fall back to torch", + result.returncode, + ) + except Exception as e: + logger.warning("nvidia-smi -L failed: %s; caller should fall back to torch", e) + return None + + +def get_primary_gpu_utilization() -> dict[str, Any]: + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=utilization.gpu,temperature.gpu," + "memory.used,memory.total,power.draw,power.limit", + "--format=csv,noheader,nounits", + ], + capture_output = True, + text = True, + timeout = 5, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("nvidia-smi query failed in get_primary_gpu_utilization: %s", e) + return {"available": False} + if result.returncode != 0 or not result.stdout.strip(): + return {"available": False} + + first_line = result.stdout.strip().splitlines()[0] + parts = [p.strip() for p in first_line.split(",")] + if len(parts) < 6: + return {"available": False} + + return _build_gpu_metrics( + vram_used_mb = _parse_smi_value(parts[2]), + vram_total_mb = _parse_smi_value(parts[3]), + power_draw = _parse_smi_value(parts[4]), + power_limit = _parse_smi_value(parts[5]), + available = True, + gpu_utilization_pct = _parse_smi_value(parts[0]), + temperature_c = _parse_smi_value(parts[1]), + ) + + +def get_visible_gpu_utilization( + parent_visible_ids: Optional[list[int]], + parent_cuda_visible_devices: Optional[str] = None, +) -> dict[str, Any]: + # When parent_visible_ids is None (UUID/MIG mask), we cannot safely + # map nvidia-smi rows to the process's visible devices. Return empty + # instead of exposing all physical GPUs. + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + visible_ordinals = _visible_ordinal_map(parent_visible_ids) + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,utilization.gpu,temperature.gpu," + "memory.used,memory.total,power.draw,power.limit", + "--format=csv,noheader,nounits", + ], + capture_output = True, + text = True, + timeout = 5, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("nvidia-smi query failed in get_visible_gpu_utilization: %s", e) + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + if result.returncode != 0 or not result.stdout.strip(): + return { + "available": False, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + devices = [] + for line in result.stdout.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) < 7: + continue + + try: + idx = int(parts[0]) + except (ValueError, TypeError): + continue + + if visible_ordinals is not None and idx not in visible_ordinals: + continue + + devices.append( + _build_gpu_metrics( + vram_used_mb = _parse_smi_value(parts[3]), + vram_total_mb = _parse_smi_value(parts[4]), + power_draw = _parse_smi_value(parts[5]), + power_limit = _parse_smi_value(parts[6]), + index = idx, + index_kind = "physical", + visible_ordinal = ( + visible_ordinals[idx] + if visible_ordinals is not None + else len(devices) + ), + gpu_utilization_pct = _parse_smi_value(parts[1]), + temperature_c = _parse_smi_value(parts[2]), + ) + ) + + return { + "available": len(devices) > 0, + "backend_cuda_visible_devices": parent_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": devices, + "index_kind": "physical", + } + + +def get_backend_visible_gpu_info( + parent_visible_ids: Optional[list[int]], + backend_cuda_visible_devices: Optional[str], +) -> dict[str, Any]: + # When parent_visible_ids is None (UUID/MIG mask), we cannot safely + # map nvidia-smi rows to the process's visible devices. + if parent_visible_ids is None: + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": [], + "devices": [], + "index_kind": "unresolved", + } + visible_ordinals = _visible_ordinal_map(parent_visible_ids) + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,name,memory.total", + "--format=csv,noheader,nounits", + ], + capture_output = True, + text = True, + timeout = 10, + ) + except (OSError, subprocess.TimeoutExpired) as e: + logger.warning("nvidia-smi query failed in get_backend_visible_gpu_info: %s", e) + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + if result.returncode != 0: + return { + "available": False, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": [], + "index_kind": "physical", + } + + devices = [] + for line in result.stdout.strip().splitlines(): + parts = [p.strip() for p in line.split(",")] + if len(parts) < 3: + continue + try: + idx = int(parts[0]) + except (ValueError, TypeError): + continue + if visible_ordinals is not None and idx not in visible_ordinals: + continue + # Use split with limit to handle GPU names containing commas + name = parts[1] if len(parts) == 3 else ", ".join(parts[1:-1]) + try: + mem_total_mb = int(parts[-1]) + except (ValueError, TypeError): + continue + devices.append( + { + "index": idx, + "index_kind": "physical", + "visible_ordinal": ( + visible_ordinals[idx] + if visible_ordinals is not None + else len(devices) + ), + "name": name, + "memory_total_gb": round(mem_total_mb / 1024, 2), + } + ) + + return { + "available": len(devices) > 0, + "backend_cuda_visible_devices": backend_cuda_visible_devices, + "parent_visible_gpu_ids": parent_visible_ids or [], + "devices": devices, + "index_kind": "physical", + } diff --git a/studio/backend/utils/hardware/vram_estimation.py b/studio/backend/utils/hardware/vram_estimation.py new file mode 100644 index 0000000000..e03665374d --- /dev/null +++ b/studio/backend/utils/hardware/vram_estimation.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +""" +Training VRAM estimation. + +Total VRAM = weights + LoRA adapters + optimizer states + gradients + + activations + CUDA overhead. +Activation formula from unsloth_zoo/vllm_utils.py. +All constants empirically calibrated against Llama-3.2-1B on B200. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, Optional + +QUANT_4BIT_FACTOR = 16 / 5 +CUDA_OVERHEAD_BYTES = int(1.4 * 1024**3) # calibrated on RTX 5070 Ti + +DEFAULT_TARGET_MODULES = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +] + +# Empirically calibrated bytes/param — see VRAM_ESTIMATION.md for rationale. +OPTIMIZER_BYTES_PER_PARAM: Dict[str, int] = { + "adamw_8bit": 4, # BNB upcasts to fp32 during step + "paged_adamw_8bit": 4, + "adamw_bnb_8bit": 4, + "paged_adamw_32bit": 8, + "adamw_torch": 6, # fused, no master copy + "adamw_torch_fused": 6, + "sgd": 4, +} + +# (full_ft_multiplier, lora_multiplier) — fraction of num_layers. +# LoRA: frozen base layers skip activation storage, but you always need +# at least ~1 layer in flight during backprop recomputation. +GC_LAYER_MULTIPLIERS = { + "none": (None, None), + "true": (2.0, 1.0), + "unsloth": (1.5, 1.0), +} + + +@dataclass +class ModelArchConfig: + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + intermediate_size: int + vocab_size: int + tie_word_embeddings: bool = True + num_experts: Optional[int] = None + moe_intermediate_size: Optional[int] = None + n_shared_experts: int = 0 + num_dense_layers: int = 0 + q_lora_rank: Optional[int] = None + kv_lora_rank: Optional[int] = None + qk_nope_head_dim: Optional[int] = None + qk_rope_head_dim: Optional[int] = None + v_head_dim: Optional[int] = None + + +@dataclass +class TrainingVramConfig: + training_method: str = "qlora" + batch_size: int = 4 + max_seq_length: int = 2048 + lora_rank: int = 16 + target_modules: list = field(default_factory = lambda: list(DEFAULT_TARGET_MODULES)) + gradient_checkpointing: str = "unsloth" + optimizer: str = "adamw_8bit" + load_in_4bit: bool = True + + +@dataclass +class VramBreakdown: + model_weights: int + lora_adapters: int + optimizer_states: int + gradients: int + activations: int + cuda_overhead: int + # The computed (formula-based) activation cost before floors. + # This is the true per-layer cost that doesn't shard across GPUs. + activations_computed: int = 0 + + @property + def total(self) -> int: + return ( + self.model_weights + + self.lora_adapters + + self.optimizer_states + + self.gradients + + self.activations + + self.cuda_overhead + ) + + def min_gpu_vram(self, n_gpus: int) -> int: + """Minimum VRAM a single GPU needs: its shard + non-shardable costs. + + Weights/LoRA/optimizer/gradients shard across GPUs. + The computed activation cost does NOT shard (one GPU runs the layer). + The floor portion (activations - computed) is overhead that shards. + """ + shardable = ( + self.model_weights + + self.lora_adapters + + self.optimizer_states + + self.gradients + + (self.activations - self.activations_computed) # floor overhead shards + ) + per_gpu_fixed = self.activations_computed + self.cuda_overhead + return shardable // max(n_gpus, 1) + per_gpu_fixed + + def to_gb_dict(self) -> Dict[str, float]: + return { + "model_weights_gb": round(self.model_weights / (1024**3), 3), + "lora_adapters_gb": round(self.lora_adapters / (1024**3), 3), + "optimizer_states_gb": round(self.optimizer_states / (1024**3), 3), + "gradients_gb": round(self.gradients / (1024**3), 3), + "activations_gb": round(self.activations / (1024**3), 3), + "cuda_overhead_gb": round(self.cuda_overhead / (1024**3), 3), + "total_gb": round(self.total / (1024**3), 3), + } + + +def _compute_num_dense_layers(text_config, total_layers: int) -> int: + """Count how many layers use dense MLP instead of MoE.""" + first_k = getattr(text_config, "first_k_dense_replace", None) + if first_k is not None: + return min(int(first_k), total_layers) + + sparse_step = getattr(text_config, "decoder_sparse_step", None) + mlp_only = getattr(text_config, "mlp_only_layers", None) or [] + if sparse_step is not None and sparse_step > 0: + mlp_only_set = set(mlp_only) + moe_count = sum( + 1 + for i in range(total_layers) + if i not in mlp_only_set and (i + 1) % sparse_step == 0 + ) + return total_layers - moe_count + + return 0 + + +def extract_arch_config(hf_config) -> Optional[ModelArchConfig]: + text_config = getattr(hf_config, "text_config", None) or hf_config + + hidden_size = getattr(text_config, "hidden_size", None) + num_layers = getattr(text_config, "num_hidden_layers", None) + num_heads = getattr(text_config, "num_attention_heads", None) + intermediate_size = getattr(text_config, "intermediate_size", None) + vocab_size = getattr(text_config, "vocab_size", None) + + if isinstance(intermediate_size, (list, tuple)): + intermediate_size = intermediate_size[0] if intermediate_size else None + if intermediate_size is None and hidden_size is not None: + intermediate_size = hidden_size * 4 + + if not all( + v is not None + for v in (hidden_size, num_layers, num_heads, intermediate_size, vocab_size) + ): + return None + if num_heads <= 0: + return None + + num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads) + + num_experts = None + for attr in ("num_local_experts", "num_experts", "n_routed_experts"): + num_experts = getattr(text_config, attr, None) + if num_experts is not None: + break + + moe_intermediate = getattr(text_config, "moe_intermediate_size", None) + n_shared_experts = getattr(text_config, "n_shared_experts", None) or 0 + + num_dense_layers = 0 + if num_experts is not None and num_experts > 1: + num_dense_layers = _compute_num_dense_layers(text_config, num_layers) + + q_lora_rank = getattr(text_config, "q_lora_rank", None) + kv_lora_rank = getattr(text_config, "kv_lora_rank", None) + qk_nope_head_dim = getattr(text_config, "qk_nope_head_dim", None) + qk_rope_head_dim = getattr(text_config, "qk_rope_head_dim", None) + v_head_dim = getattr(text_config, "v_head_dim", None) + + return ModelArchConfig( + hidden_size = hidden_size, + num_hidden_layers = num_layers, + num_attention_heads = num_heads, + num_key_value_heads = num_kv_heads, + intermediate_size = intermediate_size, + vocab_size = vocab_size, + tie_word_embeddings = getattr(text_config, "tie_word_embeddings", True), + num_experts = num_experts, + moe_intermediate_size = moe_intermediate, + n_shared_experts = n_shared_experts, + num_dense_layers = num_dense_layers, + q_lora_rank = q_lora_rank, + kv_lora_rank = kv_lora_rank, + qk_nope_head_dim = qk_nope_head_dim, + qk_rope_head_dim = qk_rope_head_dim, + v_head_dim = v_head_dim, + ) + + +def _get_kv_size(arch: ModelArchConfig) -> int: + return (arch.hidden_size // arch.num_attention_heads) * arch.num_key_value_heads + + +def _get_mlp_size(arch: ModelArchConfig) -> int: + if arch.moe_intermediate_size is not None: + return arch.moe_intermediate_size + return arch.intermediate_size + + +def _get_num_experts(arch: ModelArchConfig) -> int: + return arch.num_experts if arch.num_experts and arch.num_experts > 1 else 1 + + +def _compute_attn_elements(arch: ModelArchConfig) -> int: + """Attention weight elements per layer.""" + hd = arch.hidden_size + if arch.q_lora_rank is not None: + nh = arch.num_attention_heads + qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim + q_a = hd * arch.q_lora_rank + q_b = arch.q_lora_rank * (nh * qk_head) + kv_a = hd * (arch.kv_lora_rank + arch.qk_rope_head_dim) + kv_b = arch.kv_lora_rank * (nh * (arch.qk_nope_head_dim + arch.v_head_dim)) + o = (nh * arch.v_head_dim) * hd + norms = arch.q_lora_rank + arch.kv_lora_rank + return q_a + q_b + kv_a + kv_b + o + norms + kv_size = _get_kv_size(arch) + return (hd + kv_size + kv_size + hd) * hd + + +def _compute_dense_mlp_elements(arch: ModelArchConfig) -> int: + return arch.hidden_size * arch.intermediate_size * 3 + + +def _compute_moe_mlp_elements(arch: ModelArchConfig) -> int: + hd = arch.hidden_size + mlp_size = _get_mlp_size(arch) + n_experts = _get_num_experts(arch) + return hd * mlp_size * 3 * (n_experts + arch.n_shared_experts) + n_experts * hd + + +def _compute_layer_elements(arch: ModelArchConfig): + """Return (total_quantizable, layernorms_per_layer, embed, lm_head) element counts. + + total_quantizable is summed across ALL layers (not per-layer). + """ + hd = arch.hidden_size + n_layers = arch.num_hidden_layers + n_experts = _get_num_experts(arch) + + attn_total = _compute_attn_elements(arch) * n_layers + + if n_experts > 1: + n_dense = arch.num_dense_layers + n_moe = n_layers - n_dense + mlp_total = ( + _compute_moe_mlp_elements(arch) * n_moe + + _compute_dense_mlp_elements(arch) * n_dense + ) + else: + mlp_total = _compute_dense_mlp_elements(arch) * n_layers + + layernorms = 2 * hd + embed_tokens = arch.vocab_size * hd + lm_head = 0 if arch.tie_word_embeddings else arch.vocab_size * hd + return attn_total + mlp_total, layernorms, embed_tokens, lm_head + + +def compute_model_weights_bytes( + arch: ModelArchConfig, + training_method: str, + load_in_4bit: bool, +) -> int: + total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch) + n_layers = arch.num_hidden_layers + non_quantizable = layernorms * n_layers + embed_tokens + lm_head + + if training_method == "qlora" and load_in_4bit: + return int(total_quantizable * 2 / QUANT_4BIT_FACTOR + non_quantizable * 2) + + return int((total_quantizable + non_quantizable) * 2) + + +def compute_total_params(arch: ModelArchConfig) -> int: + total_quantizable, layernorms, embed_tokens, lm_head = _compute_layer_elements(arch) + n_layers = arch.num_hidden_layers + return total_quantizable + layernorms * n_layers + embed_tokens + lm_head + + +def _lora_attn_elements( + arch: ModelArchConfig, + r: int, + target_modules: list, +) -> int: + hd = arch.hidden_size + if arch.q_lora_rank is not None: + # MLA: q_proj->q_b, k_proj->kv_a, v_proj->kv_b, o_proj->o + nh = arch.num_attention_heads + qk_head = arch.qk_nope_head_dim + arch.qk_rope_head_dim + kv_out = nh * (arch.qk_nope_head_dim + arch.v_head_dim) + o_in = nh * arch.v_head_dim + dims = { + "q_proj": (arch.q_lora_rank, nh * qk_head), + "k_proj": (hd, arch.kv_lora_rank + arch.qk_rope_head_dim), + "v_proj": (arch.kv_lora_rank, kv_out), + "o_proj": (o_in, hd), + } + else: + kv_size = _get_kv_size(arch) + dims = { + "q_proj": (hd, hd), + "k_proj": (hd, kv_size), + "v_proj": (hd, kv_size), + "o_proj": (hd, hd), + } + total = 0 + for name, (in_dim, out_dim) in dims.items(): + if name in target_modules: + total += in_dim * r + r * out_dim + return total + + +def _lora_mlp_elements( + hd: int, + mlp_size: int, + r: int, + target_modules: list, + expert_mult: int, +) -> int: + module_ab = { + "gate_proj": (hd * r, r * mlp_size), + "up_proj": (hd * r, r * mlp_size), + "down_proj": (mlp_size * r, r * hd), + } + total = 0 + for name, (a, b) in module_ab.items(): + if name in target_modules: + total += (a + b) * expert_mult + return total + + +def compute_lora_params( + arch: ModelArchConfig, + lora_rank: int, + target_modules: list, +) -> int: + hd = arch.hidden_size + r = lora_rank + n_layers = arch.num_hidden_layers + n_experts = _get_num_experts(arch) + + attn_total = _lora_attn_elements(arch, r, target_modules) * n_layers + + if n_experts > 1: + n_dense = arch.num_dense_layers + n_moe = n_layers - n_dense + # Include shared experts alongside routed experts + moe_expert_mult = n_experts + arch.n_shared_experts + moe_mlp = _lora_mlp_elements( + hd, + _get_mlp_size(arch), + r, + target_modules, + moe_expert_mult, + ) + dense_mlp = _lora_mlp_elements( + hd, + arch.intermediate_size, + r, + target_modules, + 1, + ) + mlp_total = moe_mlp * n_moe + dense_mlp * n_dense + else: + mlp_total = ( + _lora_mlp_elements( + hd, + arch.intermediate_size, + r, + target_modules, + 1, + ) + * n_layers + ) + + return attn_total + mlp_total + + +def compute_lora_adapter_bytes(lora_params: int) -> int: + return lora_params * 2 + + +def compute_optimizer_bytes(trainable_params: int, optimizer: str) -> int: + optimizer_key = optimizer.lower().replace("-", "_") + bytes_per_param = OPTIMIZER_BYTES_PER_PARAM.get(optimizer_key, 4) + return trainable_params * bytes_per_param + + +def compute_gradient_bytes(trainable_params: int) -> int: + return trainable_params * 2 + + +def compute_activation_bytes( + arch: ModelArchConfig, + batch_size: int, + seq_len: int, + gradient_checkpointing: str, + is_lora: bool = False, +) -> int: + hd = arch.hidden_size + kv_size = _get_kv_size(arch) + mlp_size = _get_mlp_size(arch) + bsz = batch_size + n_layers = arch.num_hidden_layers + + activation_qkv = seq_len * bsz * (hd + kv_size + kv_size) + residual_memory = (seq_len * bsz) * 2 + activation_mlp = seq_len * bsz * (mlp_size + mlp_size) + + per_layer_bytes = (activation_qkv + residual_memory + activation_mlp) * 2 + per_layer_bytes = int(per_layer_bytes * 1.25) + + gc_key = gradient_checkpointing.lower() + gc_entry = GC_LAYER_MULTIPLIERS.get(gc_key, (None, None)) + full_ft_mult, lora_mult = gc_entry + gc_multiplier = lora_mult if is_lora else full_ft_mult + + if gc_multiplier is None: + effective_layers = n_layers + else: + effective_layers = gc_multiplier + + return int(per_layer_bytes * effective_layers) + + +def estimate_training_vram( + arch: ModelArchConfig, + config: TrainingVramConfig, +) -> VramBreakdown: + method = config.training_method.lower() + is_lora = method in ("qlora", "lora") + load_in_4bit = config.load_in_4bit or method == "qlora" + + model_weights = compute_model_weights_bytes(arch, method, load_in_4bit) + + lora_params = 0 + lora_adapter_bytes = 0 + if is_lora: + lora_params = compute_lora_params( + arch, + config.lora_rank, + config.target_modules, + ) + lora_adapter_bytes = compute_lora_adapter_bytes(lora_params) + + trainable_params = lora_params if is_lora else compute_total_params(arch) + optimizer_bytes = compute_optimizer_bytes(trainable_params, config.optimizer) + gradient_bytes = max( + compute_gradient_bytes(trainable_params), + int(model_weights * 0.15), + ) + activations_computed = compute_activation_bytes( + arch, + config.batch_size, + config.max_seq_length, + config.gradient_checkpointing, + is_lora = is_lora, + ) + activation_bytes = max( + activations_computed, + int(model_weights * 0.15 * (config.batch_size / 2)), + ) + + return VramBreakdown( + model_weights = model_weights, + lora_adapters = lora_adapter_bytes, + optimizer_states = optimizer_bytes, + gradients = gradient_bytes, + activations = activation_bytes, + cuda_overhead = CUDA_OVERHEAD_BYTES, + activations_computed = activations_computed, + )