Custom checkpoint : torchscript and ONNX export

[pfeatherstone]

I would like a custom checkpoint callback which:

• Converts model to torchscript using tracing, and saves
• Converts torchscript model to ONNX and saves

So far i have:

class JitCheckpoint(Callback):
    def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"):
        if trainer.is_global_zero:
            print("Validation ended")
            .... Do stuff ... here

How do i convert pl_module to normal pytorch torch.nn.Module ?
Huggingface Accelerate library has accelerator.unwrap_model(). Does lightning have something similar?
How do I synchronise all the nodes. Huggingface Accelerate has accelerator.wait_for_everyone(), which presumably I would need since checkpointing only happens on one node.

[pfeatherstone]

So far I have this:

class JitCheckpoint(Checkpoint):
    def __init__(self):
        self.lastLoss = 100000

    def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"):
        if trainer.is_global_zero:
            assert "loss/val" in trainer.callback_metrics, "monitor not in callback_metrics"
            assert "acc/val" in trainer.callback_metrics, "monitor not in callback_metrics"
            loss        = trainer.callback_metrics["loss/val"].item()
            acc         = trainer.callback_metrics["acc/val"].item()
            if loss < self.lastLoss:
                print("\nSaving checkpoint last loss {} new loss {}".format(self.lastLoss, loss))
                self.lastLoss = loss
                ckpdir      = trainer.log_dir + "/checkpoints/"
                filePath    = ckpdir + "epoch={}_loss={:.4f}_acc={:.2f}".format(trainer.current_epoch, loss, acc)
                jitPath     = filePath + '.pt'
                onnxPath    = filePath + '.onnx'
                os.makedirs(ckpdir, exist_ok=True)

                input = torch.randn(2, 20*1024, 2)
                pl_module.to_torchscript(file_path=jitPath, method='trace', example_inputs=[input])
                pl_module.to_onnx(file_path=onnxPath, input_sample=input)
        trainer.strategy.barrier()

But when I train on multi-GPU i get a deadlock the first time it calls this checkpoint, it hangs for ever spinning both my GPUs at 100% doing nothing. And that's despite calling trainer.strategy.barrier(). Can someone help?

[pfeatherstone]

@williamFalcon

[pfeatherstone]

Note, on single GPU or CPU it works fine.