Enable `backward` on `_foreach_zero_` #101149

crcrpar · 2023-05-11T02:07:35Z

Currently torchgen cannot find an appropriate DifferentiabilityInfo for _foreach_zero_ because gen_foreach_derivativeinfo doesn't correctly make use of functional_info_by_signature and differentiability_infos, and is_reference_for_foreach a bit too strict to _foreach_zero_.

Generated code in VariableType

void _foreach_zero_(c10::DispatchKeySet ks, at::TensorList self) {
  auto self_ = unpack(self, "self", 0);
  [[maybe_unused]] auto _any_requires_grad = compute_requires_grad( self );

  std::vector<c10::optional<at::Tensor>> original_selfs(self.size());
  std::vector<std::shared_ptr<ZeroBackward0>> grad_fns;
  if (_any_requires_grad) {
    for (const auto& i : c10::irange( self.size() )) {
      const auto ith_requires_grad = compute_requires_grad(self[i]);
      check_inplace(self[i], ith_requires_grad);
      grad_fns.push_back([&]() -> std::shared_ptr<ZeroBackward0> {
          if (!ith_requires_grad) {
              return nullptr;
          } else {
              auto grad_fn = std::shared_ptr<ZeroBackward0>(new ZeroBackward0(), deleteNode);
              grad_fn->set_next_edges(collect_next_edges( self[i] ));
              return grad_fn;
          }
      }());
    }
  }
  #ifndef NDEBUG
  std::vector<c10::optional<Storage>> self__storage_saved(self_.size());
  for (const Tensor& tensor : self_)
    self__storage_saved.push_back(
      tensor.has_storage() ? c10::optional<Storage>(tensor.storage()) : c10::nullopt);
  std::vector<c10::intrusive_ptr<TensorImpl>> self__impl_saved(self_.size());
  for (size_t i=0; i<self_.size(); i++)
    if (self_[i].defined()) self__impl_saved[i] = self_[i].getIntrusivePtr();
  #endif
  {
    at::AutoDispatchBelowAutograd guard;
    at::redispatch::_foreach_zero_(ks & c10::after_autograd_keyset, self_);
  }
  #ifndef NDEBUG
  for (size_t i=0; i<self_.size() && !at::impl::dispatch_mode_enabled(); i++) {
    if (self__storage_saved[i].has_value() && !at::impl::tensorlist_has_dispatch(self_))
      TORCH_INTERNAL_ASSERT(self__storage_saved[i].value().is_alias_of(self_[i].storage()));
  }
  for (size_t i=0; i<self_.size() && !at::impl::dispatch_mode_enabled(); i++) {
    if (self__impl_saved[i] && !at::impl::tensorlist_has_dispatch(self_))
      TORCH_INTERNAL_ASSERT(self__impl_saved[i] == self_[i].getIntrusivePtr());
  }
  #endif
  if (!grad_fns.empty()) {
      auto differentiable_outputs = flatten_tensor_args( self );
      TORCH_INTERNAL_ASSERT(differentiable_outputs.size() == grad_fns.size());
      for (const auto& i : c10::irange(grad_fns.size())) {
          auto grad_fn = grad_fns[i];
          if (grad_fn != nullptr) {
              rebase_history(differentiable_outputs[i], grad_fns[i]);
          }
      }
  }
}

Rel:

pytorch-bot · 2023-05-11T02:07:37Z

🔗 Helpful Links

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

📄 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 New Failure

As of commit c83cb65:

NEW FAILURE - The following job has failed:

Check Labels (gh)

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

torchgen/api/autograd.py

soulitzer

Thanks for the fix!

torchgen/api/autograd.py

test/test_foreach.py

crcrpar · 2023-05-12T01:03:27Z

@pytorchbot label suppress-bc-linter

soulitzer

OK!

crcrpar · 2023-05-16T17:56:56Z

@pytorchbot label ciflow/trunk

crcrpar · 2023-05-16T21:42:46Z

The failure refers to https://github.com/pytorch/pytorch/issues/100209🤔

crcrpar · 2023-05-17T00:00:16Z

@pytorchbot merge -r

pytorchmergebot · 2023-05-17T00:02:19Z

@pytorchbot started a rebase job onto refs/remotes/origin/viable/strict. Check the current status here

as especially `_foreach_zero_`'s reference `DifferentiabilityInfo` seems to be only available in it, not `functional_info_by_signature`. Also, made `gen_foreach_derivativeinfo` a bit cleaner as passing both `differentiability_infors` and `functional_info_by_signature` wouldn't make less sense than passing `functional_info_by_signature` and `non_functional_info_by_signature`. Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

by creating allowlist of foreach functions whose reference function is somehow inplace Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

with a fake `op` as arg to OpInfo dunder init Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

pytorchmergebot · 2023-05-17T00:02:26Z

Successfully rebased inplace_foreach_zero_backward onto refs/remotes/origin/viable/strict, please pull locally before adding more changes (for example, via git checkout inplace_foreach_zero_backward && git pull --rebase)

pytorchmergebot · 2023-05-17T00:03:32Z

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

Currently torchgen cannot find an appropriate `DifferentiabilityInfo` for `_foreach_zero_` because `gen_foreach_derivativeinfo` doesn't correctly make use of `functional_info_by_signature` and `differentiability_infos`, and `is_reference_for_foreach` a bit too strict to `_foreach_zero_`. Generated code in `VariableType` ```c++ void _foreach_zero_(c10::DispatchKeySet ks, at::TensorList self) { auto self_ = unpack(self, "self", 0); [[maybe_unused]] auto _any_requires_grad = compute_requires_grad( self ); std::vector<c10::optional<at::Tensor>> original_selfs(self.size()); std::vector<std::shared_ptr<ZeroBackward0>> grad_fns; if (_any_requires_grad) { for (const auto& i : c10::irange( self.size() )) { const auto ith_requires_grad = compute_requires_grad(self[i]); check_inplace(self[i], ith_requires_grad); grad_fns.push_back([&]() -> std::shared_ptr<ZeroBackward0> { if (!ith_requires_grad) { return nullptr; } else { auto grad_fn = std::shared_ptr<ZeroBackward0>(new ZeroBackward0(), deleteNode); grad_fn->set_next_edges(collect_next_edges( self[i] )); return grad_fn; } }()); } } #ifndef NDEBUG std::vector<c10::optional<Storage>> self__storage_saved(self_.size()); for (const Tensor& tensor : self_) self__storage_saved.push_back( tensor.has_storage() ? c10::optional<Storage>(tensor.storage()) : c10::nullopt); std::vector<c10::intrusive_ptr<TensorImpl>> self__impl_saved(self_.size()); for (size_t i=0; i<self_.size(); i++) if (self_[i].defined()) self__impl_saved[i] = self_[i].getIntrusivePtr(); #endif { at::AutoDispatchBelowAutograd guard; at::redispatch::_foreach_zero_(ks & c10::after_autograd_keyset, self_); } #ifndef NDEBUG for (size_t i=0; i<self_.size() && !at::impl::dispatch_mode_enabled(); i++) { if (self__storage_saved[i].has_value() && !at::impl::tensorlist_has_dispatch(self_)) TORCH_INTERNAL_ASSERT(self__storage_saved[i].value().is_alias_of(self_[i].storage())); } for (size_t i=0; i<self_.size() && !at::impl::dispatch_mode_enabled(); i++) { if (self__impl_saved[i] && !at::impl::tensorlist_has_dispatch(self_)) TORCH_INTERNAL_ASSERT(self__impl_saved[i] == self_[i].getIntrusivePtr()); } #endif if (!grad_fns.empty()) { auto differentiable_outputs = flatten_tensor_args( self ); TORCH_INTERNAL_ASSERT(differentiable_outputs.size() == grad_fns.size()); for (const auto& i : c10::irange(grad_fns.size())) { auto grad_fn = grad_fns[i]; if (grad_fn != nullptr) { rebase_history(differentiable_outputs[i], grad_fns[i]); } } } } ``` Rel: - #58833 - #96405 Pull Request resolved: #101149 Approved by: https://github.com/soulitzer

crcrpar requested review from albanD and soulitzer as code owners May 11, 2023 02:07

pytorch-bot bot added the release notes: foreach_frontend release notes category label May 11, 2023

pytorchbot added the open source label May 11, 2023

mikaylagawarecki added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label May 11, 2023

crcrpar commented May 11, 2023

View reviewed changes

torchgen/api/autograd.py Show resolved Hide resolved

soulitzer reviewed May 11, 2023

View reviewed changes

torchgen/api/autograd.py Outdated Show resolved Hide resolved

torchgen/api/autograd.py Show resolved Hide resolved

test/test_foreach.py Outdated Show resolved Hide resolved

pytorch-bot bot added the suppress-bc-linter Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) label May 12, 2023

crcrpar force-pushed the inplace_foreach_zero_backward branch from ef7e9c6 to 355eee6 Compare May 12, 2023 07:55

crcrpar requested review from mruberry and ngimel as code owners May 15, 2023 21:19

soulitzer approved these changes May 16, 2023

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label May 16, 2023

crcrpar added 5 commits May 17, 2023 00:02

add _foreach_zero_ test

d5ba1ec

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

remove print from test

791dea1

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

stricter is_reference_for_foreach func

6ec8fc2

by creating allowlist of foreach functions whose reference function is somehow inplace Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

ForeachFuncInfo("zero")

c83cb65

with a fake `op` as arg to OpInfo dunder init Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

pytorchmergebot force-pushed the inplace_foreach_zero_backward branch from a96cb2d to c83cb65 Compare May 17, 2023 00:02

pytorchmergebot added the merging label May 17, 2023

pytorchmergebot added Merged and removed merging labels May 17, 2023

pytorchmergebot closed this in ba2bc7d May 17, 2023

crcrpar deleted the inplace_foreach_zero_backward branch May 17, 2023 05:56

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Enable `backward` on `_foreach_zero_` #101149

Enable `backward` on `_foreach_zero_` #101149

crcrpar commented May 11, 2023 •

edited

pytorch-bot bot commented May 11, 2023 •

edited

soulitzer left a comment

crcrpar commented May 12, 2023

soulitzer left a comment

crcrpar commented May 16, 2023

crcrpar commented May 16, 2023

crcrpar commented May 17, 2023

pytorchmergebot commented May 17, 2023

pytorchmergebot commented May 17, 2023

pytorchmergebot commented May 17, 2023

Enable backward on _foreach_zero_ #101149

Enable backward on _foreach_zero_ #101149

Conversation

crcrpar commented May 11, 2023 • edited

pytorch-bot bot commented May 11, 2023 • edited

🔗 Helpful Links

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

❌ 1 New Failure

soulitzer left a comment

Choose a reason for hiding this comment

crcrpar commented May 12, 2023

soulitzer left a comment

Choose a reason for hiding this comment

crcrpar commented May 16, 2023

crcrpar commented May 16, 2023

crcrpar commented May 17, 2023

pytorchmergebot commented May 17, 2023

pytorchmergebot commented May 17, 2023

pytorchmergebot commented May 17, 2023

Merge started

Enable `backward` on `_foreach_zero_` #101149

Enable `backward` on `_foreach_zero_` #101149

crcrpar commented May 11, 2023 •

edited

pytorch-bot bot commented May 11, 2023 •

edited