Fix activation checkpoint (NVIDIA#7334)

* Fix activation checkpoint Signed-off-by: Cheng-Ping Hsieh <[email protected]> * Remove inference_mode checkpoint Signed-off-by: Cheng-Ping Hsieh <[email protected]> --------- Signed-off-by: Cheng-Ping Hsieh <[email protected]>
rohitrango · Aug 29, 2023 · 22e61ca · 22e61ca
1 parent 62c5f11
commit 22e61ca
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/attention.py b/nemo/collections/nlp/modules/common/megatron/attention.py
@@ -245,24 +245,22 @@ def _checkpointed_attention_forward(
         """Forward method with activation checkpointing."""
 
         def custom_forward(*inputs):
-            if len(inputs) == 8:
+            if len(inputs) == 7:
                 query_layer = inputs[0]
                 key_layer = inputs[1]
                 value_layer = inputs[2]
                 attention_mask = inputs[3]
                 rotary_pos_emb = inputs[4]
                 relative_position_bias = inputs[5]
                 headscale_tensor = inputs[6]
-                inference_mode = inputs[7]
-            elif len(inputs) == 9:
+            elif len(inputs) == 8:
                 query_layer = inputs[0]
                 key_layer = inputs[1]
                 value_layer = inputs[2]
                 attention_mask = inputs[3]
                 rotary_pos_emb = (inputs[4], inputs[5])
                 relative_position_bias = inputs[6]
                 headscale_tensor = inputs[7]
-                inference_mode = inputs[8]
             else:
                 raise ValueError('unexpected number of inputs')
             output_ = self.core_attention(