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

[Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602)

Browse source 
* [WIP] balanced device map for studio

* gpus as a request parameter

* API for multi GPU stuff

* return multi gpu util in new API

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

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

* Use balanced_low0 instead of balanced

* Use balanced_low0 instead of balanced

* Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests

- balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string)
- get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES
  gracefully instead of crashing on int() parse
- _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so
  CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs
- test_gpu_selection.py: add missing get_visible_gpu_utilization import and
  add required job_id arg to start_training() calls

* Smart GPU determinism using estimates

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

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

* disallow gpu selection for gguf for now

* cleanup

* Slightly larger baseline

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

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

* Treat empty list as auto

* Verbose logging/debug

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

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

* Cleanup and revert unnecessary deletions

* Cleanup excessive logs and guard against disk/cpu offload

* auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate

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

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

* support for non cuda gpus

* Fix multi-GPU auto-selection memory accounting

The multi_gpu_factor was applied uniformly to all GPUs including the
first one, which unfairly penalizes single-GPU capacity when
transitioning to multi-GPU. This created a discontinuity where a model
that barely fits 1 GPU would suddenly require 2 GPUs because the first
GPU's free memory was discounted by 20%.

Now the first GPU keeps its full free memory, and only additional GPUs
have an overhead factor (0.85) applied to account for inter-GPU
communication and sharding overhead. This gives more accurate
auto-selection and avoids unnecessary multi-GPU for models that
comfortably fit on one device.

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

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

* Add sandbox tests for multi-GPU selection logic

24 tests covering model size estimation, memory requirements, automatic
GPU selection, device map generation, GPU ID validation, and multi-GPU
overhead accounting. All tests use mocks so they run without GPUs on
Linux, macOS, and Windows.

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

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

* Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry

1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer)
   instead of the FP16 1.3x multiplier. This prevents over-sharding
   quantized models across unnecessary GPUs.

2. When model size estimation fails, auto_select_gpu_ids now falls back to
   all visible GPUs instead of returning None (which could default to
   single-GPU loading for an unknown-size model).

3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as
   None) instead of rejecting it as an unsupported explicit request.

4. Training retry path for "could not get source code" now preserves the
   gpu_ids parameter so the retry lands on the same GPUs.

5. Updated sandbox tests to cover the new 4-bit inference estimate branch.

* Remove accidentally added unsloth-zoo submodule

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

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

* Fix UUID/MIG visibility and update test expectations

1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the
   visibility APIs now return "unresolved" with empty device lists instead
   of exposing all physical GPUs. This prevents the UI from showing GPUs
   that the backend process cannot actually use.

2. test_gpu_selection.py: Updated test expectations to match the new
   multi-GPU overhead accounting (first GPU at full capacity, 0.85x for
   additional GPUs) and 4-bit inference memory estimation formula.
   All 60 tests now pass.

* Add CPU/disk offload guard to audio inference path

The audio model loading branch returned before the common
get_offloaded_device_map_entries() check, so audio models loaded with a
multi-GPU device_map that spilled layers to CPU/disk would be accepted
instead of rejected. Now audio loads also verify no modules are offloaded.

* Improve VRAM requirement estimates

* Replace balanced_low_0 with balanced

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

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

* refine calculations for slightly easier nums

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

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

* adjust estimates

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

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

* Use nums instead of obj to avoid seralisation error

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

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

* Harden nvidia-smi parsing and fix fallback GPU list

1. nvidia.py: Wrap int() casts for GPU index and memory in try/except
   so MIG slices, N/A values, or unexpected nvidia-smi output skip the
   unparseable row instead of aborting the entire GPU list.

2. nvidia.py: Handle GPU names containing commas by using the last
   field as memory instead of a fixed positional index.

3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified
   VRAM data) instead of raw devices list, which could include GPUs
   with null VRAM that were excluded from the ranking.

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

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

* cleanup

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

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

* consolidate raise_if_offload

* Improve MoE support. Guard against nvidia-smi failures

* Improve MoE support. Guard against nvidia-smi failures

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

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

* Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case

1. vram_estimation.py: compute_lora_params now includes shared experts
   (n_shared_experts) alongside routed experts when computing MoE LoRA
   adapter parameters. Previously only n_experts were counted, causing
   the estimator to undercount adapter, optimizer, and gradient memory
   for DeepSeek/GLM-style models with shared experts.

2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which
   reports system-wide VRAM usage) instead of memory_allocated (which
   only reports this process's PyTorch allocations). This prevents
   auto-selection from treating a GPU as mostly free when another
   process is consuming VRAM. Falls back to memory_allocated when
   mem_get_info is unavailable.

3. hardware.py: apply_gpu_ids([]) now returns early instead of setting
   CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty
   list inherits the parent visibility, same as None.

4. hardware.py: Upgraded fallback_all GPU selection log from debug to
   warning so operators are notified when the model likely will not fit
   in available VRAM.

* Guard nvidia-smi subprocess calls against OSError and TimeoutExpired

get_visible_gpu_utilization and get_backend_visible_gpu_info now catch
OSError (nvidia-smi not found) and TimeoutExpired internally instead
of relying on callers to wrap every invocation. Returns the standard
available=False sentinel on failure so the torch-based fallback in
hardware.py can take over.

* Guard get_primary_gpu_utilization and reset GPU caches between tests

1. nvidia.py: get_primary_gpu_utilization now catches OSError and
   TimeoutExpired internally, matching the pattern already used in
   get_visible_gpu_utilization and get_backend_visible_gpu_info. All
   three nvidia-smi callers are now self-contained.

2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the
   module-level _physical_gpu_count and _visible_gpu_count caches in
   tearDown. Applied to all test classes that exercise GPU selection,
   device map, or visibility functions. This prevents stale cache
   values from leaking between tests and causing flaky results on
   machines with real GPUs.

* Fix nvidia-smi fallback regression and physical GPU count validation

1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and
   get_backend_visible_gpu_info now check result.get("available") before
   returning the nvidia-smi result. When nvidia-smi is unavailable or
   returns no data (e.g., containers without nvidia-smi, UUID/MIG masks),
   the functions fall through to the torch-based fallback instead of
   returning an empty result. This fixes a regression where the internal
   exception handling in nvidia.py prevented the caller's except block
   from triggering the fallback.

2. hardware.py: resolve_requested_gpu_ids now separates negative-ID
   validation from physical upper-bound validation. The physical count
   check is only enforced when it is plausibly a true physical count
   (i.e., higher than the largest parent-visible ID), since
   torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the
   visible count, not the physical total. The parent-visible-set check
   remains authoritative in all cases. This prevents valid physical IDs
   like [2, 3] from being rejected as "out of range" when nvidia-smi is
   unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only
   2 devices.

* Fix UUID/MIG torch fallback to enumerate devices by ordinal

When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers,
get_parent_visible_gpu_ids() returns [] because the tokens are
non-numeric. The torch fallback in get_visible_gpu_utilization() and
get_backend_visible_gpu_info() previously passed that empty list to
_torch_get_per_device_info(), getting nothing back.

Now both functions detect the empty-list case and fall back to
enumerating torch-visible ordinals (0..device_count-1) with
index_kind="relative". This means the UI and auto-selection still
see real device data in Kubernetes, MIG, and Slurm-style UUID
environments where nvidia-smi output cannot be mapped to physical
indices.

Updated test_uuid_parent_visibility to verify the new torch fallback
path returns available=True with relative ordinals.

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

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

* Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Datta Nimmaturi 2026-03-30 15:03:15 +05:30 • committed by GitHub
commit 9311df2b29
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 4534 additions and 220 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

23
studio/backend/core/inference/llama_cpp.py
View file

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

11
studio/backend/core/inference/orchestrator.py
View file

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

4
studio/backend/core/inference/worker.py
View file

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

26
studio/backend/core/training/trainer.py
View file

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

18
studio/backend/core/training/training.py
View file

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

4
studio/backend/core/training/worker.py
View file

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

81
studio/backend/main.py
View file

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

4
studio/backend/models/inference.py
View file

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

6
studio/backend/models/training.py
View file

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

13
studio/backend/routes/inference.py
View file

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

16
studio/backend/routes/training.py
View file

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

1125
studio/backend/tests/test_gpu_selection.py Normal file
View file

File diff suppressed because it is too large Load diff

544
studio/backend/tests/test_gpu_selection_sandbox.py Normal file
View file

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

32
studio/backend/tests/test_utils.py
View file

@ -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() ==========

695
studio/backend/tests/test_vram_estimation.py Normal file
View file

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

161
studio/backend/utils/hardware/VRAM_ESTIMATION.md Normal file
View file

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

37
studio/backend/utils/hardware/__init__.py
View file

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

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

File diff suppressed because it is too large Load diff

279
studio/backend/utils/hardware/nvidia.py Normal file
View file

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

501
studio/backend/utils/hardware/vram_estimation.py Normal file
View file

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