SBERT微调始终提前终止，无法完成全部epoch的解决方案

项目背景 #

我正在开展一个文本分类项目，使用SBERT预训练模型MiniLM-L6-H384-uncased，涉及995个分类类别。大部分步骤遵循官方文档，整体运行无异常。

问题现象 #

无论设置何种训练参数，训练总是提前结束，无法完成全部批次：

• 设置num_train_epochs=1，仅能运行到0.49个epoch；
• 设置num_train_epochs=4，则始终在3.49个epoch时终止。

训练代码 #

from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    SentenceTransformerModelCardData,
)
from sentence_transformers.losses import BatchAllTripletLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import TripletEvaluator

model = SentenceTransformer(
    "nreimers/MiniLM-L6-H384-uncased",
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="apache-2.0",
        model_name="all-MiniLM-L6-v2",
    )
)

loss = BatchAllTripletLoss(model)

args = SentenceTransformerTrainingArguments(
    output_dir="finetune/model20240924",
    num_train_epochs=1,
    max_steps = -1,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    learning_rate=1e-5,
    warmup_ratio=0.1,
    fp16=True,
    bf16=False,
    batch_sampler=BatchSamplers.GROUP_BY_LABEL,
    eval_strategy="no",
    eval_steps=100,
    save_strategy="epoch",
    save_total_limit=2,
    logging_steps=100,
    run_name="miniLm-triplet",
)

trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=trainDataset,
    eval_dataset=devDataset,
    loss=loss,
)
trainer.train()

数据集信息 #

训练集结构：

Dataset({
    features: ['Title', 'Body', 'label'],
    num_rows: 23961
})

验证集结构相同，仅行数更少。

训练输出示例 #

[1473/2996 57:06 < 59:07, 0.43 it/s, Epoch 0/1]
Step    Training Loss
100     1.265600
200     0.702700
300     0.633900
400     0.505200
500     0.481900
600     0.306800
700     0.535600
800     0.369800
900     0.265400
1000    0.345300
1100    0.516700
1200    0.372600
1300    0.392300
1400    0.421900

TrainOutput(global_step=1473, training_loss=0.5003972503496366, metrics={'train_runtime': 3427.9198, 'train_samples_per_second': 6.99, 'train_steps_per_second': 0.874, 'total_flos': 0.0, 'train_loss': 0.5003972503496366, 'epoch': 0.4916555407209613})

解决方案建议 #

• 更换采样器：当前使用的GROUP_BY_LABEL采样器可能因为部分标签样本量过少，导致无法凑够完整epoch的批次。换成默认的SEQUENTIAL采样器，观察是否能完成训练。
• 检查标签分布：统计数据集中每个标签的样本数量，看是否存在大量标签只有1-2个样本的情况。这类样本会被分组采样器过滤，导致实际参与训练的样本数远少于总样本数，进而提前终止训练。
• 关闭混合精度：暂时将fp16设为False，排除混合精度训练带来的隐性停止问题。
• 验证数据集完整性：遍历训练集，确认所有样本的Title、Body和label都有效，无空值或异常值。
• 强制设置max_steps：手动计算完成1个epoch所需的steps（公式：总样本数 // (per_device_train_batch_size * 设备数量)），将max_steps设为该值，替代num_train_epochs参数，强制跑完指定步数。

内容的提问来源于stack exchange，提问作者SohmOuse