From 2a9133bfecf7f5145c6b4ffe5a849bf21fdc24b7 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Fri, 14 Jul 2023 17:33:24 +0300 Subject: [PATCH] IPEX rework --- installer.py | 2 +- modules/api/api.py | 18 +----- modules/devices.py | 23 +++++++- modules/hypernetworks/hypernetwork.py | 15 +---- modules/memmon.py | 25 ++------- modules/memstats.py | 13 ----- modules/sd_hijack_optimizations.py | 55 +++++-------------- modules/sd_hijack_unet.py | 2 +- modules/sd_models.py | 2 +- .../textual_inversion/textual_inversion.py | 5 +- 10 files changed, 47 insertions(+), 113 deletions(-) diff --git a/installer.py b/installer.py index eeaf50942..01d5bda2d 100644 --- a/installer.py +++ b/installer.py @@ -315,7 +315,7 @@ def check_torch(): log.error('Intel OneAPI Toolkit is not activated! Start the WebUI with --use-ipex or activate OneAPI manually') os.environ.setdefault('NEOReadDebugKeys', '1') os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100') - torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0 torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu') + torch_command = os.environ.get('TORCH_COMMAND', 'torch==1.13.0a0+git6c9b55e torchvision==0.14.1a0 intel_extension_for_pytorch==1.13.120+xpu -f https://developer.intel.com/ipex-whl-stable-xpu') else: machine = platform.machine() if sys.platform == 'darwin': diff --git a/modules/api/api.py b/modules/api/api.py index eda9612a4..f208ba664 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -603,23 +603,7 @@ class Api: ram = { 'error': f'{err}' } try: import torch - if devices.backend == 'ipex': - system = { 'free': (torch.xpu.get_device_properties(shared.device).total_memory - torch.xpu.memory_allocated()), 'used': torch.xpu.memory_allocated(), 'total': torch.xpu.get_device_properties(shared.device).total_memory } - s = dict(torch.xpu.memory_stats()) - allocated = { 'current': s['allocated_bytes.all.current'], 'peak': s['allocated_bytes.all.peak'] } - reserved = { 'current': s['reserved_bytes.all.current'], 'peak': s['reserved_bytes.all.peak'] } - active = { 'current': s['active_bytes.all.current'], 'peak': s['active_bytes.all.peak'] } - inactive = { 'current': s['inactive_split_bytes.all.current'], 'peak': s['inactive_split_bytes.all.peak'] } - warnings = { 'retries': s['num_alloc_retries'], 'oom': s['num_ooms'] } - cuda = { - 'system': system, - 'active': active, - 'allocated': allocated, - 'reserved': reserved, - 'inactive': inactive, - 'events': warnings, - } - elif torch.cuda.is_available(): + if torch.cuda.is_available(): s = torch.cuda.mem_get_info() system = { 'free': s[0], 'used': s[1] - s[0], 'total': s[1] } s = dict(torch.cuda.memory_stats(shared.device)) diff --git a/modules/devices.py b/modules/devices.py index 8dd32a964..ba91b8a46 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -8,7 +8,6 @@ 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() previous_oom = 0 @@ -192,10 +191,26 @@ else: backend = 'cpu' if backend == 'ipex': - #Fix broken functions with ipex - from modules.sd_hijack_utils import CondFunc + #Fix functions with ipex + torch.cuda.is_available = torch.xpu.is_available + torch.cuda.current_device = torch.xpu.current_device + torch.cuda.get_device_properties = torch.xpu.get_device_properties torch.cuda.empty_cache = torch_gc + torch.cuda.memory_stats = torch.xpu.memory_stats + torch.cuda.mem_get_info = lambda device: [(torch.xpu.get_device_properties(device).total_memory - torch.xpu.memory_allocated(device)), torch.xpu.get_device_properties(device).total_memory] + torch.cuda.memory_allocated = torch.xpu.memory_allocated + torch.cuda.max_memory_allocated = torch.xpu.max_memory_allocated + torch.cuda.reset_peak_memory_stats = torch.xpu.reset_peak_memory_stats + + torch.cuda.get_rng_state_all = torch.xpu.get_rng_state_all + torch.cuda.set_rng_state_all = torch.xpu.set_rng_state_all + try: + torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler + except Exception: + pass + + from modules.sd_hijack_utils import CondFunc #Functions with dtype errors: CondFunc('torch.nn.modules.GroupNorm.forward', lambda orig_func, *args, **kwargs: orig_func(args[0], args[1].to(args[0].weight.data.dtype)), @@ -207,6 +222,7 @@ if backend == 'ipex': CondFunc('torch.nn.modules.Conv2d._conv_forward', lambda orig_func, *args, **kwargs: orig_func(args[0], args[1].to(args[2].data.dtype), args[2], args[3]), lambda *args, **kwargs: args[2].dtype != args[3].data.dtype) + #Functions that does not work with the XPU: #UniPC: CondFunc('torch.linalg.solve', @@ -238,6 +254,7 @@ if backend == 'ipex': args[5], args[6], args[7], args[8]).to(get_cuda_device_string()), lambda *args, **kwargs: args[1].device != torch.device("cpu")) +cuda_ok = torch.cuda.is_available() and not backend == 'ipex' cpu = torch.device("cpu") device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None dtype = torch.float16 diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py index 471e75322..07d566a3e 100644 --- a/modules/hypernetworks/hypernetwork.py +++ b/modules/hypernetworks/hypernetwork.py @@ -589,10 +589,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi print("Cannot resume from saved optimizer!") print(e) - if devices.backend == 'ipex': - scaler = torch.xpu.amp.GradScaler() - else: - scaler = torch.cuda.amp.GradScaler() + scaler = torch.cuda.amp.GradScaler() batch_size = ds.batch_size gradient_step = ds.gradient_step @@ -706,10 +703,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi hypernetwork.eval() rng_state = torch.get_rng_state() cuda_rng_state = None - if devices.backend == 'ipex': - cuda_rng_state = torch.xpu.get_rng_state_all() - elif torch.cuda.is_available(): - cuda_rng_state = torch.cuda.get_rng_state_all() + cuda_rng_state = torch.cuda.get_rng_state_all() shared.sd_model.cond_stage_model.to(devices.device) shared.sd_model.first_stage_model.to(devices.device) @@ -745,10 +739,7 @@ def train_hypernetwork(id_task, hypernetwork_name, learn_rate, batch_size, gradi shared.sd_model.cond_stage_model.to(devices.cpu) shared.sd_model.first_stage_model.to(devices.cpu) torch.set_rng_state(rng_state) - if devices.backend == 'ipex': - torch.xpu.set_rng_state_all(cuda_rng_state) - elif torch.cuda.is_available(): - torch.cuda.set_rng_state_all(cuda_rng_state) + torch.cuda.set_rng_state_all(cuda_rng_state) hypernetwork.train() if image is not None: shared.state.assign_current_image(image) diff --git a/modules/memmon.py b/modules/memmon.py index 929451875..1e426b9cb 100644 --- a/modules/memmon.py +++ b/modules/memmon.py @@ -21,12 +21,7 @@ class MemUsageMonitor(threading.Thread): self.run_flag = threading.Event() self.data = defaultdict(int) if not torch.cuda.is_available(): - #torch.cuda.is_available() reports False when using IPEX. - if devices.backend == 'ipex': - self.cuda_mem_get_info() - torch.xpu.memory_stats(self.device) - else: - self.disabled = True + self.disabled = True else: try: self.cuda_mem_get_info() @@ -35,22 +30,15 @@ class MemUsageMonitor(threading.Thread): self.disabled = True def cuda_mem_get_info(self): - if devices.backend == 'ipex': - index = self.device.index if self.device.index is not None else torch.xpu.current_device() - return [(torch.xpu.get_device_properties(index).total_memory - torch.xpu.memory_allocated(index)), torch.xpu.get_device_properties(index).total_memory] - else: - index = self.device.index if self.device.index is not None else torch.cuda.current_device() - return torch.cuda.mem_get_info(index) + index = self.device.index if self.device.index is not None else torch.cuda.current_device() + return torch.cuda.mem_get_info(index) def run(self): if self.disabled: return while True: self.run_flag.wait() - if devices.backend == 'ipex': - torch.xpu.reset_peak_memory_stats() - else: - torch.cuda.reset_peak_memory_stats() + torch.cuda.reset_peak_memory_stats() self.data.clear() if self.opts.memmon_poll_rate <= 0: self.run_flag.clear() @@ -70,10 +58,7 @@ class MemUsageMonitor(threading.Thread): self.data["free"] = free self.data["total"] = total try: - if devices.backend == 'ipex': - torch_stats = torch.xpu.memory_stats(self.device) - else: - torch_stats = torch.cuda.memory_stats(self.device) + torch_stats = torch.cuda.memory_stats(self.device) self.data["active"] = torch_stats["active.all.current"] self.data["active_peak"] = torch_stats["active_bytes.all.peak"] self.data["reserved"] = torch_stats["reserved_bytes.all.current"] diff --git a/modules/memstats.py b/modules/memstats.py index ea3a772f7..1d2c21162 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -29,17 +29,4 @@ def memory_stats(): return mem except Exception: pass - try: - gpu = { 'used': gb(torch.xpu.memory_allocated()), 'total': gb(torch.xpu.get_device_properties(shared.device).total_memory) } - s = dict(torch.xpu.memory_stats()) - mem.update({ - 'gpu': gpu, - 'retries': s['num_alloc_retries'], - 'oom': s['num_ooms'] - }) - if s['num_ooms'] > 0: - shared.state.oom = True - return mem - except Exception: - pass return mem diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 730f0f09c..96516d8a3 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -29,15 +29,7 @@ else: def get_available_vram(): - if devices.backend == 'ipex': - stats = torch.xpu.memory_stats(shared.device) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_xpu = torch.xpu.get_device_properties(shared.device).total_memory - torch.xpu.memory_allocated(shared.device) - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_xpu + mem_free_torch - return mem_free_total - elif shared.device.type == 'cuda': + if shared.device.type == 'cuda' or shared.device.type == 'xpu': try: stats = torch.cuda.memory_stats(shared.device) mem_active = stats['active_bytes.all.current'] @@ -47,7 +39,6 @@ def get_available_vram(): mem_free_total = mem_free_cuda + mem_free_torch except Exception: mem_free_total = 1024 * 1024 * 1024 - return mem_free_total elif shared.device.type == 'privateuseone': mem_total, mem_active = torch.dml.memory_stats(shared.device) @@ -190,27 +181,17 @@ def einsum_op_tensor_mem(q, k, v, max_tensor_mb): return einsum_op_slice_1(q, k, v, max(q.shape[1] // div, 1)) def einsum_op_cuda(q, k, v): - if devices.backend == 'ipex': - stats = torch.xpu.memory_stats(q.device) + try: + stats = torch.cuda.memory_stats(q.device) mem_active = stats['active_bytes.all.current'] mem_reserved = stats['reserved_bytes.all.current'] - mem_free_xpu = torch.xpu.get_device_properties(q.device).total_memory - torch.xpu.memory_allocated(q.device) + mem_free_cuda, _ = torch.cuda.mem_get_info(q.device) mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_xpu + mem_free_torch - # Divide factor of safety as there's copying and fragmentation - return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) - else: - try: - stats = torch.cuda.memory_stats(q.device) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] - mem_free_cuda, _ = torch.cuda.mem_get_info(q.device) - mem_free_torch = mem_reserved - mem_active - mem_free_total = mem_free_cuda + mem_free_torch - except Exception: - mem_free_total = 1024 * 1024 * 1024 - # Divide factor of safety as there's copying and fragmentation - return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) + mem_free_total = mem_free_cuda + mem_free_torch + except Exception: + mem_free_total = 1024 * 1024 * 1024 + # Divide factor of safety as there's copying and fragmentation + return einsum_op_tensor_mem(q, k, v, mem_free_total / 3.3 / (1 << 20)) def einsum_op_dml(q, k, v): mem_total, mem_active = torch.dml.memory_stats(q.device) @@ -218,7 +199,7 @@ def einsum_op_dml(q, k, v): return einsum_op_tensor_mem(q, k, v, (mem_reserved - mem_active) if mem_reserved > mem_active else 1) def einsum_op(q, k, v): - if q.device.type == 'cuda' or devices.backend == 'ipex': + if q.device.type == 'cuda' or q.device.type == 'xpu': return einsum_op_cuda(q, k, v) if q.device.type == 'mps': @@ -410,12 +391,8 @@ def scaled_dot_product_attention_forward(self, x, context=None, mask=None): return hidden_states def scaled_dot_product_no_mem_attention_forward(self, x, context=None, mask=None): - if devices.backend == 'ipex': - with torch.backends.xpu.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): - return scaled_dot_product_attention_forward(self, x, context, mask) - else: - with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): - return scaled_dot_product_attention_forward(self, x, context, mask) + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): + return scaled_dot_product_attention_forward(self, x, context, mask) def cross_attention_attnblock_forward(self, x): h_ = x @@ -523,12 +500,8 @@ def sdp_attnblock_forward(self, x): return x + out def sdp_no_mem_attnblock_forward(self, x): - if devices.backend == 'ipex': - with torch.backends.xpu.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): - return sdp_attnblock_forward(self, x) - else: - with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): - return sdp_attnblock_forward(self, x) + with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=False): + return sdp_attnblock_forward(self, x) def sub_quad_attnblock_forward(self, x): h_ = x diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py index e5e9ded1d..3234d33b6 100644 --- a/modules/sd_hijack_unet.py +++ b/modules/sd_hijack_unet.py @@ -67,7 +67,7 @@ def hijack_ddpm_edit(): unet_needs_upcast = lambda *args, **kwargs: devices.unet_needs_upcast # pylint: disable=unnecessary-lambda-assignment CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast) CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast) -if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available() or devices.backend == 'ipex': +if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available(): CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast) CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast) CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU) diff --git a/modules/sd_models.py b/modules/sd_models.py index f8011f9e5..134a6c5c4 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -737,7 +737,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No sd_model.unet.to(memory_format=torch.channels_last) base_sent_to_cpu=False - if shared.opts.cuda_compile and (torch.cuda.is_available() or devices.backend == 'ipex'): + if shared.opts.cuda_compile and torch.cuda.is_available(): if op == 'refiner': gpu_vram = memory_stats().get('gpu', {}) free_vram = gpu_vram.get('total', 0) - gpu_vram.get('used', 0) diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index c13e027c4..8529caa0f 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -443,10 +443,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st else: shared.log.info("No saved optimizer exists in checkpoint") - if devices.backend == 'ipex': - scaler = torch.xpu.amp.GradScaler() - else: - scaler = torch.cuda.amp.GradScaler() + scaler = torch.cuda.amp.GradScaler() batch_size = ds.batch_size gradient_step = ds.gradient_step