fix: multi-GPU inference crash for bnb 4-bit/8-bit models (#5068)
* fix: multi-GPU inference crash for bnb 4-bit/8-bit models When load_in_4bit or load_in_8bit is used with device_map="sequential" and max_memory constraints that place weights across multiple GPUs (or entirely on a non-default GPU like cuda:1), the bitsandbytes loading path in transformers never calls dispatch_model. No AlignDevicesHook is installed, and the first forward/generate call crashes with: RuntimeError: Expected all tensors to be on the same device This adds _attach_bnb_multidevice_hooks() which is called after from_pretrained returns. It infers a device map from actual parameter placements and calls dispatch_model(force_hooks=True) to install the missing hooks. The function is a complete no-op for the common single-GPU cuda:0 case. Call sites: FastBaseModel.from_pretrained (vision.py) and FastLlamaModel.from_pretrained (llama.py). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: align with PR #5053 final review improvements - Add hook call to the bnb quantized loading branch in llama.py (the primary load_in_4bit path), not just the non-fast-inference fallback - Expand bnb detection: also check model.is_loaded_in_4bit, model.is_loaded_in_8bit, model.quantization_method - Pass explicit main_device and skip_keys to dispatch_model - Use logger.info instead of print for the success message - Use kwargs.get("load_in_8bit", False) at llama.py call sites * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
ee86530e55
commit
d56f980452
2 changed files with 167 additions and 0 deletions
|
|
@ -2478,6 +2478,16 @@ class FastLlamaModel:
|
|||
and not _head.weight.is_floating_point()
|
||||
):
|
||||
_head.to(dtype)
|
||||
# Attach dispatch hooks for bnb multi-device loads.
|
||||
from unsloth.models.vision import _attach_bnb_multidevice_hooks
|
||||
|
||||
_attach_bnb_multidevice_hooks(
|
||||
model,
|
||||
load_in_4bit = load_in_4bit,
|
||||
load_in_8bit = kwargs.get("load_in_8bit", False),
|
||||
offload_embedding = False,
|
||||
fast_inference = fast_inference,
|
||||
)
|
||||
elif not fast_inference:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_name,
|
||||
|
|
@ -2490,6 +2500,16 @@ class FastLlamaModel:
|
|||
attn_implementation = preferred_attn_impl,
|
||||
**kwargs,
|
||||
)
|
||||
# Attach dispatch hooks for bnb multi-device loads.
|
||||
from unsloth.models.vision import _attach_bnb_multidevice_hooks
|
||||
|
||||
_attach_bnb_multidevice_hooks(
|
||||
model,
|
||||
load_in_4bit = load_in_4bit,
|
||||
load_in_8bit = kwargs.get("load_in_8bit", False),
|
||||
offload_embedding = False,
|
||||
fast_inference = False,
|
||||
)
|
||||
model.fast_generate = make_fast_generate_wrapper(model.generate)
|
||||
model.fast_generate_batches = None
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -75,6 +75,7 @@ import functools
|
|||
import os
|
||||
import gc
|
||||
import math
|
||||
import warnings
|
||||
from typing import Optional, Tuple, List, Union
|
||||
import re, inspect, sys
|
||||
import contextlib
|
||||
|
|
@ -97,6 +98,144 @@ __all__ = [
|
|||
"FastBaseModel",
|
||||
]
|
||||
|
||||
|
||||
def _infer_device_map_from_loaded_model(model):
|
||||
"""Build a compact device_map by inspecting actual parameter placements."""
|
||||
device_map = {}
|
||||
|
||||
def _assign(module, prefix):
|
||||
params = list(module.named_parameters(remove_duplicate = False))
|
||||
if not params:
|
||||
bufs = list(module.named_buffers())
|
||||
if bufs:
|
||||
device_map[prefix] = bufs[0][1].device
|
||||
return
|
||||
devices = {p.device for _, p in params}
|
||||
if len(devices) == 1:
|
||||
device_map[prefix] = next(iter(devices))
|
||||
else:
|
||||
for child_name, child in module.named_children():
|
||||
child_prefix = f"{prefix}.{child_name}" if prefix else child_name
|
||||
_assign(child, child_prefix)
|
||||
for pname, param in module.named_parameters(remove_duplicate = False):
|
||||
if "." not in pname:
|
||||
full = f"{prefix}.{pname}" if prefix else pname
|
||||
if not any(
|
||||
full == k or full.startswith(k + ".") for k in device_map
|
||||
):
|
||||
device_map[full] = param.device
|
||||
|
||||
_assign(model, "")
|
||||
if "" in device_map and len(device_map) > 1:
|
||||
device_map.pop("")
|
||||
return device_map
|
||||
|
||||
|
||||
def _attach_bnb_multidevice_hooks(
|
||||
model, load_in_4bit, load_in_8bit, offload_embedding, fast_inference
|
||||
):
|
||||
"""
|
||||
Attach accelerate AlignDevicesHook on a bnb model loaded across multiple
|
||||
devices (or a non-default device). No-op for single-GPU cuda:0, non-bnb,
|
||||
vLLM, or already-dispatched models.
|
||||
"""
|
||||
if fast_inference:
|
||||
return
|
||||
is_bnb = (
|
||||
load_in_4bit
|
||||
or load_in_8bit
|
||||
or getattr(model, "is_loaded_in_4bit", False)
|
||||
or getattr(model, "is_loaded_in_8bit", False)
|
||||
or getattr(model, "quantization_method", None) == "bitsandbytes"
|
||||
)
|
||||
if not is_bnb:
|
||||
return
|
||||
if offload_embedding:
|
||||
return
|
||||
if getattr(model, "hf_device_map", None) is not None:
|
||||
return # already dispatched
|
||||
|
||||
try:
|
||||
all_devs = {p.device for p in model.parameters()}
|
||||
except Exception as exc:
|
||||
warnings.warn(
|
||||
"Unsloth: Failed to determine device placement from model parameters, "
|
||||
f"so multi-GPU hooks cannot be attached. ({type(exc).__name__}: {exc})",
|
||||
RuntimeWarning,
|
||||
stacklevel = 2,
|
||||
)
|
||||
return
|
||||
|
||||
cuda_devs = {d for d in all_devs if d.type == "cuda"}
|
||||
if not cuda_devs:
|
||||
return
|
||||
|
||||
default_cuda = torch.device("cuda", 0)
|
||||
if all_devs == {default_cuda}:
|
||||
return
|
||||
|
||||
try:
|
||||
from accelerate import dispatch_model
|
||||
except ImportError:
|
||||
return # accelerate not available
|
||||
|
||||
try:
|
||||
inferred_map = _infer_device_map_from_loaded_model(model)
|
||||
if not inferred_map:
|
||||
return
|
||||
|
||||
# bnb constructors reject _is_hf_initialized; strip before dispatch.
|
||||
_extra_keys = ("_is_hf_initialized",)
|
||||
_stripped = []
|
||||
for _, param in model.named_parameters():
|
||||
for key in _extra_keys:
|
||||
if key in param.__dict__:
|
||||
_stripped.append((param, key, param.__dict__.pop(key)))
|
||||
|
||||
try:
|
||||
# CUDA -> int index, non-CUDA -> type string ("cpu", "meta").
|
||||
device_map_int = {
|
||||
k: (v.index if v.type == "cuda" else v.type)
|
||||
if isinstance(v, torch.device)
|
||||
else v
|
||||
for k, v in inferred_map.items()
|
||||
}
|
||||
|
||||
# force_hooks=True: install hooks even for single-device maps.
|
||||
main_device = device_map_int.get("")
|
||||
if main_device in (None, "cpu", "disk"):
|
||||
main_device = next(
|
||||
(d for d in device_map_int.values() if d not in ("cpu", "disk")),
|
||||
None,
|
||||
)
|
||||
dispatch_model(
|
||||
model,
|
||||
device_map = device_map_int,
|
||||
main_device = main_device,
|
||||
skip_keys = getattr(model, "_skip_keys_device_placement", None),
|
||||
force_hooks = True,
|
||||
)
|
||||
desc = f"{len(inferred_map)} block(s) across {len(cuda_devs)} device(s)"
|
||||
finally:
|
||||
# Restore stripped keys.
|
||||
for param, key, val in _stripped:
|
||||
param.__dict__[key] = val
|
||||
|
||||
logger.info(
|
||||
f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) "
|
||||
f"for bnb multi-GPU inference."
|
||||
)
|
||||
except Exception as exc:
|
||||
warnings.warn(
|
||||
f"Unsloth: Could not attach multi-device dispatch hooks automatically "
|
||||
f"({type(exc).__name__}: {exc}). "
|
||||
"Cross-device inference may fail. Consider using a single GPU or "
|
||||
"calling accelerate.dispatch_model() manually.",
|
||||
RuntimeWarning,
|
||||
stacklevel = 2,
|
||||
)
|
||||
|
||||
|
||||
global NUM_LOGITS_TO_KEEP
|
||||
NUM_LOGITS_TO_KEEP = dict()
|
||||
|
||||
|
|
@ -796,6 +935,14 @@ class FastBaseModel:
|
|||
# attn_implementation = attn_implementation,
|
||||
**kwargs,
|
||||
)
|
||||
# Attach dispatch hooks for bnb multi-device loads.
|
||||
_attach_bnb_multidevice_hooks(
|
||||
model,
|
||||
load_in_4bit = load_in_4bit,
|
||||
load_in_8bit = load_in_8bit,
|
||||
offload_embedding = offload_embedding,
|
||||
fast_inference = fast_inference,
|
||||
)
|
||||
if hasattr(model, "generate"):
|
||||
model.fast_generate = make_fast_generate_wrapper(model.generate)
|
||||
model.fast_generate_batches = error_out_no_vllm
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue