diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index b54ceaf842..463f80003a 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -38,6 +38,7 @@ from .loader_utils import ( _tag_model_with_fp8_torchao_config, get_model_name, prepare_device_map, + resolve_optimal_device_map, ) import os, contextlib, sys @@ -288,6 +289,16 @@ class FastLanguageModel(FastLlamaModel): if isinstance(bnb_compute_dtype, torch.dtype): dtype = bnb_compute_dtype + # Resolve device_map="optimal" -- auto-detect best GPU allocation + if isinstance(device_map, str) and device_map == "optimal": + device_map, cuda_override = resolve_optimal_device_map( + model_name, + load_in_4bit = load_in_4bit, + token = token, + ) + if cuda_override is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_override + # Distributed-safe device placement for quantized models. # In multi-GPU (torchrun), each rank must load the model on its own device # to avoid Accelerate device relocation errors with quantized weights. @@ -952,6 +963,16 @@ class FastModel(FastBaseModel): if qat_scheme == "phone-deployment": qat_scheme = "int8-int4" + # Resolve device_map="optimal" -- auto-detect best GPU allocation + if isinstance(device_map, str) and device_map == "optimal": + device_map, cuda_override = resolve_optimal_device_map( + model_name, + load_in_4bit = load_in_4bit, + token = token, + ) + if cuda_override is not None: + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_override + # Distributed-safe device placement for quantized models. # In multi-GPU (torchrun), each rank must load the model on its own device # to avoid Accelerate device relocation errors with quantized weights. diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index cf5af983a6..bf1cf5a51c 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -99,6 +99,214 @@ def prepare_device_map(): return device_map, True +def _get_gpu_free_memory(): + """Return a list of (gpu_index, free_memory_bytes) for all visible GPUs. + + Uses torch.cuda if available. Falls back to returning an empty list + on non-CUDA platforms or import failures. + """ + if DEVICE_TYPE_TORCH != "cuda": + return [] + try: + count = torch.cuda.device_count() + result = [] + for i in range(count): + free, total = torch.cuda.mem_get_info(i) + result.append((i, free)) + return result + except Exception: + return [] + + +def _estimate_model_bytes(model_name, load_in_4bit = False, token = None): + """Estimate model weight size in bytes using HF safetensors metadata or config. + + Returns the estimated size or None if unavailable. For 4bit models + (bnb-4bit in name or load_in_4bit=True), we estimate the compressed size. + """ + try: + from huggingface_hub import model_info as hf_model_info + + info = hf_model_info(model_name, token = token) + safetensors = getattr(info, "safetensors", None) + if isinstance(safetensors, dict): + total_params = safetensors.get("total") + if total_params: + total_params = int(total_params) + # fp16 = 2 bytes per param + fp16_bytes = total_params * 2 + if load_in_4bit or "-bnb-4bit" in model_name.lower(): + # ~5 bits per param effective (4bit + scales) + return int(fp16_bytes * 5 / 16) + return fp16_bytes + except Exception: + pass + + # Fallback: try to estimate from config + try: + from transformers import AutoConfig + + config = AutoConfig.from_pretrained( + model_name, + token = token, + trust_remote_code = model_name.lower().startswith("unsloth/"), + ) + text_config = getattr(config, "text_config", None) or config + + vocab_size = getattr(text_config, "vocab_size", None) + hidden_size = getattr(text_config, "hidden_size", None) + intermediate_size = getattr(text_config, "intermediate_size", None) + num_layers = getattr(text_config, "num_hidden_layers", None) + num_heads = getattr(text_config, "num_attention_heads", 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 (vocab_size, hidden_size, intermediate_size, num_layers, num_heads) + ): + return None + if num_heads <= 0: + return None + + num_kv_heads = getattr(text_config, "num_key_value_heads", num_heads) + kv_size = (hidden_size // num_heads) * num_kv_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) + if moe_intermediate is not None: + intermediate_size = moe_intermediate + + qkvo = (hidden_size + kv_size + kv_size + hidden_size) * hidden_size + if num_experts and num_experts > 1: + mlp = (hidden_size * intermediate_size) * 3 * num_experts + mlp += num_experts * hidden_size # router + else: + mlp = (hidden_size * intermediate_size) * 3 + embed_tokens = vocab_size * hidden_size + lm_head = ( + 0 + if getattr(text_config, "tie_word_embeddings", True) + else vocab_size * hidden_size + ) + + total_elements = (qkvo + mlp) * num_layers + embed_tokens + lm_head + fp16_bytes = int(total_elements * 2) + + if load_in_4bit or "-bnb-4bit" in model_name.lower(): + return int(fp16_bytes * 5 / 16) + return fp16_bytes + except Exception: + pass + return None + + +def resolve_optimal_device_map( + model_name, + load_in_4bit = False, + token = None, + overhead_factor = 1.5, +): + """Resolve device_map='optimal' to a concrete device map string. + + Logic: + 1. If only 1 GPU visible, return 'sequential'. + 2. Estimate model weight size in bytes. + 3. If model * overhead_factor fits on the single GPU with most free + memory, restrict to that GPU and return 'sequential'. + 4. If it needs 2+ GPUs, find the minimum set and return + 'balanced' (4bit) or 'balanced_low_0' (non-4bit). + 5. If it cannot fit at all, raise ValueError. + + Falls back to 'sequential' if estimation fails. + + Returns: + tuple: (device_map_str, cuda_visible_override) + device_map_str: 'sequential', 'balanced', or 'balanced_low_0' + cuda_visible_override: None or a CUDA_VISIBLE_DEVICES string + to restrict GPU visibility. Caller should set env var. + """ + gpus = _get_gpu_free_memory() + if len(gpus) <= 1: + return "sequential", None + + model_bytes = _estimate_model_bytes( + model_name, load_in_4bit = load_in_4bit, token = token + ) + if model_bytes is None: + # Cannot estimate -- fall back to sequential and let accelerate handle it + print( + f"Unsloth: Could not estimate model size for '{model_name}'. " + f"Using device_map='sequential' with all {len(gpus)} GPUs." + ) + return "sequential", None + + required_bytes = int(model_bytes * overhead_factor) + required_gb = required_bytes / (1024**3) + model_gb = model_bytes / (1024**3) + + # Sort GPUs by free memory descending + ranked = sorted(gpus, key = lambda x: -x[1]) + + # Check if single GPU suffices + best_gpu_idx, best_free = ranked[0] + if best_free >= required_bytes: + print( + f"Unsloth: Model ~{model_gb:.1f}GB fits on GPU {best_gpu_idx} " + f"({best_free / (1024**3):.1f}GB free). Using single GPU." + ) + return "sequential", str(best_gpu_idx) + + # Find minimum number of GPUs needed + multi_gpu_overhead = 0.85 # Each additional GPU has sharding overhead + accumulated = 0 + selected_indices = [] + for gpu_idx, free_mem in ranked: + selected_indices.append(gpu_idx) + if len(selected_indices) == 1: + accumulated = free_mem + else: + accumulated = ranked[0][1] + sum( + mem * multi_gpu_overhead for _, mem in ranked[1 : len(selected_indices)] + ) + if accumulated >= required_bytes: + break + + if accumulated < required_bytes: + total_available = sum(free for _, free in ranked) + raise ValueError( + f"Model '{model_name}' requires ~{required_gb:.1f}GB but only " + f"{total_available / (1024**3):.1f}GB total free across " + f"{len(gpus)} GPUs. Cannot fit the model." + ) + + # Build CUDA_VISIBLE_DEVICES override + selected_str = ",".join(str(idx) for idx in selected_indices) + n_gpus = len(selected_indices) + + if n_gpus == 1: + device_map = "sequential" + elif load_in_4bit or "-bnb-4bit" in model_name.lower(): + device_map = "balanced" + else: + device_map = "balanced_low_0" + + print( + f"Unsloth: Model ~{model_gb:.1f}GB needs {n_gpus} GPU(s) " + f"({accumulated / (1024**3):.1f}GB usable). " + f"Using device_map='{device_map}' on GPUs [{selected_str}]." + ) + return device_map, selected_str + + def __get_model_name( model_name, load_in_4bit = True,