ModernBERT for Sequence Classification - issues with finetuning
Author: 98MM
Pull request: no
Labels: name: bug

### System Info

transformers version 4.53.0.dev0
flash_attn version 2.7.4.post1 (via pip)

Using offical transformers image from: https://hub.docker.com/r/huggingface/transformers-pytorch-gpu

### Who can help?

@ArthurZucker
@zach-huggingface

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

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

### Reproduction

Standard BERT-like finetuning approach for NLI using AutoModelForSequenceCassification:

```python

nli = load_dataset('kiddothe2b/contract-nli', 'contractnli_a') # also tested with subset of sentence-transformer/all-nli

def tokenize_function(examples, tokenizer, task_inputs):
    inps = [examples[inp] for inp in task_inputs]
    tokenized = tokenizer(*inps, truncation=True)
    return tokenized

# tokenize the datasets ...

train_dataset = nli['train'].map(
    tokenize_function,
    batched=True,
    remove_columns=['premise', 'hypothesis'],
    fn_kwargs={"tokenizer":tokenizer, "task_inputs":['premise', 'hypothesis']})

tokenizer=AutoTokenizer.from_pretrained('answerdotai/ModernBERT-base')
data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding='longest')

# compute metrics fn ...

model = AutoModelForSequenceClassification.from_pretrained(
    'answerdotai/ModernBERT-base',
    num_labels=3,
    classifier_pooling='cls',
    torch_dtype='float16',
    id2label=id2label,
    label2id=label2id,
    reference_compile=False # also tested without this setting
).to('cuda') # tested this on L4, A100, A100_80GB, H100

train_bsz, val_bsz = 16, 16
lr = 2e-5
betas = (0.9, 0.98)
n_epochs = 10
eps = 1e-6
wd = 8e-6

training_args = TrainingArguments(
    output_dir='/out',
    learning_rate=lr,
    weight_decay=wd,
    per_device_train_batch_size=train_bsz,
    per_device_eval_batch_size=val_bsz,
    num_train_epochs=n_epochs,
    lr_scheduler_type="linear",
    optim="adamw_torch",
    adam_beta1=betas[0],
    adam_beta2=betas[1],
    adam_epsilon=eps,
    logging_strategy="epoch",
    eval_strategy="epoch",
    save_strategy="epoch",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    processing_class=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

trainer.train()
```

### Expected behavior

I want to preface this that I first encountered this problem on my own pretrained ModernBERT implementation. After extensive testing on my version I decided to test the official  version (answerdotai/ModernBERT-base), where I encountered the same issue. 

My base model works well on fill-mask tasks (via pipeline), outputting expected, correct words. After pretraining it achieved an accuracy of 0.86 and perplexity of 2.1 on an eval dataset. These results are why I decided to test the offical implementation as well.

This issue I encountered when attempting to use AutoModelForSequenceClassification / ModernBertForSequenceClassification on both my and the official version.


Issue behaviour:

Training begins and runs. However, the prediction accuracy and F1 scorse indicate random guessing/very unstable training (with three classes macro accuracy and f1 hower around 0.33, while per-class values vary wildly from 0.00 up to even 0.91 per epoch, seemingly randomly).

Looking at the embeddings themselves, they seem fine - no issues with NaN values (as raised [here](https://github.com/huggingface/transformers/issues/35574)). I've also tested masked word prediction (via fill-mask pipeline), which also output correct/expected predictions.

Expected behaviour would be at least stable training, if not similar results as claimed in the original ModernBERT paper.

Comments:
Upon further testing I believe the issue may be related to a mismatch in hiddenstate outputs, highlighted in https://github.com/huggingface/transformers/issues/38499


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Can you confirm whether this is fixed or not? 🤗 

Hi! Yes - I tested with the latest container (the one from 9h ago), works fine! I also found that with the previous version (4.53.0.dev0), switching to 'mean' pooling fixed the issue as well. Thank you!

Closing!