diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index 47b9d1343b..86b458e229 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -779,21 +779,17 @@ def native_batch_norm(input, weight, bias, running_mean, running_var, training, @torch.library.impl("aten::native_batch_norm_backward", "privateuseone") def native_batch_norm_backward(grad_out, input, weight, running_mean, running_var, save_mean, save_invstd, train, eps, output_mask): grad_out_t, input_t = unwrap(grad_out), unwrap(input) - weight_t = unwrap(weight) if weight is not None else None - save_mean_t = unwrap(save_mean) - save_invstd_t = unwrap(save_invstd) - out = input_t.batchnorm(weight_t, None, save_mean_t, save_invstd_t) - targets = [t for t, m in zip([input_t, weight_t], output_mask[:2]) if t is not None and m] - if targets: - grads = out.gradient(*targets, gradient=grad_out_t) - grad_input = grads.pop(0) if output_mask[0] else None - grad_weight = grads.pop(0) if output_mask[1] and weight_t is not None else None - else: - grad_input, grad_weight = None, None - grad_bias = grad_out_t.sum(axis=tuple(x for x in range(grad_out_t.ndim) if x != 1)) if output_mask[2] else None - return (wrap(grad_input) if grad_input is not None else None, - wrap(grad_weight) if grad_weight is not None else None, - wrap(grad_bias) if grad_bias is not None else None) + dims, shape = tuple(x for x in range(input_t.ndim) if x != 1), (1, -1) + (1,)*(input_t.ndim-2) + # training differentiates the batch stats it was given, eval treats the running stats as constants + if train: mean, invstd = unwrap(save_mean), unwrap(save_invstd) + else: mean, invstd = unwrap(running_mean), unwrap(running_var).add(eps).rsqrt() + xhat = (input_t - mean.reshape(shape)) * invstd.reshape(shape) + grad_bias, grad_weight = grad_out_t.sum(axis=dims), (grad_out_t * xhat).sum(axis=dims) + grad_input = grad_out_t if not train else \ + grad_out_t - (grad_bias.reshape(shape) + xhat * grad_weight.reshape(shape)) / (input_t.numel() // input_t.shape[1]) + grad_input = grad_input * invstd.reshape(shape) * (unwrap(weight).reshape(shape) if weight is not None else 1) + return (wrap(grad_input) if output_mask[0] else None, wrap(grad_weight) if output_mask[1] else None, + wrap(grad_bias) if output_mask[2] else None) # _pad_circular is not CompositeImplicitAutograd (unlike reflect/replicate pad) # we need torch.autograd.Function with explicit AutogradPrivateUse1 registration diff --git a/extra/torch_backend/test.py b/extra/torch_backend/test.py index 306cd42f8f..1aad2157a2 100644 --- a/extra/torch_backend/test.py +++ b/extra/torch_backend/test.py @@ -343,6 +343,21 @@ class TestTorchBackend(unittest.TestCase): assert b.shape == (4, 2, 3) np.testing.assert_equal(b.cpu().numpy(), a.cpu().numpy().transpose(2, 0, 1)) + def test_batchnorm_backward_realized_stats(self): + # the saved stats are a function of input in training, so grad_input must flow through them even when handed in realized. + # the backward eps is unused in training: torch differentiates the save_invstd it was given + x0, g0 = torch.randn(8, 4, 3, 3), torch.randn(8, 4, 3, 3) + def run(dev, bwd_eps): + x, go = x0.to(dev), g0.to(dev) + w, b = torch.linspace(0.5, 2.0, 4).to(dev), torch.zeros(4, device=dev) + rm, rv = torch.zeros(4, device=dev), torch.ones(4, device=dev) + out, sm, si = torch.ops.aten.native_batch_norm(x, w, b, rm, rv, True, 0.1, 1e-5) + grads = torch.ops.aten.native_batch_norm_backward(go, x, w, rm, rv, sm.clone().detach(), si.clone().detach(), + True, bwd_eps, [True,True,True]) + return [t.cpu().numpy() for t in grads] + for bwd_eps in [1e-5, 0.3]: + for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3) + def test_batchnorm_unsqueeze(self): bn = torch.nn.BatchNorm2d(4).to(device) x = torch.randn(8, 4, 3, 3, device=device) @@ -773,6 +788,12 @@ class TestBackendHelpers(unittest.TestCase): # unwrap casts to the tiny impl, so a tensor from another backend must be refused rather than reinterpreted with self.assertRaises(RuntimeError): extra.torch_backend.backend.unwrap(torch.ones(4)) + def test_update_metadata_rejects_foreign_tensor(self): + # resizing a tensor we don't own would expose memory past its allocation + t = torch.ones(4) + with self.assertRaises(RuntimeError): extra.torch_backend.backend.mod.update_metadata(t, [8], [1], 0) + self.assertEqual(t.shape, (4,)) + def test_unwrap_parameter_and_detached(self): # nn.Parameter and detach rebuild the base OpaqueTensorImpl, which unwrap still has to accept extra.torch_backend.backend.unwrap(torch.nn.Parameter(torch.ones(4, device="tiny"))) diff --git a/extra/torch_backend/wrapped_tensor.cpp b/extra/torch_backend/wrapped_tensor.cpp index 120137fab6..82f5c7a58f 100644 --- a/extra/torch_backend/wrapped_tensor.cpp +++ b/extra/torch_backend/wrapped_tensor.cpp @@ -124,18 +124,21 @@ at::Tensor wrap_tensor(py::object &py_obj, c10::ScalarType dtype, c10::DeviceInd sizes, strides, storage_offset); } +// shallow_copy_and_detach (nn.Parameter, aten.detach) rebuilds the base OpaqueTensorImpl, so that is the type every tiny tensor has +at::OpaqueTensorImpl> *tiny_impl(const at::Tensor &tensor) { + auto* impl = dynamic_cast>*>(tensor.unsafeGetTensorImpl()); + TORCH_CHECK(impl != nullptr, "expected a tiny tensor, got a ", tensor.device().str(), " one. move it with .to(\"tiny\") first"); + return impl; +} + py::object unwrap_tensor(const at::Tensor &tensor) { - auto* impl = tensor.unsafeGetTensorImpl(); - // shallow_copy_and_detach (nn.Parameter, aten.detach) rebuilds the base OpaqueTensorImpl, so that is the type every tiny tensor has - auto* opaque_impl = dynamic_cast>*>(impl); - TORCH_CHECK(opaque_impl != nullptr, "expected a tiny tensor, got a ", tensor.device().str(), " one. move it with .to(\"tiny\") first"); - std::shared_ptr tiny = opaque_impl->opaque_handle(); + std::shared_ptr tiny = tiny_impl(tensor)->opaque_handle(); return py::reinterpret_borrow(tiny->ptr(getPyInterpreter())); } void update_metadata(const at::Tensor &tensor, const std::vector &sizes, const std::vector &strides, int64_t storage_offset) { - auto* impl = tensor.unsafeGetTensorImpl(); + auto* impl = tiny_impl(tensor); impl->set_allow_tensor_metadata_change(true); impl->set_sizes_and_strides(sizes, strides, storage_offset); }