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:
Daniel Han 2026-04-16 14:41:44 +00:00
commit d6b758035a
2 changed files with 21 additions and 23 deletions

View file

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

View file

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