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.
Architecture: GitHub - arghosh/AKT → akt.py. I reproduced this repository and I am trying to train this with privacy. 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”. 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 and I am wondering why is this happening with opacus.

FULL TRACEBACK:

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. warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes " Traceback (most recent call last): File “main.py”, line 283, in best_epoch = train_one_dataset( File “main.py”, line 73, in train_one_dataset train_loss, train_accuracy, train_auc = train( File “/media/OS/code-exp/KT/knowledge-tracing-bert-transformers/context-aware-transformer-priv (grad-sampler)/run.py”, line 96, in train loss.backward() File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_tensor.py”, line 522, in backward torch.autograd.backward( File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/autograd/init.py”, line 266, in backward Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 71, in call return self.hook(module, *args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/opacus/grad_sample/grad_sample_module.py”, line 337, in capture_backprops_hook grad_samples = grad_sampler_fn(module, activations, backprops) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/opacus/grad_sample/functorch.py”, line 58, in ft_compute_per_sample_gradient per_sample_grads = layer.ft_compute_sample_grad( File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/apis.py”, line 188, in wrapped return vmap_impl(func, in_dims, out_dims, randomness, chunk_size, *args, *kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 278, in vmap_impl return _flat_vmap( File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 44, in fn return f(args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 391, in _flat_vmap batched_outputs = func(*batched_inputs, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/apis.py”, line 363, in wrapper return eager_transforms.grad_impl(func, argnums, has_aux, args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py”, line 1295, in grad_impl results = grad_and_value(func, argnums, has_aux=has_aux)(*args, *kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/vmap.py”, line 44, in fn return f(args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py”, line 1256, in wrapper output = func(*args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py”, line 489, in _fn return fn(*args, *kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/opacus/grad_sample/functorch.py”, line 34, in compute_loss_stateless_model output = flayer(params, batched_activations) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1511, in _wrapped_call_impl return self._call_impl(args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1520, in _call_impl return forward_call(*args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/_functorch/make_functional.py”, line 342, in forward return self.stateless_model(*args, *kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1511, in _wrapped_call_impl return self._call_impl(args, kwargs) File “/home/miniconda3/envs/kt-transformers/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1520, in _call_impl return forward_call(*args, **kwargs) TypeError: forward() missing 4 required positional arguments: ‘k’, ‘v’, ‘mask’, and ‘zero_pad’

It originates at this return statement in torchnn module,py:

def _call_impl(self, *args, **kwargs): forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)

If we don't have any hooks, we want to skip the rest of the logic in

    # this function, and just call forward.
    if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
            or _global_backward_pre_hooks or _global_backward_hooks
            or _global_forward_hooks or _global_forward_pre_hooks):
        return forward_call(*args, **kwargs)

Changes I made to main.py: I reproduced it exactly as it is and added the following lines in main.py: model = load_model(params)

optimizer = torch.optim.Adam(model.parameters(), lr=params.lr, betas=(0.9, 0.999), eps=1e-8)

optimizer = torch.optim.SGD(model.parameters(), lr=params.lr, momentum=0.9) print(“\n”)

opacus_utils.register_grad_sampler(MonotonicMultiHeadAttention)(compute_monotonic_attention_grad_sample)

batch_size = 24 accountant = RDPAccountant() dp_model = GradSampleModule(model) dp_optimizer = DPOptimizer( optimizer=optimizer, noise_multiplier=1.0, # same as make_private arguments max_grad_norm=1.0, # same as make_private arguments expected_batch_size=batch_size # if you’re averaging your gradients, you need to know the denominator ) sample_rate = batch_size/len(train_q_data) print(sample_rate) dp_optimizer.attach_step_hook( accountant.get_optimizer_hook_fn(

sample_rate=sample_rate ) )

training

all_train_loss = {} all_train_accuracy = {} all_train_auc = {} all_valid_loss = {} all_valid_accuracy = {} all_valid_auc = {} best_valid_auc = 0

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' )

**I am proving all the details to reproduce this: opacus version: 1.4.0 pytorch: 2.2.1+CUDA 11.8 pandas, tqdm all latest versions work fine

pytorch / opacus

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

If we don't have any hooks, we want to skip the rest of the logic in

optimizer = torch.optim.Adam(model.parameters(), lr=params.lr, betas=(0.9, 0.999), eps=1e-8)

opacus_utils.register_grad_sampler(MonotonicMultiHeadAttention)(compute_monotonic_attention_grad_sample)

training

Train Model