Support serialized checkpoint loading [redo #9605] #10141
Conversation
jjenniferdai
requested review from
awaelchli,
Borda,
carmocca,
justusschock,
kaushikb11,
rohitgr7,
SeanNaren,
tchaton and
williamFalcon
as code owners
| del checkpoint["state_dict"] | ||
| self._has_loaded_state_dict = True | ||
| log.info(f"Rank {self.global_rank}: done loading model states from {checkpoint_path}.") | ||
| self.barrier() |
There was a problem hiding this comment.
Will this hang if the user runs any collective op in their on_load_checkpoint hook?
There was a problem hiding this comment.
good point I think yes ... is this resolvable though, how could we support collective op in model.on_load_checkpoint(checkpoint) if it's impossible to concurrently load then pass checkpoint for large models that OOM?
If it's not possible (please correct me if i'm wrong) - should we control serialized loading option with a boolean, and note that if serialized loading is set, model.on_load_checkpoint would not support collective ops?
There was a problem hiding this comment.
@carmocca good catch. Yes this will hang if the user on_load_checkpoint runs any collective.
There was a problem hiding this comment.
for large models that OOM?
I guess it's more of a problem for models that don't need this serialized loading. Because this PR makes serialization the default mechanism with no opt-out.
should we control serialized loading option with a boolean
This would probably be too ugly to implement.
This weak point seems very niche, as you would generally not run a collective op inside these hooks. Some possibilities that come to mind:
- IIRC, FB has mentioned in the past that they call
model.setupfrommodel.on_load_checkpointin some cases (sharded only?). I'd saysetupis more likely to run such ops. - Naive property accesses (like
self.trainer.log_dir) can runbroadcast.
For this PR, I'd suggest adding a comment here about this weak point. But it's something to remember if we see such hangs.
Perhaps this boolean would be easier to implement and manage after the collective refactor.
cc @awaelchli
There was a problem hiding this comment.
IIRC, FB has mentioned in the past that they call model.setup from model.on_load_checkpoint in some cases (sharded only?). I'd say setup is more likely to run such ops.
It's not allowed to call setup() in this hook. With the change we made to the loading logic in #????, we removed the need to call setup() in that hook.
But generally I agree, it is going to be hard to debug hangs as this mechanism may not be obvious. Adding a comment would help minimally.
There was a problem hiding this comment.
That sounds great @yifuwang !!
There was a problem hiding this comment.
This could also resolve the current bug where the checkpoint doesn't exist on some nodes when the filesystem isn't shared.
There was a problem hiding this comment.
@yifuwan, Would you mind opening an issue for this proposal?
There was a problem hiding this comment.
Should we close this PR and use the proposed logic ?
There was a problem hiding this comment.
@yifuwang @jjenniferdai Who would make a PR with the above proposition?
| return | ||
| self.lightning_module.load_state_dict(checkpoint["state_dict"]) | ||
|
|
||
| def load_checkpoint(self, checkpoint_path: _PATH) -> Dict[str, Any]: |
There was a problem hiding this comment.
can we test this by mocking the hooks called and asserting the order? also the warning/info message
|
@jjenniferdai What is the status with this PR ? |
|
if everyone's onboard I think we can close this PR, I'll work with Yifu on the proposed alternative |
What does this PR do?
Fixes #9406: redo #9605 now that #10061 resolves #10057.
Does your PR introduce any breaking changes? If yes, please list them.
Before submitting
PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:
Did you have fun?
Make sure you had fun coding 🙃