Improve CI tests #87
Description
The quality of CI tests needs to be improved
- Fine-grained tests for pipeline parallelism. Current test is very coarse grained and checks the training loss at the end of 10 epochs. We should have more fine-grained tests that check the outputs of an individual forward and backward pass.
- Some tests for the new lightning strategy.
- Do a sweep over pytorch and python versions.
- Move intra-layer tests to torchrun