[ROCm] return an empty state dict for TE DPA module

ROCm · Nov 6, 2024 · 871c161 · 871c161
1 parent a2d0bdf
commit 871c161
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
@@ -706,6 +706,15 @@ def forward(
         else:
             return core_attn_out
 
+    def sharded_state_dict(self, prefix='', sharded_offsets=(), metadata=None):
+        state_dict = self.state_dict(prefix='', keep_vars=True)
+        # TE with version>=1.9 introduces an extra state in DotProductAttention Module
+        if is_te_min_version("1.9.0.dev0") and ('_extra_state' in state_dict):
+            state_dict.pop('_extra_state')
+        return make_sharded_tensors_for_checkpoint(
+            state_dict, prefix, {}, sharded_offsets
+        )
+
 
 if is_te_min_version("1.9.0.dev0"):