Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(backend): revert non-blocking device transfer

Browse Source 
In #6490 we enabled non-blocking torch device transfers throughout the model manager's memory management code. When using this torch feature, torch attempts to wait until the tensor transfer has completed before allowing any access to the tensor. Theoretically, that should make this a safe feature to use.

This provides a small performance improvement but causes race conditions in some situations. Specific platforms/systems are affected, and complicated data dependencies can make this unsafe.

- Intermittent black images on MPS devices - reported on discord and #6545, fixed with special handling in #6549.
- Intermittent OOMs and black images on a P4000 GPU on Windows - reported in #6613, fixed in this commit.

On my system, I haven't experience any issues with generation, but targeted testing of non-blocking ops did expose a race condition when moving tensors from CUDA to CPU.

One workaround is to use torch streams with manual sync points. Our application logic is complicated enough that this would be a lot of work and feels ripe for edge cases and missed spots.

Much safer is to fully revert non-locking - which is what this change does.
This commit is contained in:
psychedelicious 2024-07-16 07:05:29 +10:00
parent 5a0c99816c
commit 38343917f8
8 changed files with 43 additions and 115 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
invokeai/backend/ip_adapter/ip_adapter.py
View File

@ -124,16 +124,14 @@ class IPAdapter(RawModel):
self.device, dtype=self.dtype self.device, dtype=self.dtype
) )
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, non_blocking: bool = False
):
if device is not None: if device is not None:
self.device = device self.device = device
if dtype is not None: if dtype is not None:
self.dtype = dtype self.dtype = dtype
self._image_proj_model.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking) self._image_proj_model.to(device=self.device, dtype=self.dtype)
self.attn_weights.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking) self.attn_weights.to(device=self.device, dtype=self.dtype)
def calc_size(self) -> int: def calc_size(self) -> int:
# HACK(ryand): Fix this issue with circular imports. # HACK(ryand): Fix this issue with circular imports.

94
invokeai/backend/lora.py
View File

@ -11,7 +11,6 @@ from typing_extensions import Self
from invokeai.backend.model_manager import BaseModelType from invokeai.backend.model_manager import BaseModelType
from invokeai.backend.raw_model import RawModel from invokeai.backend.raw_model import RawModel
from invokeai.backend.util.devices import TorchDevice
class LoRALayerBase: class LoRALayerBase:
@ -57,14 +56,9 @@ class LoRALayerBase:
model_size += val.nelement() * val.element_size() model_size += val.nelement() * val.element_size()
return model_size return model_size
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
if self.bias is not None: if self.bias is not None:
self.bias = self.bias.to(device=device, dtype=dtype, non_blocking=non_blocking) self.bias = self.bias.to(device=device, dtype=dtype)
# TODO: find and debug lora/locon with bias # TODO: find and debug lora/locon with bias
@ -106,19 +100,14 @@ class LoRALayer(LoRALayerBase):
model_size += val.nelement() * val.element_size() model_size += val.nelement() * val.element_size()
return model_size return model_size
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self, super().to(device=device, dtype=dtype)
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
super().to(device=device, dtype=dtype, non_blocking=non_blocking)
self.up = self.up.to(device=device, dtype=dtype, non_blocking=non_blocking) self.up = self.up.to(device=device, dtype=dtype)
self.down = self.down.to(device=device, dtype=dtype, non_blocking=non_blocking) self.down = self.down.to(device=device, dtype=dtype)
if self.mid is not None: if self.mid is not None:
self.mid = self.mid.to(device=device, dtype=dtype, non_blocking=non_blocking) self.mid = self.mid.to(device=device, dtype=dtype)
class LoHALayer(LoRALayerBase): class LoHALayer(LoRALayerBase):
@ -167,23 +156,18 @@ class LoHALayer(LoRALayerBase):
model_size += val.nelement() * val.element_size() model_size += val.nelement() * val.element_size()
return model_size return model_size
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
super().to(device=device, dtype=dtype) super().to(device=device, dtype=dtype)
self.w1_a = self.w1_a.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w1_a = self.w1_a.to(device=device, dtype=dtype)
self.w1_b = self.w1_b.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w1_b = self.w1_b.to(device=device, dtype=dtype)
if self.t1 is not None: if self.t1 is not None:
self.t1 = self.t1.to(device=device, dtype=dtype, non_blocking=non_blocking) self.t1 = self.t1.to(device=device, dtype=dtype)
self.w2_a = self.w2_a.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w2_a = self.w2_a.to(device=device, dtype=dtype)
self.w2_b = self.w2_b.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w2_b = self.w2_b.to(device=device, dtype=dtype)
if self.t2 is not None: if self.t2 is not None:
self.t2 = self.t2.to(device=device, dtype=dtype, non_blocking=non_blocking) self.t2 = self.t2.to(device=device, dtype=dtype)
class LoKRLayer(LoRALayerBase): class LoKRLayer(LoRALayerBase):
@ -264,12 +248,7 @@ class LoKRLayer(LoRALayerBase):
model_size += val.nelement() * val.element_size() model_size += val.nelement() * val.element_size()
return model_size return model_size
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
super().to(device=device, dtype=dtype) super().to(device=device, dtype=dtype)
if self.w1 is not None: if self.w1 is not None:
@ -277,19 +256,19 @@ class LoKRLayer(LoRALayerBase):
else: else:
assert self.w1_a is not None assert self.w1_a is not None
assert self.w1_b is not None assert self.w1_b is not None
self.w1_a = self.w1_a.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w1_a = self.w1_a.to(device=device, dtype=dtype)
self.w1_b = self.w1_b.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w1_b = self.w1_b.to(device=device, dtype=dtype)
if self.w2 is not None: if self.w2 is not None:
self.w2 = self.w2.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w2 = self.w2.to(device=device, dtype=dtype)
else: else:
assert self.w2_a is not None assert self.w2_a is not None
assert self.w2_b is not None assert self.w2_b is not None
self.w2_a = self.w2_a.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w2_a = self.w2_a.to(device=device, dtype=dtype)
self.w2_b = self.w2_b.to(device=device, dtype=dtype, non_blocking=non_blocking) self.w2_b = self.w2_b.to(device=device, dtype=dtype)
if self.t2 is not None: if self.t2 is not None:
self.t2 = self.t2.to(device=device, dtype=dtype, non_blocking=non_blocking) self.t2 = self.t2.to(device=device, dtype=dtype)
class FullLayer(LoRALayerBase): class FullLayer(LoRALayerBase):
@ -319,15 +298,10 @@ class FullLayer(LoRALayerBase):
model_size += self.weight.nelement() * self.weight.element_size() model_size += self.weight.nelement() * self.weight.element_size()
return model_size return model_size
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
super().to(device=device, dtype=dtype) super().to(device=device, dtype=dtype)
self.weight = self.weight.to(device=device, dtype=dtype, non_blocking=non_blocking) self.weight = self.weight.to(device=device, dtype=dtype)
class IA3Layer(LoRALayerBase): class IA3Layer(LoRALayerBase):
@ -359,16 +333,11 @@ class IA3Layer(LoRALayerBase):
model_size += self.on_input.nelement() * self.on_input.element_size() model_size += self.on_input.nelement() * self.on_input.element_size()
return model_size return model_size
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None):
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
):
super().to(device=device, dtype=dtype) super().to(device=device, dtype=dtype)
self.weight = self.weight.to(device=device, dtype=dtype, non_blocking=non_blocking) self.weight = self.weight.to(device=device, dtype=dtype)
self.on_input = self.on_input.to(device=device, dtype=dtype, non_blocking=non_blocking) self.on_input = self.on_input.to(device=device, dtype=dtype)
AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer] AnyLoRALayer = Union[LoRALayer, LoHALayer, LoKRLayer, FullLayer, IA3Layer]
@ -390,15 +359,10 @@ class LoRAModelRaw(RawModel): # (torch.nn.Module):
def name(self) -> str: def name(self) -> str:
return self._name return self._name
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
# TODO: try revert if exception? # TODO: try revert if exception?
for _key, layer in self.layers.items(): for _key, layer in self.layers.items():
layer.to(device=device, dtype=dtype, non_blocking=non_blocking) layer.to(device=device, dtype=dtype)
def calc_size(self) -> int: def calc_size(self) -> int:
model_size = 0 model_size = 0
@ -521,7 +485,7 @@ class LoRAModelRaw(RawModel): # (torch.nn.Module):
# lower memory consumption by removing already parsed layer values # lower memory consumption by removing already parsed layer values
state_dict[layer_key].clear() state_dict[layer_key].clear()
layer.to(device=device, dtype=dtype, non_blocking=TorchDevice.get_non_blocking(device)) layer.to(device=device, dtype=dtype)
model.layers[layer_key] = layer model.layers[layer_key] = layer
return model return model

6
invokeai/backend/model_manager/load/model_cache/model_cache_default.py
View File

@ -289,11 +289,9 @@ class ModelCache(ModelCacheBase[AnyModel]):
else: else:
new_dict: Dict[str, torch.Tensor] = {} new_dict: Dict[str, torch.Tensor] = {}
for k, v in cache_entry.state_dict.items(): for k, v in cache_entry.state_dict.items():
new_dict[k] = v.to( new_dict[k] = v.to(target_device, copy=True)
target_device, copy=True, non_blocking=TorchDevice.get_non_blocking(target_device)
)
cache_entry.model.load_state_dict(new_dict, assign=True) cache_entry.model.load_state_dict(new_dict, assign=True)
cache_entry.model.to(target_device, non_blocking=TorchDevice.get_non_blocking(target_device)) cache_entry.model.to(target_device)
cache_entry.device = target_device cache_entry.device = target_device
except Exception as e: # blow away cache entry except Exception as e: # blow away cache entry
self._delete_cache_entry(cache_entry) self._delete_cache_entry(cache_entry)

15
invokeai/backend/model_patcher.py
View File

@ -139,15 +139,12 @@ class ModelPatcher:
# We intentionally move to the target device first, then cast. Experimentally, this was found to # We intentionally move to the target device first, then cast. Experimentally, this was found to
# be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the # be significantly faster for 16-bit CPU tensors being moved to a CUDA device than doing the
# same thing in a single call to '.to(...)'. # same thing in a single call to '.to(...)'.
layer.to(device=device, non_blocking=TorchDevice.get_non_blocking(device)) layer.to(device=device)
layer.to(dtype=torch.float32, non_blocking=TorchDevice.get_non_blocking(device)) layer.to(dtype=torch.float32)
# TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA # TODO(ryand): Using torch.autocast(...) over explicit casting may offer a speed benefit on CUDA
# devices here. Experimentally, it was found to be very slow on CPU. More investigation needed. # devices here. Experimentally, it was found to be very slow on CPU. More investigation needed.
layer_weight = layer.get_weight(module.weight) * (lora_weight * layer_scale) layer_weight = layer.get_weight(module.weight) * (lora_weight * layer_scale)
layer.to( layer.to(device=TorchDevice.CPU_DEVICE)
device=TorchDevice.CPU_DEVICE,
non_blocking=TorchDevice.get_non_blocking(TorchDevice.CPU_DEVICE),
)
assert isinstance(layer_weight, torch.Tensor) # mypy thinks layer_weight is a float|Any ??! assert isinstance(layer_weight, torch.Tensor) # mypy thinks layer_weight is a float|Any ??!
if module.weight.shape != layer_weight.shape: if module.weight.shape != layer_weight.shape:
@ -156,7 +153,7 @@ class ModelPatcher:
layer_weight = layer_weight.reshape(module.weight.shape) layer_weight = layer_weight.reshape(module.weight.shape)
assert isinstance(layer_weight, torch.Tensor) # mypy thinks layer_weight is a float|Any ??! assert isinstance(layer_weight, torch.Tensor) # mypy thinks layer_weight is a float|Any ??!
module.weight += layer_weight.to(dtype=dtype, non_blocking=TorchDevice.get_non_blocking(device)) module.weight += layer_weight.to(dtype=dtype)
yield # wait for context manager exit yield # wait for context manager exit
@ -164,9 +161,7 @@ class ModelPatcher:
assert hasattr(model, "get_submodule") # mypy not picking up fact that torch.nn.Module has get_submodule() assert hasattr(model, "get_submodule") # mypy not picking up fact that torch.nn.Module has get_submodule()
with torch.no_grad(): with torch.no_grad():
for module_key, weight in original_weights.items(): for module_key, weight in original_weights.items():
model.get_submodule(module_key).weight.copy_( model.get_submodule(module_key).weight.copy_(weight)
weight, non_blocking=TorchDevice.get_non_blocking(weight.device)
)
@classmethod @classmethod
@contextmanager @contextmanager

7
invokeai/backend/onnx/onnx_runtime.py
View File

@ -190,12 +190,7 @@ class IAIOnnxRuntimeModel(RawModel):
return self.session.run(None, inputs) return self.session.run(None, inputs)
# compatability with RawModel ABC # compatability with RawModel ABC
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
pass pass
# compatability with diffusers load code # compatability with diffusers load code

7
invokeai/backend/raw_model.py
View File

@ -20,10 +20,5 @@ class RawModel(ABC):
"""Abstract base class for 'Raw' model wrappers.""" """Abstract base class for 'Raw' model wrappers."""
@abstractmethod @abstractmethod
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
pass pass

9
invokeai/backend/textual_inversion.py
View File

@ -65,17 +65,12 @@ class TextualInversionModelRaw(RawModel):
return result return result
def to( def to(self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None) -> None:
self,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
non_blocking: bool = False,
) -> None:
if not torch.cuda.is_available(): if not torch.cuda.is_available():
return return
for emb in [self.embedding, self.embedding_2]: for emb in [self.embedding, self.embedding_2]:
if emb is not None: if emb is not None:
emb.to(device=device, dtype=dtype, non_blocking=non_blocking) emb.to(device=device, dtype=dtype)
def calc_size(self) -> int: def calc_size(self) -> int:
"""Get the size of this model in bytes.""" """Get the size of this model in bytes."""

12
invokeai/backend/util/devices.py
View File

@ -112,15 +112,3 @@ class TorchDevice:
@classmethod @classmethod
def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype: def _to_dtype(cls, precision_name: TorchPrecisionNames) -> torch.dtype:
return NAME_TO_PRECISION[precision_name] return NAME_TO_PRECISION[precision_name]
@staticmethod
def get_non_blocking(to_device: torch.device) -> bool:
"""Return the non_blocking flag to be used when moving a tensor to a given device.
MPS may have unexpected errors with non-blocking operations - we should not use non-blocking when moving _to_ MPS.
When moving _from_ MPS, we can use non-blocking operations.
See:
- https://github.com/pytorch/pytorch/issues/107455
- https://discuss.pytorch.org/t/should-we-set-non-blocking-to-true/38234/28
"""
return False if to_device.type == "mps" else True