Repository: pytorch
Issue: test_layer_norm_backward and test_layer_norm_backward_5d run OOM in slow gradcheck
Labels: triaged,module: nestedtensor

### 🐛 Describe the bug

The two new tests test_layer_norm_backward and test_layer_norm_backward_5d in TestNestedTensorAutograd added by https://github.com/pytorch/pytorch/pull/94781 consistently run OOM in slow gradcheck mode. This occurs even after the job has been updated to run the test sequentially in https://github.com/pytorch/pytorch/pull/95494.

Here is an example failure https://github.com/pytorch/pytorch/actions/runs/4268550203/jobs/7431239137:

```
  test_layer_norm_backward_5d_size_128_cuda (__main__.TestNestedTensorAutogradCUDA) ...     test_layer_norm_backward_5d_size_128_cuda errored - num_retries_left: 3
Traceback (most recent call last):
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 2081, in wrapper
    method(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 2081, in wrapper
    method(*args, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_device_type.py", line 414, in instantiated_test
    raise rte
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_device_type.py", line 401, in instantiated_test
    result = test(self, **param_kwargs)
  File "/var/lib/jenkins/workspace/test/test_nestedtensor.py", line 2405, in test_layer_norm_backward_5d
    assert gradcheck(grad_test_func, inputs=data, check_batched_grad=False)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 3804, in gradcheck
    return torch.autograd.gradcheck(fn, inputs, **kwargs)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 1476, in gradcheck
    return _gradcheck_helper(**args)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 1490, in _gradcheck_helper
    _gradcheck_real_imag(gradcheck_fn, func, func_out, tupled_inputs, outputs, eps,
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 1113, in _gradcheck_real_imag
    gradcheck_fn(func, func_out, tupled_inputs, outputs, eps,
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 1150, in _slow_gradcheck
    numerical = _transpose(_get_numerical_jacobian(func, tupled_inputs, func_out, eps=eps, is_forward_ad=use_forward_ad))
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 187, in _get_numerical_jacobian
    jacobians += [get_numerical_jacobian_wrt_specific_input(fn, inp_idx, inputs, outputs, eps,
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 330, in get_numerical_jacobian_wrt_specific_input
    jacobian_cols[d_idx] = _compute_numerical_jvps_wrt_specific_input(jvp_fn, eps, x.is_complex(), is_forward_ad)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 251, in _compute_numerical_jvps_wrt_specific_input
    ds_dx_tup = jvp_fn(delta[0] if isinstance(delta, tuple) else delta)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 446, in jvp_fn
    return _compute_numerical_gradient(wrapped_fn, input_to_perturb, delta, eps, nbhd_checks_fn)
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 240, in _compute_numerical_gradient
    return tuple(compute(a, b) for (a, b) in zip(outa, outb))
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 240, in <genexpr>
    return tuple(compute(a, b) for (a, b) in zip(outa, outb))
  File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/autograd/gradcheck.py", line 237, in compute
    ret = (b - a) / (2 * norm_v)
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.00 MiB (GPU 0; 7.44 GiB total capacity; 6.98 GiB already allocated; 10.81 MiB free; 7.02 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

ERROR (2.276s)
```

### Versions

CI periodic slow gradcheck test

cc @seemethere @malfet @pytorch/pytorch-dev-infra @cpuhrsch @jbschlosser @bhosmer @drisspg @mikaylagawarecki