Refactored recomputable ops a bit and added a bunch more ops #85993
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/85993
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 FailuresAs of commit 20a413b: 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,
anjali411,
antoniojkim,
ezyang,
Krovatkin,
miladm and
wconstab
[ghstack-poisoned]
| pointwise_ops += [prims.div, prims.convert_element_type, aten.sign, aten.clone, aten._to_copy] # noqa: E501 | ||
| misc_ops = [aten.to, aten.type_as, operator.getitem] | ||
| recomputable_ops += [prims.div, prims.convert_element_type, aten.sign, aten.clone, aten._to_copy, aten.full_like, prims.var, prims.sum, aten.var, aten.std, prims.broadcast_in_dim, aten.select, aten.permute, aten._unsafe_view, aten.view, aten.expand, aten.slice, aten.reshape, aten.broadcast_tensors, aten.scalar_tensor, aten.ones, aten.new_zeros, aten.lift_fresh_copy, aten.minimum, aten.arange, aten.bitwise_and, aten.triu, aten.var_mean, aten.isinf, aten.any, aten.isnan, aten.full, aten.as_strided, aten.zeros, aten.argmax, aten.maximum, aten.bitwise_or, aten.logical_and, aten.logical_or] # noqa: E501 | ||
| # Natalia said that we should allow recomputing indexing :) |
There was a problem hiding this comment.
It probably has to have a slightly higher multiplier on top of i/o bytes, but generally yes, it's a regular bw bound op, so recomputing it may be beneficial, subject to comparison of bytes read/written that mincut takes care of
There was a problem hiding this comment.
Yeah, we might need to do something smarter when it comes to "limiting amount of recomputation".
|
@pytorchbot merge -f "failing test is flaky OOM" |
|
@pytorchbot successfully started a merge job. Check the current status here. |
|
Hey @Chillee. |
Pull Request resolved: #85993 Approved by: https://github.com/ngimel
Stack from ghstack (oldest at bottom):