Stuck on model forward with 100% GPU-Util

Nobody has claimed this yet.

When I run finetune/adapter.py with my dataset with almost no modification (with devices=2), the code stuck on `logits = model(input_ids, lm_head_chunk_size=128)`(https://github.com/Lightning-AI/lit-gpt/blob/main/finetune/adapter.py#L157C13-L157C62), the `model ()` call nerver finish but GPU-Util is always 100%.

The process even can't be terminated with `Ctrl-C`

After replacing `validate(fabric, model, val_data, tokenizer, max_iters=2)`(https://github.com/Lightning-AI/lit-gpt/blob/main/finetune/adapter.py#L137) with `model.train()`, output error log with following:
```
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: A different number of tensors was saved during the original forward and recomputation.
Number of tensors saved during forward: 49
Number of tensors saved during recomputation: 47
raise CheckpointError(
torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: A different number of tensors was saved during the original forward and recomputation.
Number of tensors saved during forward: 49
Number of tensors saved during recomputation: 47
```

The script param I used to run is:
```
python finetune/adapter.py \
--data_dir data/xxx \
--checkpoint_dir checkpoints/meta-llama/Llama-2-7b-chat-hf \
--out_dir out/adapter/xxx
```

Any tips will be appreciated.

Start with finetune/adapter.py at the validation call around line 137 and the model forward at line 157, using the provided two-device command and dataset setup. Reproduce the hang, then compare it with the reported torch.utils.checkpoint.CheckpointError; done means validation completes normally and the process can be interrupted without the forward remaining stuck.

Written by the indexing model from the issue text.