[PyTorch] Distributed intermediate/activation tensors for FSDP #687
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
denera
force-pushed
the
databricks/distribute-fp8-weights-fsdp
branch
from
d18b49f
to
71f696b
Compare
denera
changed the title
rahul003
reviewed
Mar 6, 2024
| fsdp_states, fsdp_modules = _get_fsdp_states_with_modules(fsdp_root) | ||
| for state, module in zip(fsdp_states, fsdp_modules): | ||
| if _is_te_module(module): | ||
| setattr(module, "fsdp_wrapepd", True) |
denera
force-pushed
the
databricks/distribute-fp8-weights-fsdp
branch
from
1b822f4
to
9d8a7f5
Compare
denera
force-pushed
the
databricks/distribute-fp8-weights-fsdp
branch
2 times, most recently
from
0267b08
to
d62330d
Compare
|
/te-ci pytorch |
denera
force-pushed
the
databricks/distribute-fp8-weights-fsdp
branch
2 times, most recently
from
b260d50
to
8c6e9b7
Compare
|
/te-ci pytorch |
timmoon10
reviewed
May 23, 2024
…s distribute their activations after the forward pass and gather them before the backward pass Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
…or base TE modules Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
…ass and gathered before forward Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
denera
force-pushed
the
databricks/distribute-fp8-weights-fsdp
branch
from
7cb919b
to
518df99
Compare
|
/te-ci pytorch |
|
/te-ci pytorch |
ksivaman
reviewed
Jun 5, 2024
| @@ -856,3 +865,110 @@ def allreduce( | |||
| handle = torch.distributed.all_reduce(input_, group=tp_group, async_op=async_op) | |||
|
|
|||
| return input_, handle | |||
|
|
|||
|
|
|||
| def _fsdp_scatter_tensors( | |||
There was a problem hiding this comment.
Interesting that the linter doesn't complain about the missing docstring, is this due to the __all__ decl or just using internal function convention with _func_name? Either way I think this is good practice going forward as well instead of adding filler docs!
There was a problem hiding this comment.
Yes, I believe PyLint ignores some warnings by default for internal functions designated with the leading underscore.
|
/te-ci pytorch |
|
te-ci |
|
/te-ci |
|
/te-ci pytorch |
|
Pipeline 15637774 |
phu0ngng
pushed a commit
to phu0ngng/TransformerEngine
that referenced
this pull request
Jun 11, 2024
…A#687) * New TE wrapper for PyTorch FullyShardedDataParallel to make TE modules distribute their activations after the forward pass and gather them before the backward pass Signed-off-by: Alp Dener <adener@nvidia.com> * simplified TE module setup for FSDP comms Signed-off-by: Alp Dener <adener@nvidia.com> * FSDP scatter/gather for tensors saved into autograd ctx now working for base TE modules Signed-off-by: Alp Dener <adener@nvidia.com> * make sure activation recompute disables FSDP scatter/gather Signed-off-by: Alp Dener <adener@nvidia.com> * make sure Fp8 weight buffers are sharded at the end of the backward pass and gathered before forward Signed-off-by: Alp Dener <adener@nvidia.com> * Fixed typo in attribute name Signed-off-by: Alp Dener <adener@nvidia.com> * fixed bug in finding FSDP-wrapped TE modules Signed-off-by: Alp Dener <adener@nvidia.com> * fixed typo in fp8 weight tensor name Signed-off-by: Alp Dener <adener@nvidia.com> * fixed incorrect # of gradients Signed-off-by: Alp Dener <adener@nvidia.com> * Added fp8 amax gradient hook tensor to the parameter reset Signed-off-by: Alp Dener <adener@nvidia.com> * get rid of erroneous dummy tensor leftover from incorrect rebase Signed-off-by: Alp Dener <adener@nvidia.com> * Linting fixes Signed-off-by: Alp Dener <adener@nvidia.com> * fixing git snafu and removing debug statements Signed-off-by: Alp Dener <adener@nvidia.com> --------- Signed-off-by: Alp Dener <adener@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
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.
torch.distributed.fsdp.FullyShardedDataParallelcannot scatter/gather the intermediate/activation tensors that TE modules pack into the autograd context at the end of their forward passes, resulting in globally sized activation and Fp8 weight tensors staying in memory.This PR provides a
te.distributed.prepare_te_modules_for_fsdp(fsdp_root)API that inserts references to the correct FSDP process group into FSDP-wrapped TE modules in a given model. The TE modules then use these process groups to scatter the intermediate/activation tensors at the end of the forward pass before packing them into the autograd context. The same tensors are gathered in the beginning of the backward pass before compute.Using
te.distributed.checkpoint()turns off the scatters/gathers to avoid unnecessary comm for tensors that need to be recomputed anyway.nn.Sequential( 3 x te.LayerNormMLP )before Fp8/intermediate sharding:nn.Sequential( 3 x te.LayerNormMLP )after Fp8/intermediate sharding: