Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 16:54:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

IPEX improve autocast hijack

Browse Source
This commit is contained in:
Disty0
2024-01-28 02:20:43 +03:00
parent ebdc2144a8
commit 4b30d6aeb7
2 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/intel/ipex/__init__.py
+4
View File
@@ -125,9 +125,13 @@ def ipex_init(): # pylint: disable=too-many-statements
# AMP:
torch.cuda.amp = torch.xpu.amp
torch.is_autocast_enabled = torch.xpu.is_autocast_xpu_enabled
torch.get_autocast_gpu_dtype = torch.xpu.get_autocast_xpu_dtype
if not hasattr(torch.cuda.amp, "common"):
torch.cuda.amp.common = contextlib.nullcontext()
torch.cuda.amp.common.amp_definitely_not_available = lambda: False
try:
torch.cuda.amp.GradScaler = torch.xpu.amp.GradScaler
except Exception: # pylint: disable=broad-exception-caught
modules/intel/ipex/hijacks.py
+9 -10
View File
@@ -29,16 +29,15 @@ def return_xpu(device):
# Autocast
original_autocast = torch.autocast
@wraps(torch.autocast)
def ipex_autocast(*args, **kwargs):
if len(args) > 0 and (args[0] == "cuda" or args[0] == "xpu"):
if "dtype" in kwargs:
return original_autocast("xpu", *args[1:], **kwargs)
else:
return original_autocast("xpu", *args[1:], dtype=devices.dtype, **kwargs)
original_autocast_init = torch.amp.autocast_mode.autocast.__init__
@wraps(torch.amp.autocast_mode.autocast.__init__)
def autocast_init(self, device_type, dtype=None, enabled=True, cache_enabled=None):
if device_type == "cuda" or device_type == "xpu":
if dtype is None:
dtype = devices.dtype
return original_autocast_init(self, device_type="xpu", dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)
else:
return original_autocast(*args, **kwargs)
return original_autocast_init(self, device_type=device_type, dtype=dtype, enabled=enabled, cache_enabled=cache_enabled)
# Latent Antialias CPU Offload:
original_interpolate = torch.nn.functional.interpolate
@@ -275,7 +274,7 @@ def ipex_hijacks():
torch.backends.cuda.sdp_kernel = return_null_context
torch.nn.DataParallel = DummyDataParallel
torch.UntypedStorage.is_cuda = is_cuda
torch.autocast = ipex_autocast
torch.amp.autocast_mode.autocast.__init__ = autocast_init
torch.nn.functional.scaled_dot_product_attention = scaled_dot_product_attention
torch.nn.functional.group_norm = functional_group_norm