Notebook crashes before training

Nobody has claimed this yet.

### Bug description

When training a T5 finetuner, model.fit() ends without any output. Attempts to run any other cells hang until the notebook is restarted, so I can assume that the notebook has crashed.

### What version are you seeing the problem on?

master

### How to reproduce the bug

```python
from transformers import (
T5ForConditionalGeneration,
T5Tokenizer
)
import torch
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader

tokenizer = T5Tokenizer.from_pretrained('t5-base')

input_sequences = []
rewrite_prompts = [f"{prompt}" for prompt in data["transformation"]]

i = 0
for item in data["rewrite"]:
original = data['original'][i]
# format: original|new
line = f'{original}|{item}'
input_sequences.append(line)
i += 1

input_train, input_test, prompts_train, prompts_test = train_test_split(input_sequences, rewrite_prompts, test_size=0.1, random_state=42)

train_encodings = tokenizer(input_train, padding=True, truncation=True, return_tensors="pt", max_length=384)
train_labels = tokenizer(prompts_train, padding=True, truncation=True, return_tensors="pt", max_length=384).input_ids

# Tokenize testing data
test_encodings = tokenizer(input_test, padding=True, truncation=True, return_tensors="pt", max_length=384)
test_labels = tokenizer(prompts_test, padding=True, truncation=True, return_tensors="pt", max_length=384).input_ids

class TextDataset(Dataset):
def __init__(self, input_ids, labels):
self.input_ids = input_ids
self.labels = labels

def __len__(self):
return len(self.input_ids)

def __getitem__(self, idx):
item = {"input_ids": self.input_ids[idx], "labels": self.labels[idx]}
return item

train_dataset = TextDataset(input_ids=train_encodings["input_ids"], labels=train_labels)
validation_dataset = TextDataset(input_ids=test_encodings["input_ids"], labels=test_labels)

import pytorch_lightning as pl
import torch
from torch.utils.data import DataLoader
from transformers import T5ForConditionalGeneration, AdamW

class T5Tuner(pl.LightningModule):
def __init__(self, batchsize, t5model, t5tokenizer):
super(T5Tuner, self).__init__()
self.batch_size = batchsize
self.model = t5model
self.tokenizer = t5tokenizer

def forward(self, input_ids, labels=None):
return self.model(input_ids=input_ids, labels=labels)

def training_step(self, batch, batch_idx):
outputs = self.forward(batch['input_ids'], batch['labels'])
loss = outputs.loss
print('completed train step')
self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, rank_zero_only=True)
print('logged')
return loss

def validation_step(self, batch, batch_idx):
outputs = self.forward(batch['input_ids'], batch['labels'])

loss = outputs.loss
print('completed val step')
self.log("val_loss", loss, on_step=True, on_epoch=True, prog_bar=True, rank_zero_only=True)
print('logged (2)')
return loss

def configure_optimizers(self):
optimizer = AdamW(self.parameters(), lr=3e-4, eps=1e-8)
return optimizer

def train_dataloader(self):
return DataLoader(train_dataset, batch_size=self.batch_size,
num_workers=4)

def val_dataloader(self):
return DataLoader(validation_dataset,
batch_size=self.batch_size,
num_workers=4)

# Model Fine-Tuning
t5_model = T5ForConditionalGeneration.from_pretrained('t5-base')

model = T5Tuner(16, t5_model, tokenizer)
trainer = pl.Trainer(max_epochs=3, accelerator="tpu", devices=1)
trainer.fit(model)
```

### Error messages and logs

```
INFO:pytorch_lightning.utilities.rank_zero:GPU available: False, used: False
INFO:pytorch_lightning.utilities.rank_zero:TPU available: True, using: 1 TPU cores
INFO:pytorch_lightning.utilities.rank_zero:IPU available: False, using: 0 IPUs
INFO:pytorch_lightning.utilities.rank_zero:HPU available: False, using: 0 HPUs
```

### Environment

Current environment

```
#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): Trainer
#- PyTorch Lightning Version (e.g., 1.5.0): latest
#- Lightning App Version (e.g., 0.5.2): n/a
#- PyTorch Version (e.g., 2.0): latest. torch-xla version is 1.13
#- Python version (e.g., 3.9): unknown (kaggle tpu)
#- OS (e.g., Linux): Linux
#- CUDA/cuDNN version: unknown
#- GPU models and configuration: TPU x1
#- How you installed Lightning(`conda`, `pip`, source): !pip install lightning
#- Running environment of LightningApp (e.g. local, cloud): n/a

Current environment

* CUDA:
- GPU: None
- available: False
- version: 12.1
* Lightning:
- torch: 2.1.0
- torch-xla: 2.1.0+libtpu
- torchaudio: 2.1.0
- torchdata: 0.7.0
- torchtext: 0.16.0
- torchvision: 0.16.0
* Packages:
- absl-py: 1.4.0
- accelerate: 0.27.2
- aiofiles: 22.1.0
- aiosqlite: 0.20.0
- anyio: 4.3.0
- argon2-cffi: 23.1.0
- argon2-cffi-bindings: 21.2.0
- array-record: 0.5.0
- arrow: 1.3.0
- astroid: 3.0.3
- asttokens: 2.4.1
- astunparse: 1.6.3
- attrs: 23.2.0
- audioread: 3.0.1
- autopep8: 2.0.4
- babel: 2.14.0
- beautifulsoup4: 4.12.3
- bleach: 6.1.0
- cachetools: 5.3.2
- certifi: 2024.2.2
- cffi: 1.16.0
- charset-normalizer: 3.3.2
- chex: 0.1.85
- click: 8.1.7
- cloud-tpu-client: 0.10
- cloudpickle: 3.0.0
- comm: 0.2.1
- contourpy: 1.2.0
- cycler: 0.12.1
- debugpy: 1.8.1
- decorator: 5.1.1
- defusedxml: 0.7.1
- diffusers: 0.26.3
- dill: 0.3.8
- distrax: 0.1.5
- dm-haiku: 0.0.12.dev0
- dm-tree: 0.1.8
- docstring-to-markdown: 0.15
- entrypoints: 0.4
- etils: 1.7.0
- exceptiongroup: 1.2.0
- executing: 2.0.1
- fastjsonschema: 2.19.1
- filelock: 3.13.1
- flake8: 7.0.0
- flatbuffers: 23.5.26
- flax: 0.8.1
- fonttools: 4.49.0
- fqdn: 1.5.1
- fsspec: 2024.2.0
- funcsigs: 1.0.2
- gast: 0.5.4
- gin-config: 0.5.0
- google-api-core: 1.34.1
- google-api-python-client: 1.8.0
- google-auth: 2.28.1
- google-auth-httplib2: 0.2.0
- google-auth-oauthlib: 1.2.0
- google-pasta: 0.2.0
- googleapis-common-protos: 1.62.0
- grpcio: 1.62.0
- gym: 0.26.2
- gym-notices: 0.0.8
- h5py: 3.10.0
- httplib2: 0.22.0
- huggingface-hub: 0.20.3
- idna: 3.6
- importlib-metadata: 7.0.1
- importlib-resources: 6.1.1
- ipykernel: 6.29.2
- ipython: 8.22.0
- ipython-genutils: 0.2.0
- isoduration: 20.11.0
- isort: 5.13.2
- jax: 0.4.23
- jaxlib: 0.4.23
- jedi: 0.19.1
- jinja2: 3.1.3
- jmp: 0.0.4
- joblib: 1.3.2
- jraph: 0.0.6.dev0
- json5: 0.9.17
- jsonpointer: 2.4
- jsonschema: 4.21.1
- jsonschema-specifications: 2023.12.1
- jupyter-client: 7.4.9
- jupyter-core: 5.7.1
- jupyter-events: 0.9.0
- jupyter-lsp: 1.5.1
- jupyter-server: 2.12.5
- jupyter-server-fileid: 0.9.1
- jupyter-server-terminals: 0.5.2
- jupyter-server-ydoc: 0.8.0
- jupyter-ydoc: 0.2.5
- jupyterlab: 3.6.7
- jupyterlab-pygments: 0.3.0
- jupyterlab-server: 2.25.3
- kagglehub: 0.1.9
- keras: 3.0.5
- keras-cv: 0.8.2
- keras-nlp: 0.8.1
- kiwisolver: 1.4.5
- lazy-loader: 0.3
- libclang: 16.0.6
- librosa: 0.10.1
- libtpu-nightly: 0.1.dev20231213
- llvmlite: 0.42.0
- markdown: 3.5.2
- markdown-it-py: 3.0.0
- markupsafe: 2.1.5
- matplotlib: 3.8.3
- matplotlib-inline: 0.1.6
- mccabe: 0.7.0
- mdurl: 0.1.2
- mistune: 3.0.2
- ml-dtypes: 0.2.0
- mpmath: 1.3.0
- msgpack: 1.0.7
- nbclassic: 1.0.0
- nbclient: 0.9.0
- nbconvert: 7.16.1
- nbformat: 5.9.2
- nest-asyncio: 1.6.0
- networkx: 3.2.1
- notebook: 6.5.6
- notebook-shim: 0.2.4
- numba: 0.59.0
- numpy: 1.26.4
- nvidia-cublas-cu12: 12.1.3.1
- nvidia-cuda-cupti-cu12: 12.1.105
- nvidia-cuda-nvrtc-cu12: 12.1.105
- nvidia-cuda-runtime-cu12: 12.1.105
- nvidia-cudnn-cu12: 8.9.2.26
- nvidia-cufft-cu12: 11.0.2.54
- nvidia-curand-cu12: 10.3.2.106
- nvidia-cusolver-cu12: 11.4.5.107
- nvidia-cusparse-cu12: 12.1.0.106
- nvidia-nccl-cu12: 2.18.1
- nvidia-nvjitlink-cu12: 12.3.101
- nvidia-nvtx-cu12: 12.1.105
- oauth2client: 4.1.3
- oauthlib: 3.2.2
- opencv-python-headless: 4.9.0.80
- opt-einsum: 3.3.0
- optax: 0.1.9
- orbax-checkpoint: 0.4.4
- overrides: 7.7.0
- packaging: 23.2
- pandas: 2.2.0
- pandocfilters: 1.5.1
- papermill: 2.5.0
- parso: 0.8.3
- pexpect: 4.9.0
- pillow: 10.2.0
- pip: 23.0.1
- platformdirs: 4.2.0
- pluggy: 1.4.0
- pooch: 1.8.1
- prometheus-client: 0.20.0
- promise: 2.3
- prompt-toolkit: 3.0.43
- protobuf: 3.20.3
- psutil: 5.9.8
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pyasn1: 0.5.1
- pyasn1-modules: 0.3.0
- pycodestyle: 2.11.1
- pycparser: 2.21
- pydocstyle: 6.3.0
- pyflakes: 3.2.0
- pygments: 2.17.2
- pylint: 3.0.3
- pyparsing: 3.1.1
- python-dateutil: 2.8.2
- python-json-logger: 2.0.7
- python-lsp-jsonrpc: 1.1.2
- python-lsp-server: 1.10.0
- pytoolconfig: 1.3.1
- pytz: 2024.1
- pyyaml: 6.0.1
- pyzmq: 24.0.1
- referencing: 0.33.0
- regex: 2023.12.25
- requests: 2.31.0
- requests-oauthlib: 1.3.1
- rfc3339-validator: 0.1.4
- rfc3986-validator: 0.1.1
- rich: 13.7.0
- rope: 1.12.0
- rpds-py: 0.18.0
- rsa: 4.9
- safetensors: 0.4.2
- scikit-learn: 1.4.1.post1
- scipy: 1.12.0
- send2trash: 1.8.2
- setuptools: 65.5.1
- six: 1.16.0
- sniffio: 1.3.0
- snowballstemmer: 2.2.0
- soundfile: 0.12.1
- soupsieve: 2.5
- soxr: 0.3.7
- stack-data: 0.6.3
- sympy: 1.12
- tabulate: 0.9.0
- tenacity: 8.2.3
- tensorboard: 2.15.2
- tensorboard-data-server: 0.7.2
- tensorflow: 2.15.0
- tensorflow-datasets: 4.9.4
- tensorflow-estimator: 2.15.0
- tensorflow-hub: 0.16.1
- tensorflow-io: 0.36.0
- tensorflow-io-gcs-filesystem: 0.36.0
- tensorflow-metadata: 1.14.0
- tensorflow-probability: 0.23.0
- tensorflow-text: 2.15.0
- tensorstore: 0.1.45
- termcolor: 2.4.0
- terminado: 0.18.0
- tf-keras: 2.15.0
- threadpoolctl: 3.3.0
- tinycss2: 1.2.1
- tokenizers: 0.15.2
- toml: 0.10.2
- tomli: 2.0.1
- tomlkit: 0.12.3
- toolz: 0.12.1
- torch: 2.1.0
- torch-xla: 2.1.0+libtpu
- torchaudio: 2.1.0
- torchdata: 0.7.0
- torchtext: 0.16.0
- torchvision: 0.16.0
- tornado: 6.4
- tqdm: 4.66.2
- traitlets: 5.14.1
- transformers: 4.38.1
- trax: 1.4.1
- triton: 2.1.0
- types-python-dateutil: 2.8.19.20240106
- typing-extensions: 4.9.0
- tzdata: 2024.1
- ujson: 5.9.0
- uri-template: 1.3.0
- uritemplate: 3.0.1
- urllib3: 2.2.1
- wcwidth: 0.2.13
- webcolors: 1.13
- webencodings: 0.5.1
- websocket-client: 1.7.0
- werkzeug: 3.0.1
- whatthepatch: 1.0.5
- wheel: 0.42.0
- wrapt: 1.14.1
- y-py: 0.6.2
- yapf: 0.40.2
- ypy-websocket: 0.8.4
- zipp: 3.17.0
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor:
- python: 3.10.13
- release: 6.1.75+
- version: #1 SMP PREEMPT_DYNAMIC Fri Mar 1 15:14:26 UTC 2024

```

### More info

_No response_

cc @JackCaoG @Liyang90 @gkroiz

Start by reproducing the supplied notebook at trainer.fit(model), using the TPU configuration and the shown training_step and validation_step methods. Compare the behavior with the reported logs; done means training completes or reports an actionable error without hanging the notebook.

Written by the indexing model from the issue text.