Repository navigation
[PyTorch] Deferred Initialization via device='meta' option - #596
Merged
Merged
Conversation
denera
force-pushed
the
databricks/fp8-fsdp-oom-fix
branch
from
January 8, 2024 17:55
908f1a5 to
70a90a3
Compare
Collaborator
Author
|
/te-ci pytorch |
denera
force-pushed
the
databricks/fp8-fsdp-oom-fix
branch
from
January 9, 2024 20:01
0ed0550 to
a12bb2b
Compare
Collaborator
Author
|
/te-ci pytorch |
ptrendx
reviewed
Jan 9, 2024
ptrendx
reviewed
Jan 9, 2024
ptrendx
reviewed
Jan 9, 2024
Collaborator
Author
|
/te-ci pytorch |
ksivaman
reviewed
Jan 10, 2024
ksivaman
reviewed
Jan 10, 2024
ptrendx
reviewed
Jan 11, 2024
denera
force-pushed
the
databricks/fp8-fsdp-oom-fix
branch
from
January 12, 2024 18:17
cb055e5 to
05b0e12
Compare
Collaborator
Author
|
/te-ci pytorch |
….Linear and added new PyTorch example to demonstrate its use with FullyShardedDataParallel execution. 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>
…ally Signed-off-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Alp Dener <adener@nvidia.com>
denera
force-pushed
the
databricks/fp8-fsdp-oom-fix
branch
from
January 12, 2024 22:43
05b0e12 to
1cc1be6
Compare
Collaborator
Author
|
/te-ci pytorch |
denera
changed the title
te.Linear via device='meta' optiondevice='meta' option
Collaborator
Author
Signed-off-by: Alp Dener <adener@nvidia.com>
Collaborator
Author
|
/te-ci pytorch |
ksivaman
reviewed
Jan 17, 2024
ksivaman
reviewed
Jan 17, 2024
…d clarifying comments in parameter reset Signed-off-by: Alp Dener <adener@nvidia.com>
Marks101
added a commit
to Marks101/TransformerEngine
that referenced
this pull request
Jan 22, 2024
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
ptrendx
pushed a commit
that referenced
this pull request
Jan 22, 2024
* Implemented deferred initialization via `device='meta'` option for te.Linear and added new PyTorch example to demonstrate its use with FullyShardedDataParallel execution. Signed-off-by: Alp Dener <adener@nvidia.com> * correcting Float8Tensor initialization and fixing linting errors Signed-off-by: Alp Dener <adener@nvidia.com> * removed duplicate code from upstream rebase, local tests passing Signed-off-by: Alp Dener <adener@nvidia.com> * improved comments/documentation for FSDP example Signed-off-by: Alp Dener <adener@nvidia.com> * converted reset_parameters() into a base module function Signed-off-by: Alp Dener <adener@nvidia.com> * fixed Float8Tensor creation with deferred init, all tests passing locally Signed-off-by: Alp Dener <adener@nvidia.com> * extended deferred initialization to all TE modules Signed-off-by: Alp Dener <adener@nvidia.com> * fixed linting errors Signed-off-by: Alp Dener <adener@nvidia.com> * removed unnecessary reference to the parent module of parameter, added clarifying comments in parameter reset Signed-off-by: Alp Dener <adener@nvidia.com> --------- Signed-off-by: Alp Dener <adener@nvidia.com>
ksivaman
pushed a commit
that referenced
this pull request
Jan 22, 2024
ptrendx
pushed a commit
that referenced
this pull request
Jan 22, 2024
Closed
xiaoteng-cui-mt
pushed a commit
to MooreThreads/MT-TransformerEngine
that referenced
this pull request
Oct 9, 2026
) * Implemented deferred initialization via `device='meta'` option for te.Linear and added new PyTorch example to demonstrate its use with FullyShardedDataParallel execution. Signed-off-by: Alp Dener <adener@nvidia.com> * correcting Float8Tensor initialization and fixing linting errors Signed-off-by: Alp Dener <adener@nvidia.com> * removed duplicate code from upstream rebase, local tests passing Signed-off-by: Alp Dener <adener@nvidia.com> * improved comments/documentation for FSDP example Signed-off-by: Alp Dener <adener@nvidia.com> * converted reset_parameters() into a base module function Signed-off-by: Alp Dener <adener@nvidia.com> * fixed Float8Tensor creation with deferred init, all tests passing locally Signed-off-by: Alp Dener <adener@nvidia.com> * extended deferred initialization to all TE modules Signed-off-by: Alp Dener <adener@nvidia.com> * fixed linting errors Signed-off-by: Alp Dener <adener@nvidia.com> * removed unnecessary reference to the parent module of parameter, added clarifying comments in parameter reset Signed-off-by: Alp Dener <adener@nvidia.com> --------- Signed-off-by: Alp Dener <adener@nvidia.com>
xiaoteng-cui-mt
pushed a commit
to MooreThreads/MT-TransformerEngine
that referenced
this pull request
Oct 9, 2026
Signed-off-by: Markus Schnoes <markus.schnoes@gmx.de>
This file contains hidden or 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
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.
This is primarily meant to support training very large models in
torch.distributed.fsdp.FullyShardedDataParallelmode.FSDP execution flow requires models to be initialized in serial mode before being passed into an FSDP wrapper that shards their parameters. Without
device='meta'support, this materializes a copy of every module in memory for every host/device and causes OOM errors for large models. Deferring initialization withdevice='meta'materializes the memory only when the FSDP wrapper recursively executesmodule.reset_parameters()on already sharded modules, thus avoiding unsharded duplicates.Databricks has requested us to accelerate upstreaming this feature forte.Linearin support of their immediate work instead of waiting to upstream it until the same feature is implemented for other TransformerEngine modules. I will be filing additional PRs to extend this feature to other modules in the coming days/weeks.The PR also includes an example for demonstrating
device='meta'use with FSDP.I am planning to convert this into a distributed/FSDP regression test once the meta device feature is implemented for other TE modules.Edit: Extending this to all TE modules turned out to be easy enough to roll into this PR instead postponing to a follow-up.