Error when using packed sequence and gradient checkpoint: save_for_backward can only save variables, but argument 5 is of type PakedSeqParams
This issue is stale because it has been open for 30 days with no activity. Remove stale label or comment or this will be closed in 7 days. This issue was closed because it has been inactive for 7 days since being marked as stale. @xingyaoww did you solve this problem
**Describe the bug**

When trying to train the model using the latest NeMo code (`32e630220d2c24550a869f43ce618949ccf1a1a5`) with gradient checkpoint and packed sequence, getting the following error:

```
...
  File "/scratch/bcsi/xingyao6/nemo-util/Megatron-LM/megatron/core/transformer/tran
sformer_block.py", line 371, in forward
    hidden_states = self._checkpointed_forward(
  File "/scratch/bcsi/xingyao6/nemo-util/Megatron-LM/megatron/core/transformer/tran
sformer_block.py", line 257, in _checkpointed_forward
    hidden_states, context = checkpoint_handler(
  File "/scratch/bcsi/xingyao6/nemo-util/Megatron-LM/megatron/core/transformer/tran
sformer_block.py", line 240, in checkpoint_handler
    return tensor_parallel.checkpoint(
  File "/scratch/bcsi/xingyao6/nemo-util/Megatron-LM/megatron/core/tensor_parallel$
random.py", line 301, in checkpoint                                                    return CheckpointFunction.apply(function, distribute_saved_activations, *args)
  File "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py", line $
51, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
TypeError: save_for_backward can only save variables, but argument 5 is of type PakedSeqParams
```

**Environment overview (please complete the following information)**

Using commit `32e630220d2c24550a869f43ce618949ccf1a1a5` from the `main` branch with a docker environment build from `Dockerfile`.
