Fix the gradient checkpointing bug of the llama model (huggingface#22270

) fix grad ckpt bug of llama
sangethmathewjohn · Mar 20, 2023 · 89f0fda · 89f0fda
1 parent cf0af9a
commit 89f0fda
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py
@@ -372,7 +372,7 @@ def _init_weights(self, module):
                 module.weight.data[module.padding_idx].zero_()
 
     def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, (LlamaDecoderLayer)):
+        if isinstance(module, LlamaModel):
             module.gradient_checkpointing = value