[Roadmap] GraphGym via PyTorch Lightning 🚀

[aniketmaurya]

🚀 The feature, motivation and pitch

PyTorch Lightning Integration

GraphGym training experience can be improved for scalability, mixed precision support, logging and checkpoints with PyTorch Lightning integration.

• Write LightningModule
• Refactor load_ckpt and save_ckpt with PL checkpoint save and load method
• Update train method with PL Trainer and the LightningModule implementation
• Integrate LightningDataset, LightningNodeData and LightningLinkData modules
• ...

[rusty1s]

Thanks for this effort. I think the roadmap looks pretty good.

• LightningModule: Ideally, all we need to do is to let create_model return a LightningModule (first PR). This shouldn't break anything as far as I can tell. We can then start adding train_step, val_step and test_step, and configure_optimizers as part of this newly created LightningModule (second PR).
• As far as I can tell, load_ckpt and save_ckpt could be completely dropped by using PL in-house implementations. Correct me if I am wrong, but I think this has to be done after the LightningModule and its training pipeline have been properly integrated.
• We can also think about wrapping datasets inside our LightningDataModule.

cc: @JiaxuanYou

[rusty1s]

Closing this issue in favor of #5132. Thanks for your contributions @aniketmaurya. Hope you are also on board on future refactors of GraphGym :)