Can't vmap autograd.grad over outputs

username: valerian.rey
cooked: <p>Hi! I’m trying to compute per-sample gradients efficiently. I know from the <a href="https://docs.pytorch.org/tutorials/intermediate/per_sample_grads.html" rel="noopener nofollow ugc">tutorial about per-sample gradients</a> that I can vmap a call to <code>torch.func.grad</code>.</p>
<p>However, the function returned by <code>torch.func.grad</code> does both a forward and a backward pass. In my specific case, I have to do a forward pass anyway, which creates the autograd graph, so it feels like a waste to not re-use this autograd graph.</p>
<p><strong>I would like to know if there is a way to vmap a call to <code>torch.autograd.grad</code> instead of <code>torch.func.grad</code>, so that the already-existing autograd graph can be reused, and so that no extra forward pass is done.</strong></p>
<p>I tried the following:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

batch_size = 64

inputs = torch.randn((batch_size, 5))
targets = torch.randn((batch_size, 1))

model = nn.Linear(5, 1)
params = list(model.parameters())

outputs = model(inputs)  # shape: [batch_size, 1]
losses = F.mse_loss(outputs, targets, reduction="none").squeeze()  # shape: [batch_size]

def compute_one_gradient(loss: Tensor) -&gt; tuple[Tensor, ...]:
    return torch.autograd.grad(loss, params)

grads = torch.vmap(compute_one_gradient)(losses)
</code></pre>
<p>But I end up with the following error:</p>
<pre><code class="lang-auto">tests/unit/autogram/test_per_sample_grads.py:65 (test_per_sample_grads)
def test_per_sample_grads():
        import torch
        import torch.nn as nn
        import torch.nn.functional as F
        from torch import Tensor
    
        batch_size = 64
    
        inputs = torch.randn((batch_size, 5))
        targets = torch.randn((batch_size, 1))
    
        model = nn.Linear(5, 1)
        params = list(model.parameters())

        outputs = model(inputs)  # shape: [batch_size, 1]
        losses = F.mse_loss(outputs, targets, reduction="none").squeeze()  # shape: [batch_size]
    
        def compute_one_gradient(loss: Tensor) -&gt; tuple[Tensor, ...]:
            return torch.autograd.grad(loss, params)
    
&gt;       grads = torch.vmap(compute_one_gradient)(losses)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

unit/autogram/test_per_sample_grads.py:86: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../.venv/lib/python3.13/site-packages/torch/_functorch/apis.py:202: in wrapped
    return vmap_impl(
../.venv/lib/python3.13/site-packages/torch/_functorch/vmap.py:334: in vmap_impl
    return _flat_vmap(
../.venv/lib/python3.13/site-packages/torch/_functorch/vmap.py:484: in _flat_vmap
    batched_outputs = func(*batched_inputs, **kwargs)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
unit/autogram/test_per_sample_grads.py:84: in compute_one_gradient
    return torch.autograd.grad(loss, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
../.venv/lib/python3.13/site-packages/torch/autograd/__init__.py:502: in grad
    result = _engine_run_backward(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

t_outputs = (BatchedTensor(lvl=1, bdim=0, value=
    tensor([4.2187e-01, 7.2812e-04, 3.7193e-02, 3.6601e+00, 4.0343e+00, 6.0330e-0...03, 7.4318e-01,
            2.8548e-01, 4.4032e-02, 7.8910e-04, 7.2651e-01],
           grad_fn=&lt;SqueezeBackward0&gt;)
),)
args = ((None,), False, False, (Parameter containing:
tensor([[ 0.1479,  0.4471,  0.2320,  0.2780, -0.1565]], requires_grad=True), Parameter containing:
tensor([0.2146], requires_grad=True)), False)
kwargs = {'accumulate_grad': False}, attach_logging_hooks = False

    def _engine_run_backward(
        t_outputs: Sequence[Union[torch.Tensor, GradientEdge]],
        *args: Any,
        **kwargs: Any,
    ) -&gt; tuple[torch.Tensor, ...]:
        attach_logging_hooks = log.getEffectiveLevel() &lt;= logging.DEBUG
        if attach_logging_hooks:
            unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
        try:
&gt;           return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
                t_outputs, *args, **kwargs
            )  # Calls into the C++ engine to run the backward pass
E           RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

../.venv/lib/python3.13/site-packages/torch/autograd/graph.py:824: RuntimeError
</code></pre>
<p>I think the reason is that when vmap splits the <code>losses</code> tensor to parallelize over it, it does so without creating a <code>grad_fn</code>, so the <code>loss</code> tensor provided to <code>compute_one_gradient</code> does not require grad anymore.</p>
<p>Is there any proper way, or even any trick, to do what I want to do in parallel and without any extra forward pass? (I think <code>torch.func.grad</code>, <code>torch.func.vjp</code>, <code>torch.func.jacobian</code>, <code>torch.func.jacrev</code>, etc. all require an extra forward pass). It almost seems like what I need is the <code>is_grads_batched</code> parameter of <code>torch.autograd.grad</code>, but this is for batched <code>grad_outputs</code>, not for batched <code>outputs</code>.</p>
<p>Thanks in advance for the help!</p>

username: Soumya_Kundu
cooked: <p>Would <a href="https://toiaydcdyywlhzvlob.github.io/backpack/" rel="noopener nofollow ugc">backpack</a> help here?</p>

username: valerian.rey
cooked: <p>Thanks for the suggestion! I just tried with backpack’s <a href="https://docs.backpack.pt/en/master/extensions.html#backpack.extensions.BatchGrad" rel="noopener nofollow ugc">BatchGrad</a>, but it’s actually slightly slower than my current implementation with the extra forward pass <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=14" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"> It also seems to come with a few bugs that I had to work around, so I would rather not depend on it.</p>

username: Soumya_Kundu
cooked: <p>Ah okay. Sorry, but i have never stress tested backpack to be honest.</p>
<p>I found this bit</p>
<ul>
<li><strong>is_grads_batched</strong> (<a href="https://docs.python.org/3/library/functions.html#bool" rel="noopener nofollow ugc"><em>bool</em></a><em>,</em> <em>optional</em>) – If <code>True</code>, the first dimension of each tensor in <code>grad_outputs</code> will be interpreted as the batch dimension. Instead of computing a single vector-Jacobian product, we compute a batch of vector-Jacobian products for each “vector” in the batch. We use the vmap prototype feature as the backend to vectorize calls to the autograd engine so that this computation can be performed in a single call. This should lead to performance improvements when compared to manually looping and performing backward multiple times. Note that due to this feature being experimental, there may be performance cliffs. Please use <code>torch._C._debug_only_display_vmap_fallback_warnings(True)</code> to show any performance warnings and file an issue on github if warnings exist for your use case. Defaults to <code>False</code>.</li>
</ul>
<p>From: <a href="https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad.html" class="inline-onebox" rel="noopener nofollow ugc">torch.autograd.grad — PyTorch 2.7 documentation</a></p>
<p>Still trying to fully understand it, but is this something you could potentially use? Did not know of this option till now.</p>

username: valerian.rey
cooked: <p>I also tried this, but it really is for batched <code>grad_outputs</code> and not <code>outputs</code>.</p>
<p>For instance, you can do:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.autograd.grad(
    outputs,  # One tensor of shape: [10]
    inputs,  # any shapes
    grad_outputs,  # One batch of grad_outputs, of shape: [batch_size, 10]
    is_grads_batched=True
)
</code></pre>
<p>My case is similar to this, except that both <code>outputs</code> and <code>grad_outputs</code> are batched (e.g. the shape of <code>outputs</code> is <code>[batch_size, 10]</code>), which makes the call to <code>torch.autograd.grad</code> fail, even with <code>is_grads_batched=True</code>.</p>

username: Soumya_Kundu
cooked: <p>I will have a deeper read over the weekend and try to get back to with something reasonable hopefully. Thanks for the responses.</p>

username: Dhia-naouali
cooked: <p>hii guys, I think I have a few ways to solve that</p>
<p>first and most similar to your original code is using a  <code>vmap</code>ped <code>autograd.grad</code><br>
in this solution I’m vmapping a call to a function that computes the grads of a single loss w.r.t to the params using autograd.grad<br>
vmapping over an identity matrix that I’m using to select a single loss each time by pasisng it as <code>grad_outputs</code> + retaining the graph so that once the forward pass is computed we’re parallelizing the backward pass using the same graph without reconstruction</p>
<pre data-code-wrap="python"><code class="lang-python">
import torch
import torch.nn as nn
import torch.nn.functional as F

batch_size = 64

inputs = torch.randn((batch_size, 5))
targets = torch.randn((batch_size, 1))

model = nn.Linear(5, 1)
params = list(model.parameters())

outputs = model(inputs)  # shape: [batch_size, 1]
losses = F.mse_loss(outputs, targets, reduction="none").squeeze()  # shape: [batch_size]

bs = losses.shape[0]
identity = torch.eye(bs)
def compute_sample_grad(loss_selector):
    grads = torch.autograd.grad(
        outputs=losses,
        inputs=params,
        grad_outputs=loss_selector,
        retain_graph=True
    )
    
    return {
        name: grad for (name, _), grad in zip(model.named_parameters(), grads)
    }
    
per_sample_grads = torch.vmap(compute_sample_grad)(identity)
</code></pre>
<p>a second solution would be using just using jvp (I believe it runs the forward pas just once, then reuses the graph it computed)</p>
<p>I’ve also found this tutorial that suggest using a functional call <a href="https://docs.pytorch.org/tutorials/intermediate/per_sample_grads.html" rel="noopener nofollow ugc">link</a></p>
<p>lastly, if you’re running pytorch 2.7 +, I think this is the wrapper api is what you’re looking for <a href="https://docs.pytorch.org/tutorials/intermediate/per_sample_grads.html" rel="noopener nofollow ugc">docs</a></p>

username: AlphaBetaGamma96
cooked: <p>Hi <a class="mention" href="/u/valerian.rey">@valerian.rey</a>, why don’t you use <code>torch.func.grad</code> with the <code>has_aux</code> flag? So, you can compute the forward and backward in a single call and return the forward with the backward via a single<code>torch.func.grad</code> call? (rather than only using <code>torch.func.grad</code> for the backward?</p>

username: valerian.rey
cooked: <p>Hi! Thanks for your suggestions!</p>
<h3><a name="p-472852-first-solution-1" class="anchor" href="#p-472852-first-solution-1"></a>First solution</h3>
<p>Your first solution is quite good in the general case (in fact it’s very similar to <a href="https://github.com/TorchJD/torchjd/blob/main/src/torchjd/_autojac/_transform/_jac.py" rel="noopener nofollow ugc">what we use internally in torchjd</a>) but it is also quite inefficient (in terms of memory) when the number of losses grows.</p>
<p>The thing is that internally, this will compute the gradient of each loss with respect to all layer outputs (as these are needed to compute the gradients wrt the model parameters). In particular, it will compute the gradient of loss1 (the loss associated to batch element 1) with respect to output1_2 (the output of the first layer for batch element 2). But if your model operates on each batch element independently (which is the case unless you have BatchNorm layers or other weird things mixing elements of the batch), this gradient will always be 0: the loss for batch element 1 does not depend at all on batch element 2 (or on any of its activations throughout the network).</p>
<p>So while this method works, it has to compute internal jacobian matrices that are very big and sparse, which takes a ton of memory.</p>
<p>What we currently do is to just compute the gradient of the sum of the losses with respect to each layer’s output tensor, and then use it to compute the jacobian of the losses with respect to that layer’s parameters.<br>
For a given layer, the output will have shape [batch_size, …], and so the gradient (call it g) of the sum of losses wrt this output will also have shape [batch_size, …]. Further, the gradient of the i’th loss wrt this output also has shape [batch_size, …], but all of its rows j != i will be 0 and its i’th row will be g[i]. All of this means that g contains all the information we care about: despite being a simple gradient, it has the same information as a jacobian.</p>
<p>The problem arises when we want to use this g as the grad_outputs to compute the jacobian of the losses wrt the layer’s parameters. We can tell <code>torch.autograd.grad</code> that the <code>grad_outputs</code> are batched, using <code>is_grads_batched=True</code>. However, we cannot tell it that the outputs are also batched, so it doesn’t work (btw I opened an <a href="https://github.com/pytorch/pytorch/issues/156616" rel="noopener nofollow ugc">issue</a> to add an <code>is_outputs_batched</code> parameter to <code>torch.autograd.grad</code>).</p>
<p>Your first solution is still applicable here, in a slightly modified way: instead of giving a simple identity matrix as grad_output, we could give a “diagonalized” version of the gradient wrt to the output of the layer: its first element would be the gradient of the first loss wrt the output of the layer (including all of its rows of zeros), its second element would be that of the second loss, and so on. This works, but such a tensor is extremely large, which blows memory usage up.</p>
<h3><a name="p-472852-second-solution-2" class="anchor" href="#p-472852-second-solution-2"></a>Second solution</h3>
<p><code>jvp</code> is not really usable in my case because it does not return a <code>jvp_func</code>: it gives directly the output of the model and the grad. But (vmapped) <code>vjp</code> is a viable solution to compute both the batched output of a layer and its (batched) <code>vjp_func</code>, which we can store and then use during the backward pass to obtain the jacobians wrt the parameters. See the code of <a href="https://github.com/pytorch/pytorch/issues/156616#issuecomment-3041385352" rel="noopener nofollow ugc">this comment</a> for the implementation. This is not the solution we ended up using because despite being better in theory (single forward pass), it is a bit slower.</p>
<h3><a name="p-472852-tutorial-about-per-sample-gradients-3" class="anchor" href="#p-472852-tutorial-about-per-sample-gradients-3"></a>Tutorial about per-sample gradients</h3>
<p>Sadly, we can’t directly use what is given in this tutorial: to save memory, we want to compute per-sample gradients on a parameter-by-parameter basis, compute and accumulate some minimally sufficient statistic about those per-sample gradients, and “drop” them (so that they don’t use memory anymore). So we really need a  fine-grained solution that doesn’t just give the per-sample gradients for the full model.</p>

username: valerian.rey
cooked: <p>Well, with <code>torch.func.vjp</code>, we can directly obtain both the output and the <code>vjp_func</code>, which is, I think, what you’re suggesting. We can then <code>vmap</code> this <code>torch.func.vjp</code> to do that for a whole batch in parallel.</p>
<p>In our case, however, we want to be able to do some extra operations after obtaining the per-sample gradients for each parameter, and we want to “drop” them right after that (to save memory). So we would need to apply the forward pass of each layer through a different call to <code>vmap(torch.func.vjp(...))</code>, and store all of the <code>vjp_funcs</code> obtained along the way. This is a viable solution, but it seems a bit slower than not storing anything and just re-doing the call to <code>vmap(torch.func.vjp(...))</code> each time we need it, which does both a (extra) forward pass and a backward pass.</p>
<p>So yeah, in theory you’re right that using a <code>torch.func</code> function transform to do both the forward and the backward makes sense, and it’s actually doable, but it’s somehow a bit slower <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=14" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"></p>
<p>For reference, the implementation for this is available in <a href="https://github.com/pytorch/pytorch/issues/156616#issuecomment-3041385352" rel="noopener nofollow ugc">this comment</a>.</p>

username: qq-me
cooked: <p>With <code>is_grads_batched</code> you could do this:</p>
<pre data-code-wrap="py"><code class="lang-py">def jacobian_batched(output: list[torch.Tensor], wrt: list[torch.Tensor], create_graph=False):
    flat_input = torch.cat([i.reshape(-1) for i in output])
    return torch.autograd.grad(
        flat_input,
        wrt,
        torch.eye(len(flat_input), device=output[0].device, dtype=output[0].dtype),
        create_graph=create_graph,
        is_grads_batched=True,
    )

jacobian_batched([losses], params)
</code></pre>
<p>This is batched vjp, I haven’t done much testing, but for a hessian seems similar in performance to torch.func.hessian and torch.autograd.functional.hessian</p>