[RFC] Integration with DCP - Benchmark Results #443
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
✅ Deploy Preview for torchtune-preview ready!
To edit notification comments on pull requests, go to your Netlify site configuration. |
rohan-varma
reviewed
Mar 4, 2024
rohan-varma
reviewed
Mar 4, 2024
LucasLLC
changed the title
fegin
reviewed
Mar 12, 2024
| self.log(msg=f"Waited for {time.monotonic() - t0} seconds.") | ||
|
|
||
| t0 = time.monotonic() | ||
| with FSDP.state_dict_type( |
There was a problem hiding this comment.
We should use get_state_dict(), which will avoid future migration when FSDP2 is used.
| ckpt_dict[MODEL_KEY] = self._model.state_dict() | ||
|
|
||
| if OPT_KEY in ckpt_dict: | ||
| ckpt_dict[OPT_KEY] = FSDP.optim_state_dict(self._model, self._optimizer) |
There was a problem hiding this comment.
ditto, we should use get_optimizer_state_dict().
| self.log(msg=f"creating sharded state_dict {time.monotonic() - t0} seconds.") | ||
|
|
||
| t0 = time.monotonic() | ||
| self._checkpoint_future = async_save( |
There was a problem hiding this comment.
nice! Good to see async_save is actually being used.
|
Superseded by #2006 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Context
We've pivoted on this PR from an RFC on torchtune + dcp UX, to a PR which showcases async saving benchmarks.
If consensus is reached that we want an
async_checkpointingintegration in torchtune, I'll pull the latest GPU recipes and implement this code in a less-hacky way there. Please don't review this PR for code cleanliness as it's a bit of a dumpster fire in the current state.Results:
full_finetune:
What's shown here is the blocking portion of saving -- or specifically the portion that blocks the training loop. LLM training is a great case for async checkpointing since the training loops are fairly long. This means by the time we reach the end of the epoch, serialization has finished and we never have to wait on the previous checkpoint before staging another one.
In the
full_tunecase, we will also notice an inverse relation to main in terms of ranks. As opposed to the current implementation, which increases latency with number of ranks, we expect the sharded state dict to be smaller per rank as we increase world_size, a fact that is taken advantage of in DCP. Meaning: "more GPUs, faster save"lora_finetune:
In the Lora case, we'll notice performance is actually at parity or worst for sharded state dict + DCP. There's a couple reasons for this:
Most of the time spent saving the lora checkpoints is actually spent in FSDP's
state_dictcalls. Corresponding to the graph above, on the main branch with 2 ranks we see full state dict takes ~0.624seconds. Sharded state dict can take slightly longer to materialize.The file sizes are really small, looks like 16MB - 48MB. Essentially async overhead will be almost equivalent to the serialization time of just using
torch.save( at least that's what I'm seeing in this case)Proposal:
For intermediary checkpoints on multi-gpu recipes, we should implement DCP.async_save as the checkpointing solution.
Some decisions points (comments are welcome):
We can go about converting to an un-sharded
torch.savefile in two ways. We can either :a) Go back to the slower
torch.save+ full state dict implementation on the last call to save in the last epochb) Save in the DCP format always, and ask users to convert themselves using
torch.distributed.checkpoint.format_utilsUsing the DCP checkpointer for
lora_finetunewould keep things consistent, but could incur some (pretty minimal) performance hits. Any thoughts?Changelog
Test plan