From d4aa840a77de06eee4a26f013a0858d47788ab6b Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Mon, 24 Jul 2023 16:10:01 +0900 Subject: [PATCH] DirectML update. DirectML reuses GPU memory instead of returning it. So prints "practical" GPU memory utilization too. --- modules/devices.py | 5 ++++- modules/dml/__init__.py | 23 ++++++++++++++++++++++- modules/dml/backend.py | 31 +++++++++++++++++++++++++------ modules/dml/device_properties.py | 7 +++++++ modules/dml/hijack/torch.py | 8 ++++++++ modules/shared.py | 4 ++-- 6 files changed, 68 insertions(+), 10 deletions(-) diff --git a/modules/devices.py b/modules/devices.py index 207ddbac3..b1e03c017 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -33,7 +33,7 @@ def get_cuda_device_string(): elif backend == 'directml' and torch.dml.is_available(): if shared.cmd_opts.device_id is not None: return f"privateuseone:{shared.cmd_opts.device_id}" - return torch.dml.get_default_device_string() + return torch.dml.get_device_string(torch.dml.default_device().index) else: if shared.cmd_opts.device_id is not None: return f"cuda:{shared.cmd_opts.device_id}" @@ -70,6 +70,9 @@ def torch_gc(force=False): if used > 95: shared.log.warning(f'GPU high memory utilization: {used}% {mem}') force = True + if backend == "directml": + practical_used = round(100 * torch.cuda.memory_allocated() / (1 << 30) / gpu.get('total', 1)) + shared.log.info(f'Practical GPU memory utilization: {practical_used}%') if shared.opts.disable_gc and not force: return diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py index 9ce75598e..a656b970f 100644 --- a/modules/dml/__init__.py +++ b/modules/dml/__init__.py @@ -1,5 +1,10 @@ +import os import torch +from modules.sd_hijack_utils import CondFunc + +do_nothing = lambda: None + def directml_init(): from modules.dml.backend import DirectML # pylint: disable=ungrouped-imports # Alternative of torch.cuda for DirectML. @@ -7,18 +12,34 @@ def directml_init(): torch.cuda.is_available = lambda: False torch.cuda.device = torch.dml.device + torch.cuda.device_count = torch.dml.device_count torch.cuda.current_device = torch.dml.current_device torch.cuda.get_device_name = torch.dml.get_device_name torch.cuda.get_device_properties = torch.dml.get_device_properties + torch.cuda.empty_cache = do_nothing + torch.cuda.ipc_collect = do_nothing torch.cuda.memory_stats = torch.dml.memory_stats torch.cuda.mem_get_info = torch.dml.mem_get_info torch.cuda.memory_allocated = torch.dml.memory_allocated torch.cuda.max_memory_allocated = torch.dml.max_memory_allocated torch.cuda.reset_peak_memory_stats = torch.dml.reset_peak_memory_stats + torch.cuda.utilization = lambda: 0 -def directml_hijack_init(): + torch.Tensor.directml = lambda self: self.to(torch.dml.current_device()) + + mem_bound = os.environ.get("DML_GPU_MEMORY_BOUND", None) + if mem_bound is not None: + torch.dml.set_gpu_memory_bound(int(mem_bound)) + +def directml_do_hijack(): import modules.dml.hijack + from modules.devices import device + + if not torch.dml.has_float64_support(device): + CondFunc('torch.from_numpy', + lambda orig_func, *args, **kwargs: orig_func(args[0].astype('float32')), + lambda *args, **kwargs: args[1].dtype == float) def directml_override_opts(): from modules import shared diff --git a/modules/dml/backend.py b/modules/dml/backend.py index b6f4d7e53..ed26da886 100644 --- a/modules/dml/backend.py +++ b/modules/dml/backend.py @@ -15,6 +15,8 @@ class DirectML: context_device: torch.device | None = None + __gpu_memory_bound: int | None = None + is_autocast_enabled = False autocast_gpu_dtype = torch.float16 @@ -33,18 +35,30 @@ class DirectML: return memory_control except Exception: return UnknownMemoryControl + + def set_gpu_memory_bound(bound: int | None): + DirectML.__gpu_memory_bound = bound def is_available() -> bool: return torch_directml.is_available() + def is_directml_device(device: torch.device) -> bool: + return device.type == "privateuseone" + + def has_float64_support(device: Optional[rDevice]=None) -> bool: + return torch_directml.has_float64_support(get_device(device).index) + + def device_count() -> int: + return torch_directml.device_count() + def current_device() -> torch.device: return DirectML.context_device or DirectML.default_device() def default_device() -> torch.device: return torch_directml.device(torch_directml.default_device()) - def get_default_device_string() -> str: - return f"privateuseone:{torch_directml.default_device()}" + def get_device_string(device: Optional[rDevice]=None) -> str: + return f"privateuseone:{get_device(device).index}" def get_device_name(device: Optional[rDevice]=None) -> str: return torch_directml.device_name(get_device(device)) @@ -59,14 +73,19 @@ class DirectML: "num_alloc_retries": mem_stat_fill, } - def mem_get_info(device: Optional[rDevice]=None): + def mem_get_info(device: Optional[rDevice]=None) -> tuple[int, int]: device = get_device(device) memory_control = DirectML.__get_memory_control(device) - return memory_control.mem_get_info(device.index) + mem_info = memory_control.mem_get_info(device.index) + if DirectML.__gpu_memory_bound is None: + return mem_info + used = mem_info[1] - mem_info[0] + available = DirectML.__gpu_memory_bound - used + return (0 if available < 0 else available, DirectML.__gpu_memory_bound) - def memory_allocated(device: Optional[rDevice]=None): + 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(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/device_properties.py b/modules/dml/device_properties.py index 5b5239a9d..abd146653 100644 --- a/modules/dml/device_properties.py +++ b/modules/dml/device_properties.py @@ -7,6 +7,13 @@ class DeviceProperties: minor: int = 0 total_memory: int multi_processor_count: int = 1 + def __init__(self, device: torch.device): self.name = torch.dml.get_device_name(device) self.total_memory = torch.dml.mem_get_info(device)[0] + + def __str__(self): + return f"DeviceProperties(name='{self.name}', total_memory='{self.total_memory}')" + + def __repr__(self): + return f"DeviceProperties(name='{self.name}', total_memory='{self.total_memory}')" diff --git a/modules/dml/hijack/torch.py b/modules/dml/hijack/torch.py index 04e777258..2400e6f6a 100644 --- a/modules/dml/hijack/torch.py +++ b/modules/dml/hijack/torch.py @@ -3,3 +3,11 @@ import torch from modules.sd_hijack_utils import CondFunc CondFunc('torchsde._brownian.brownian_interval._randn', lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=torch.device("cpu"), generator=torch.Generator(torch.device("cpu")).manual_seed(int(seed))).to(device), lambda _, size, dtype, device, seed: device.type == 'privateuseone') + +_new = torch.Tensor.new +def new(self: torch.Tensor, *args, **kwargs): + if torch.dml.is_directml_device(self.device): + return _new(self.cpu(), *args, **kwargs).to(self.device) + return _new(self, *args, **kwargs) + +torch.Tensor.new = new diff --git a/modules/shared.py b/modules/shared.py index 1ad2d54f1..b9c21f61b 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -12,7 +12,7 @@ import requests import fasteners from modules import errors, ui_components, shared_items, cmd_args from modules.paths_internal import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611 -from modules.dml import directml_hijack_init, directml_override_opts +from modules.dml import directml_do_hijack, directml_override_opts import modules.interrogate import modules.memmon import modules.styles @@ -803,7 +803,7 @@ parallel_processing_allowed = not cmd_opts.lowvram mem_mon = modules.memmon.MemUsageMonitor("MemMon", device, opts) mem_mon.start() if devices.backend == "directml": - directml_hijack_init() + directml_do_hijack() directml_override_opts()