diff --git a/modules/devices.py b/modules/devices.py index b1e03c017..cf3e5011f 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -135,9 +135,6 @@ def set_cuda_params(): except Exception: pass global dtype, dtype_vae, dtype_unet, unet_needs_upcast # pylint: disable=global-statement - if shared.cmd_opts.use_directml and not shared.cmd_opts.experimental: # TODO DirectML does not have full autocast capabilities - shared.opts.no_half = True - shared.opts.no_half_vae = True if shared.opts.cuda_dtype == 'FP32': dtype = torch.float32 dtype_vae = torch.float32 diff --git a/modules/dml/amp/autocast_mode.py b/modules/dml/amp/autocast_mode.py index f04728bf5..a5766dd32 100644 --- a/modules/dml/amp/autocast_mode.py +++ b/modules/dml/amp/autocast_mode.py @@ -2,18 +2,25 @@ import importlib from typing import Any, Optional import torch -ops = ["torch.Tensor.__matmul__", "torch.addbmm", "torch.addmm", "torch.addmv", "torch.addr", "torch.baddbmm", "torch.bmm", "torch.chain_matmul", "torch.linalg.multi_dot", "torch.nn.functional.conv1d", "torch.nn.functional.conv2d", "torch.nn.functional.conv3d", "torch.nn.functional.conv_transpose1d", "torch.nn.functional.conv_transpose2d", "torch.nn.functional.conv_transpose3d", "torch.nn.GRUCell", "torch.nn.functional.linear", "torch.nn.LSTMCell", "torch.matmul", "torch.mm", "torch.mv", "torch.prelu", "torch.nn.RNNCell"] +ops = ["torch.Tensor.__matmul__", "torch.addbmm", "torch.addmm", "torch.addmv", "torch.addr", "torch.baddbmm", "torch.bmm", "torch.chain_matmul", "torch.linalg.multi_dot", "torch.nn.functional.conv1d", "torch.nn.functional.conv2d", "torch.nn.functional.conv3d", "torch.nn.functional.conv_transpose1d", "torch.nn.functional.conv_transpose2d", "torch.nn.functional.conv_transpose3d", "torch.nn.GRUCell", "torch.nn.functional.linear", "torch.nn.LSTMCell", "torch.matmul", "torch.mm", "torch.mv", "torch.prelu", "torch.nn.RNNCell", "torch.embedding"] +supported_cast_pairs = { + torch.float16: (torch.float32,), + torch.float32: (torch.float16,), +} -def pre_forward(forward, args, kwargs): +def forward(op, args: tuple, kwargs: dict): if not torch.dml.is_autocast_enabled: - return forward(*args, **kwargs) + return op(*args, **kwargs) args = list(map(cast, args)) - for keyword in kwargs: - kwargs[keyword] = cast(kwargs[keyword]) - return forward(*args, **kwargs) + for kwarg in kwargs: + kwargs[kwarg] = cast(kwargs[kwarg]) + return op(*args, **kwargs) -def cast(tensor): - if not isinstance(tensor, torch.Tensor) or torch.dml.autocast_gpu_dtype == tensor.dtype: +def cast(tensor: torch.Tensor): + if not torch.is_tensor(tensor): + return tensor + dtype: torch.dtype = tensor.dtype + if dtype not in supported_cast_pairs or (torch.dml.autocast_gpu_dtype != dtype and torch.dml.autocast_gpu_dtype not in supported_cast_pairs[dtype]): return tensor return tensor.type(torch.dml.autocast_gpu_dtype) @@ -29,21 +36,25 @@ def cond(op: str): for attr_name in func_path[i:-1]: resolved_obj = getattr(resolved_obj, attr_name) op = getattr(resolved_obj, func_path[-1]) - setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: pre_forward(op, args, kwargs)) + setattr(resolved_obj, func_path[-1], lambda *args, **kwargs: forward(op, args, kwargs)) for op in ops: cond(op) class autocast: - def __init__(self, dtype: Optional[torch.dtype] = None): - self.fast_dtype = dtype or torch.dml.autocast_gpu_dtype + prev: bool + + fast_dtype: torch.dtype = torch.float16 + prev_fast_dtype: torch.dtype + def __init__(self, dtype: Optional[torch.dtype] = torch.float16): + self.fast_dtype = dtype def __enter__(self): self.prev = torch.dml.is_autocast_enabled - self.prev_fastdtype = torch.dml.autocast_gpu_dtype + self.prev_fast_dtype = torch.dml.autocast_gpu_dtype torch.dml.is_autocast_enabled = True torch.dml.autocast_gpu_dtype = self.fast_dtype def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): torch.dml.is_autocast_enabled = self.prev - torch.dml.autocast_gpu_dtype = self.prev_fastdtype + torch.dml.autocast_gpu_dtype = self.prev_fast_dtype diff --git a/modules/dml/backend.py b/modules/dml/backend.py index ed26da886..87ab9d043 100644 --- a/modules/dml/backend.py +++ b/modules/dml/backend.py @@ -61,7 +61,7 @@ class DirectML: return f"privateuseone:{get_device(device).index}" def get_device_name(device: Optional[rDevice]=None) -> str: - return torch_directml.device_name(get_device(device)) + return torch_directml.device_name(get_device(device).index) def get_device_properties(device: Optional[rDevice]=None) -> DeviceProperties: return DeviceProperties(get_device(device)) @@ -84,8 +84,7 @@ class DirectML: return (0 if available < 0 else available, DirectML.__gpu_memory_bound) def memory_allocated(device: Optional[rDevice]=None) -> int: - device = get_device(device) - return sum(torch_directml.gpu_memory(device.index)) * (1 << 20) + return sum(torch_directml.gpu_memory(get_device(device).index)) * (1 << 20) def max_memory_allocated(device: Optional[rDevice]=None): return DirectML.memory_allocated(device) # DirectML does not empty GPU memory diff --git a/modules/dml/hijack/transformers.py b/modules/dml/hijack/transformers.py index 712d5281e..6ae6bfb50 100644 --- a/modules/dml/hijack/transformers.py +++ b/modules/dml/hijack/transformers.py @@ -1,4 +1,5 @@ import torch +from typing import Optional import transformers.models.clip.modeling_clip # Copied from transformers.models.bart.modeling_bart._make_causal_mask @@ -19,4 +20,25 @@ def _make_causal_mask( mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1) return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length) +def CLIPTextEmbeddings_forward( + self: transformers.models.clip.modeling_clip.CLIPTextEmbeddings, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, +) -> torch.Tensor: + from modules.devices import dtype + seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if inputs_embeds is None: + inputs_embeds = self.token_embedding(input_ids).type(dtype) # Type correction. + + position_embeddings = self.position_embedding(position_ids) + embeddings = inputs_embeds + position_embeddings + + return embeddings + transformers.models.clip.modeling_clip._make_causal_mask = _make_causal_mask +transformers.models.clip.modeling_clip.CLIPTextEmbeddings.forward = CLIPTextEmbeddings_forward