aot autograd refactor: make all synthetic base logic layered in a single location #96235
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
…gle location [ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/96235
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 FailuresAs of commit 6f19e89: NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
github-actions
bot
requested review from
albanD,
antoniojkim,
jbschlosser,
miladm,
SherlockNoMad,
voznesenskym and
wconstab
bdhirsh
added a commit
that referenced
this pull request
Mar 7, 2023
…gle location ghstack-source-id: 330eaa87797f205ba5456c17ce1349a1510975ce Pull Request resolved: #96235
…ed in a single location" This refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts). The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it. Specifically, I added a new function `aot_wrapper_synthetic_base`, similar to the existing `aot_wrapper_dedupe`. The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places). The downsides: there are two. (1) `aot_wrapper_synthetic_base()` needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation. Ordinarily, metadata mutations are handled by the runtime epilogue, in `create_runtime_wrapper`. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we **need** to apply the metadata mutation directly to the original inputs. The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs. `aot_wrapper_synthetic_base` will know to grab these extra outputs and use them to apply the metadata mutations. This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in `aot_wrapper_synthetic_base()`. (2) input mutations are now performed on the synthetic base instead of the individual aliases. You can see the original code comment [here](https://github.com/pytorch/pytorch/blob/b0b5f3c6c681896febbd9ff7ad7649b13def345d/torch/_functorch/aot_autograd.py#L1131) for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias). To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think. [ghstack-poisoned]
bdhirsh
added a commit
that referenced
this pull request
Mar 8, 2023
…gle location ghstack-source-id: 529ecb436f17194feddf3a88aada7946a6875571 Pull Request resolved: #96235
…ed in a single location" This refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts). The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it. Specifically, I added a new function `aot_wrapper_synthetic_base`, similar to the existing `aot_wrapper_dedupe`. The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places). The downsides: there are two. (1) `aot_wrapper_synthetic_base()` needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation. Ordinarily, metadata mutations are handled by the runtime epilogue, in `create_runtime_wrapper`. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we **need** to apply the metadata mutation directly to the original inputs. The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs. `aot_wrapper_synthetic_base` will know to grab these extra outputs and use them to apply the metadata mutations. This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in `aot_wrapper_synthetic_base()`. (2) input mutations are now performed on the synthetic base instead of the individual aliases. You can see the original code comment [here](https://github.com/pytorch/pytorch/blob/b0b5f3c6c681896febbd9ff7ad7649b13def345d/torch/_functorch/aot_autograd.py#L1131) for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias). To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think. [ghstack-poisoned]
bdhirsh
added a commit
that referenced
this pull request
Mar 8, 2023
…gle location ghstack-source-id: d77ac663d86b61f4ae94f1cdd502fd37dd0140ad Pull Request resolved: #96235
This was referenced Mar 8, 2023
ezyang
reviewed
Mar 9, 2023
torch/_functorch/aot_autograd.py
Outdated
| return NotImplemented | ||
| return (self.output_type == other.output_type and | ||
| self.base_idx == other.base_idx and | ||
| (issubclass(self.raw_type, other.raw_type) or issubclass(other.raw_type, self.raw_type))) |
There was a problem hiding this comment.
Yeah... I'll try to fix this for real (should have at least left a comment here). What I should be doing is grabbing the raw_type off the the original input to get the synthetics base's raw type. I spent 10 minutes trying to do that and it was painful, so I ended up doing this instead. I'm trying to remember exactly where my metadata check failed, but I think one had a type of FakeTensor and the other was torch.Tensor.
There was a problem hiding this comment.
yeaa i was wrong, not necessary
ezyang
reviewed
Mar 9, 2023
ezyang
reviewed
Mar 9, 2023
| args.clear() | ||
| outs = compiled_fn(args_with_synthetic_bases) | ||
| if num_aliased_args_with_metadata_mutations > 0: | ||
| # This code does not handle **all** input metadata mutations. |
There was a problem hiding this comment.
It's a bit interesting that we have to handle metadata mutation here. I wonder if it's possible to have another layer that just gets rid of all metadata mutation upfront
ezyang
reviewed
Mar 9, 2023
| compiler_fn, | ||
| ): | ||
| is_inference = not needs_autograd | ||
| flat_args_with_synthetic_bases, synthetic_base_info = merge_view_inputs( |
There was a problem hiding this comment.
It's cool that all of this logic was already factored out, definitely was nice for review here!
ezyang
reviewed
Mar 9, 2023
| requires_grad_info = mutated_inp_require_grad_info + output_grad_info + input_metadata_mutation_grad_info | ||
| output_info = existing_output_infos + input_metadata_output_info | ||
| # Regenerate traced tangents to include mutated inputs including synthetic bases | ||
| traced_tangents = inner_mutated_data_inps + m.traced_tangents[num_outer_mutated_data_inps:] |
There was a problem hiding this comment.
Remark: we don't have a "length" comment for traced_tangents in ViewAndMutationMeta
ezyang
reviewed
Mar 9, 2023
torch/_functorch/aot_autograd.py
Outdated
|
|
||
| # For each arg in the inner calling convention, map it to its idx in the outer calling convention | ||
| # (there might be multiple indices, if the inner arg is a synthetic base) | ||
| synthetic_base_to_indices: List[List[int]] = [] |
There was a problem hiding this comment.
mypy lets you define newtypes, and I think two newtypes here would be really useful. One for referring to "inner" indexes, and the other for "outer" indexes. Inner/outer is ambiguous when we have multiple layers, so probably need some prefix (maybe SB or S for synthetic base?)
(I don't remember: you said you were going to type the PR after? So you can do this later.)
There was a problem hiding this comment.
So, it's not List[List[int]], but Map[S_Inner, List[S_Outer]] (here Map is just List but with the index explicitly typed) which nearly makes the comment above unnecessary
There was a problem hiding this comment.
nice, I'll add that (although I'll save the mypy-ifying aot_autograd.py for a future PR)
ezyang
reviewed
Mar 9, 2023
torch/_functorch/aot_autograd.py
Outdated
| # For each arg in the inner calling convention, map it to its idx in the outer calling convention | ||
| # (there might be multiple indices, if the inner arg is a synthetic base) | ||
| synthetic_base_to_indices: List[List[int]] = [] | ||
| for idx in range(len(flat_args_with_synthetic_bases)): |
There was a problem hiding this comment.
If you like inner/outer, we should probably name this inner_args (and then flat_args is outer_args). Not sure if I like the name best but it is certainly more compact.
There was a problem hiding this comment.
And then, you don't say idx, you say inner_idx (?)
There was a problem hiding this comment.
yeah i don't love it either but ill stick with it for consistency
ezyang
reviewed
Mar 9, 2023
torch/_functorch/aot_autograd.py
Outdated
| synthetic_base_to_indices: List[List[int]] = [] | ||
| for idx in range(len(flat_args_with_synthetic_bases)): | ||
| outer_aliased_indices_of_current_base_arg = [ | ||
| i for i, outer_idx_or_tuple in enumerate(synthetic_base_info) |
There was a problem hiding this comment.
outer_idx_or_tuple seems misnamed here. Isn't this the inner idx? Per the documentation on synthetic_base_info.
ezyang
reviewed
Mar 9, 2023
| # and compare it to this output for sanity. | ||
| # | ||
| # In addition to the updated metadata, also return the list of input indices | ||
| # that will need to be updated in the synthetic base epilogue |
ezyang
reviewed
Mar 9, 2023
torch/_functorch/aot_autograd.py
Outdated
| # mutations, they will be hidden from the rest of aot autograd. | ||
| mutates_data=True if len(indices) > 1 else m.input_info[indices[0]].mutates_data, | ||
| mutates_metadata=False if len(indices) > 1 else m.input_info[indices[0]].mutates_metadata, | ||
| # leaf-ness should be all-or-nothing for aliased tensor. |
There was a problem hiding this comment.
so what, some sort of assert that all(m.input_info[x].is_leaf) == any(...)?
There was a problem hiding this comment.
sure, ill add a check
ezyang
reviewed
Mar 9, 2023
| ] | ||
|
|
||
| # grab the original requires grad info on the outputs, except the ones from the mutated inputs | ||
| num_original_input_data_mutations = len([x for x in m.input_info if x.mutates_data or x.mutates_metadata]) |
There was a problem hiding this comment.
yeah I would really really love if we could assume there's no metadata mutation at this point anymore
There was a problem hiding this comment.
for a future refactor :) (discussed offline and I think we agreed this would be nice, but not necessary for now)
ezyang
approved these changes
Mar 9, 2023
…ed in a single location" This refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts). The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it. Specifically, I added a new function `aot_wrapper_synthetic_base`, similar to the existing `aot_wrapper_dedupe`. The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places). The downsides: there are two. (1) `aot_wrapper_synthetic_base()` needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation. Ordinarily, metadata mutations are handled by the runtime epilogue, in `create_runtime_wrapper`. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we **need** to apply the metadata mutation directly to the original inputs. The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs. `aot_wrapper_synthetic_base` will know to grab these extra outputs and use them to apply the metadata mutations. This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in `aot_wrapper_synthetic_base()`. (2) input mutations are now performed on the synthetic base instead of the individual aliases. You can see the original code comment [here](https://github.com/pytorch/pytorch/blob/b0b5f3c6c681896febbd9ff7ad7649b13def345d/torch/_functorch/aot_autograd.py#L1131) for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias). To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think. [ghstack-poisoned]
…ed in a single location" This refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts). The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it. Specifically, I added a new function `aot_wrapper_synthetic_base`, similar to the existing `aot_wrapper_dedupe`. The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places). The downsides: there are two. (1) `aot_wrapper_synthetic_base()` needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation. Ordinarily, metadata mutations are handled by the runtime epilogue, in `create_runtime_wrapper`. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we **need** to apply the metadata mutation directly to the original inputs. The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs. `aot_wrapper_synthetic_base` will know to grab these extra outputs and use them to apply the metadata mutations. This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in `aot_wrapper_synthetic_base()`. (2) input mutations are now performed on the synthetic base instead of the individual aliases. You can see the original code comment [here](https://github.com/pytorch/pytorch/blob/b0b5f3c6c681896febbd9ff7ad7649b13def345d/torch/_functorch/aot_autograd.py#L1131) for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias). To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think. [ghstack-poisoned]
cyyever
pushed a commit
to cyyever/pytorch_private
that referenced
this pull request
Mar 23, 2023
…gle location (#96235) This refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts). The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it. Specifically, I added a new function `aot_wrapper_synthetic_base`, similar to the existing `aot_wrapper_dedupe`. The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places). The downsides: there are two. (1) `aot_wrapper_synthetic_base()` needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation. Ordinarily, metadata mutations are handled by the runtime epilogue, in `create_runtime_wrapper`. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we **need** to apply the metadata mutation directly to the original inputs. The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs. `aot_wrapper_synthetic_base` will know to grab these extra outputs and use them to apply the metadata mutations. This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in `aot_wrapper_synthetic_base()`. (2) input mutations are now performed on the synthetic base instead of the individual aliases. You can see the original code comment [here](https://github.com/pytorch/pytorch/blob/b0b5f3c6c681896febbd9ff7ad7649b13def345d/torch/_functorch/aot_autograd.py#L1131) for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias). To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think. Pull Request resolved: pytorch/pytorch#96235 Approved by: https://github.com/ezyang
cyyever
pushed a commit
to cyyever/pytorch_private
that referenced
this pull request
Mar 27, 2023
…gle location (#96235) This refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts). The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it. Specifically, I added a new function `aot_wrapper_synthetic_base`, similar to the existing `aot_wrapper_dedupe`. The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places). The downsides: there are two. (1) `aot_wrapper_synthetic_base()` needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation. Ordinarily, metadata mutations are handled by the runtime epilogue, in `create_runtime_wrapper`. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we **need** to apply the metadata mutation directly to the original inputs. The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs. `aot_wrapper_synthetic_base` will know to grab these extra outputs and use them to apply the metadata mutations. This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in `aot_wrapper_synthetic_base()`. (2) input mutations are now performed on the synthetic base instead of the individual aliases. You can see the original code comment [here](https://github.com/pytorch/pytorch/blob/b0b5f3c6c681896febbd9ff7ad7649b13def345d/torch/_functorch/aot_autograd.py#L1131) for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias). To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think. Pull Request resolved: pytorch/pytorch#96235 Approved by: https://github.com/ezyang
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 refactor doesn't significantly change LoC in aot autograd, but I think this nets out to making it clearer (interested in peoples' thoughts).
The idea is that I tried to re-write the part of aot autograd that deals with synthetic bases in a layered way, similar to how Ed wrote the logic for dedup'ing inputs: it happens in one place, and all of the downstream transformation in aot autograd don't have to worry about it.
Specifically, I added a new function
aot_wrapper_synthetic_base, similar to the existingaot_wrapper_dedupe.The benefit: none of the other code in aot autograd needs to think about synthetic bases (previously, synthetic base code was intertwined in several places).
The downsides: there are two.
(1)
aot_wrapper_synthetic_base()needs to have its own epilogue. There is one particularly hairy case, where factoring the synthetic base logic to a single location was painful: If you have two inputs that alias each other, where one gets a data mutation, and the other gets a metadata mutation.Ordinarily, metadata mutations are handled by the runtime epilogue, in
create_runtime_wrapper. However, now that things are factored this way, the runtime wrapper operates only on synthetic bases instead of operating on the original inputs. For data mutations, it is fine to apply the data mutation to the synthetic base instead of the original input alias. But for metadata mutations, we need to apply the metadata mutation directly to the original inputs.The way that I handled this was by tracking which inputs slot into this specific case (part of a synthetic base, and get metadata mutations), and updateing the flat_fn() that we pass downstream to return these updated inputs as extra outputs. From the perspective of downstream logic, these are real user outputs, that it can treat like any other user outputs.
aot_wrapper_synthetic_basewill know to grab these extra outputs and use them to apply the metadata mutations.This was pretty annoying, but has the benefit that all of that logic is encapsulated entirely in
aot_wrapper_synthetic_base().(2) input mutations are now performed on the synthetic base instead of the individual aliases.
You can see the original code comment here for details. We used to do the optimized thing in this case, and now we do the less optimized thing (copying the entire synthetic base, instead of the potentially smaller alias).
To be fair, we had no data showing that this optimization was showing improvements on any models in practice. I also think that the main reason anyone would ever run across this problem is because of a graph break - so if you care about perf, you probably want to avoid the extra graph breaks to begin with. I haven't added any warnings for this, but we probably could depending on what people think.
Stack from ghstack (oldest at bottom):