This commit is contained in:
Disty0
2025-06-25 23:37:41 +03:00
parent e43d1d2ba7
commit c9f49720c5
3 changed files with 4 additions and 4 deletions
+1 -1
View File
@@ -665,6 +665,6 @@ def normalize_device(dev):
def same_device(d1, d2):
if d1.type != d2.type:
if torch.device(d1).type != torch.device(d2).type:
return False
return normalize_device(d1) == normalize_device(d2)
+2 -2
View File
@@ -36,7 +36,7 @@ def hijack_set_module_tensor(
# note: majority of time is spent on .to(old_value.dtype)
if tensor_name in module._buffers: # pylint: disable=protected-access
module._buffers[tensor_name] = value.to(device, old_value.dtype) # pylint: disable=protected-access
elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access
elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device): # pylint: disable=protected-access
param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype) # pylint: disable=protected-access
t1 = time.time()
@@ -64,7 +64,7 @@ def hijack_set_module_tensor_simple(
with devices.inference_context():
if tensor_name in module._buffers: # pylint: disable=protected-access
module._buffers[tensor_name] = value.to(device) # pylint: disable=protected-access
elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access
elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device): # pylint: disable=protected-access
param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device) # pylint: disable=protected-access
t1 = time.time()
+1 -1
View File
@@ -171,7 +171,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
return module
def pre_forward(self, module, *args, **kwargs):
if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
if not devices.same_device(module.device, devices.device):
device_index = torch.device(devices.device).index
if device_index is None:
device_index = 0