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

fix grad ckpt bug of llama

src/transformers/models/llama/modeling_llama.py

@@ -372,7 +372,7 @@ class LlamaPreTrainedModel(PreTrainedModel):
                 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