Support for Transformers

[KastanDay]

Requesting a new feature (I can't seem to add a label myself).

Description

Do you currently support Huggingface Transformers? In particular, I'd like to debug T5.

I tried an editable install of Transformers, and modifying the PyTorch implementation underneath, but I can't quite get the two different loss functions to play together properly.
For example, I'm trying to steal this implementation from the Basic Examples, and add it to Transformers:

loss_fn = extend(torch.nn.CrossEntropyLoss(reduction="mean"))
individual_loss_fn = torch.nn.CrossEntropyLoss(reduction="none")

After 30 minutes of tinkering, I can't get arround this error. Thanks for any assistance.

AssertionError: BackPACK extension expects a backpropagation quantity but it is None. Module: Linear(in_features=768, out_features=32128, bias=False), Extension: <backpack.extensions.secondorder.diag_hessian.DiagHessian object at 0x7f6a72f102b0>.

During handling of the above exception, another exception occurred:

AttributeError                            Traceback (most recent call last)
Cell In[6], line 80
     69 # loss.sum().backward()
     70 # backward pass
     71 with cockpit(
     72     global_step,
     73     info={
   (...)
     78     },
     79 ):
---> 80     loss.sum().backward(create_graph=cockpit.create_graph(global_step))
     82 # optimizer step
     83 optimizer.step()

File ~/githubs/cockpit/cockpit/context.py:155, in BackwardCTX.__exit__(self, type, value, traceback)
    152 for ctx in self.contexts:
    153     ctx.__exit__(type, value, traceback)
--> 155 self.cp.track(self.global_step, protected_savefields=self.protected_savefields)
    157 CockpitCTX.erase()

File ~/githubs/cockpit/cockpit/cockpit.py:195, in Cockpit.track(self, global_step, protected_savefields)
    190 before_cleanup = [
    191     q for q in self.quantities if not isinstance(q, quantities.HessMaxEV)
    192 ]
    194 for q in before_cleanup:
--> 195     q.track(global_step, self.params, batch_loss)
    197 self._free_backpack_buffers(global_step, protected_savefields)
    199 after_cleanup = [
    200     q for q in self.quantities if isinstance(q, quantities.HessMaxEV)
    201 ]

File ~/githubs/cockpit/cockpit/quantities/quantity.py:101, in Quantity.track(self, global_step, params, batch_loss)
     92 """Perform scheduled computations and store result.
     93 
     94 Args:
   (...)
...
--> 335 grad_dict = {id(p): p.grad.data.clone().detach() for p in params}
    336 self.save_to_cache(global_step, f"grad_{point}", grad_dict, block_fn)
    338 # L = ¹/ₙ ∑ᵢ ℓᵢ, BackPACK's BatchGrad computes ¹/ₙ ∇ℓᵢ, we have to rescale

AttributeError: 'NoneType' object has no attribute 'data'

Update: after 2 more hours of tinkering, I'm pretty sure I am exactly matching your examples, but I'm still facing the exact same issue.

BTW, I think I found this project at NeurIPS. Great work on it!

[f-dangel]

Hi @KastanDay, thanks for giving Cockpit a try!

From the title and description, I think you've already accurately boiled down the problem. BackPACK, which is used by Cockpit to compute some of its quantities, currently does not support transformer architectures.

I am maintaining BackPACK and can add a feature request to support transformers. A first step would probably be to support the computation of per-sample gradients, which would already enable a large number of Cockpit instruments.

However, since I am quite unfamiliar with these architectures, it would be helpful to see some code how they are composed in practice. Do you have a code reference to a 'basic' transformer architecture that could serve as a starting point?

Best,
Felix

[b11z]

Hi there! If I understand correctly, this issue is blocked on the ability to compute the per-training-sample gradient of a transformer model.

That problem has been studied in the context of Differential Privacy -- the popular DP-SGD algorithm requires per-sample gradients. Opacus is a library that's relatively mature and implements it. There's also https://github.com/lxuechen/private-transformers which has some transformer-specific optimizations (specifically ghost clipping).

Have you looked at whether those approaches could be used for backpack?

[f-dangel]

Hi,

so I didn't dig into the internals of Opacus and the other source you mentioned, but for making BackPACK support per-sample gradients for attention layers, there are two options:

1. Rewrite the attention layers in terms of nn.Linear layers. Then, you can extend these layers and BackPACK will be able to compute their per-sample gradients.
2. Add support for nn.MultiheadAttention in BackPACK's BatchGrad extension. There is an example in the docs showing how to do that.

Best,
Felix