Gradients returning None in huggingface module

user_249341 0
Viewed 593

I want to get the gradient of an embedding layer from a pytorch/huggingface model. Here's a minimal working example:

from transformers import pipeline

nlp = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

responses = ["I'm having a great day!!"]
hypothesis_template = 'This person feels {}'
candidate_labels = ['happy', 'sad']
nlp(responses, candidate_labels, hypothesis_template=hypothesis_template)

I can extract the logits just fine,

inputs = nlp._parse_and_tokenize(responses, candidate_labels, hypothesis_template)
predictions = nlp.model(**inputs, return_dict=True, output_hidden_states=True)
predictions['logits']

and the model returns a layer I'm interested in. I tried to retain the gradient and backprop with respect to a single logit I'm interested in:

layer = predictions['encoder_hidden_states'][0]
layer.retain_grad()
predictions['logits'][0][2].backward(retain_graph=True)

However, layer.grad == None no matter what I try. The other named parameters of the model have their gradients computed, so I'm not sure what I'm doing wrong. How do I get the grad of the encoder_hidden_states?

huggingface-transformers pytorch nlp python
1 Answers

I was also very surprised of this issue. Although I have never used the library I went down and did some debugging and found out that the issue is coming from the library transformers. The problem is comming from from this line :

encoder_states = tuple(hidden_state.transpose(0, 1) for hidden_state in encoder_states)

If you comment it out, you will get the gradient just with some dimensions transposed. This issue is related to the fact that Pytorch Autograd does not do very well on inplace operations as mentioned here.

So to recap the solution is to comment line 382 in modeling_bart.py.

You will get the gradient with this shape T x B x C instead of B x T x C, but you can reshape it as you want later.

user_28476990
Related