Integrating Opacus for a custom pytorch model raises errors but works fine on its own

Nahian: <p>I am trying to run a Transformer model with opacus, but running into the following error that relates to the custom layers. This does not arise when I train the model without privacy.<br>
In general I understand that sometimes a separate grad sampler needs to be registered in case opacus cannot handle certain layers. But that does not seem to be the case because I do not have the “Layer incompatible/not recognized error”.<br>
I also know that opacus has grad samplers for basic layers i.e. nn.linear, layernorm etc. which the architecture is composed of. But could it still be something opacus can’t handle? Or is it the way the architecture was implemented? What are some things to note when implementing custom models like this using opacus? Looking at the traceback, I think its related to backward pass at loss.backward().</p>
<p>FULL TRACEBACK:</p>
<p>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 "<br>
Traceback (most recent call last):<br>
File “main.py”, line 283, in <br>
best_epoch = train_one_dataset(<br>
File “main.py”, line 73, in train_one_dataset<br>
train_loss, train_accuracy, train_auc = train(<br>
File “/media/OS/code-exp/KT/knowledge-tracing-bert-transformers/context-aware-transformer-priv (grad-sampler)/run.py”, line 96, in train<br>
loss.backward()<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_tensor.py”, line 522, in backward<br>
torch.autograd.backward(<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/autograd/<strong>init</strong>.py”, line 266, in backward<br>
Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 71, in <strong>call</strong><br>
return self.hook(module, *args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/opacus/grad_sample/grad_sample_module.py”, line 337, in capture_backprops_hook<br>
grad_samples = grad_sampler_fn(module, activations, backprops)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/opacus/grad_sample/functorch.py”, line 58, in ft_compute_per_sample_gradient<br>
per_sample_grads = layer.ft_compute_sample_grad(<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/apis.py”, line 188, in wrapped<br>
return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 278, in vmap_impl<br>
return _flat_vmap(<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 44, in fn<br>
return f(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 391, in _flat_vmap<br>
batched_outputs = func(*batched_inputs, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/apis.py”, line 363, in wrapper<br>
return eager_transforms.grad_impl(func, argnums, has_aux, args, kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py”, line 1295, in grad_impl<br>
results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 44, in fn<br>
return f(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py”, line 1256, in wrapper<br>
output = func(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 489, in _fn<br>
return fn(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/opacus/grad_sample/functorch.py”, line 34, in compute_loss_stateless_model<br>
output = flayer(params, batched_activations)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1511, in _wrapped_call_impl<br>
return self._call_impl(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1520, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/make_functional.py”, line 342, in forward<br>
return self.stateless_model(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1511, in _wrapped_call_impl<br>
return self._call_impl(*args, **kwargs)<br>
File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1520, in _call_impl<br>
return forward_call(*args, **kwargs)<br>
TypeError: forward() missing 4 required positional arguments: ‘k’, ‘v’, ‘mask’, and ‘zero_pad’</p>
<p>architecture:<br>
<a href="https://github.com/arghosh/AKT" class="inline-onebox" rel="noopener nofollow ugc">GitHub - arghosh/AKT</a> → akt.py</p>
<p>Changes I made:<br>
I reproduced it exactly as it is and added the following lines in main.py:<br>
model = load_model(params)<br>
<span class="hashtag-raw">#optimizer</span> = torch.optim.Adam(model.parameters(), lr=params.lr, betas=(0.9, 0.999), eps=1e-8)<br>
optimizer = torch.optim.SGD(model.parameters(), lr=params.lr, momentum=0.9)<br>
print(“\n”)<br>
<span class="hashtag-raw">#opacus_utils</span>.register_grad_sampler(MonotonicMultiHeadAttention)(compute_monotonic_attention_grad_sample)<br>
batch_size = 24<br>
accountant = RDPAccountant()<br>
dp_model = GradSampleModule(model)<br>
dp_optimizer = DPOptimizer(<br>
optimizer=optimizer,<br>
noise_multiplier=1.0, # same as make_private arguments<br>
max_grad_norm=1.0, # same as make_private arguments<br>
expected_batch_size=batch_size # if you’re averaging your gradients, you need to know the denominator<br>
)<br>
sample_rate = batch_size/len(train_q_data)<br>
print(sample_rate)<br>
dp_optimizer.attach_step_hook(<br>
accountant.get_optimizer_hook_fn(<br>
# this is an important parameter for privacy accounting. Should be equal to batch_size / len(dataset)<br>
sample_rate=sample_rate<br>
)<br>
)<br>
<span class="hashtag-raw">#training</span><br>
all_train_loss = {}<br>
all_train_accuracy = {}<br>
all_train_auc = {}<br>
all_valid_loss = {}<br>
all_valid_accuracy = {}<br>
all_valid_auc = {}<br>
best_valid_auc = 0</p>
<pre><code>for idx in range(params.max_iter):
    # Train Model
    train_loss, train_accuracy, train_auc = train(
         dp_model, params, dp_optimizer, train_q_data, train_qa_data, train_pid, accountant,  label='Train')
    # Validation step
    valid_loss, valid_accuracy, valid_auc = test(
        dp_model,  params, dp_optimizer, valid_q_data, valid_qa_data, valid_pid, label='Valid' )
</code></pre>
<p>**I am providing all the details to reproduce this:<br>
opacus version: 1.4.0<br>
pytorch: 2.2.1+CUDA 11.8<br>
pandas, tqdm all latest versions work fine</p>

Nahian: <p>could multiple forward passes in the model be an issue?</p>