[Autograd] Use in-place input accumulation fast path for dense Tensors.

[robieta]

Stack from ghstack (oldest at bottom):

• -> [Autograd] Use in-place input accumulation fast path for dense Tensors. #88339

There is a fast path in InputBuffer to steal memory when use count is zero, however it is only used for sparse Tensors. According to Natalia, this is just because it wasn't obvious that there would be a benefit for dense Tensors so there was no reason to live dangerously. However I've noticed large Tensors in internal models which would benefit from this optimization as well.

Differential Revision: D40946601

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 our office hours

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

❌ 1 Failures

As of commit 835e532:

The following jobs have failed:

• linux-focal-py3.7-clang7-asan / test (default, 3, 5, linux.2xlarge)

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

[ngimel]

I don't think we want to take this path when both var and old_var are sparse, who even knows if we have working inplace sparse-sparse addition, and in any case it's a conceptually weird operation, so you might want to push this case to the else branch.

[robieta]

If old_var is sparse it will get hooked by the first part of the conditional. (Which has a !var.is_sparse() on its steal check.) Maybe I should add a cautionary comment?

[ngimel]

Ah yeah that would be good, thanks for checking!

[robieta]

Looks like I also need a !old_var._is_zerotensor() check. @anjali411 is there a more general mutability check that I should use?

[albanD]

I think there are a few things wrong here that need updating. In particular looking at the logic for AccumulateGrad show a lot of possible failure cases here:

pytorch/torch/csrc/autograd/functions/accumulate_grad.h

Line 56 in 0fa2366

// "Gradient Layout Contract"

• We need to check that grad mode is disabled. This inplace might be invalid if old_var is saved by autograd
• You need to guard for .is_sparse_csr() which is not covered by .is_sparse()
• We want is_non_overlapping_and_dense on top of is_contiguous as we really don't want weird memory overlapping Tensors to be written inplace.
• Why the special case on zero_tensor here?

[robieta]

We need to check that grad mode is disabled. This inplace might be invalid if old_var is saved by autograd

Should that be GradMode::is_enabled() or a property of old_var?

SGTM on is_sparse_csr and is_non_overlapping_and_dense.

Why the special case on zero_tensor here?

ZeroTensors are immutable

[ngimel]

is_non_overlapping_and_dense on top of is_contiguous

instead of is contiguous? Yeah that's reasonable. IIRC it errors out on sparse tensors though, so the logic will need to become more tortured

[robieta]

FWIW I'm also updating the path where var is used, so there will be sparse checks before any other checks.

[albanD]

Should that be GradMode::is_enabled() or a property of old_var?

GradMode::is_enabled() because it might not be ok to modify old_var even if it doesn't require gradients.
Note that it could still be invalid when grad mode is disabled but that should be unlikely enough here that we're ok.

[robieta]

Updated with more rigorous checks and comments.

[robieta]

The latest failure is torch.autograd.gradcheck.GradcheckError: While computing batched gradients, got: Cannot access storage of BatchedTensorImpl. Nested Tensor, sparse Tensor, ZeroTensor, BatchedTensor etc. will all blow up if various methods are called and it's virtually impossible to get everything right, much less future proof it. I'll fix the unit test failures, but I'm increasingly coming to the conclusion that we need to add something like is_plain_old_tensor() to Tensor.

[zou3519]

there's an isTensorSubclassLike check for this, though it is not completely comprehensive yet:

pytorch/aten/src/ATen/TensorSubclassLikeUtils.h

Line 51 in ef4ce6d

inline bool isTensorSubclassLike(const Tensor& tensor) {

We probably want to exclude tensors that check isTensorSubclassLike True from this fast path

[robieta]

there's an isTensorSubclassLike check for this, though it is not completely comprehensive yet:

pytorch/aten/src/ATen/TensorSubclassLikeUtils.h

Line 51 in ef4ce6d

inline bool isTensorSubclassLike(const Tensor& tensor) {

We probably want to exclude tensors that check isTensorSubclassLike True from this fast path

This is AWESOME! (CC @chaekit, @aaronenyeshi, @slgong-fb we should use this in profiler too.) AFAICT it doesn't check nested or any other C++ subclass? (And if not, do you think it should?) Plus it includes sparse and sparse_csr, so we shouldn't need those explicit checks.

[robieta]

I also had to gate on var not being a subclass because of a forward mode AD test. In theory I could have made the test check grad mode, but given how varied Tensor subclasses can be it seems like a better idea to just exclude them from the fast path. PTAL

[robieta]

@pytorchbot merge -f "test failure is unrelated. (failed to install triton)"

[pytorchmergebot]

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

[mehtanirav]

@pytorchbot revert -m "Internal test failures" -c ghfirst

[pytorchmergebot]

@pytorchbot successfully started a revert job. Check the current status here.
Questions? Feedback? Please reach out to the PyTorch DevX Team

[pytorchmergebot]

@robieta your PR has been successfully reverted.