fix: address review findings -- force_hooks, CPU device conversion, llama.py path
Three fixes from 13-reviewer consensus:
1. Use dispatch_model(force_hooks=True) instead of manual AlignDevicesHook
fallback. This properly handles single-device non-default GPU maps
(e.g. all weights on cuda:1) with io_same_device routing, and
eliminates the need for the separate add_hook_to_module call.
2. Convert non-CUDA devices (cpu, meta) to type strings in device_map_int.
dispatch_model uses string equality (device == "cpu") internally, so
torch.device("cpu") must become "cpu" not stay as an object.
3. Add _attach_bnb_multidevice_hooks call to FastLlamaModel.from_pretrained
in llama.py. The default 4-bit FastLanguageModel path routes through
llama.py:2482, not FastBaseModel, so the hook was never called for the
most common use case (Llama, Gemma, Mistral, etc.).
Also removes redundant retie_parameters call (dispatch_model already
calls it internally) and unused imports.
This commit is contained in:
parent
42b29a35ac
commit
d6b758035a
2 changed files with 21 additions and 23 deletions
|
|
@ -2490,6 +2490,17 @@ class FastLlamaModel:
|
|||
attn_implementation = preferred_attn_impl,
|
||||
**kwargs,
|
||||
)
|
||||
# Attach AlignDevicesHook for bnb multi-device / non-default-device
|
||||
# loads. The bnb loading path places weights on CUDA devices but
|
||||
# never calls dispatch_model, so no hooks are installed.
|
||||
from unsloth.models.vision import _attach_bnb_multidevice_hooks
|
||||
_attach_bnb_multidevice_hooks(
|
||||
model,
|
||||
load_in_4bit = load_in_4bit,
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -209,8 +209,6 @@ def _attach_bnb_multidevice_hooks(
|
|||
|
||||
try:
|
||||
from accelerate import dispatch_model
|
||||
from accelerate.hooks import AlignDevicesHook, add_hook_to_module
|
||||
from accelerate.utils import find_tied_parameters, retie_parameters
|
||||
except ImportError:
|
||||
return # accelerate not available
|
||||
|
||||
|
|
@ -219,9 +217,6 @@ def _attach_bnb_multidevice_hooks(
|
|||
if not inferred_map:
|
||||
return
|
||||
|
||||
# Preserve tied-parameter references before dispatching.
|
||||
tied_params = find_tied_parameters(model)
|
||||
|
||||
# accelerate's set_module_tensor_to_device grabs param.__dict__ and
|
||||
# passes it as **kwargs to the parameter class constructor. HF
|
||||
# Transformers adds _is_hf_initialized to parameter __dict__ which
|
||||
|
|
@ -235,36 +230,28 @@ def _attach_bnb_multidevice_hooks(
|
|||
_stripped.append((param, key, param.__dict__.pop(key)))
|
||||
|
||||
try:
|
||||
# Convert device_map values from torch.device to int (HF convention)
|
||||
# Convert device_map values to the format dispatch_model expects:
|
||||
# CUDA devices -> int index, non-CUDA devices -> type string.
|
||||
# dispatch_model uses string equality (device == "cpu") internally,
|
||||
# so torch.device("cpu") must become "cpu", not stay as an object.
|
||||
device_map_int = {
|
||||
k: v.index if isinstance(v, torch.device) and v.type == "cuda" else v
|
||||
k: (v.index if v.type == "cuda" else v.type)
|
||||
if isinstance(v, torch.device) else v
|
||||
for k, v in inferred_map.items()
|
||||
}
|
||||
|
||||
# dispatch_model installs AlignDevicesHook on every module and
|
||||
# sub-module, exactly matching HF's own loading behavior.
|
||||
dispatch_model(model, device_map = device_map_int)
|
||||
|
||||
# dispatch_model skips hook installation for single-device maps
|
||||
# (accelerate fast-path). When all weights sit on a non-default
|
||||
# GPU (e.g. cuda:1) the caller's inputs may still arrive on
|
||||
# cuda:0, causing a device mismatch. Detect this and manually
|
||||
# add a root-level AlignDevicesHook to route inputs.
|
||||
if len(cuda_devs) == 1 and not hasattr(model, "_hf_hook"):
|
||||
main_device = next(iter(cuda_devs))
|
||||
add_hook_to_module(
|
||||
model, AlignDevicesHook(execution_device = main_device)
|
||||
)
|
||||
|
||||
# force_hooks=True ensures hooks are installed even for single-
|
||||
# device maps (e.g. all weights on cuda:1), where the default
|
||||
# fast-path would skip hook installation entirely.
|
||||
dispatch_model(model, device_map = device_map_int, 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
|
||||
|
||||
# Retie parameters that hook installation may have unlinked.
|
||||
retie_parameters(model, tied_params)
|
||||
|
||||
print(
|
||||
f"Unsloth: Attached accelerate AlignDevicesHook ({desc}) "
|
||||
f"for bnb multi-GPU inference."
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue