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

🐛 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

# %%
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.

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