update

ppfleetx/models/language_model/gpt/dygraph/single_model.py

@@ -33,8 +33,7 @@
     ForcedBOSTokenLogitsProcessor, ForcedEOSTokenLogitsProcessor)
 
 from ppfleetx.distributed.moe import MoELayer
-import numpy as np
-
+from ppfleetx.utils import matmul as multi_device_matmul
 def get_attr(layer, name):
     if getattr(layer, name, None) is not None:
         return getattr(layer, name, None)
@@ -232,7 +231,7 @@ def core_attn(self, q, k, v, attn_mask=None):
                 mode="upscale_in_train")
 
 
-        out = _multi_device_matmul(weights, v)
+        out = multi_device_matmul(weights, v)
 
         # combine heads
         out = tensor.transpose(out, perm=[0, 2, 1, 3])

ppfleetx/utils/__init__.py

@@ -12,5 +12,3 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from .patches_for_multi_device import matmul
-if paddle.matmul != matmul:
-    paddle.matmul = matmul

ppfleetx/utils/patches_for_multi_device.py

@@ -1,4 +1,5 @@
 import paddle
+import numpy as np
 # There is a bug in matmul case 13 (batchsize>1) on DCU in fp16 mode, (maybe too large).
 # Here split the big matrix to few small matrices with batchsize=1 (matmul case 8) to compute the result.
 # TODO: fix the bug in matmul kernel