25% Performance regression from v0.1.1 to 0.2.0 when calculating hessian

[yueyericardo]

Hi developers,

After I upgraded functorch from v0.1.1 to 0.2.0, I noticed a 25% performance regression when calculating hessian, please check the following benchmark result and the attached benchmark script.

Please let me know if I did anything wrong, and also whether the perf regression could be fixed.
Thanks!

Benchmark result

Benchmark result on NVIDIA A100

# torch 111 and functorch 0.1.1
===== benchmark without backward =====
max pred       error: functorch: 0.00e+00
max hessian    error: functorch: 0.00e+00
reference_hessian: 61.837 ms
functorch_hessian: 29.474 ms

# torch 112 and functorch 0.2.0
===== benchmark without backward =====
max pred       error: functorch: 1.49e-08
max hessian    error: functorch: 0.00e+00
reference_hessian: 62.519 ms
functorch_hessian: 39.666 ms  (0.75 X)

Benchmark result on NVIDIA A6000

# torch 111 and functorch 0.1.1
===== benchmark without backward =====
max pred       error: functorch: 1.49e-08
max hessian    error: functorch: 0.00e+00
reference_hessian: 65.984 ms
functorch_hessian: 33.662 ms

# torch 112 and functorch 0.2.0
===== benchmark without backward =====
max pred       error: functorch: 1.86e-08
max hessian    error: functorch: 0.00e+00
reference_hessian: 67.285 ms
functorch_hessian: 49.723 ms (0.68 X)

benchmark script

benchmark.py

import time
import argparse
from functorch import vmap, jacrev, jacfwd
import torch
import torch.nn as nn

torch.backends.cuda.matmul.allow_tf32 = False


_ = torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
D1 = 2  # x, y
D2 = 3  # u, v, p
B = 10000
x = torch.randn(B, D1).to(device)
run_backward = False

model = nn.Sequential(
    nn.Linear(D1, 512),
    nn.ReLU(),
    nn.Linear(512, 512),
    nn.ReLU(),
    nn.Linear(512, 512),
    nn.ReLU(),
    nn.Linear(512, 512),
    nn.ReLU(),
    nn.Linear(512, 512),
    nn.ReLU(),
    nn.Linear(512, 512),
    nn.ReLU(),
    nn.Linear(512, D2),
).to(device)


def predict(x):
    torch.cuda.nvtx.range_push("forward")
    out = model(x)
    torch.cuda.nvtx.range_pop()
    return out, out  # return two outputs is needed for jacrev auxiliary object


def reference_hessian():
    x_ = x.clone().requires_grad_()
    ones = torch.ones(B, device=x.device)
    pred, _ = predict(x_)
    jacobian_rows = [None] * D2
    hessian_rows = [None] * (D2 * D1)
    for i in range(D2):
        torch.cuda.nvtx.range_push("autograd jacobian")
        jacobian_rows[i] = torch.autograd.grad(pred[:, i], x_, ones, create_graph=True)[
            0
        ]
        torch.cuda.nvtx.range_pop()

    for i in range(D2):
        for j in range(D1):
            torch.cuda.nvtx.range_push("autograd hesian")
            hessian_rows[i * D1 + j] = torch.autograd.grad(
                jacobian_rows[i][:, j], x_, ones, create_graph=True
            )[0]
            torch.cuda.nvtx.range_pop()

    jacobian = torch.stack(jacobian_rows)  # [D2, B, D1]
    hessian = torch.stack(hessian_rows)  # [D2 * D1, B, D1]
    if run_backward:
        l = hessian.sum()
        l.backward()
    return hessian.transpose(0, 1), pred


def functorch_hessian():
    x_ = x.clone().requires_grad_()
    hessian, pred = vmap(
        jacfwd(jacrev(predict, argnums=0, has_aux=True), argnums=0, has_aux=True),
        in_dims=0,
    )(
        x_
    )  # [B, D2, D1, D1]
    if run_backward:
        l = hessian.sum()
        l.backward()
    return hessian, pred


def validate_result():
    # test functorch result
    ref_hes, ref_pred = reference_hessian()
    ft_hes, ft_pred = functorch_hessian()
    ref_hes = ref_hes.view_as(ft_hes)
    print(f"max pred       error: functorch: {(ref_pred - ft_pred).max():.2e}")
    print(f"max hessian    error: functorch: {(ref_hes - ft_hes).max():.2e}")


def benchmark(func):
    N = 20

    torch.cuda.synchronize()
    start = time.time()

    for i in range(N):
        torch.cuda.nvtx.range_push(func.__name__)
        _ = func()
        torch.cuda.nvtx.range_pop()

    torch.cuda.synchronize()
    time_ms = ((time.time() - start) / N) * 1000
    print(f"{func.__name__}: {time_ms:.3f} ms")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--backward", default=False, action="store_true")
    args = parser.parse_args()
    if args.backward:
        run_backward = True
        print("===== benchmark with backward =====")
    else:
        print("===== benchmark without backward =====")

    validate_result()

    # warm up
    for i in range(10):
        reference_hessian()
        functorch_hessian()

    # benchmark hessian
    benchmark(reference_hessian)
    benchmark(functorch_hessian)

[yueyericardo]

ping @samdow @zou3519

[zou3519]

Thanks for the report, we'll take a look soon

[zou3519]

Bisected to https://github.com/pytorch/pytorch/pull/75195/files. https://github.com/pytorch/pytorch/pull/75195/files by itself may not be a problem, perhaps the problem is our batching rule for mv.

@yueyericardo is the repro you provided the entire model, or is it a subset of some model that you're running?

[zou3519]

cc @lezcano @ezyang for pytorch/pytorch#75195 -- this led to a performance regression in functorch. I'm not sure what the original intent of the PR is (there are no tests). I'm still trying to root cause this, but it is a bit difficult to visualize. What are the chances we could revert that PR?

(I've confirmed that reverting that single PR on pytorch/pytorch master makes the performance regression go away)

[yueyericardo]

Great thanks to @zou3519 for the quick debugging!!
I'm working on NVIDIA Modulus project, we are using functorch because it provides a lot of perfs for the Jacobian and Hessian calculations.
The minimal repro I provided is only a subset of our model to demonstrate the performance regression.

Thanks again!

[zou3519]

@yueyericardo - for the original model itself, is the performance regression also 25%, or is it a smaller number? Is the original model public? One thing we can do to prevent future regressions is to check the original model into https://github.com/pytorch/benchmark.

I've noticed a lot of other similar models where folks have a nn.Sequential that is just made up of nn.Linear and activations and need to compute a vmap(jacrev or vmap(hessian of the quantity, so we could also potentially just check your script into torchbench otherwise.

[yueyericardo]

Hi @zou3519 Our source code is free to download from the website, but it is not developed on GitHub. And our code base is also might too large to put into pytorch/benchmark.

Yes, exactly! I believe the minimal repro I provided is enough to prevent future regression for our model.
Thanks!

[lezcano]

I think that rather than blindly reverting, we should get to the root of the problem, as it is very weird to get such a regression when dispatching from a more general function to a more concrete (that was the reason for that PR).

Things that come to mind are:

• Is this regression architecture / cublas version - dependent?
• Is this regresion also happening on regular matmul for that codepath?

If the answer to the above two is no, then this performance issue is likely on the functorch end and should be fixed. Otherwise, it's on the cuBLAS end and should be reported to NVIDIA

cc @ngimel @xwang233 @IvanYashchuk

[ezyang]

@lezcano It's fair to submit the upstream bugs, but if we know that our upstream library's general kernel has better perf than a specialized one, we might as well use it.

[zou3519]

@lezcano @ezyang let's say that we did revert the PR (because we're trying to release PyTorch 1.12.1 as soon as possible). Would it cause any other problems?

Because the motivation was "dispatching from a more general function to a more concrete [function]", it sounds like this wouldn't change very much else.

[ezyang]

No, I don't think so. The PR is supposed to make the kernel run faster.

[lezcano]

fwiw, I think this may be related to the open PR I have to avoid copies in matmul. Could you check whether pytorch/pytorch#76828 fixes this?

In any case, I'm fine with reverting, but we should investigate what's causing this regardless

[ngimel]

@lezcano @ezyang let's say that we did revert the PR (because we're trying to release PyTorch 1.12.1 as soon as possible). Would it cause any other problems?

Because the motivation was "dispatching from a more general function to a more concrete [function]", it sounds like this wouldn't change very much else.

That PR is fixing spurious resize warnings that were previously generated, and by itself is supposed to speedup things by avoiding squeeze/unsqueeze calls which are not free (and especially not free when autograd is needed). As for more general/more concrete function performance, we should investigate this, but I doubt that's the case.

[ngimel]

@zou3519 can you by any chance collect profiling results for old and new versions?

[zou3519]

In any case, I'm fine with reverting, but we should investigate what's causing this regardless

I agree this warrants more investigation. We've got a problem in that there is a timeline for 1.12.1, and I am not sure how long it is going to take to actually get to the bottom of this.

@zou3519 can you by any chance collect profiling results for old and new versions?

I can try but it's been a long time since I touched nvprof or nsight profiler, so I will need to relearn the magic invocations.

but we should investigate what's causing this regardless

Since we changed the mm to mv, functorch generates different code for the vmap(jacrev(jacfwd(mm)) as opposed to vmap(jacrev(jacfwd(mv)). It's plausible that the problem is that "functorch should generate better code"; we're still digging into it

[ngimel]

Isn't 1.12.1 done already?
No nsight needed, just torch profiler should be enough. (with torch.profiler.profile as p(): and print key averages and export chrome trace in the end).

[yueyericardo]

FYI, the nsight profiling result
before

Time	Total Time	Instances	Avg	Med	Min	Max	StdDev	Name
64.7%	928.624 ms	1000	928.623 μs	867.621 μs	297.345 μs	1.686 ms	493.406 μs	ampere_sgemm_128x64_nn
10.3%	147.304 ms	250	589.217 μs	589.187 μs	588.643 μs	590.435 μs	275 ns	ampere_sgemm_128x128_tn
9.4%	134.966 ms	1300	103.819 μs	70.432 μs	5.344 μs	194.209 μs	75.470 μs	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::BinaryFunctor<float, float, float, at::native::AddFunctor<float>>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
6.8%	98.146 ms	900	109.050 μs	94.848 μs	70.081 μs	177.921 μs	43.651 μs	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::BinaryFunctor<float, float, float, void at::native::threshold_kernel_impl<float>(at::TensorIteratorBase &, T1, T1)::[lambda(float, float) (instance 1)]>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
5.4%	77.550 ms	350	221.572 μs	300.321 μs	18.080 μs	302.050 μs	124.914 μs	ampere_sgemm_128x64_tn
0.8%	10.981 ms	1450	7.572 μs	3.552 μs	2.880 μs	26.112 μs	7.973 μs	void at::native::vectorized_elementwise_kernel<(int)4, at::native::FillFunctor<float>, at::detail::Array<char *, (int)1>>(int, T2, T3)
0.7%	9.949 ms	150	66.324 μs	82.928 μs	27.968 μs	88.672 μs	26.863 μs	ampere_sgemm_32x32_sliced1x4_nn
0.6%	8.074 ms	300	26.912 μs	26.976 μs	25.696 μs	27.776 μs	390 ns	void at::native::vectorized_elementwise_kernel<(int)4, at::native::<unnamed>::launch_clamp_scalar(at::TensorIteratorBase &, c10::Scalar, c10::Scalar, at::native::detail::ClampLimits)::[lambda() (instance 1)]::operator ()() const::[lambda() (instance 8)]::operator ()() const::[lambda(float) (instance 1)], at::detail::Array<char *, (int)2>>(int, T2, T3)
0.5%	7.713 ms	50	154.257 μs	154.177 μs	153.184 μs	155.905 μs	634 ns	ampere_sgemm_32x128_nn
0.3%	4.847 ms	100	48.468 μs	47.664 μs	34.688 μs	63.393 μs	13.397 μs	ampere_sgemm_32x32_sliced1x4_tn
0.2%	3.342 ms	650	5.141 μs	5.344 μs	4.256 μs	5.728 μs	423 ns	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::<unnamed>::direct_copy_kernel_cuda(at::TensorIteratorBase &)::[lambda() (instance 2)]::operator ()() const::[lambda() (instance 8)]::operator ()() const::[lambda(float) (instance 1)]>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
0.1%	2.028 ms	50	40.569 μs	40.512 μs	40.032 μs	41.536 μs	272 ns	ampere_sgemm_64x64_nn
0.1%	994.305 μs	50	19.886 μs	19.856 μs	19.649 μs	20.896 μs	169 ns	ampere_sgemm_128x32_nn
0.0%	435.235 μs	100	4.352 μs	4.384 μs	4.160 μs	4.576 μs	100 ns	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::FillFunctor<float>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)

after (the third row 18.4% on a copy kernel)

Time	Total Time	Instances	Avg	Med	Min	Max	StdDev	Name
21.8%	433.409 ms	500	866.818 μs	866.756 μs	865.028 μs	870.116 μs	881 ns	ampere_sgemm_32x128_nt
21.2%	420.235 ms	250	1.681 ms	1.681 ms	1.678 ms	1.685 ms	813 ns	ampere_sgemm_128x64_nn
18.4%	365.797 ms	2850	128.349 μs	83.809 μs	4.096 μs	608.931 μs	173.023 μs	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::<unnamed>::direct_copy_kernel_cuda(at::TensorIteratorBase &)::[lambda() (instance 2)]::operator ()() const::[lambda() (instance 14)]::operator ()() const::[lambda(float) (instance 1)]>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
10.8%	213.750 ms	300	712.501 μs	588.802 μs	588.195 μs	1.335 ms	277.174 μs	ampere_sgemm_128x128_tn
8.1%	160.586 ms	1300	123.527 μs	93.057 μs	5.089 μs	266.497 μs	92.181 μs	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::CUDAFunctor_add<float>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
7.9%	156.152 ms	500	312.304 μs	312.290 μs	310.145 μs	314.689 μs	781 ns	ampere_sgemm_128x32_nn
5.9%	117.864 ms	900	130.960 μs	129.824 μs	77.504 μs	203.681 μs	49.220 μs	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::BinaryFunctor<float, float, float, void at::native::<unnamed>::threshold_kernel_impl<float>(at::TensorIteratorBase &, T1, T1)::[lambda(float, float) (instance 1)]>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)
3.3%	66.223 ms	50	1.324 ms	1.324 ms	1.324 ms	1.326 ms	397 ns	ampere_sgemm_128x128_tt
1.1%	21.094 ms	1850	11.402 μs	3.584 μs	2.816 μs	65.857 μs	12.662 μs	void at::native::vectorized_elementwise_kernel<(int)4, at::native::FillFunctor<float>, at::detail::Array<char *, (int)1>>(int, T2, T3)
0.4%	8.071 ms	300	26.904 μs	26.976 μs	25.792 μs	28.128 μs	392 ns	void at::native::vectorized_elementwise_kernel<(int)4, at::native::<unnamed>::launch_clamp_scalar(at::TensorIteratorBase &, c10::Scalar, c10::Scalar, at::native::detail::ClampLimits)::[lambda() (instance 1)]::operator ()() const::[lambda() (instance 14)]::operator ()() const::[lambda(float) (instance 1)], at::detail::Array<char *, (int)2>>(int, T2, T3)
0.4%	7.903 ms	50	158.063 μs	158.049 μs	156.609 μs	159.617 μs	634 ns	ampere_sgemm_32x128_nn
0.4%	7.359 ms	100	73.594 μs	73.600 μs	72.000 μs	78.368 μs	1.223 μs	ampere_sgemm_32x32_sliced1x4_nt
0.2%	3.242 ms	50	64.844 μs	64.928 μs	62.464 μs	65.761 μs	528 ns	ampere_sgemm_32x32_sliced1x4_tn
0.1%	2.255 ms	100	22.551 μs	22.688 μs	20.865 μs	24.192 μs	1.106 μs	void gemmSN_NN_kernel<float, (int)256, (int)4, (int)2, (int)8, (int)3, (int)4, (bool)0, cublasGemvTensorStridedBatched<const float>, cublasGemvTensorStridedBatched<const float>, cublasGemvTensorStridedBatched<float>>(cublasGemmSmallNParams<T9, T10, T11, T1>)
0.1%	2.043 ms	100	20.427 μs	20.416 μs	20.160 μs	22.656 μs	257 ns	ampere_sgemm_128x32_tn
0.0%	432.292 μs	100	4.322 μs	4.336 μs	4.128 μs	4.640 μs	97 ns	void at::native::elementwise_kernel<(int)128, (int)2, void at::native::gpu_kernel_impl<at::native::FillFunctor<float>>(at::TensorIteratorBase &, const T1 &)::[lambda(int) (instance 1)]>(int, T3)

[lezcano]

@ngimel the PR that fixed the warnings was already merged, this one is just concerned about avoiding copies. One of the cases where it elides a copy is when you multiply a matrix by a batch of matrices. This is exactly the batched version of the vector-matrix product that"s causing the regression. That's why I think it may fix it

[ngimel]

That must be coming from functorch, as that PR doesn't introduce any additional copies.

[lezcano]

yup, I think pytorch/pytorch#76828 would fix the regression

[zou3519]

I patched in pytorch/pytorch#76828 and the above script ends up OOM-ing :/

[zou3519]

Isn't 1.12.1 done already?

Not yet, we have a chance to change it (or request a 1.12.2 if necessary since this regression is large and these types of models are typical functorch usage)

[lezcano]

Re. OOM. Wow, that's certainly unexpected. I'm not sure what's the best way to follow up on that. Probably @ngimel has a better idea how to proceed.

Regardless, landing that PR (the avoid copies matmul...) will be tricky due to some discrepancies in the accuracy of mm vs mv that we found for float16. As such, I think the most realistic way forward would be to revert the offending PR and investigate what's causing that OMM after 1.12.1 is released

[ngimel]

the PR that fixed the warnings was already merged

I don't think that's true, #75195 itself is fixing a warning (otherwise a user-supplied correct 1d out was sent to mm, and mm complained about resizing it to 2d).

[lezcano]

I think the one that fixed the out version was pytorch/pytorch#75197 and a previous one in that stack, but it may be the case that 75195 also fixed an out warning, I don't remember now.

[zou3519]

@yueyericardo I got access to the modulus repo -- could you point me to which model in the modulus repo contains the nn.Sequential above please? We're figuring out how to check the above code into torchbench and I'm looking for some more context -- is there a name for computing the hessian for the nn.Sequential? Do you generally run .backward() after computing the hessian? What are some representative input sizes (are the sizes in the benchmark script representative?)

[yueyericardo]

Hi @zou3519, thanks for the following up! We are having some internal discussions regarding this and will come back to you tomorrow.

[samdow]

From looking at this a bit, I think what happened is:

• pytorch changes unsqueeze + mm -> mv in case where first element is vector
• this example hits that codepath under vmap => functorch's mv kernel
• vector is batched and matrix is not => functorch moves the vector's batch dim to the end and matmuls. Batch dim is now the 1st dimension, not the 0th
• relu's output is now BatchedTensor([512, 10000], lvl=0, dim=1). This needs to be saved as an activation for the backwards

During backwards,

• grad_output is a BatchedTensor(BatchedTensor([10000, 512, 6], lvl=0, dim =0), lvl=1, dim=2)
• threshold_backward (relu's backward) takes in grad_output and relu (binary pointwise batch rule). Starts with the lvl 1 batch rule
  • grad_output is the only value batched at the highest level, we move its batch dim to the front with a movedim. grad_output is no longer contiguous
  • redispatch on threshold_backward for lvl 0
  • move relu's bdim to front with a movedim and pad to be 3D. relu is no longer contiguous grad_output does not need to since it's lvl 0 bdim is 0. It is still contiguous
  • the view on relu is not contiguous, the view on grad_output is not contiguous => threshold_backward triggers copy
    Locally I've tested where only one of the inputs is not contiguous and haven't seen the same slowdowns.

We can also validate that this is our issue since we hit pre-regression performance numbers by changing functorch's mv batch rule here:

  auto other_ = moveBatchDimToFront(other, other_bdim);
  auto self_ = at::movedim(self, 0, 1);
  auto result = at::matmul(other_, self_);
  return std::make_tuple( std::move(result), 0 );

This doesn't trigger the copy since the batch dimension for the saved relu activation is now the first dimension. However, this may hit other perf issue from the transposing both self and result

---

This is the smallest subset of @zou3519's trace where I can see the perf differences

mm_2 = torch.randn(60000, 512)
relu = torch.randn(10000, 512)
def old_linear_faster():
    mm_2_view = mm_2.view([10000, 6, 1, 512])
    mm_2_view_squeeze = mm_2_view.squeeze(2)

    relu_view = relu.view([10000, 1, 512])

    threshold_backward = torch.ops.aten.threshold_backward(mm_2_view_squeeze, relu_view, 0)
    return threshold_backward

bmm = torch.randn(10000, 512, 6)
other_relu = torch.randn(512, 10000)
def new_linear_faster():
    bmm_view = bmm.view([10000, 512, 6])
    bmm_view_permute = bmm_view.permute([0, 2, 1])

    other_relu_permute = other_relu.permute([1, 0])
    other_relu_permute_view = other_relu_permute.view([10000, 1, 512])

    threshold_backward = torch.ops.aten.threshold_backward(bmm_view_permute, other_relu_permute_view, 0)
    return threshold_backward

notably, if If I change the views on either bmm or other_relu to be contiguous constants, it has much faster performance. So it seems like threshold_backwards doesn't copy if only one of its inputs is not contiguous but does if both are not

[ngimel]

Thanks @samdow, copy from relu definitely seems to affect perf, however there's also another copy coming from MvBackward

Also, why does threshold_backward on discontiguous inputs trigger a copy? In eager threshold_backward should be able to handle them via tensorIterator

[lezcano]

As discussed, I think that mv_backward copy_ would be fixed by pytorch/pytorch#76828. Now, Richard reports in #989 (comment) that patching in this fix results in OOM, so it's not clear what's going on there.

In my opinion, the path forward to fix this regression would be to:

1. Figure out why is it OOM. I'm not sure how to tackle this, but perhaps @ngimel can help here
2. Unblock and land Avoid copies in matmul pytorch#76828. I will try to get a small repro tomorrow and pass it on to the NVIDIA folks for them to investigate further.

[zou3519]

@yueyericardo I got access to the modulus repo -- could you point me to which model in the modulus repo contains the nn.Sequential above please? We're figuring out how to check the above code into torchbench and I'm looking for some more context -- is there a name for computing the hessian for the nn.Sequential? Do you generally run .backward() after computing the hessian? What are some representative input sizes (are the sizes in the benchmark script representative?)

@yueyericardo quick bump on the above. We're looking to merge some form of the above code into our benchmark suite and additional information would be very helpful

[yueyericardo]

@yueyericardo quick bump on the above. We're looking to merge some form of the above code into our benchmark suite and additional information would be very helpful

@zou3519 Sorry, I already finished my internship there. I believe the next release of Modulus (this month) will include the functorch integration and will be available on GitLab.
Including @NickGeneva from Modulus team for further communications.

[lezcano]

For what is worth, the fix pytorch/pytorch#76828 is completely stalled on some errors on NestedTensor on some GPU architectures. If any one would like to chime in on how to tackle it, you are more than welcome.

[akshaysubr]

@zou3519 The modulus version with functorch support was made public last week. Here is the repo.

could you point me to which model in the modulus repo contains the nn.Sequential above please?

The specific place where functorch is used is in this wrapper class.

Do you generally run .backward() after computing the hessian? What are some representative input sizes (are the sizes in the benchmark script representative?)

Yes, after computing the hessian, we need to run .backward() to actually get the weight gradients. In terms of the representative input sizes, they are typically O(1)-O(10) and same for the outputs. The sizes in the benchmark script are representative for some of the simpler cases and they might go up from there but the hidden dimensions are usually the same as in the benchmark.

[IvanYashchuk]

@akshaysubr provided links do not work, maybe project access settings are not correctly set.

[akshaysubr]

@IvanYashchuk Sorry, should've mentioned that you have to apply for access to that repo here and after that those links should work.

[zou3519]

I was able to add a version of the original benchmark to pytorch/benchmark so we now can prevent regressions in this model, so I'm lowering the priority of this issue. Leaving it open to discuss the changes to matmul above though.

[lezcano]

fwiw, the matmul PR is as stalled as it was before. Neither Ivan or me have been able to find the time to put together a standalone C repro for the cublas team.

[lezcano]

After the patching in the one-liner pytorch/pytorch#75195 on top of the stack pytorch/pytorch#76828, I still get an OOM using the script in the OP.
The issue happens to be that, somehow, we end up with a tensor of shape [10000, 512, 512] and strides equal to zero within bmm_out_cuda. Then, we try to materialize this tensor and we OOM. Interestingly enough, if you print the sizes/strides of the tensors that go through matmul you get

size t1
[512, 512]     
strides t1
[512, 1]  
size t2                                                                                                                                                                                                           
[10000, 512, 3]               
strides t2                                                                                                                                                                                       
[3, 30000, 1] 

so that tensor that gets to bmm looks like a batched version of the one that gets to the last if-else within matmul. Perhaps this is a feature of functorch, but I don't really understand it. Note that this would be fine if this tensor rather than strides equal to zero, it had strides equal to, say, [0, 512, 1], as we would not need to materialise it there. Now perhaps my question is, is this tensor with strides equal to zero reasonable, or is it a bug / optimisation opportunity in functorch, or whether this follows from some functorch feature, and I may be doing something wrong here.

cc @kshitij12345