[Hotfix] Fix OPT gradient checkpointing forward

Co-authored-by: Edenzzzz <[email protected]>
hpcaitech · Jul 3, 2024 · eb24fcd · eb24fcd
1 parent ea94c07
commit eb24fcd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/colossalai/shardformer/modeling/opt.py b/colossalai/shardformer/modeling/opt.py
@@ -221,7 +221,7 @@ def opt_model_forward(
             past_key_value = past_key_values[idx] if past_key_values is not None else None
 
             if decoder.gradient_checkpointing and decoder.training:
-                layer_outputs = self._gradient_checkpointing_func(
+                layer_outputs = self.decoder._gradient_checkpointing_func(
                     decoder_layer.__call__,
                     hidden_states,
                     causal_attention_mask,