Reland Simplify handle indexing (#105006) #106357

voznesenskym · 2023-08-01T02:25:11Z

Stack from ghstack (oldest at bottom):

This reverts commit a9a3c45.

This PR changes the following:

_ExecOrderData.handle_to_handle_index -> FlatParamHandle._handle_index
_ExecOrderData.handles_to_pre_forward_order_index -> FlatParamHandle._pre_forward_order_index
_ExecOrderData.handles_to_post_forward_order_index -> FlatParamHandle._post_forward_index
_FSDPState._needs_pre_forward_unshard -> FlatParamHandle._needs_pre_forward_unshard
_FSDPState._needs_pre_backward_unshard -> FlatParamHandle._needs_pre_backward_unshard
_FSDPState._handles_prefetched -> FlatParamHandle._prefetched

This reverts commit a9a3c45. fixes [ghstack-poisoned]

pytorch-bot · 2023-08-01T02:25:13Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106357

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ 1 Unrelated Failure

As of commit be4c9a8:

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

linux-focal-rocm5.6-py3.8 / test (default, 1, 3, linux.rocm.gpu, unstable) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

This reverts commit a9a3c45. fixes ghstack-source-id: a025b0c Pull Request resolved: #106357

This reverts commit a9a3c45. fixes [ghstack-poisoned]

voznesenskym · 2023-08-03T00:27:26Z

Fixed pre backward fetch allgather

voznesenskym · 2023-08-03T00:39:06Z

torch/distributed/fsdp/_exec_order_utils.py

        if not handle:
            return
        # Only record the first usage of a handles key
-        if handle in self.handles_to_post_forward_order_index:


How it looked before, in the reverted PR. That would cause a None to show up.

awgu · 2023-08-03T02:36:40Z

It looks like the test failures are real :/

voznesenskym · 2023-08-03T05:57:46Z

It looks like the test failures are real :/

yeah they are

This reverts commit a9a3c45. fixes [ghstack-poisoned]

This reverts commit a9a3c45. fixes ghstack-source-id: 0342a7c Pull Request resolved: #106357 fix fix

awgu · 2023-08-03T16:42:19Z

torch/distributed/fsdp/_exec_order_utils.py

-        self.handles_post_forward_order: List[FlatParamHandle] = []
-        # Maps each handles key to its index in `handles_post_forward_order`
-        self.handles_to_post_forward_order_index: Dict[FlatParamHandle, int] = {}
+        self.handles_post_forward_order: List[Optional[FlatParamHandle]] = []


nit: IIUC, this is actually a List[FlatParamHandle] now because we never append None into the list?

awgu

LGTM!

awgu · 2023-08-03T18:00:54Z

torch/distributed/fsdp/_exec_order_utils.py

        Returns the handle indices (i.e. indices into ``self.all_handles``)
        corresponding to the handles in ``handle``. An entry in the
        returned tuple is ``None`` if the handle is invalid.
        """


Before this PR (making note to myself):

pytorch/torch/distributed/fsdp/_exec_order_utils.py

Lines 193 to 201 in be4c9a8

if self.is_first_iter:

msg_prefix = "Forward order differs across ranks:"

optional_local_indices: Tuple[

Optional[int], ...

] = self._get_handle_indices(handle)

device = handle.device # guaranteed to be non-CPU

num_valid_indices = sum(

(index is not None) for index in optional_local_indices

)

pytorch/torch/distributed/fsdp/_exec_order_utils.py

Line 235 in be4c9a8

local_indices = torch.tensor(optional_local_indices, **tensor_kwargs) # type: ignore[arg-type]

The above looks buggy. If optional_local_indices has a None in it, we should not be able to pass it to torch.tensor().

voznesenskym · 2023-08-03T18:54:42Z

@pytorchbot merge - Once more unto the breach

pytorch-bot · 2023-08-03T18:54:44Z

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: unrecognized arguments: - Once more unto the breach

usage: @pytorchbot [-h] {merge,revert,rebase,label,drci} ...

Try @pytorchbot --help for more info.

voznesenskym · 2023-08-03T19:15:31Z

@pytorchbot merge

pytorchmergebot · 2023-08-03T19:17:27Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Revert "Revert "Simplify handle indexing (#105006)" (#105984)"

0ee4822

This reverts commit a9a3c45. fixes [ghstack-poisoned]

voznesenskym requested review from H-Huang, awgu, d4l3k, fduwjj, fegin, kiukchung, kwen2501, mrshenli, rohan-varma, wanchaol and zhaojuanmao as code owners August 1, 2023 02:25

pytorch-bot bot added release notes: distributed (fsdp) release notes category labels Aug 1, 2023

voznesenskym added a commit that referenced this pull request Aug 1, 2023

Revert "Revert "Simplify handle indexing (#105006)" (#105984)"

92f0a0a

This reverts commit a9a3c45. fixes ghstack-source-id: a025b0c Pull Request resolved: #106357

github-actions bot requested review from SherlockNoMad, albanD, antoniojkim, bdhirsh, ezyang, jbschlosser, miladm and wconstab August 1, 2023 02:25

voznesenskym changed the title ~~Revert "Revert "Simplify handle indexing (#105006)" (#105984)"~~ WIP - reland Simplify handle indexing (#105006) Aug 1, 2023

albanD removed their request for review August 1, 2023 16:02

Update on "WIP - reland Simplify handle indexing (#105006)"

5e64db0

This reverts commit a9a3c45. fixes [ghstack-poisoned]

voznesenskym commented Aug 3, 2023

View reviewed changes

voznesenskym added ciflow/trunk Trigger trunk jobs on your pull request topic: not user facing topic category labels Aug 3, 2023

Update on "WIP - reland Simplify handle indexing (#105006)"

be4c9a8

This reverts commit a9a3c45. fixes [ghstack-poisoned]

voznesenskym added a commit that referenced this pull request Aug 3, 2023

Revert "Revert "Simplify handle indexing (#105006)" (#105984)"

6ae189e

This reverts commit a9a3c45. fixes ghstack-source-id: 0342a7c Pull Request resolved: #106357 fix fix

awgu self-assigned this Aug 3, 2023

awgu reviewed Aug 3, 2023

View reviewed changes

voznesenskym mentioned this pull request Aug 3, 2023

Better dynamo dict support via SetVariable keys #106559

Closed

awgu approved these changes Aug 3, 2023

View reviewed changes

voznesenskym changed the title ~~WIP - reland Simplify handle indexing (#105006)~~ Reland Simplify handle indexing (#105006) Aug 3, 2023

pytorchmergebot added the merging label Aug 3, 2023

pytorchmergebot added Merged and removed merging labels Aug 3, 2023

pytorchmergebot closed this in d1a99a0 Aug 3, 2023

	if self.is_first_iter:
	msg_prefix = "Forward order differs across ranks:"
	optional_local_indices: Tuple[
	Optional[int], ...
	] = self._get_handle_indices(handle)
	device = handle.device # guaranteed to be non-CPU
	num_valid_indices = sum(
	(index is not None) for index in optional_local_indices
	)

Reland Simplify handle indexing (#105006) #106357

Reland Simplify handle indexing (#105006) #106357

Uh oh!

Conversation

voznesenskym commented Aug 1, 2023 • edited by awgu Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 1, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106357

✅ 1 Unrelated Failure

Uh oh!

voznesenskym commented Aug 3, 2023

Uh oh!

voznesenskym Aug 3, 2023

Choose a reason for hiding this comment

Uh oh!

awgu commented Aug 3, 2023

Uh oh!

voznesenskym commented Aug 3, 2023

Uh oh!

awgu Aug 3, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

awgu left a comment

Choose a reason for hiding this comment

Uh oh!

awgu Aug 3, 2023

Choose a reason for hiding this comment

Uh oh!

voznesenskym commented Aug 3, 2023

Uh oh!

pytorch-bot bot commented Aug 3, 2023

Uh oh!

voznesenskym commented Aug 3, 2023

Uh oh!

pytorchmergebot commented Aug 3, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

voznesenskym commented Aug 1, 2023 •

edited by awgu

Loading

pytorch-bot bot commented Aug 1, 2023 •

edited

Loading

awgu Aug 3, 2023 •

edited

Loading