add gradient checkpointing for the final_layernorm module.

modeling_chatglm.py

@@ -1012,7 +1012,10 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
                 all_self_attentions = all_self_attentions + (layer_ret[2 if use_cache else 1],)
 
         # Final layer norm.
-        hidden_states = self.final_layernorm(hidden_states)
+        if self.gradient_checkpointing and self.training:
+            hidden_states = torch.utils.checkpoint.checkpoint(self.final_layernorm, hidden_states)
+        else:
+            hidden_states = self.final_layernorm(hidden_states)
 
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)