Modified multihead attention and its compatibility with opacus

Nahian: <p>This is a continuation of the first issue: <a href="https://discuss.pytorch.org/t/making-a-custom-transformer-architecture-work-with-opacus/200316" class="inline-onebox">Making a custom transformer architecture work with opacus</a><br>
This is another notebook to reproduce and understand the problem with monotonic multihead attention: <a href="https://colab.research.google.com/drive/103GoZl-o1fJ0mZrg727-Womxd4Xa1Sfe?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a><br>
(See MultiheadAttention class)</p>
<p>def attention(q, k, v, d_k, mask, dropout, zero_pad, gamma=None):<br>
“”"<br>
This is called by Multi-head atention object to find the values.<br>
“”"<br>
scores = torch.matmul(q, k.transpose(-2, -1)) / <br>
math.sqrt(d_k)  # BS, 8, seqlen, seqlen<br>
bs, head, seqlen = scores.size(0), scores.size(1), scores.size(2)</p>
<pre><code>x1 = torch.arange(seqlen).expand(seqlen, -1).to(device)
x2 = x1.transpose(0, 1).contiguous()

with torch.no_grad():
    scores_ = scores.masked_fill(mask == 0, -1e32)
    scores_ = F.softmax(scores_, dim=-1)  # BS,8,seqlen,seqlen
    scores_ = scores_ * mask.float().to(device)
    distcum_scores = torch.cumsum(scores_, dim=-1)  # bs, 8, sl, sl
    disttotal_scores = torch.sum(
        scores_, dim=-1, keepdim=True)  # bs, 8, sl, 1
    position_effect = torch.abs(
        x1-x2)[None, None, :, :].type(torch.FloatTensor).to(device)  # 1, 1, seqlen, seqlen
    # bs, 8, sl, sl positive distance
    dist_scores = torch.clamp(
        (disttotal_scores-distcum_scores)*position_effect, min=0.)
    dist_scores = dist_scores.sqrt().detach()
m = nn.Softplus()
gamma = -1. * m(gamma).unsqueeze(0)  # 1,8,1,1
# Now after do exp(gamma*distance) and then clamp to 1e-5 to 1e5
total_effect = torch.clamp(torch.clamp(
    (dist_scores*gamma).exp(), min=1e-5), max=1e5)
scores = scores * total_effect

scores.masked_fill_(mask == 0, -1e32)
scores = F.softmax(scores, dim=-1)  # BS,8,seqlen,seqlen
if zero_pad:
    pad_zero = torch.zeros(bs, head, 1, seqlen).to(device)
    scores = torch.cat([pad_zero, scores[:, :, 1:, :]], dim=2)
scores = dropout(scores)
output = torch.matmul(scores, v)
return output
</code></pre>
<p>To me it seemed more like the model parameters are being used to calculate terms that are applied to the distance. In this case, could this be an issue that prevents back propagation?</p>
<h2><a name="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="#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>/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>TypeError                                 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>21 frames<br>
 in train_one_dataset(train_q_data, train_qa_data, train_pid, valid_q_data, valid_qa_data, valid_pid)<br>
38     for idx in range(max_iter):<br>
39         # Train Model<br>
—&gt; 40         train_loss, train_accuracy, train_auc = train(<br>
41             dp_model, dp_optimizer, train_q_data, train_qa_data, train_pid, accountant, label=‘Train’)<br>
42         # Validation step</p>
<p> in train(net, optimizer, q_data, qa_data, pid_data, accountant, label)<br>
82             loss, pred, true_ct = net(input_q, input_qa, target)<br>
83         pred = pred.detach().cpu().numpy()  # (seqlen * batch_size, 1)<br>
—&gt; 84         loss.backward()<br>
85         true_el += true_ct.cpu().numpy()<br>
86</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_tensor.py in backward(self, gradient, retain_graph, create_graph, inputs)<br>
520                 inputs=inputs,<br>
521             )<br>
 → 522         torch.autograd.backward(<br>
523             self, gradient, retain_graph, create_graph, inputs=inputs<br>
524         )</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/autograd/<strong>init</strong>.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)<br>
264     # some Python versions print out the first line of a multi-line function<br>
265     # calls in the traceback and some print out the last line<br>
 → 266     Variable.<em>execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
267         tensors,<br>
268         grad_tensors</em>,</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in <strong>call</strong>(self, *args, **kwargs)<br>
69             if module is None:<br>
70                 raise RuntimeError(“You are trying to call the hook of a dead Module!”)<br>
—&gt; 71             return self.hook(module, *args, **kwargs)<br>
72         return self.hook(*args, **kwargs)<br>
73</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/grad_sample/grad_sample_module.py in capture_backprops_hook(self, module, _forward_input, forward_output, loss_reduction, batch_first)<br>
336             grad_sampler_fn = ft_compute_per_sample_gradient<br>
337<br>
 → 338         grad_samples = grad_sampler_fn(module, activations, backprops)<br>
339         for param, gs in grad_samples.items():<br>
340             create_or_accumulate_grad_sample(</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/grad_sample/functorch.py in ft_compute_per_sample_gradient(layer, activations, backprops)<br>
92         prepare_layer(layer)<br>
93<br>
—&gt; 94     per_sample_grads = layer.ft_compute_sample_grad(<br>
95         parameters, activations[0], backprops<br>
96     )</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/apis.py in wrapped(*args, **kwargs)<br>
186     # <span class="mention">@functools.wraps</span>(func)<br>
187     def wrapped(*args, **kwargs):<br>
 → 188         return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)<br>
189<br>
190     return wrapped</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)<br>
276<br>
277     # If chunk_size is not specified.<br>
 → 278     return _flat_vmap(<br>
279         func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs<br>
280     )</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in fn(*args, **kwargs)<br>
42     def fn(*args, **kwargs):<br>
43         with torch.autograd.graph.disable_saved_tensors_hooks(message):<br>
—&gt; 44             return f(*args, **kwargs)<br>
45     return fn<br>
46</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in _flat_vmap(func, batch_size, flat_in_dims, flat_args, args_spec, out_dims, randomness, **kwargs)<br>
389     try:<br>
390         batched_inputs = _create_batched_inputs(flat_in_dims, flat_args, vmap_level, args_spec)<br>
 → 391         batched_outputs = func(*batched_inputs, **kwargs)<br>
392         return _unwrap_batched(batched_outputs, out_dims, vmap_level, batch_size, func)<br>
393     finally:</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/apis.py in wrapper(*args, **kwargs)<br>
361     <span class="mention">@functools.wraps</span>(func)<br>
362     def wrapper(*args, **kwargs):<br>
 → 363         return eager_transforms.grad_impl(func, argnums, has_aux, args, kwargs)<br>
364     return wrapper</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/<em>functorch/eager_transforms.py in grad_impl(func, argnums, has_aux, args, kwargs)<br>
1293 def grad_impl(func: Callable, argnums: argnums_t, has_aux: bool, args, kwargs):<br>
1294     func = lazy_dynamo_disable(func)<br>
 → 1295     results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)<br>
1296     if has_aux:<br>
1297         grad, (</em>, aux) = results</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/vmap.py in fn(*args, **kwargs)<br>
42     def fn(*args, **kwargs):<br>
43         with torch.autograd.graph.disable_saved_tensors_hooks(message):<br>
—&gt; 44             return f(*args, **kwargs)<br>
45     return fn<br>
46</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/<em>functorch/eager_transforms.py in wrapper(*args, **kwargs)<br>
1254                 tree_map</em>(partial(_create_differentiable, level=level), diff_args)<br>
1255<br>
 → 1256                 output = func(*args, **kwargs)<br>
1257                 if has_aux:<br>
1258                     if not (isinstance(output, tuple) and len(output) == 2):</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py in _fn(*args, **kwargs)<br>
487                 dynamo_config_ctx.<strong>enter</strong>()<br>
488             try:<br>
 → 489                 return fn(*args, **kwargs)<br>
490             finally:<br>
491                 set_eval_frame(prior)</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/grad_sample/functorch.py in compute_loss_stateless_model(params, activations, backprops)<br>
69             batched_backprops = backprops.unsqueeze(1)<br>
70<br>
—&gt; 71         output = flayer(params, batched_activations)<br>
72         loss = (output * batched_backprops).sum()<br>
73         return loss</p>
<p>/usr/local/lib/python3.10/dist-packages/opacus/grad_sample/functorch.py in fmodel(new_params_values, *args, **kwargs)<br>
34             name: value for name, value in zip(params_names, new_params_values)<br>
35         }<br>
—&gt; 36         return torch.func.functional_call(stateless_mod, new_params_dict, args, kwargs)<br>
37<br>
38     if disable_autograd_tracking:</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/_functorch/functional_call.py in functional_call(module, parameter_and_buffer_dicts, args, kwargs, tie_weights, strict)<br>
141         )<br>
142<br>
 → 143     return nn.utils.stateless._functional_call(<br>
144         module,<br>
145         parameters_and_buffers,</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/utils/stateless.py in _functional_call(module, parameters_and_buffers, args, kwargs, tie_weights, strict)<br>
261         module, parameters_and_buffers, tie_weights=tie_weights, strict=strict<br>
262     ):<br>
 → 263         return module(*args, **kwargs)</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _wrapped_call_impl(self, *args, **kwargs)<br>
1509             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]<br>
1510         else:<br>
 → 1511             return self._call_impl(*args, **kwargs)<br>
1512<br>
1513     def _call_impl(self, *args, **kwargs):</p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py in _call_impl(self, *args, **kwargs)<br>
1518                 or _global_backward_pre_hooks or _global_backward_hooks<br>
1519                 or _global_forward_hooks or _global_forward_pre_hooks):<br>
 → 1520             return forward_call(*args, **kwargs)<br>
1521<br>
1522         try:</p>
<p>TypeError: MultiHeadAttention.forward() missing 4 required positional arguments: ‘k’, ‘v’, ‘mask’, and ‘zero_pad’</p>
<p>here is a notebook of the model running without opacus: <a href="https://colab.research.google.com/drive/1jg-ygK7Vfou-IaJqNaujk-CHfdMruup3?usp=sharing" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a></p>