已合并
Register npuBasicAutogradNotImplementedFallbackImpl for AutogradPrivateUse1. #18042
AtomGit-Bot创建于 2025年2月18日
Register npuBasicAutogradNotImplementedFallbackImpl for AutogradPrivateUse1. #18042
已合并
从refs/pull/18042/head合入到master
共 2 个文件变更+228-1
| @@ -0,0 +1,30 @@ | |||
| 1 | +import torch | ||
| 2 | +from torch.testing._internal.common_utils import ( | ||
| 3 | + run_tests, | ||
| 4 | + TestCase, | ||
| 5 | +) | ||
| 6 | +import torch_npu | ||
| 7 | + | ||
| 8 | +class TestAutogradFallback(TestCase): | ||
| 9 | + | ||
| 10 | + def test_pad_backward_warn(self): | ||
| 11 | + | ||
| 12 | + def _exec_npu_pad(): | ||
| 13 | + npu_input = torch.randn(2, 3).npu() | ||
| 14 | + npu_input.requires_grad = True | ||
| 15 | + pads = (1, 1, 1, 1) | ||
| 16 | + output = torch_npu.npu_pad(npu_input, pads) | ||
| 17 | + output.backward(torch.ones_like(output)) | ||
| 18 | + | ||
| 19 | + # When set to "nothing," calling the reverse function directly causes an error. | ||
| 20 | + torch._C._set_autograd_fallback_mode("nothing") | ||
| 21 | + with self.assertRaisesRegex(RuntimeError, "does not require grad"): | ||
| 22 | + _exec_npu_pad() | ||
| 23 | + | ||
| 24 | + # When set to "warn," calling the print function emits a warning. | ||
| 25 | + torch._C._set_autograd_fallback_mode("warn") | ||
| 26 | + _exec_npu_pad() | ||
| 27 | + | ||
| 28 | + | ||
| 29 | +if __name__ == "__main__": | ||
| 30 | + run_tests() | ||
| @@ -2,6 +2,11 @@ | |||
| 2 | 2 | ||
| 3 | 3 | ||
| 4 | 4 | ||
| 5 | + | ||
| 6 | + | ||
| 7 | + | ||
| 8 | + | ||
| 9 | + | ||
| 5 | 10 | ||
| 6 | 11 | ||
| 7 | 12 | ||
| @@ -25,16 +30,208 @@ using c10::DispatchKey; | |||
| 25 | using c10::DispatchKeySet; | 30 | using c10::DispatchKeySet; |
| 26 | using c10::Dispatcher; | 31 | using c10::Dispatcher; |
| 27 | using c10::KernelFunction; | 32 | using c10::KernelFunction; |
| 33 | +using torch::autograd::edge_list; | ||
| 34 | +using torch::autograd::Node; | ||
| 35 | +using torch::autograd::variable_list; | ||
| 28 | 36 | ||
| 29 | namespace { | 37 | namespace { |
| 30 | 38 | ||
| 39 | +template <typename F> | ||
| 40 | +void _foreach_tensor( | ||
| 41 | + F fn, | ||
| 42 | + torch::jit::Stack* stack, | ||
| 43 | + size_t stack_start, | ||
| 44 | + size_t size) | ||
| 45 | +{ | ||
| 46 | + // Enumerate over tensors in a stack, including ones in TensorLists | ||
| 47 | + int idx_tensor = 0; | ||
| 48 | + for (const auto idx_arg : c10::irange(size)) { | ||
| 49 | + auto& ivalue = (*stack)[stack_start + idx_arg]; | ||
| 50 | + if (ivalue.isTensor()) { // true for optional tensor that has value | ||
| 51 | + const auto& tensor = ivalue.toTensor(); | ||
| 52 | + fn(idx_tensor, idx_arg, tensor); | ||
| 53 | + idx_tensor++; | ||
| 54 | + } else if (ivalue.isTensorList()) { | ||
| 55 | + for (const auto& iv : ivalue.toListRef()) { | ||
| 56 | + const auto& tensor = iv.toTensor(); | ||
| 57 | + fn(idx_tensor, idx_arg, tensor); | ||
| 58 | + idx_tensor++; | ||
| 59 | + } | ||
| 60 | + } | ||
| 61 | + } | ||
| 62 | +} | ||
| 63 | + | ||
| 64 | + | ||
| 65 | +static void warnAutogradNotImplemented(const std::string& op_name) | ||
| 66 | +{ | ||
| 67 | + TORCH_NPU_WARN_ONCE( | ||
| 68 | + op_name, | ||
| 69 | + ": an autograd kernel was not registered to the Autograd key(s) ", | ||
| 70 | + "but we are trying to backprop through it. This may lead to silently incorrect behavior. ", | ||
| 71 | + "This behavior is deprecated and will be removed in a future version of PyTorch. ", | ||
| 72 | + "If your operator is differentiable, please ensure you have registered an " | ||
| 73 | + "autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, " | ||
| 74 | + "DispatchKey::CompositeImplicitAutograd). If your operator is not " | ||
| 75 | + "differentiable, or to squash this warning and use the previous behavior, " | ||
| 76 | + "please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd."); | ||
| 77 | +} | ||
| 78 | + | ||
| 79 | + | ||
| 80 | +struct WarnNotImplemented : public Node { | ||
| 81 | + WarnNotImplemented( | ||
| 82 | + std::string op_name, | ||
| 83 | + int64_t num_outputs, | ||
| 84 | + edge_list&& next_edges) | ||
| 85 | + : Node(std::move(next_edges)), op_name(std::move(op_name)), num_outputs(num_outputs) {} | ||
| 86 | + | ||
| 87 | + WarnNotImplemented(std::string op_name, int64_t num_outputs) | ||
| 88 | + : op_name(std::move(op_name)), num_outputs(num_outputs) {} | ||
| 89 | + | ||
| 90 | + variable_list apply(variable_list&& inputs) override; | ||
| 91 | + | ||
| 92 | + std::string op_name; | ||
| 93 | + int64_t num_outputs; | ||
| 94 | +}; | ||
| 95 | + | ||
| 96 | +auto WarnNotImplemented::apply(variable_list&& inputs) -> variable_list | ||
| 97 | +{ | ||
| 98 | + warnAutogradNotImplemented(op_name); | ||
| 99 | + std::vector<at::Tensor> output(num_outputs); | ||
| 100 | + return output; | ||
| 101 | +} | ||
| 102 | + | ||
| 103 | +static void npuBasicAutogradNotImplementedFallbackImpl( | ||
| 104 | + const c10::OperatorHandle& op, | ||
| 105 | + c10::DispatchKeySet dispatch_keys, | ||
| 106 | + torch::jit::Stack* stack) | ||
| 107 | +{ | ||
| 108 | + const auto& schema = op.schema(); | ||
| 109 | + const auto& op_name = schema.operator_name().name; | ||
| 110 | + const auto num_arguments = schema.arguments().size(); | ||
| 111 | + const auto num_returns = schema.returns().size(); | ||
| 112 | + const auto stack_start = stack->size() - num_arguments; | ||
| 113 | + | ||
| 114 | + if (torch::autograd::getAutogradFallbackMode() == torch::autograd::AutogradFallbackMode::Nothing) { | ||
| 115 | + op.redispatchBoxed(dispatch_keys & c10::after_autograd_keyset, stack); | ||
| 116 | + return; | ||
| 117 | + } | ||
| 118 | + TORCH_INTERNAL_ASSERT( | ||
| 119 | + torch::autograd::getAutogradFallbackMode() == torch::autograd::AutogradFallbackMode::Warn); | ||
| 120 | + | ||
| 121 | + bool any_input_requires_grad = false; | ||
| 122 | + _foreach_tensor( | ||
| 123 | + [&](size_t _, size_t idx_arg, const at::Tensor& t) { | ||
| 124 | + if (t.requires_grad()) { | ||
| 125 | + any_input_requires_grad = true; | ||
| 126 | + } | ||
| 127 | + }, | ||
| 128 | + stack, | ||
| 129 | + stack_start, | ||
| 130 | + num_arguments); | ||
| 131 | + // Optimization: TLS access can be slow. So we only check if it necessary | ||
| 132 | + // by putting it after the requires_grad checks. | ||
| 133 | + any_input_requires_grad = any_input_requires_grad && at::GradMode::is_enabled(); | ||
| 134 | + | ||
| 135 | + std::shared_ptr<WarnNotImplemented> grad_fn; | ||
| 136 | + if (any_input_requires_grad) { | ||
| 137 | + // NB: It is standard to collect edges from all tensors | ||
| 138 | + // (see generated/VariableTypeEverything.cpp for examples) | ||
| 139 | + std::vector<const at::Tensor*> all_tensors_on_stack; | ||
| 140 | + _foreach_tensor( | ||
| 141 | + [&](size_t _, size_t idx_arg, const at::Tensor& t) { | ||
| 142 | + all_tensors_on_stack.push_back(&t); | ||
| 143 | + }, | ||
| 144 | + stack, | ||
| 145 | + stack_start, | ||
| 146 | + num_arguments); | ||
| 147 | + grad_fn = std::shared_ptr<WarnNotImplemented>( | ||
| 148 | + new WarnNotImplemented(op_name, all_tensors_on_stack.size()), | ||
| 149 | + torch::autograd::deleteNode); | ||
| 150 | + grad_fn->set_next_edges(torch::autograd::collect_next_edges(all_tensors_on_stack)); | ||
| 151 | + } | ||
| 152 | + | ||
| 153 | + op.redispatchBoxed(dispatch_keys & c10::after_autograd_keyset, stack); | ||
| 154 | + | ||
| 155 | + if (any_input_requires_grad) { | ||
| 156 | + // NB: if the operator mutates any inputs in-place and does not return them | ||
| 157 | + // as outputs, we are unable to lazily raise a warning. This is OK because | ||
| 158 | + // we don't expect many existing operators to do this because of the amount | ||
| 159 | + // of technical expertise necessary (you would need to manually register an | ||
| 160 | + // autograd kernel without using autograd.Function) | ||
| 161 | + _foreach_tensor( | ||
| 162 | + [&](size_t _, size_t idx_ret, const at::Tensor& t) { | ||
| 163 | + if (!torch::autograd::isDifferentiableType(t.scalar_type())) { | ||
| 164 | + return; | ||
| 165 | + } | ||
| 166 | + const bool is_mutable_output = | ||
| 167 | + schema.is_aliasing({c10::SchemaArgType::output, idx_ret}) && | ||
| 168 | + schema.is_mutable({c10::SchemaArgType::output, idx_ret}); | ||
| 169 | + | ||
| 170 | + // If the post-autograd implementation returns Tensors that require | ||
| 171 | + // grad, then we install a hook that will warn during the backwards. | ||
| 172 | + // | ||
| 173 | + // NB: If the operation is inplace and the inputs were views, | ||
| 174 | + // it is possible that the history was rebased and the hook will | ||
| 175 | + // not warn in all places where it should. That is, the following | ||
| 176 | + // won't warn: | ||
| 177 | + // >>> x = torch.randn(3, 3, requires_grad=True) | ||
| 178 | + // >>> z = x.clone() | ||
| 179 | + // >>> w = z[0] | ||
| 180 | + // >>> k = w[0] | ||
| 181 | + // >>> y = op(k) | ||
| 182 | + // >>> torch.autograd.grad(z.sum(), w) | ||
| 183 | + if (t.requires_grad()) { | ||
| 184 | + t.register_hook([op_name](const at::Tensor& grad) { | ||
| 185 | + warnAutogradNotImplemented(op_name); | ||
| 186 | + }); | ||
| 187 | + // If history is rebased, then we will attempt to warn | ||
| 188 | + // on the view's base. This will catch most cases (because | ||
| 189 | + // users typically call .backward() and backprop through | ||
| 190 | + // the entire program). | ||
| 191 | + if (t.is_view() && is_mutable_output) { | ||
| 192 | + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) | ||
| 193 | + auto& base = const_cast<at::TensorBase&>(t._base()); | ||
| 194 | + if (base.requires_grad()) { | ||
| 195 | + // Can only register_hook on tensors that require grad. | ||
| 196 | + base.register_hook([op_name](const at::TensorBase& grad) { | ||
| 197 | + warnAutogradNotImplemented(op_name); | ||
| 198 | + }); | ||
| 199 | + } | ||
| 200 | + } | ||
| 201 | + return; | ||
| 202 | + } | ||
| 203 | + | ||
| 204 | + // If the post-autograd implementation returns any Tensors that | ||
| 205 | + // don't require grad, then we install the WarnNotImplemented grad_fn. | ||
| 206 | + // This grad_fn warns in backward and returns undefined tensor | ||
| 207 | + // gradients. | ||
| 208 | + // | ||
| 209 | + // NOTE [autograd fallback and in-place operations] | ||
| 210 | + // If the schema says the output is mutable, and the output | ||
| 211 | + // is an input, and the input is a view Tensor, then... | ||
| 212 | + // we're not sure if set_history is OK to do, so we just skip | ||
| 213 | + // adding the grad_fn. Builtin operators do rebase_history here, | ||
| 214 | + // but custom operators may have multiple Tensor(a!) returns, | ||
| 215 | + // rebase_history assumes single Tensor(a!) return, and in general | ||
| 216 | + // custom ops don't have a good in-place story. | ||
| 217 | + if (!is_mutable_output) { | ||
| 218 | + // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast) | ||
| 219 | + torch::autograd::set_history(const_cast<at::Tensor&>(t), grad_fn); | ||
| 220 | + } | ||
| 221 | + }, | ||
| 222 | + stack, | ||
| 223 | + stack->size() - num_returns, | ||
| 224 | + num_returns); | ||
| 225 | + } | ||
| 226 | +} | ||
| 227 | + | ||
| 31 | // Register fallthrough for Autograd backends dispatch keys | 228 | // Register fallthrough for Autograd backends dispatch keys |
| 32 | // NB: But not the private use ones; maybe the extension wants | 229 | // NB: But not the private use ones; maybe the extension wants |
| 33 | // to override it themselves! | 230 | // to override it themselves! |
| 34 | 231 | ||
| 35 | // (Ascend) TORCH_LIBRARY_IMPL | 232 | // (Ascend) TORCH_LIBRARY_IMPL |
| 36 | TORCH_LIBRARY_IMPL(_, AutogradPrivateUse1, m) { | 233 | TORCH_LIBRARY_IMPL(_, AutogradPrivateUse1, m) { |
| 37 | - m.fallback(torch::CppFunction::makeFallthrough()); | 234 | + m.fallback(torch::CppFunction::makeFromBoxedFunction<&npuBasicAutogradNotImplementedFallbackImpl>()); |
| 38 | } | 235 | } |
| 39 | 236 | ||
| 40 | bool has_op_name_warned(const std::string& op_name) | 237 | bool has_op_name_warned(const std::string& op_name) |