Nan/0 logits when finetuning ModernBERT with flash attention enabled.
Author: anunknowperson
Pull request: no
Labels: name: bug

### System Info

Output without fa:
, attn_implementation="sdpa"

{'loss': 0.3681, 'grad_norm': 5.589271545410156, 'learning_rate': 4e-05, 'epoch': 1.0}
{'eval_loss': 0.2998541593551636, 'eval_runtime': 2.4136, 'eval_samples_per_second': 20.716, 'eval_steps_per_second': 5.386, 'epoch': 1.0}
{'loss': 0.1703, 'grad_norm': 1.1856054067611694, 'learning_rate': 0.0, 'epoch': 2.0}
{'eval_loss': 0.21692198514938354, 'eval_runtime': 1.0645, 'eval_samples_per_second': 46.97, 'eval_steps_per_second': 12.212, 'epoch': 2.0}
{'train_runtime': 86.2749, 'train_samples_per_second': 10.432, 'train_steps_per_second': 2.62, 'train_loss': 0.2692194153777266, 'epoch': 2.0}
100%|████████████████████████████████████████████████████████████████████████████████| 226/226 [01:26<00:00, 2.62it/s]
Model saved successfully.

(loss != nan/0 and resulting model inference works ok)

With fa enabled:

trainer = Trainer(
{'loss': 0.4619, 'grad_norm': nan, 'learning_rate': 4e-05, 'epoch': 1.0}
{'eval_loss': nan, 'eval_runtime': 1.04, 'eval_samples_per_second': 48.079, 'eval_steps_per_second': 12.501, 'epoch': 1.0}
{'loss': 0.0, 'grad_norm': nan, 'learning_rate': 0.0, 'epoch': 2.0}
{'eval_loss': nan, 'eval_runtime': 0.7291, 'eval_samples_per_second': 68.58, 'eval_steps_per_second': 17.831, 'epoch': 2.0}
{'train_runtime': 74.6236, 'train_samples_per_second': 12.061, 'train_steps_per_second': 3.029, 'train_loss': 0.23094445625237658, 'epoch': 2.0}
100%|████████████████████████████████████████████████████████████████████████████████| 226/226 [01:14<00:00, 3.03it/s]
Model saved successfully.

Model produses nans on evals and resulting model is unusable too because of nans during inference.

Speed is improved 10x compared to disabled fa, so i can't train without it.


code snippet:




torch 2.5.1+cu124
transformers 4.48.2 (same problem with latest 4.49 built from git)
triton 3.1.0
flash_attn         2.7.1.post1

windows

```
(ve) C:\Users\Admin\Desktop\bert>nvidia-smi
Fri Jan 31 15:28:53 2025
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 560.94                 Driver Version: 560.94         CUDA Version: 12.6     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                  Driver-Model | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  Tesla P40                    TCC   |   00000000:03:00.0 Off |                  Off |
| N/A   28C    P8              9W /  250W |       9MiB /  24576MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA GeForce RTX 3090      WDDM  |   00000000:04:00.0 Off |                  N/A |
|  0%   47C    P8             30W /  350W |     316MiB /  24576MiB |      2%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
```

(ve) C:\Users\Admin\Desktop\bert>echo %CUDA_VISIBLE_DEVICES%
00000000:04:00.0

(seems like tesla is unused, as it should be)

mirror: https://huggingface.co/answerdotai/ModernBERT-base/discussions/59

### Who can help?

@muellerzr @tomaarsen @ArthurZucker

### Information

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

### Tasks

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

### Reproduction

This is my testing code. Any dataset should be loaded and adjusted a bit (fields) to get nans:

```python
# Split dataset
train_texts, test_texts, train_labels, test_labels = train_test_split(
    df_sampled["prompt"].tolist(), 
    binary_labels,
    test_size=0.1,
    random_state=42
)

# Create Hugging Face datasets
def create_hf_dataset(texts, labels):
    return Dataset.from_dict({
        "text": texts,
        "labels": labels.astype(np.float32)   # Now contains float32 values
    })
dataset_train = create_hf_dataset(train_texts, train_labels)
print(dataset_train[0]["labels"])  # Should show [1.0, 0.0, ...] instead of [1, 0, ...]
dataset_test = create_hf_dataset(test_texts, test_labels)

# Load tokenizer
checkpoint = "answerdotai/ModernBERT-large"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

# Tokenization function
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True, padding=True)

# Tokenize datasets
tokenized_train = dataset_train.map(preprocess_function, batched=True)
tokenized_test = dataset_test.map(preprocess_function, batched=True)

# Load model for multi-label classification
model = AutoModelForSequenceClassification.from_pretrained(
    checkpoint,
    num_labels=len(unique_labels),
    id2label={i: label for i, label in enumerate(unique_labels)},
    label2id={label: i for i, label in enumerate(unique_labels)},
    problem_type="multi_label_classification"
)

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

# Training setup (optimized for memory efficiency)
training_args = TrainingArguments(
    output_dir="modernbert_finetuned",
    learning_rate=lr,
    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",
    load_best_model_at_end=True,
    bf16=True,
    bf16_full_eval=True,
    push_to_hub=False,
)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    tokenizer=tokenizer,
    data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
)

# Train and save the model
trainer.train()
model.save_pretrained("modernbert_finetuned_model")
tokenizer.save_pretrained("modernbert_finetuned_model")
print("Model saved successfully.")

```



Also I have same problems with example code:

```python
# Install necessary libraries
# !pip install transformers datasets accelerate scikit-learn -Uqq

import os
import numpy as np
import pandas as pd
import torch
from functools import partial
import gc

from datasets import load_dataset
from sklearn.metrics import f1_score, accuracy_score, matthews_corrcoef
from scipy.stats import pearsonr, spearmanr
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    TrainingArguments,
    Trainer,
    TrainerCallback,
)

# Set environment variables
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_VISIBLE_DEVICES"] = "00000000:04:00.0"

# Define GLUE tasks metadata
glue_tasks = {
    "cola": {
        "abbr": "CoLA",
        "name": "Corpus of Linguistic Acceptability",
        "description": "Predict whether a sequence is a grammatical English sentence",
        "task_type": "Single-Sentence Task",
        "domain": "Misc.",
        "size": "8.5k",
        "metrics": "Matthews corr.",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["sentence"],
        "target": "label",
        "metric_funcs": [matthews_corrcoef],
        "n_labels": 2,
    },
    "sst2": {
        "abbr": "SST-2",
        "name": "Stanford Sentiment Treebank",
        "description": "Predict the sentiment of a given sentence",
        "task_type": "Single-Sentence Task",
        "domain": "Movie reviews",
        "size": "67k",
        "metrics": "Accuracy",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["sentence"],
        "target": "label",
        "metric_funcs": [accuracy_score],
        "n_labels": 2,
    },
    "mrpc": {
        "abbr": "MRPC",
        "name": "Microsoft Research Paraphrase Corpus",
        "description": "Predict whether two sentences are semantically equivalent",
        "task_type": "Similarity and Paraphrase Tasks",
        "domain": "News",
        "size": "3.7k",
        "metrics": "F1/Accuracy",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["sentence1", "sentence2"],
        "target": "label",
        "metric_funcs": [accuracy_score, f1_score],
        "n_labels": 2,
    },
    "stsb": {
        "abbr": "SST-B",
        "name": "Semantic Textual Similarity Benchmark",
        "description": "Predict the similarity score for two sentences on a scale from 1 to 5",
        "task_type": "Similarity and Paraphrase Tasks",
        "domain": "Misc.",
        "size": "7k",
        "metrics": "Pearson/Spearman corr.",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["sentence1", "sentence2"],
        "target": "label",
        "metric_funcs": [pearsonr, spearmanr],
        "n_labels": 1,
    },
    "qqp": {
        "abbr": "QQP",
        "name": "Quora question pair",
        "description": "Predict if two questions are a paraphrase of one another",
        "task_type": "Similarity and Paraphrase Tasks",
        "domain": "Social QA questions",
        "size": "364k",
        "metrics": "F1/Accuracy",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["question1", "question2"],
        "target": "label",
        "metric_funcs": [f1_score, accuracy_score],
        "n_labels": 2,
    },
    "mnli-matched": {
        "abbr": "MNLI",
        "name": "Mulit-Genre Natural Language Inference",
        "description": "Predict whether the premise entails, contradicts or is neutral to the hypothesis",
        "task_type": "Inference Tasks",
        "domain": "Misc.",
        "size": "393k",
        "metrics": "Accuracy",
        "dataset_names": {"train": "train", "valid": "validation_matched", "test": "test_matched"},
        "inputs": ["premise", "hypothesis"],
        "target": "label",
        "metric_funcs": [accuracy_score],
        "n_labels": 3,
    },
    "mnli-mismatched": {
        "abbr": "MNLI",
        "name": "Mulit-Genre Natural Language Inference",
        "description": "Predict whether the premise entails, contradicts or is neutral to the hypothesis",
        "task_type": "Inference Tasks",
        "domain": "Misc.",
        "size": "393k",
        "metrics": "Accuracy",
        "dataset_names": {"train": "train", "valid": "validation_mismatched", "test": "test_mismatched"},
        "inputs": ["premise", "hypothesis"],
        "target": "label",
        "metric_funcs": [accuracy_score],
        "n_labels": 3,
    },
    "qnli": {
        "abbr": "QNLI",
        "name": "Stanford Question Answering Dataset",
        "description": "Predict whether the context sentence contains the answer to the question",
        "task_type": "Inference Tasks",
        "domain": "Wikipedia",
        "size": "105k",
        "metrics": "Accuracy",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["question", "sentence"],
        "target": "label",
        "metric_funcs": [accuracy_score],
        "n_labels": 2,
    },
    "rte": {
        "abbr": "RTE",
        "name": "Recognize Textual Entailment",
        "description": "Predict whether one sentence entails another",
        "task_type": "Inference Tasks",
        "domain": "News, Wikipedia",
        "size": "2.5k",
        "metrics": "Accuracy",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["sentence1", "sentence2"],
        "target": "label",
        "metric_funcs": [accuracy_score],
        "n_labels": 2,
    },
    "wnli": {
        "abbr": "WNLI",
        "name": "Winograd Schema Challenge",
        "description": "Predict if the sentence with the pronoun substituted is entailed by the original sentence",
        "task_type": "Inference Tasks",
        "domain": "Fiction books",
        "size": "634",
        "metrics": "Accuracy",
        "dataset_names": {"train": "train", "valid": "validation", "test": "test"},
        "inputs": ["sentence1", "sentence2"],
        "target": "label",
        "metric_funcs": [accuracy_score],
        "n_labels": 2,
    },
}

# Function to get label maps
def get_label_maps(raw_datasets, train_ds_name):
    labels = raw_datasets[train_ds_name].features["label"]
    id2label = {idx: name.upper() for idx, name in enumerate(labels.names)} if hasattr(labels, "names") else None
    label2id = {name.upper(): idx for idx, name in enumerate(labels.names)} if hasattr(labels, "names") else None
    return id2label, label2id

# Function to preprocess data
def preprocess_function(examples, task_inputs, hf_tokenizer):
    inps = [examples[inp] for inp in task_inputs]
    tokenized = hf_tokenizer(*inps, truncation=True)
    return tokenized

# Function to compute metrics
def compute_metrics(eval_pred, task_metrics):
    predictions, labels = eval_pred
    metrics_d = {}
    for metric_func in task_metrics:
        metric_name = metric_func.__name__
        if metric_name in ["pearsonr", "spearmanr"]:
            score = metric_func(labels, np.squeeze(predictions))
        else:
            score = metric_func(np.argmax(predictions, axis=-1), labels)
        if isinstance(score, tuple):
            metrics_d[metric_func.__name__] = score[0]
        else:
            metrics_d[metric_func.__name__] = score
    return metrics_d

# Callback to capture training history
class MetricsCallback(TrainerCallback):
    def __init__(self):
        self.training_history = {"train": [], "eval": []}

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is not None:
            if "loss" in logs:  # Training logs
                self.training_history["train"].append(logs)
            elif "eval_loss" in logs:  # Evaluation logs
                self.training_history["eval"].append(logs)

# Function to fine-tune a GLUE task
def finetune_glue_task(task, checkpoint="answerdotai/ModernBERT-base", train_subset=None, do_cleanup=True):
    task_meta = glue_tasks[task]
    train_ds_name = task_meta["dataset_names"]["train"]
    valid_ds_name = task_meta["dataset_names"]["valid"]
    task_inputs = task_meta["inputs"]
    n_labels = task_meta["n_labels"]
    task_metrics = task_meta["metric_funcs"]

    # Load dataset
    raw_datasets = load_dataset("glue", task.split("-")[0] if "-" in task else task)
    if train_subset is not None and len(raw_datasets["train"]) > train_subset:
        raw_datasets["train"] = raw_datasets["train"].shuffle(seed=42).select(range(train_subset))

    id2label, label2id = get_label_maps(raw_datasets, train_ds_name)

    # Load tokenizer
    hf_tokenizer = AutoTokenizer.from_pretrained(checkpoint)
    tokenized_datasets = raw_datasets.map(partial(preprocess_function, task_inputs=task_inputs, hf_tokenizer=hf_tokenizer), batched=True)

    # Define compute metrics function
    task_compute_metrics = partial(compute_metrics, task_metrics=task_metrics)

    # Load model
    model_additional_kwargs = {"id2label": id2label, "label2id": label2id} if id2label and label2id else {}
    hf_model = AutoModelForSequenceClassification.from_pretrained(
        checkpoint, num_labels=n_labels, **model_additional_kwargs
    )

    # Data collator
    hf_data_collator = DataCollatorWithPadding(tokenizer=hf_tokenizer)

    # Training arguments
    training_args = TrainingArguments(
        output_dir=f"aai_ModernBERT_{task}_ft",
        learning_rate=8e-5,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=32,
        num_train_epochs=2,
        lr_scheduler_type="linear",
        optim="adamw_torch",
        adam_beta1=0.9,
        adam_beta2=0.98,
        adam_epsilon=1e-6,
        logging_strategy="epoch",
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        bf16=True,
        bf16_full_eval=True,
        push_to_hub=False,
    )

    # Trainer
    trainer = Trainer(
        model=hf_model,
        args=training_args,
        train_dataset=tokenized_datasets[train_ds_name],
        eval_dataset=tokenized_datasets[valid_ds_name],
        tokenizer=hf_tokenizer,
        data_collator=hf_data_collator,
        compute_metrics=task_compute_metrics,
    )

    # Add callback to trainer
    metrics_callback = MetricsCallback()
    trainer.add_callback(metrics_callback)

    # Train
    trainer.train()

    # Get training results and hyperparameters
    train_history_df = pd.DataFrame(metrics_callback.training_history["train"]).add_prefix("train_")
    eval_history_df = pd.DataFrame(metrics_callback.training_history["eval"])
    train_res_df = pd.concat([train_history_df, eval_history_df], axis=1)
    args_df = pd.DataFrame([training_args.to_dict()])

    # Cleanup
    if do_cleanup:
        cleanup(things_to_delete=[trainer, hf_model, hf_tokenizer, tokenized_datasets, raw_datasets])

    return train_res_df, args_df, hf_model, hf_tokenizer

# Function to cleanup GPU memory
def cleanup(things_to_delete=None):
    if things_to_delete is not None:
        for thing in things_to_delete:
            if thing is not None:
                del thing
    gc.collect()
    torch.cuda.empty_cache()

# Fine-tune all GLUE tasks
# Fine-tune all GLUE tasks
for task in glue_tasks.keys():
    print(f"----- Finetuning {task} -----")
    train_res_df, args_df, hf_model, hf_tokenizer = finetune_glue_task(
        task, checkpoint="answerdotai/ModernBERT-base", train_subset=1_000, do_cleanup=True
    )
    print(":: Results ::")
    print(train_res_df)  # Print training results DataFrame
    print("\n:: Hyperparameters ::")
    print(args_df)  # Print hyperparameters DataFrame
```

Outputs:
{'eval_loss': nan, 'eval_matthews_corrcoef': 0.0, 'eval_runtime': 5.3788, 'eval_samples_per_second': 193.908, 'eval_steps_per_second': 6.135, 'epoch': 1.0}
{'loss': 0.0, 'grad_norm': nan, 'learning_rate': 0.0, 'epoch': 2.0}
{'eval_loss': nan, 'eval_matthews_corrcoef': 0.0, 'eval_runtime': 1.2161, 'eval_samples_per_second': 857.625, 'eval_steps_per_second': 27.135, 'epoch': 2.0}




### Expected behavior

It should work the same way it works without flash attention enabled (e.g. sdpa).

Comments:
Maybe similar problem: https://github.com/AnswerDotAI/ModernBERT/issues/163

Someone on https://huggingface.co/answerdotai/ModernBERT-base/discussions/59 also has nans with example training code.

cc @arthurzucker for FlashAttention (and sorry for more pings!)

Did you use the FA2 build from https://github.com/bdashore3/flash-attention/releases?
The FA2 build here is only suitable for inference and cannot be used for training, which is exactly as you described.
You can use the script from https://huggingface.co/lldacing/flash-attention-windows-wheel to compile a working Windows FA2 yourself.
The speed boost you saw is also a fake illusion :(

The correctly compiled Windows FA2 wheel is approximately 178MB, not the 58MB as indicated in the link.

![Image](https://github.com/user-attachments/assets/f72658bb-778d-4ec4-9469-6888aab42d9b)

Oh, ok. Sad that there are no full windows builds. It takes more than an hour to build fa locally.

Happy to hear that the issue is fixed! 🤗 

> Did you use the FA2 build from https://github.com/bdashore3/flash-attention/releases? The FA2 build here is only suitable for inference and cannot be used for training, which is exactly as you described. You can use the script from https://huggingface.co/lldacing/flash-attention-windows-wheel to compile a working Windows FA2 yourself. The speed boost you saw is also a fake illusion :(

Hey neovo! May I know where do you get the info that "The FA2 build here is only suitable for inference and cannot be used for training"? I plan to install fa2 with the source compile way mentioned on README：python setup.py install,  hope it will work.