Computing Jacobian Vector Product (jvp) using nn.DataParallel

MitchellKT: <p>Hello <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"><br>
I have an application which requires computing the derivative of the model’s logits (more accurately - the log softmax of them) with respect to the model’s parameters, multiplied by a constant vector.<br>
Theoretically, this can be done exactly by a single forward pass, as demonstrated in the following PyTorch tutorial: <a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" class="inline-onebox" rel="noopener nofollow ugc">Forward-mode Automatic Differentiation (Beta) — PyTorch Tutorials 1.13.1+cu117 documentation</a>.</p>
<p>Following this tutorial I had managed to write a version of this code for my use case:</p>
<pre><code class="lang-auto">def jacobian_vector_product(model, inputs, vector):
    params = {name: p for name, p in model.named_parameters()}
    tangents = {name: vec_pt for name, vec_pt in zip(params.keys(), vector)}

    with torch.autograd.forward_ad.dual_level():
        for name, p in params.items():
            del_attr(model, name.split("."))
            set_attr(model, name.split("."), torch.autograd.forward_ad.make_dual(p, tangents[name]))

        out = model(inputs)
        lsm = F.log_softmax(out, dim=1)
        jvp = torch.autograd.forward_ad.unpack_dual(lsm).tangent
        return jvp
</code></pre>
<p>(Note here that the original <code>setattr</code> and <code>delattr</code> were replaced by a nested version of them).</p>
<p>Indeed, the above code works perfectly when I don’t use any kind of data parallelism.</p>
<p>However, when I wrap the model with nn.DataParallel(), the forward pass in this function yields an error indicating that the model and the inputs are not on the same device.<br>
After a quick debugging I found out that the batch indeed splits to multiple chunks, and that the input is distributed correctly, however <strong>all the forward passes are ocuuring on cuda:0 instead of being distributed correspondingly to the inputs</strong>.</p>
<p>I think the issue has to do something with the attribute editing loop.<br>
This is because when I tried to comment out the attribute editing loop - forward pass worked properly (although, of course, the jvp equals None).<br>
I have also tried to wrap the model with DataParallel after editing the attributes but I get the same error.</p>
<p>I would appreciate any assistance on this,<br>
Thanks!</p>

ptrblck: <p>I guess the newly added features such as forward-mode AD might not be supported and tested in e.g. <code>nn.DataParallel</code> as it’s planned to be deprecated. We generally recommend to use <code>DistributedDataParallel</code> instead which should also perform better than <code>DataParallel</code>.</p>

MitchellKT: <p>Hi, and thanks for the very quick reply!</p>
<p>After changing <code>DataParallel</code> to <code>DDP</code> I get the following error in my function:</p>
<pre><code class="lang-auto">NotImplementedError: You must implement the jvp function for custom autograd.Function to use it with forward mode AD.
</code></pre>
<p>Tracing the error, seems that the Synchronized Batch Normalization (SyncBN) in my model is causing the problem.</p>
<pre><code class="lang-auto">File "/home/mitchell/DivideMix/meta_utils.py", line 241, in jacobian_vector_product
    out = model(inputs)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1008, in forward
    output = self._run_ddp_forward(*inputs, **kwargs)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 969, in _run_ddp_forward
    return module_to_run(*inputs[0], **kwargs[0])
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torchvision/models/resnet.py", line 285, in forward
    return self._forward_impl(x)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torchvision/models/resnet.py", line 269, in _forward_impl
    x = self.bn1(x)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/nn/modules/batchnorm.py", line 748, in forward
    return sync_batch_norm.apply(
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/autograd/function.py", line 257, in apply_jvp
    return self._forward_cls.jvp(self, *args)  # type: ignore[attr-defined]
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/autograd/function.py", line 391, in jvp
    raise NotImplementedError("You must implement the jvp function for custom "
NotImplementedError: You must implement the jvp function for custom autograd.Function to use it with forward mode AD.
</code></pre>
<p>What can I do to resolve it?</p>
<p>Thanks in advance!</p>

ptrblck: <p>It seems the static <code>jvp</code> method is missing in an <code>autograd.Function</code> and you would need to implement it as described <a href="https://discuss.pytorch.org/t/notimplementederror-you-must-implement-the-jvp-function-for-custom-autograd-function-to-use-it-with-forward-mode-ad/138245/2">here</a>.</p>

MitchellKT: <p>As I undersrand, problem is that <code>sync_batch_norm</code> (which inherits from <code>autograd.Function</code>) is the one that is missing the <code>jvp</code> method (this can be observed here: <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/_functions.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch/_functions.py at master · pytorch/pytorch · GitHub</a>, note that <code>SyncBatchNorm</code> implements forward and backward, but not jvp).</p>
<p>Am I expected to implement the <code>jvp</code> call for <code>SyncBatchNorm</code> myself? It seems too complicated…<br>
Am I missing something?<br>
Are there any workarounds? (I.e using normal batch norm or the double grad trick(?) <a href="https://discuss.pytorch.org/t/forward-mode-ad-vs-double-grad-trick-for-jacobian-vector-product/159037" class="inline-onebox">Forward Mode AD vs double grad trick for jacobian vector product</a>)</p>

ptrblck: <p>Yes, you would either need to implement <code>jvp</code> for <code>SyncBatchNorm</code> or fall back to the standard <code>nn.BatchNormXd</code> layers (assuming they are providing this method), which would then not synchronize the running stats but would hopefully work otherwise.</p>

MitchellKT: <p>Hi and thanks for your answers!</p>
<p>I have tried to implement the double-back trick for computing <code>jvp</code>s as in:<br>
<a href="https://j-towns.github.io/2017/06/12/A-new-trick.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://j-towns.github.io/2017/06/12/A-new-trick.html</a><br>
and here:</p><aside class="quote quote-modified" data-post="1" data-topic="159037">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/youyoun/48/40540_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/forward-mode-ad-vs-double-grad-trick-for-jacobian-vector-product/159037">Forward Mode AD vs double grad trick for jacobian vector product</a> <a class="badge-category__wrapper " href="/c/autograd/7"><span data-category-id="7" data-drop-close="true" class="badge-category " title="A category of posts relating to the autograd engine itself."><span class="badge-category__name">autograd</span></span></a>
  </div>
  <blockquote>
    Hello, 
I stumbled on this <a href="https://pytorch.org/tutorials/intermediate/forward_ad_usage.html" rel="noopener nofollow ugc">page</a> of the pytorch doc, and I had a few questions about the use of this method. 
First of all, I’m not really comfortable with auto-diff, and I’ve had a hard time understanding the difference between reverse mode AD and forward mode AD. The notable difference that I seem to have understood is that one will be run alongside the forward pass, in order to minimize the numbers of operations used to compute a JVP. 
If this understanding is correct, I’d expect the forward m…
  </blockquote>
</aside>
<p>
and is also in the PyTorch’s original implementation of <code>jvp</code>:<br>
<a href="https://pytorch.org/docs/stable/generated/torch.autograd.functional.jvp.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.autograd.functional.jvp.html</a></p>
<p>However, it seems that syncBN continues to make problems.<br>
I get the following error:</p>
<pre><code class="lang-auto">File "/home/mitchell/DivideMix/meta_utils.py", line 261, in jacobian_vector_product
    jvp = torch.autograd.grad(g, v, vector)[0]
  File "/home/mitchell/miniconda3/envs/deep_learn/lib/python3.9/site-packages/torch/autograd/__init__.py", line 276, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: derivative for aten::batch_norm_backward_elemt is not implemented
</code></pre>
<p>Here is my implementation:</p>
<pre><code class="lang-auto">def jacobian_vector_product(model, inputs, vector, method='double-back-trick'):
    
    if method == 'forward':
        '''
        jvp products using forward mode AD as demonstrated in:
        https://pytorch.org/tutorials/intermediate/forward_ad_usage.html
        '''
        params = {name: p for name, p in model.named_parameters()}
        tangents = {name: vec_pt for name, vec_pt in zip(params.keys(), vector)}

        with torch.autograd.forward_ad.dual_level():
            for name, p in params.items():
                del_attr(model, name.split("."))
                set_attr(model, name.split("."), torch.autograd.forward_ad.make_dual(p, tangents[name]))

            out = model(inputs)
            lsm = F.log_softmax(out, dim=1)
            jvp = torch.autograd.forward_ad.unpack_dual(lsm).tangent
            return jvp
    elif method == 'double-back-trick':
        '''
        jvp products using double backward as demonstrated in:
        https://j-towns.github.io/2017/06/12/A-new-trick.html
        https://discuss.pytorch.org/t/forward-mode-ad-vs-double-grad-trick-for-jacobian-vector-product/159037
        '''
        out = model(inputs)
        lsm = F.log_softmax(out, dim=1)
        v = torch.zeros_like(lsm, requires_grad=True)
        g = torch.autograd.grad(lsm, model.parameters(), v, create_graph=True)
        jvp = torch.autograd.grad(g, v, vector)[0]
        return jvp
    else:
        raise NotImplementedError
</code></pre>
<p>Am I doing a mistake or is it another issue with syncBn?</p>