Just to confirm that the weight decay may have been the culprit… Added the following snippet and things are training smoothly…
def group_weight(module):
group_decay = []
group_no_decay = []
for n,m in module.named_parameters():
if ‘input_embedding’ in n:
group_no_decay.append(m)
else:
group_decay.append(m)
assert len(list(module.parameters())) == len(group_decay) + len(group_no_decay)
groups = [dict(params=group_decay), dict(params=group_no_decay, weight_decay=.0)]
return groups