Colab TPU training stuck at end of epoch if checkpoint_callback=True

[ktrapeznikov]

I am seeing training stuck at the last step of the 1st epoch when checkpoint_callback=True is enabled. I am passing tpu_cores=8 to the trainer.

If it's False training is still slow for the first few steps and then it speeds up (I guess that this is expected since it takes a few steps for xla to compile stuff).

Here is a link to the colab notebook example.

### Setup Model
class GPT2Tuned(pl.LightningModule):
    def __init__(self, hparams: argparse.Namespace):
        super().__init__()
        self.hparams = hparams
        self.model = AutoModelForCausalLM.from_pretrained(self.hparams.model_name_or_path)

    def forward(self, **inputs):
        return self.model(**inputs)
        
    def _step(self,batch):
        inputs = {"input_ids": batch[0]}
        inputs["labels"] = batch[0]
        outputs = self(**inputs)
        loss = outputs[0]
        return dict(loss=loss)

    def training_step(self, batch, batch_idx):
        out_dict = self._step(batch)
        return dict(loss=out_dict["loss"], log=out_dict)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

    def train_dataloader(self):
      # fake data
        dataset = TensorDataset(torch.randint(10000,(10000,128)).long()) 
        return  DataLoader(
            dataset,
            batch_size=self.hparams.batch_size,
            shuffle=True
        )

    @staticmethod
    def add_model_specific_args(parser):
        parser.add_argument("--model_name_or_path",type=str)
        parser.add_argument("--batch_size", default=4, type=int)
        parser.add_argument("--learning_rate", default=5e-5, type=float)
        return parser

#### Training 
parser = argparse.ArgumentParser()
parser = GPT2Tuned.add_model_specific_args(parser)
parser = pl.Trainer.add_argparse_args(parser)
args = parser.parse_args(["--gpus","0"])

args.__dict__.update(dict(model_name_or_path  = "gpt2-medium",
                          batch_size = 3,
                          precision=16,
                          tpu_cores=8))

model = GPT2Tuned(args)
trainer = pl.Trainer.from_argparse_args(args, checkpoint_callback=True)  
trainer.fit(model)

[github-actions]

Hi! thanks for your contribution!, great first issue!

[awaelchli]

Thanks for the reproduction script.
Could be related #2700
There is a PR for fixing checkpointing on TPU: #2726 (need check if this applies)

[ktrapeznikov]

Thanks. Seems related. Is there another mechanism to save checkpoints as a temporary fix?

[Borda]

@ktrapeznikov mind checking it now as Google Colab had an internal issue with TPUs in past days...

[ktrapeznikov]

Last time I tried... same issue

[evanatyourservice]

This would be a nice feature to have since I would rarely do inference on TPU, just training.

[edenlightning]

@lezwon maybe you can help revive the tpu + checkpoint fix?

[lezwon]

Been working on it. Finding it a bit tricky. I'll raise a PR with the current work I've done.

[sarmientoj24]

hi, i am experiencing being stuck on 0% at first epoch when using TPU with pytorch-lightning. any possible reasons why?

[lezwon]

@sarmientoj24 mind share the notebook?