From a22862d5c6da68c463e2b82f2a187efe858bf125 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 6 Aug 2023 02:20:45 +0300 Subject: [PATCH] IPEX fix embedding on FP32 and BF16 --- modules/ipex_specific/__init__.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/modules/ipex_specific/__init__.py b/modules/ipex_specific/__init__.py index 9206b95f2..f7af296c1 100644 --- a/modules/ipex_specific/__init__.py +++ b/modules/ipex_specific/__init__.py @@ -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)