bug fix

dcu_megatron/adaptor/megatron_adaptor.py

@@ -241,15 +241,15 @@ class LegacyAdaptation(MegatronAdaptationABC):
         self.patch_legacy_models()
 
     def patch_legacy_models(self):
-        from ..legacy.model.transformer import ParallelMLP, ParallelAttention
+        from ..legacy.model.transformer import ParallelMLPPatch, ParallelAttentionPatch
         from ..legacy.model.utils import get_norm
 
         # ParallecMLP
         MegatronAdaptation.register('megatron.legacy.model.transformer.ParallelMLP.__init__',
-                                    ParallelMLP.__init__)
+                                    ParallelMLPPatch.__init__)
 
         MegatronAdaptation.register('megatron.legacy.model.transformer.ParallelAttention.forward',
-                                    ParallelAttention.forward)
+                                    ParallelAttentionPatch.forward)
 
         # rms_norm.RMSNorm
         MegatronAdaptation.register('megatron.legacy.model.rms_norm.RMSNorm.forward',

dcu_megatron/legacy/model/transformer.py

@@ -10,7 +10,7 @@ from megatron.legacy.model.utils import (
 )
 
 
-class ParallelMLP(MegatronModule):
+class ParallelMLPPatch(MegatronModule):
     """MLP.
 
     MLP will take the input with h hidden state, project it to 4*h
@@ -74,7 +74,7 @@ class ParallelMLP(MegatronModule):
         )
 
 
-class ParallelAttention(MegatronModule):
+class ParallelAttentionPatch(MegatronModule):
     """Parallel self-attention layer abstract class.
 
     Self-attention layer takes input with size [s, b, h]