IPEX fix embedding on FP32 and BF16

This commit is contained in:
Disty0
2023-08-06 02:20:45 +03:00
parent 8aba6d8288
commit a22862d5c6
+11 -1
View File
@@ -112,15 +112,25 @@ def ipex_init():
CondFunc('torch.nn.modules.GroupNorm.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)
#FP32:
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)
#Embedding FP32:
CondFunc('torch.bmm',
lambda orig_func, input, mat2, *args, **kwargs: orig_func(input, mat2.to(input.dtype), *args, **kwargs),
lambda orig_func, input, mat2, *args, **kwargs: input.dtype != mat2.dtype)
#BF16:
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:
#Embedding BF16
CondFunc('torch.cat',
lambda orig_func, input, *args, **kwargs: orig_func([input[0].to(input[1].dtype), input[1], input[2].to(input[1].dtype)], *args, **kwargs),
lambda orig_func, input, *args, **kwargs: len(input) == 3 and (input[0].dtype != input[1].dtype or input[2].dtype != input[1].dtype))
#Diffusers BF16:
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)