Making a custom transformer architecture work with opacus

Nahian: <p>I am trying to make an architecture work with opacus . It consists of two encoders that use Self-attention and produces context embeddings x_t and y_t. “Knowledge Retriever”  is using masked attention.<br>
I suppose there are a few issues with this. It uses a modified multihead attention that uses  an exponential decay function applied to the scaled dot product and a distance adjustment factor gamma that requires no gradient. It uses the model parameters that has been already calculated to obtain the distance adjustments. This causes conflicts with opacus for which I will create a separate issue later.<br>
For simplicity, I have used just multihead attention to avoid conflicts with opacus. Here is the notebook that can be used to reproduce this:  <a href="https://colab.research.google.com/drive/1Sp3jILzB3HvizIAw3OTiQGnVq7LB5gee?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a></p>
<h2><a name="and-this-still-produces-the-following-error-usrlocallibpython310dist-packagestorchnnmodulesmodulepy1352-userwarning-using-a-non-full-backward-hook-when-the-forward-contains-multiple-autograd-nodes-is-deprecated-and-will-be-removed-in-future-versions-this-hook-will-be-missing-some-grad_input-please-use-register_full_backward_hook-to-get-the-documented-behavior-warningswarnusing-a-non-full-backward-hook-when-the-forward-contains-multiple-autograd-nodes-1" class="anchor" href="#and-this-still-produces-the-following-error-usrlocallibpython310dist-packagestorchnnmodulesmodulepy1352-userwarning-using-a-non-full-backward-hook-when-the-forward-contains-multiple-autograd-nodes-is-deprecated-and-will-be-removed-in-future-versions-this-hook-will-be-missing-some-grad_input-please-use-register_full_backward_hook-to-get-the-documented-behavior-warningswarnusing-a-non-full-backward-hook-when-the-forward-contains-multiple-autograd-nodes-1"></a>And this still produces the following error:<br>
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py:1352: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.<br>
warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "</h2>
<p>ValueError                                Traceback (most recent call last)<br>
 in &lt;cell line: 1&gt;()<br>
----&gt; 1 best_epoch = train_one_dataset(train_q_data, train_qa_data, train_pid, valid_q_data, valid_qa_data, valid_pid)<br>
2</p>
<p>5 frames<br>
 in train_one_dataset(train_q_data, train_qa_data, train_pid, valid_q_data, valid_qa_data, valid_pid)<br>
37     for idx in range(max_iter):<br>
38         # Train Model<br>
—&gt; 39         train_loss, train_accuracy, train_auc = train(<br>
40             dp_model, dp_optimizer, train_q_data, train_qa_data, train_pid, accountant, label=‘Train’)<br>
41         # Validation step</p>
<p> in train(net, optimizer, q_data, qa_data, pid_data, accountant, label)<br>
89                 net.parameters(), max_norm=maxgradnorm)<br>
90<br>
—&gt; 91         optimizer.step()<br>
92<br>
93         # correct: 1.0; wrong 0.0; padding -1.0</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/optimizers/optimizer.py in step(self, closure)<br>
516                 closure()<br>
517<br>
 → 518         if self.pre_step():<br>
519             return self.original_optimizer.step()<br>
520         else:</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/optimizers/optimizer.py in pre_step(self, closure)<br>
494         # The corner case when the optimizer has no trainable parameters.<br>
495         # Essentially the DPOptimizer act as a normal optimizer<br>
 → 496         if self.grad_samples is None or len(self.grad_samples) == 0:<br>
497             return True<br>
498</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/optimizers/optimizer.py in grad_samples(self)<br>
343         ret = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
344         for p in self.params:<br>
 → 345             ret.append(self._get_flat_grad_sample(p))<br>
346         return ret<br>
347</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/optimizers/optimizer.py in _get_flat_grad_sample(self, p)<br>
280             )<br>
281         if p.grad_sample is None:<br>
 → 282             raise ValueError(<br>
283                 “Per sample gradient is not initialized. Not updated in backward pass?”<br>
284             )</p>
<p>ValueError: Per sample gradient is not initialized. Not updated in backward pass?</p>
<p>There is also some behavior I had to note. In the architecture class, transformer layers are initialized. In the forward pass the x and y embeddings are passed into the encoders.  The flag is there to ensure when the knowledge retriever block(masked attention)  is executed. This is clearer in the forward pass of the transformer layer, where the “if statement block” is for the masked attention (knowledge retriever) and the  “else block” corresponds to the encoders on the left( see picture in notebook). All three components use the same forward pass.( see forward calls of Architecture, Transformer Layer classes)<br>
Training/ optimizer step only seems to execute when I leave out the if/else conditions and have one forward pass for all three parts of the model: two encoders and knowledge retriever that uses masked attention.<br>
Is there a way around this? Is there a way this could be reimplemented in a way which would allow per sample gradient computation?</p>
<p>here is a notebook of the model running without opacus:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1jg-ygK7Vfou-IaJqNaujk-CHfdMruup3?usp=sharing">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/e11fc3dbd9214fb6318f5e141c13f3f6/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1jg-ygK7Vfou-IaJqNaujk-CHfdMruup3?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img width="260" height="260" src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar">

<h3><a href="https://colab.research.google.com/drive/1jg-ygK7Vfou-IaJqNaujk-CHfdMruup3?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>