Fix NNCF Lora with model offload

This commit is contained in:
Disty0
2025-04-23 17:13:08 +03:00
parent d9aa9674cd
commit 75d169bc1c
+4 -1
View File
@@ -80,7 +80,10 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
try:
t0 = time.time()
if self.__class__.__name__.startswith('NNCF') and hasattr(self, "pre_ops") and len(self.pre_ops) == 1:
return_device = self.weight.data.device
self.weight.data = self.weight.data.to(devices.device) # pre_ops are always in devices.device
weight = self.pre_ops["0"](self, return_decompressed_only=True).to(devices.device)
self.weight.data = self.weight.data.to(return_device)
else:
weight = self.weight.to(devices.device) # must perform calc on gpu due to performance
updown, ex_bias = module.calc_updown(weight)
@@ -215,7 +218,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
else:
self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
if hasattr(self, "nncf_decompressor_backup"):
self.pre_ops["0"] = self.nncf_decompressor_backup.to(device)
self.pre_ops["0"] = self.nncf_decompressor_backup.to(devices.device) # model.to doesn't move pre_ops, send them here
if bias_backup is not None:
self.bias = None