diff --git a/test/test_nn.py b/test/test_nn.py index 14d4eed971d2e..e7a898f0cf22c 100644 --- a/test/test_nn.py +++ b/test/test_nn.py @@ -7437,10 +7437,15 @@ def test_layer_norm_backwards_eps(self): ln_out_cuda = ln_cuda(x_cuda) ln_out.backward(grad_output) ln_out_cuda.backward(grad_output_cuda) + atol = 1e-4 + rtol = 1e-5 + if m > 64 * 1024: + atol = 1e-3 + rtol = 1e-3 if elementwise_affine: - self.assertEqual(ln.weight.grad, ln_cuda.weight.grad, f"weight grad failed: {m=} {n=}", rtol=1e-4, atol=1e-4) + self.assertEqual(ln.weight.grad, ln_cuda.weight.grad, f"weight grad failed: {m=} {n=}", rtol=rtol, atol=atol) if bias and elementwise_affine: - self.assertEqual(ln.bias.grad, ln_cuda.bias.grad, f"bias grad failed: {m=} {n=}", rtol=1e-5, atol=1e-4) + self.assertEqual(ln.bias.grad, ln_cuda.bias.grad, f"bias grad failed: {m=} {n=}", rtol=rtol, atol=atol) @largeTensorTest("40GB", device="cuda") def test_layer_norm_large_tensor(self):