Keep getting nan-loss when using Seq2SeqTransformer - tensorflow

Question

I am trying to train a transformer model on text data. The task is to predict missing (masked) words so e.g. the input "How are you ?" gets mapped to "How [MASKED] you ?" like so:
inputs = [69, 4, 1337, 666] # How [MASK] you ?
targets = [69, 42, 1337, 666] # How are you ?
The problem is, that after a few steps, sometimes after a few hundred, sometimes after a few thousand, the loss becomes nan.
I have tried a model with just 90k parameters but also one with 10m parameters. The result is always the same.
The code below shows how I instantiate a Seq2SeqTransformer.
Using the debugger does not give me anything in the "Graph Executions" section. All I see is this:
Any idea what I could be doing wrong here? The learning rate is already rather small so I can't imagine that this is the problem.
model = Seq2SeqTransformer(
vocab_size=vocab_size, # <= 2000
embedding_width=32,
dropout_rate=0.1,
encoder_layer=TransformerEncoder(
num_layers=1, num_attention_heads=2, intermediate_size=64,
dropout_rate=0.1, intermediate_dropout=0.1, attention_dropout_rate=0.1
),
decoder_layer=TransformerDecoder(
num_layers=1, num_attention_heads=2, intermediate_size=64,
dropout_rate=0.1, intermediate_dropout=0.1, attention_dropout_rate=0.1
)
)
optimizer = Adam(
learning_rate=TransformerSchedule(
min_lr=2.5e-6,
max_lr=1.5e-4,
warmup_steps=6000,
warm_steps=30000
)
)
model.compile(
optimizer=optimizer,
loss=SmoothedSparseCategoricalCrossentropy(0.1),
)
model.fit(
train_data.repeat(),
steps_per_epoch=1000,
epochs=500,
callbacks=callbacks,
validation_data=valid_data,
validation_steps=100,
)
Batch Data
Just to verify that the data I present to the model is alright, here I print(a_batch). Since samples get bucketed, they all have the same length which is also why input_masks is all 1s.
Note: ID of [MASK] is 4.
X
{'inputs': <tf.Tensor: shape=(119, 41), dtype=int32, numpy=
array([[ 2, 192, 214, ..., 525, 7, 3],
[ 2, 57, 964, ..., 15, 7, 3],
[ 2, 4, 191, ..., 646, 7, 3],
...,
[ 2, 430, 29, ..., 675, 4, 3],
[ 2, 101, 45, ..., 15, 7, 3],
[ 2, 421, 11, ..., 15, 4, 3]], dtype=int32)>,
'input_masks': <tf.Tensor: shape=(119, 41), dtype=float32, numpy=
array([[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
...,
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.]], dtype=float32)>,
'targets': <tf.Tensor: shape=(119, 41), dtype=int32, numpy=
array([[ 2, 192, 214, ..., 525, 7, 3],
[ 2, 57, 964, ..., 15, 7, 3],
[ 2, 104, 191, ..., 646, 7, 3],
...,
[ 2, 430, 29, ..., 675, 7, 3],
[ 2, 101, 45, ..., 15, 7, 3],
[ 2, 421, 11, ..., 15, 7, 3]], dtype=int32)>}
Y
tf.Tensor(
[[ 2 192 214 ... 525 7 3]
[ 2 57 964 ... 15 7 3]
[ 2 104 191 ... 646 7 3]
...
[ 2 430 29 ... 675 7 3]
[ 2 101 45 ... 15 7 3]
[ 2 421 11 ... 15 7 3]], shape=(119, 41), dtype=int32)