Pytorch BERT input gradient

One of the problems with Bert models is that your input mostly contains token IDs rather than token embeddings, which makes getting gradient difficult since the relation between token ID and token embeddings is discontinued. To solve this issue, you can work with token embeddings.

# get your batch data: token_id, mask and labels
token_ids, mask, labels = batch
  
# get your token embeddings
token_embeds=BertModel.bert.get_input_embeddings().weight[token_ids].clone()
# track gradient of token embeddings
token_embeds.requires_grad=True
    
# get model output that contains loss value
outs = BertModel(inputs_embeds=inputs_embeds,labels=labels)
loss=outs.loss

After getting loss value, you can use torch.autograd.grad in this answer or backward function

loss.backward()
grad=token_embeds.grad

You can use torch.autograd.grad (documentation):

y_pred = BertModel(x)
out = loss_func(y_label, y_pred)  # not necessary a scalar!
grad = torch.autograd.grad(
    outputs=out,
    inputs=x,
    grad_outputs=torch.ones(out.size()).to(device), # or simply None if out is a scalar
    retain_graph=False,
    create_graph=False,
    only_inputs=True)[0]

You should pass retain_graph and create_graph to True if you want to use grad for computing a loss and apply backward (typically for computing a gradient penalty). Otherwise keep it to False to save memory and time.