answerdotai

/

ModernBERT-base

like 720

Follow

Answer.AI 376

, 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.

C:\Users\Admin\Desktop\bert\train_fin.py:152: FutureWarning: tokenizer is deprecated and will be removed in version 5.0.0 for Trainer.__init__. Use processing_class instead.
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.

# 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.")

When I enable flash attention on windows, speed is greatly improved, but resulting model is broken (nan values during eval/inference).

torch 2.5.1+cu124
transformers 4.48.2
triton 3.1.0

Any suggestions on how to fix?