add option to disable ft checkpoints #1795
Conversation
Differential Revision: D83846965
tushar00jain
requested review from
tianyu-l,
fegin,
wwwjn and
wconstab
as code owners
meta-cla
bot
added
the
CLA Signed
|
@tushar00jain has exported this pull request. If you are a Meta employee, you can view the originating Diff in D83846965. |
| enable: bool = False | ||
| """Whether to enable checkpoint""" | ||
|
|
||
| enable_ft: bool = True |
There was a problem hiding this comment.
enable_ft is too vague. I got confused that this is a similar thing as FT.enabled. Can we have a longer and more meaningful naming?
| def load_state_dict(self, state_dict: dict[str, Any]): | ||
| self.step = state_dict["step"] | ||
| self.ntokens_seen = state_dict["ntokens_seen"] | ||
| if not self.job_config.checkpoint.enable_ft and self.ft_manager is not None: |
There was a problem hiding this comment.
Setting dataloader state_dict is not safe. data.loader_state_dict() may be invoked before or after this load_state_dict().
There was a problem hiding this comment.
Setting dataloader state_dict is not safe.
Not safe to call inside this method you mean? Why is that the case? And it's safe to call after this load_state_dict returns?
The method doesn't set data loader state dict right now but we could set it manually based on step count. What's the best way to do that? Could put that in a separate PR.
Essentially with this change, users don't really have to rely on any external storage. So it reduces set up time to get things up and running. Since we also don't really need model checkpoints when we have torchft. And if checkpoint storage has issues, this can work as a killswitch to completely disable the storage so it doesn't impact training.
|
I would also suggest that you should not try to do internal code first. The internal TorchTitan is always days behind the OSS one. So it can cause some merge issues. |
|
|
||
| enable_ft: bool = True | ||
| """ | ||
| Checkpoints data loader state if enabled. Otherwise infers the data loader |
There was a problem hiding this comment.
The description is a little bit vague to me, how does this option helps checkpoint?Do I have to enable FT.enabled to use this?
Differential Revision: D83846965