From 5134471bc8a5a79ba8a75a470ff40ffbfaca3864 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sun, 14 May 2023 13:24:59 -0400 Subject: [PATCH] dml autocast --- modules/devices.py | 30 ++++++++++++------- modules/dml/__init__.py | 21 ++++++++++++-- modules/dml/amp/__init__.py | 1 + modules/dml/amp/autocast_mode.py | 49 ++++++++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 13 deletions(-) create mode 100644 modules/dml/amp/__init__.py create mode 100644 modules/dml/amp/autocast_mode.py diff --git a/modules/devices.py b/modules/devices.py index 4adf74634..bc1c8de52 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -7,6 +7,7 @@ from modules import cmd_args, shared, memstats if sys.platform == "darwin": from modules import mac_specific # pylint: disable=ungrouped-imports +cuda_ok = torch.cuda.is_available() def has_mps() -> bool: if sys.platform != "darwin": @@ -33,7 +34,7 @@ def get_cuda_device_string(): def get_optimal_device_name(): if shared.cmd_opts.use_ipex: return "xpu" - elif torch.cuda.is_available() and not shared.cmd_opts.use_directml: + elif cuda_ok and not shared.cmd_opts.use_directml: return get_cuda_device_string() if has_mps(): return "mps" @@ -69,14 +70,14 @@ def torch_gc(): torch.xpu.empty_cache() except: pass - elif torch.cuda.is_available(): + elif cuda_ok: try: with torch.cuda.device(get_cuda_device_string()): torch.cuda.empty_cache() torch.cuda.ipc_collect() except: pass - shared.log.debug(f'gc: {memstats.memory_stats()}') + shared.log.debug(f'gc: {torch.device(get_optimal_device_name())} {memstats.memory_stats()}') def test_fp16(): @@ -96,7 +97,7 @@ def test_fp16(): def set_cuda_params(): shared.log.debug('Verifying Torch settings') - if torch.cuda.is_available(): + if cuda_ok: try: torch.backends.cuda.matmul.allow_tf32 = shared.opts.cuda_allow_tf32 torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = shared.opts.cuda_allow_tf16_reduced @@ -112,10 +113,9 @@ def set_cuda_params(): except: pass global dtype, dtype_vae, dtype_unet, unet_needs_upcast # pylint: disable=global-statement - # set dtype ok = test_fp16() - if shared.cmd_opts.use_directml: - shared.opts.no_half = True + # if shared.cmd_opts.use_directml: # TODO + # shared.opts.no_half = True if ok and shared.opts.cuda_dtype == 'FP32': shared.log.info('CUDA FP16 test passed but desired mode is set to FP32') if shared.opts.cuda_dtype == 'FP16' and ok: @@ -135,12 +135,12 @@ def set_cuda_params(): unet_needs_upcast = shared.opts.upcast_sampling shared.log.debug(f'Desired Torch parameters: dtype={shared.opts.cuda_dtype} no-half={shared.opts.no_half} no-half-vae={shared.opts.no_half_vae} upscast={shared.opts.upcast_sampling}') shared.log.info(f'Setting Torch parameters: dtype={dtype} vae={dtype_vae} unet={dtype_unet}') + shared.log.debug(f'Torch default device: {torch.device(get_optimal_device_name())}') args = cmd_args.parser.parse_args() if args.use_ipex: cpu = torch.device("xpu") #Use XPU instead of CPU. %20 Perf improvement on weak CPUs. - print("Using XPU instead of CPU.") else: cpu = torch.device("cpu") device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None @@ -176,19 +176,27 @@ def autocast(disable=False): return contextlib.nullcontext() if dtype == torch.float32 or shared.cmd_opts.precision == "Full": return contextlib.nullcontext() + if shared.cmd_opts.use_directml: + return torch.dml.amp.autocast(dtype) if shared.cmd_opts.use_ipex: return torch.xpu.amp.autocast(enabled=True, dtype=dtype, cache_enabled=False) - else: + if cuda_ok: return torch.autocast("cuda") + else: + return torch.autocast("cpu") def without_autocast(disable=False): if disable: return contextlib.nullcontext() + if shared.cmd_opts.use_directml: + return torch.dml.amp.autocast(enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext() if shared.cmd_opts.use_ipex: - return torch.autocast("xpu", enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext() - else: + return torch.xpu.amp.autocast(enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext() + if cuda_ok: return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext() + else: + return torch.autocast("cpu", enabled=False) if torch.is_autocast_enabled() else contextlib.nullcontext() class NansException(Exception): diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index f613e1279..46501d9cc 100644 --- a/modules/dml/__init__.py +++ b/modules/dml/__init__.py @@ -1,12 +1,15 @@ -# pylint: disable=no-member,no-self-argument +# pylint: disable=no-member,no-self-argument,no-method-argument import torch import torch_directml # pylint: disable=import-error - import modules.dml.hijack +import modules.dml.amp as amp from .optimizer.unknown import UnknownOptimizer class DirectML(): + _is_autocast_enabled = False + _autocast_dtype = torch.float16 + def get_optimizer(device: torch.device): assert device.type == 'privateuseone' try: @@ -27,5 +30,19 @@ class DirectML(): optimizer = DirectML.get_optimizer(device) return optimizer.memory_stats(device.index) + def get_autocast_gpu_dtype(): + return DirectML._autocast_dtype + + def set_autocast_gpu_dtype(dtype): + DirectML._autocast_dtype = dtype + + def is_autocast_enabled(): + return DirectML._is_autocast_enabled + + def set_autocast_enabled(enabled: bool): + DirectML._is_autocast_enabled = enabled + + # Alternative of torch.cuda for DirectML. +DirectML.amp = amp torch.dml = DirectML diff --git a/modules/dml/amp/__init__.py b/modules/dml/amp/__init__.py new file mode 100644 index 000000000..93d038da0 --- /dev/null +++ b/modules/dml/amp/__init__.py @@ -0,0 +1 @@ +from .autocast_mode import * diff --git a/modules/dml/amp/autocast_mode.py b/modules/dml/amp/autocast_mode.py new file mode 100644 index 000000000..a8a636ee3 --- /dev/null +++ b/modules/dml/amp/autocast_mode.py @@ -0,0 +1,49 @@ +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"] + +def pre_forward(forward, args, kwargs): + if not torch.dml.is_autocast_enabled(): + return forward(*args, **kwargs) + args = list(map(cast, args)) + for keyword in kwargs: + kwargs[keyword] = cast(kwargs[keyword]) + return forward(*args, **kwargs) + +def cast(tensor): + if not isinstance(tensor, torch.Tensor): + return tensor + return tensor.type(torch.dml.get_autocast_gpu_dtype()) + +def cond(op: str): + if isinstance(op, str): + func_path = op.split('.') + for i in range(len(func_path)-1, -1, -1): + try: + resolved_obj = importlib.import_module('.'.join(func_path[:i])) + break + except ImportError: + pass + 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)) + +for op in ops: + cond(op) + +class autocast: + def __init__(self, dtype: Optional[torch.dtype] = None): + self.fast_dtype = dtype or torch.dml.get_autocast_gpu_dtype() + + def __enter__(self): + self.prev = torch.dml.is_autocast_enabled() + self.prev_fastdtype = torch.dml.get_autocast_gpu_dtype() + torch.dml.set_autocast_enabled(True) + torch.dml.set_autocast_gpu_dtype(self.fast_dtype) + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any): + torch.dml.set_autocast_enabled(self.prev) + torch.dml.set_autocast_gpu_dtype(self.prev_fastdtype)