gradient checkpointing is not implement

[xiyang-aads-lilly]

I have noticed that there is no implementation on integrating gradient checkpointing which can significantly save GPU RAM during the training. In https://huggingface.co/togethercomputer/StripedHyena-Hessian-7B/tree/main, the gradient checkpoint was implemented.

def stateless_forward(self, x, padding_mask=None):
  if type(padding_mask) == torch.Tensor:
  x = x * padding_mask[..., None]
   
  for _, block in enumerate(self.blocks):
  x, _ = block(x, inference_params=None, padding_mask=padding_mask)
  return x, None