Add lazy checkpoint loading for FSDP full-state checkpoints #18150
Conversation
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
awaelchli
changed the title
awaelchli
changed the title
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
awaelchli
changed the title
for more information, see https://pre-commit.ci
It looks like it is difficult to integrate it into the rest of the code base. Several things are standing in the way that require us to brainstorm on design decisions. Here are some raw notes from my notebook:
Happy to open an issue about this for further discussion!
There might be slight differences between the version here and in lit-gpt. I need to double check, there was something about quantization. |
| # TODO: needed for us? | ||
| # materializing with contiguous is needed for quantization | ||
| if name in {"contiguous"}: | ||
| return getattr(self._load_tensor(), name) |
There was a problem hiding this comment.
I should probably keep this here if we want to import and use the util in lit-gpt?
There was a problem hiding this comment.
It might also be a generally applicable limitation. On the other hand, we could do this manually in the quantization code directly.
What do you suggest @t-vi?
What does this PR do?
Addresses (2) in #18008 (comment)
Closes #18138
Before submitting
PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:
Reviewer checklist
cc @Borda @justusschock @awaelchli @carmocca