Add tests to test each component

[wanchaol]

The tentative tests we could add:

1. test the llama debug model init and forward/backward works
2. test checkpoint save/load works
3. metrics logging test (metrics to be added)

[tianyu-l]

TODO: test correctness of checkpointable data loading in #279

[tianyu-l]

@wconstab is going to help set up an example multi-gpu unit test. We can add more functionality unit test following the example.