[dtensor] enable foreach operators for adam optimizer #112108
Closed
Commits on Oct 26, 2023
-
[dtensor] enable foreach operators for adam optimizer
This PR enables basic foreach ops in DTensor for adam optimizer, to improve performance compare to optimizer using torch.Tensor. Currently by default optimizer won't do this for tensor subclass, we will need to enable this by default in DTensor when all ops are covered, or early enable this when exploring new FSDP, we just need to append DTensor to the optimizer allow list. cc @awgu [ghstack-poisoned]
Commits on Oct 30, 2023
-
Update on "[dtensor] enable foreach operators for adam optimizer"
This PR enables basic foreach ops in DTensor for adam optimizer, to improve performance compare to optimizer using torch.Tensor. Currently by default optimizer won't do this for tensor subclass, we will need to enable this by default in DTensor when all ops are covered, or early enable this when exploring new FSDP, we just need to append DTensor to the optimizer allow list. cc awgu [ghstack-poisoned]
-
Update on "[dtensor] enable foreach operators for adam optimizer"
This PR enables basic foreach ops in DTensor for adam optimizer, to improve performance compare to optimizer using torch.Tensor. Currently by default optimizer won't do this for tensor subclass, we will need to enable this by default in DTensor when all ops are covered, or early enable this when exploring new FSDP, we just need to append DTensor to the optimizer allow list. Some latency measurement, on a 5-layer MLP model: single tensor adam: 17ms  foreach multitensor adam: 4ms  so around 4.25x improvement cc awgu [ghstack-poisoned]
Commits on Oct 31, 2023
-
Update on "[dtensor] enable foreach operators for adam optimizer"
This PR enables basic foreach ops in DTensor for adam optimizer, to improve performance compare to optimizer using torch.Tensor. Currently by default optimizer won't do this for tensor subclass, we will need to enable this by default in DTensor when all ops are covered, or early enable this when exploring new FSDP, we just need to append DTensor to the optimizer allow list. Some latency measurement, on a 5-layer MLP model: single tensor adam: 17ms  foreach multitensor adam: 4ms  so around 4.25x improvement cc awgu [ghstack-poisoned]
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.