Add test for checkpoint loading after save

[achalddave]

We should add a test that:

1. trains a (small) model for a couple steps
2. saves it to disk
3. calls main() again with a path to the checkpoint on disk
4. trains a few steps

The test should test single process, DDP, and FSDP.

[jmercat]

I cannot assign this to myself but I can do it

[Vaishaal]

This is the error I am getting right now:

Traceback (most recent call last):  
File "/miniconda/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/miniconda/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/mnt/task_runtime/open_lm/open_lm/main.py", line 798, in <module>
    main(sys.argv[1:])
  File "/mnt/task_runtime/open_lm/open_lm/main.py", line 502, in main
    start_epoch, global_step = load_model(args, model)
  File "/mnt/task_runtime/open_lm/open_lm/main.py", line 111, in load_model
    model.module.load_state_dict(sd)
  File "/miniconda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for Transformer:
	Missing key(s) in state_dict: "_flat_param".
	Unexpected key(s) in state_dict: "tok_embeddings.weight", "norm.weight", "norm.bias", "output.weight".
	```

[achalddave]

Closing this issue as we have added a test in #148, but @Vaishaal feel free to open a new issue if you reproduce the checkpoint loading bug you have