Deadlock after training model with TPU VM 3.8 and while resuming training.

Nobody has claimed this yet.

## 🐛 Bug

As the title said, the deadlock occurs twice. After training, the process does not terminate itself. When I load checkpoint to resume training, it takes too much of time.

### To Reproduce

boring_transformer.py to produce 1st deadlock
```python
# %%
from datetime import datetime
from typing import Optional

import datasets
import torch
from pytorch_lightning import LightningDataModule, LightningModule, Trainer, seed_everything
from pytorch_lightning.loggers import CSVLogger, TensorBoardLogger
from pytorch_lightning.callbacks import RichProgressBar, ModelCheckpoint

from torch.utils.data import DataLoader
from transformers import (
AdamW,
AutoConfig,
AutoModelForSequenceClassification,
AutoTokenizer,
get_linear_schedule_with_warmup,
)

# %%
class GLUEDataModule(LightningDataModule):

task_text_field_map = {
"cola": ["sentence"],
"sst2": ["sentence"],
"mrpc": ["sentence1", "sentence2"],
"qqp": ["question1", "question2"],
"stsb": ["sentence1", "sentence2"],
"mnli": ["premise", "hypothesis"],
"qnli": ["question", "sentence"],
"rte": ["sentence1", "sentence2"],
"wnli": ["sentence1", "sentence2"],
"ax": ["premise", "hypothesis"],
}

glue_task_num_labels = {
"cola": 2,
"sst2": 2,
"mrpc": 2,
"qqp": 2,
"stsb": 1,
"mnli": 3,
"qnli": 2,
"rte": 2,
"wnli": 2,
"ax": 3,
}

loader_columns = [
"datasets_idx",
"input_ids",
"token_type_ids",
"attention_mask",
"start_positions",
"end_positions",
"labels",
]

def __init__(
self,
model_name_or_path: str,
task_name: str = "mrpc",
max_seq_length: int = 128,
train_batch_size: int = 32,
eval_batch_size: int = 32,
**kwargs,
):
super().__init__()
self.model_name_or_path = model_name_or_path
self.task_name = task_name
self.max_seq_length = max_seq_length
self.train_batch_size = train_batch_size
self.eval_batch_size = eval_batch_size

self.text_fields = self.task_text_field_map[task_name]
self.num_labels = self.glue_task_num_labels[task_name]
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)

def setup(self, stage: str):
self.dataset = datasets.load_dataset("glue", self.task_name)

for split in self.dataset.keys():
self.dataset[split] = self.dataset[split].map(
self.convert_to_features,
batched=True,
remove_columns=["label"],
)
self.columns = [c for c in self.dataset[split].column_names if c in self.loader_columns]
self.dataset[split].set_format(type="torch", columns=self.columns)

self.eval_splits = [x for x in self.dataset.keys() if "validation" in x]

def prepare_data(self):
datasets.load_dataset("glue", self.task_name)
AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)

def train_dataloader(self):
return DataLoader(self.dataset["train"], batch_size=self.train_batch_size, shuffle=True)

def val_dataloader(self):
if len(self.eval_splits) == 1:
return DataLoader(self.dataset["validation"], batch_size=self.eval_batch_size)
elif len(self.eval_splits) > 1:
return [DataLoader(self.dataset[x], batch_size=self.eval_batch_size) for x in self.eval_splits]

def test_dataloader(self):
if len(self.eval_splits) == 1:
return DataLoader(self.dataset["test"], batch_size=self.eval_batch_size)
elif len(self.eval_splits) > 1:
return [DataLoader(self.dataset[x], batch_size=self.eval_batch_size) for x in self.eval_splits]

def convert_to_features(self, example_batch, indices=None):

# Either encode single sentence or sentence pairs
if len(self.text_fields) > 1:
texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]]))
else:
texts_or_text_pairs = example_batch[self.text_fields[0]]

# Tokenize the text/text pairs
features = self.tokenizer.batch_encode_plus(
texts_or_text_pairs, max_length=self.max_seq_length, pad_to_max_length=True, truncation=True
)

# Rename label to labels to make it easier to pass to model forward
features["labels"] = example_batch["label"]

return features

# %%
class GLUETransformer(LightningModule):
def __init__(
self,
model_name_or_path: str,
num_labels: int,
task_name: str,
learning_rate: float = 2e-5,
adam_epsilon: float = 1e-8,
warmup_steps: int = 0,
weight_decay: float = 0.0,
train_batch_size: int = 32,
eval_batch_size: int = 32,
eval_splits: Optional[list] = None,
**kwargs,
):
super().__init__()

self.save_hyperparameters()

self.config = AutoConfig.from_pretrained(model_name_or_path, num_labels=num_labels)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path, config=self.config)
self.metric = datasets.load_metric(
"glue", self.hparams.task_name, experiment_id=datetime.now().strftime("%d-%m-%Y_%H-%M-%S")
)

def forward(self, **inputs):
return self.model(**inputs)

def training_step(self, batch, batch_idx):
outputs = self(**batch)
loss = outputs[0]
return loss

def validation_step(self, batch, batch_idx, dataloader_idx=0):
outputs = self(**batch)
val_loss, logits = outputs[:2]

if self.hparams.num_labels > 1:
preds = torch.argmax(logits, axis=1)
elif self.hparams.num_labels == 1:
preds = logits.squeeze()

labels = batch["labels"]

return {"loss": val_loss, "preds": preds, "labels": labels}

def validation_epoch_end(self, outputs):
if self.hparams.task_name == "mnli":
for i, output in enumerate(outputs):
# matched or mismatched
split = self.hparams.eval_splits[i].split("_")[-1]
preds = torch.cat([x["preds"] for x in output]).detach().cpu().numpy()
labels = torch.cat([x["labels"] for x in output]).detach().cpu().numpy()
loss = torch.stack([x["loss"] for x in output]).mean()
self.log(f"val_loss_{split}", loss, prog_bar=True)
split_metrics = {
f"{k}_{split}": v for k, v in self.metric.compute(predictions=preds, references=labels).items()
}
self.log_dict(split_metrics, prog_bar=True)
return loss

preds = torch.cat([x["preds"] for x in outputs]).detach().cpu().numpy()
labels = torch.cat([x["labels"] for x in outputs]).detach().cpu().numpy()
loss = torch.stack([x["loss"] for x in outputs]).mean()
self.log("val_loss", loss, prog_bar=True)
self.log_dict(self.metric.compute(predictions=preds, references=labels), prog_bar=True)

def configure_optimizers(self):
"""Prepare optimizer and schedule (linear warmup and decay)"""
model = self.model
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
{
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
"weight_decay": self.hparams.weight_decay,
},
{
"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
"weight_decay": 0.0,
},
]
optimizer = AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)

scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=self.hparams.warmup_steps,
num_training_steps=self.trainer.estimated_stepping_batches,
)
scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
return [optimizer], [scheduler]

# %%
seed_everything(42)

dm = GLUEDataModule(model_name_or_path="albert-base-v2", task_name="cola")
dm.setup("fit")
model = GLUETransformer(
model_name_or_path="albert-base-v2",
num_labels=dm.num_labels,
eval_splits=dm.eval_splits,
task_name=dm.task_name,
)

LOG_DIR = 'exp_logs'
trainer = Trainer(
default_root_dir=LOG_DIR, logger=[CSVLogger(LOG_DIR), TensorBoardLogger(LOG_DIR)],
max_epochs=1,
accelerator='tpu', devices=8,
callbacks=[RichProgressBar(), ModelCheckpoint(dirpath=LOG_DIR, every_n_train_steps=1)],
strategy="tpu_spawn_debug",
precision='bf16',
profiler='xla',
limit_train_batches=10
)
trainer.fit(model, datamodule=dm)
```

To produce second deadlock, replace last few lines of boring_transformer.py with these lines.
```python
trainer = Trainer(
default_root_dir=LOG_DIR, logger=[CSVLogger(LOG_DIR), TensorBoardLogger(LOG_DIR)],
max_epochs=1,
accelerator='tpu', devices=8,
callbacks=[RichProgressBar(), ModelCheckpoint(dirpath=LOG_DIR, every_n_train_steps=1)],
strategy="tpu_spawn_debug",
precision='bf16',
profiler='xla',
limit_train_batches=20
)
trainer.fit(model, datamodule=dm, ckpt_path='exp_logs/epoch=0-step=10.ckpt')
```

### Expected behavior

No deadlock

### Environment

```
* CUDA:
- GPU: None
- available: False
- version: 10.2
* Lightning:
- pytorch-lightning: 1.6.5
- torch: 1.11.0
- torch-tb-profiler: 0.4.0
- torch-xla: 1.11
- torchinfo: 1.7.0
- torchmetrics: 0.9.3
- torchvision: 0.12.0
* Packages:
- absl-py: 1.2.0
- aiohttp: 3.8.1
- aiosignal: 1.2.0
- argon2-cffi: 21.3.0
- argon2-cffi-bindings: 21.2.0
- astroid: 2.11.7
- asttokens: 2.0.5
- async-timeout: 4.0.2
- attrs: 22.1.0
- backcall: 0.2.0
- beautifulsoup4: 4.11.1
- bleach: 5.0.1
- cachetools: 4.2.4
- certifi: 2022.6.15
- cffi: 1.15.1
- charset-normalizer: 2.1.0
- cloud-tpu-client: 0.10
- cloud-tpu-profiler: 2.4.0
- commonmark: 0.9.1
- datasets: 2.4.0
- debugpy: 1.6.2
- decorator: 5.1.1
- defusedxml: 0.7.1
- dill: 0.3.5.1
- einops: 0.4.1
- entrypoints: 0.4
- executing: 0.9.1
- fastjsonschema: 2.16.1
- filelock: 3.7.1
- flake8: 5.0.4
- frozenlist: 1.3.1
- fsspec: 2022.7.1
- google-api-core: 1.32.0
- google-api-python-client: 1.8.0
- google-auth: 1.35.0
- google-auth-httplib2: 0.1.0
- google-auth-oauthlib: 0.4.6
- googleapis-common-protos: 1.56.4
- grpcio: 1.47.0
- httplib2: 0.20.4
- huggingface-hub: 0.8.1
- hurry.filesize: 0.9
- idna: 3.3
- importlib-metadata: 4.12.0
- importlib-resources: 5.9.0
- install: 1.3.5
- ipykernel: 6.15.1
- ipython: 8.4.0
- ipython-genutils: 0.2.0
- ipywidgets: 7.7.1
- isort: 5.10.1
- jedi: 0.18.1
- jinja2: 3.1.2
- joblib: 1.1.0
- jsonschema: 4.9.1
- jupyter-client: 7.3.4
- jupyter-core: 4.11.1
- jupyterlab-pygments: 0.2.2
- jupyterlab-widgets: 1.1.1
- lazy-object-proxy: 1.7.1
- libtpu-nightly: 0.1.dev20220303
- markdown: 3.4.1
- markupsafe: 2.1.1
- matplotlib-inline: 0.1.3
- mccabe: 0.7.0
- mistune: 0.8.4
- multidict: 6.0.2
- multiprocess: 0.70.13
- nbclient: 0.6.6
- nbconvert: 6.5.0
- nbformat: 5.4.0
- nest-asyncio: 1.5.5
- notebook: 6.4.12
- numpy: 1.23.1
- oauth2client: 4.1.3
- oauthlib: 3.2.0
- packaging: 21.3
- pandas: 1.4.3
- pandocfilters: 1.5.0
- parso: 0.8.3
- pexpect: 4.8.0
- pickleshare: 0.7.5
- pillow: 9.2.0
- pip: 22.1.2
- pkgutil-resolve-name: 1.3.10
- platformdirs: 2.5.2
- prometheus-client: 0.14.1
- prompt-toolkit: 3.0.30
- protobuf: 3.19.4
- psutil: 5.9.1
- ptyprocess: 0.7.0
- pure-eval: 0.2.2
- pyarrow: 9.0.0
- pyasn1: 0.4.8
- pyasn1-modules: 0.2.8
- pycodestyle: 2.9.1
- pycparser: 2.21
- pydeprecate: 0.3.2
- pyflakes: 2.5.0
- pygments: 2.12.0
- pylint: 2.14.5
- pyparsing: 3.0.9
- pyrsistent: 0.18.1
- python-dateutil: 2.8.2
- pytorch-lightning: 1.6.5
- pytz: 2022.1
- pyyaml: 6.0
- pyzmq: 23.2.0
- regex: 2022.7.25
- requests: 2.28.1
- requests-oauthlib: 1.3.1
- responses: 0.18.0
- rich: 12.5.1
- rsa: 4.9
- scikit-learn: 1.1.2
- scipy: 1.9.0
- send2trash: 1.8.0
- sentencepiece: 0.1.97
- setuptools: 61.2.0
- six: 1.16.0
- sklearn: 0.0
- soupsieve: 2.3.2.post1
- stack-data: 0.3.0
- tablign: 0.3.4
- tensorboard: 2.10.0
- tensorboard-data-server: 0.6.1
- tensorboard-plugin-wit: 1.8.1
- terminado: 0.15.0
- threadpoolctl: 3.1.0
- tinycss2: 1.1.1
- tokenizers: 0.12.1
- tomli: 2.0.1
- tomlkit: 0.11.1
- torch: 1.11.0
- torch-tb-profiler: 0.4.0
- torch-xla: 1.11
- torchinfo: 1.7.0
- torchmetrics: 0.9.3
- torchvision: 0.12.0
- tornado: 6.2
- tqdm: 4.64.0
- traitlets: 5.3.0
- transformers: 4.21.1
- typing-extensions: 4.3.0
- uritemplate: 3.0.1
- urllib3: 1.26.11
- wcwidth: 0.2.5
- webencodings: 0.5.1
- werkzeug: 2.2.2
- wheel: 0.37.1
- widgetsnbextension: 3.6.1
- xxhash: 3.0.0
- yarl: 1.8.1
- zipp: 3.8.1
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.8.13
- version: #46-Ubuntu SMP Mon Apr 19 19:17:04 UTC 2021
```

### Additional context

```
export TPU_LOG_DIR="disabled"
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
export PT_XLA_DEBUG=1
export USE_TORCH=ON
export TOKENIZERS_PARALLELISM=false
```

cc @carmocca @JackCaoG @Liyang90 @gkroiz @kaushikb11 @rohitgr7

Start with the boring_transformer.py reproducer and run both trainer.fit cases on a TPU VM 3.8, first without and then with ckpt_path='exp_logs/epoch=0-step=10.ckpt'. Trace the TPU/XLA strategy and checkpoint-resume paths involved in the hang. Done means training exits after the first run and resume training completes without a deadlock.

Written by the indexing model from the issue text.