Splitting optimizer graph is needed to prevent embedding table update error

[ptoulme-aws]

System Info

Top of tree of this package. Running on AWS Neuron Trn1 instance.

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Run this repository for Llama 7B pretraining on an AWS Neuron Trn1 instance. A xm.mark_step() is required before the optimizer.step() to prevent Neuron runtime errors related to the optimizer update of the embedding tables. The runtime error is an out of bound memory access related to the embedding table. Note all weights are sharded in 1D over 32 devices in tensor parallel.

I am reaching out for advice on why this would be happening?

Expected behavior

Fwd/bwd/optimizer can all run on AWS Neuron Trn1 in one graph.