How to compute gradient using BERt on inferentia

Question

Hi All,

I have to compute gradient on BERT model on inferentia. For this I guess I also need access to the hidden layers. Im currently not able to proceed because of not finding literature on the net for inferentia. There seems to be some non inferentia links but those do not work for this case.

Amongst the things which are missing are how to get positional embeddings.
Please let me know of any code snippets that may have the content Im looking for.

Thanks in advance
Ajay

Answer

Hi Ajay,

Inf1 is inference only and not intended to allow gradient calculation.  You should look at the use of torch XLA (using torch_neuronx) on inf2 or trn1 for this, this environment is intended for training tasks (as well as inference) and gradients are calculated there.  This is also likely better suited to your embedding output too.

You may want to look at https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/tutorials/training/mlp.html, but instead load the bert model of interest, using the XLA device.

How to compute gradient using BERt on inferentia

Relevant content