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, + )