[Python] Model returns NaNs when using ModernBERT instead of Roberta

I'm trying to fine-tune ModernBERT for a classification task. For this I had some old code written using PyTorch that I've used to fine-tune BERT and recently Roberta with no issues. But when I switch out Roberta with ModernBERT, I suddenly get NaNs in my output tensor from the very first batch. This means my loss is also NaNs and training can't happen. I did some research and found the problem could be exploding gradients and or NaNs in the Input, both of which is not true in my case. So, I'm looking for some help in what might be going wrong.

Here's my code for training loop. I'd also appreciate some feedback on this loop.

for epoch_i in range(0, EPOCHS):
total_train_loss = 0

model.train()

for step, batch in enumerate(train_dataloader):

b_input_ids = batch[0].to(device)
b_input_mask = batch[1].to(device)
b_labels = batch[2].to(device)

model.zero_grad()

output = model(
b_input_ids,
token_type_ids=None,
attention_mask=b_input_mask,
labels=b_labels,
)

loss = output[0]
total_train_loss += loss.item()
loss.backward()
optimizer.step()
scheduler.step()

avg_train_loss = total_train_loss / len(train_dataloader)
training_time = format_time(time.time() - t0)

t0 = time.time()
model.eval()
total_eval_accuracy = 0
total_eval_loss = 0
nb_eval_steps = 0

for batch in validation_dataloader:
b_input_ids = batch[0].to(device)
b_input_mask = batch[1].to(device)
b_labels = batch[2].to(device)

with torch.no_grad():
output = model(
b_input_ids,
token_type_ids=None,
attention_mask=b_input_mask,
labels=b_labels,
)

loss = output[0]
logits = output[1]
# loss = loss_fn(output.logits, b_labels.squeeze())
total_eval_loss += loss.item()
label_ids = b_labels.to("cpu").numpy()
total_eval_accuracy += f1_score(
torch.argmax(logits, dim=1).cpu().numpy(),
label_ids,
average="weighted",
)

avg_val_accuracy = total_eval_accuracy / len(validation_dataloader)
print(" F1 Score: {0:.2f}".format(avg_val_accuracy))

avg_val_loss = total_eval_loss / len(validation_dataloader)


Here's how I'm setting up the model:

config = AutoConfig.from_pretrained("answerdotai/ModernBERT-base")
config.num_labels = 2
tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
model = AutoModelForSequenceClassification.from_pretrained("answerdotai/ModernBERT-base", config=config)

Continue reading...