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().
|
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