Fix checkpoint loading and resume for fsdp

[achalddave]

This partially reverts #138. If we use FSDP, we should call load_model directly, instead of model.module.load_model. I tested training and resuming a model using main.py and it works with this commit.

Note that the tests in #148 pass despite this bug because of another bug that @kernelmachine discovered incidentally: Our tests are not properly using distributed functionality, because

open_lm/open_lm/distributed.py

Lines 20 to 25 in 6570b81

	def is_using_distributed():
	if "WORLD_SIZE" in os.environ:
	return int(os.environ["WORLD_SIZE"]) > 1
	if "SLURM_NTASKS" in os.environ:
	return int(os.environ["SLURM_NTASKS"]) > 1
	return False

only checks if world_size > 1 (instead of >= 1). This will be tracked in a follow up issue.

[sagadre]

can we actually get a version bump on this?