IPEX fix BF16

This commit is contained in:
Disty0
2023-08-05 19:29:52 +03:00
parent 4234555566
commit 8aba6d8288
2 changed files with 11 additions and 5 deletions
+8 -3
View File
@@ -115,10 +115,15 @@ def ipex_init():
CondFunc('torch.nn.modules.Linear.forward',
lambda orig_func, self, input: orig_func(self, input.to(self.weight.data.dtype)),
lambda orig_func, self, input: input.dtype != self.weight.data.dtype)
CondFunc('torch.nn.functional.layer_norm',
lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
orig_func(input.to(weight.data.dtype), normalized_shape, weight, *args, **kwargs),
lambda orig_func, input, normalized_shape=None, weight=None, *args, **kwargs:
input.dtype != weight.data.dtype and weight is not None)
#Diffusers bfloat16:
CondFunc('torch.nn.modules.Conv2d._conv_forward',
lambda orig_func, self, input, weight, bias=None: orig_func(self, input.to(weight.data.dtype), weight, bias=bias),
lambda orig_func, self, input, weight, bias=None: input.dtype != weight.data.dtype)
CondFunc('torch.nn.functional.conv2d',
lambda orig_func, input, weight, *args, **kwargs: orig_func(input.to(weight.data.dtype), weight, *args, **kwargs),
lambda orig_func, input, weight, *args, **kwargs: input.dtype != weight.data.dtype)
#Functions that does not work with the XPU:
#UniPC:
+3 -2
View File
@@ -52,12 +52,13 @@ class SlicedAttnProcessor:
(batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
)
block_size = (batch_size_attention * query_tokens * shape_three) / 1024 * 1.2 #MB
block_multiply = 2.4 if query.dtype == torch.float32 else 1.2
block_size = (batch_size_attention * query_tokens * shape_three) / 1024 * block_multiply #MB
split_2_slice_size = query_tokens
if block_size >= 4000:
do_split_2 = True
#Find something divisible with the query_tokens
while ((self.slice_size * split_2_slice_size * shape_three) / 1024 * 1.2) > 4000:
while ((self.slice_size * split_2_slice_size * shape_three) / 1024 * block_multiply) > 4000:
split_2_slice_size = split_2_slice_size // 2
else:
do_split_2 = False