rewrite transformer_engine

dcu_megatron/adaptor/megatron_adaptor.py

@@ -143,11 +143,11 @@ class CoreAdaptation(MegatronAdaptationABC):
     def patch_core_extentions(self):
         import transformer_engine as te
 
-        from ..core.extensions.transformer_engine import te_dot_product_attention_init
+        from ..core.extensions.transformer_engine import TEDotProductAttentionPatch
         from megatron.core.extensions.transformer_engine import TEGroupedLinear
 
         MegatronAdaptation.register('megatron.core.extensions.transformer_engine.TEDotProductAttention.__init__',
-                                    te_dot_product_attention_init)
+                                    TEDotProductAttentionPatch.__init__)
 
         if int(os.getenv("GROUPED_GEMM_BatchLinear", '0')):
             TEGroupedLinear.__bases__ = (te.pytorch.BatchLinear,)

dcu_megatron/core/extensions/transformer_engine.py

 import os
 import dataclasses
+import transformer_engine as te
 
 from typing import Any, Optional
 from packaging.version import Version as PkgVersion
@@ -19,7 +20,8 @@ from megatron.core.parallel_state import (
 )
 
 
-def te_dot_product_attention_init(
+class TEDotProductAttentionPatch(te.pytorch.DotProductAttention):
+    def __init__(
         self,
         config: TransformerConfig,
         layer_number: int,
@@ -30,7 +32,7 @@ def te_dot_product_attention_init(
         k_channels: Optional[int] = None,
         v_channels: Optional[int] = None,
         cp_comm_type: str = "p2p",
-):
+    ):
         self.config = config
         self.te_forward_mask_type = False
         self.qkv_format: str = 'sbhd'