Ettin Reranker系列模型发布
新的Reranker模型家族推出。对从事信息检索、RAG或向量搜索优化的程序员有参考价值。
新的Reranker模型家族推出。对从事信息检索、RAG或向量搜索优化的程序员有参考价值。
今天,我发布了六款新的 Sentence Transformers CrossEncoder 重排序模型。它们基于 Ettin ModernBERT 编码器构建,在各自的模型规模上均达到了当前最佳水平。同时,我也发布了用于训练这些模型的数据和完整训练方案:
cross-encoder/ettin-reranker-17m-v1
cross-encoder/ettin-reranker-32m-v1
cross-encoder/ettin-reranker-68m-v1
cross-encoder/ettin-reranker-150m-v1
cross-encoder/ettin-reranker-400m-v1
cross-encoder/ettin-reranker-1b-v1
这些模型采用知识蒸馏方案训练:在 cross-encoder/ettin-reranker-v1-data 上,针对 mixedbread-ai/mxbai-rerank-large-v2 生成的分数计算逐点 MSE。该数据集由 lightonai/embeddings-pre-training 的一个子集,与 lightonai/embeddings-fine-tuning 中经过重排序的一个子集混合而成。
我们将六款重排序模型与 google/embeddinggemma-300m 搭配,在 MTEB(eng, v2) Retrieval 上进行了评测。另请参阅“结果”部分,了解与另外五款嵌入模型搭配时的表现。
如果你刚接触重排序模型,想先了解“为什么”,请跳转至“什么是重排序模型,为什么要将它与嵌入模型搭配使用?”;如果你只想直接接入模型,请跳转至“用法”;如果你想训练自己的模型,请跳转至“训练”。
我使用 Sentence Transformers v5.5.0 中新推出的 train-sentence-transformers AI 智能体技能搭建了下文的初始训练方案。可通过 hf skills add train-sentence-transformers [--global] [--claude] 安装它,然后让你的 AI 编程智能体(Claude Code、Codex、Cursor、Gemini CLI……)使用你的数据对 SentenceTransformer、CrossEncoder 或 SparseEncoder 模型进行微调。
什么是重排序模型,为什么要将它与嵌入模型搭配使用?
用法 端到端的“检索后重排序”流水线
端到端的“检索后重排序”流水线
结果 MTEB(eng, v2) Retrieval 速度
MTEB(eng, v2) Retrieval
训练 蒸馏方案 数据集 训练参数 评估 完整训练脚本
完整训练脚本
重排序模型(又称逐点交叉编码器)是一种神经网络模型,它接收一个(查询,文档)对,并输出单个相关性分数。嵌入模型会分别对查询和文档进行编码,再根据两个嵌入向量计算相似度;与之不同,重排序模型允许两段文本在每一层 Transformer 中相互关注。这种联合编码更加准确,但成本也更高:模型必须针对每一个(查询,文档)对运行一次,而不是每段文本只运行一次。
由于在整个语料库上运行交叉编码器的成本过高,生产环境中常见的模式是“检索后重排序”:先由快速的嵌入模型检索出 top-K 个候选项(成本较低),再使用交叉编码器仅对这 K 个候选项进行高精度的重新排序。这样既能让总成本保持在可控范围内,又能使最终排序结果非常接近对所有内容执行穷举式交叉编码器推理所得到的结果。
本文会交替使用“重排序模型”和“交叉编码器”这两个术语。
本次发布的模型都是标准的 Sentence Transformers CrossEncoder 模型,因此只需 3 行代码即可使用:
from sentence_transformers import CrossEncoder
model = CrossEncoder("cross-encoder/ettin-reranker-32m-v1")
scores = model.predict([
("Where was Apple founded?", "Apple Inc. was founded in Cupertino, California in 1976 by Steve Jobs, Steve Wozniak, and Ronald Wayne."),
("Where was Apple founded?", "The Fuji apple is an apple cultivar developed in the late 1930s and brought to market in 1962."),
])
print(scores)
# [11.393298 2.968891] <- larger means more relevant
对于一个查询和一组候选项,你还可以使用 rank 获取排序后的索引和分数:
ranked = model.rank(
query="Which planet is known as the Red Planet?",
documents=[
"Venus is often called Earth's twin because of its similar size and proximity.",
"Mars, known for its reddish appearance, is often referred to as the Red Planet.",
"Jupiter, the largest planet in our solar system, has a prominent red spot.",
"Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
],
top_k=4,
return_documents=True,
)
for r in ranked:
print(f"({r['score']:.2f}): {r['text']}")
# (10.82): Mars, known for its reddish appearance, is often referred to as the Red Planet.
# (9.86): Saturn, famous for its rings, is sometimes mistaken for the Red Planet.
# (8.55): Jupiter, the largest planet in our solar system, has a prominent red spot.
# (6.21): Venus is often called Earth's twin because of its similar size and proximity.
你可以将 cross-encoder/ettin-reranker-32m-v1 替换为其他任意规模的模型,在质量和速度之间进行权衡。得益于 ModernBERT 的长上下文预训练,这六款模型均支持最多 8K token 的上下文,这对于长文档重排序非常有用。
为了获得最高吞吐量,建议安装 kernels,并设置 model_kwargs={"dtype": "bfloat16", "attn_implementation": "flash_attention_2"}。更多详情请参阅下文的“速度”部分。总体而言,根据模型规模和序列长度的不同,与默认加载方式相比,你可以获得 1.7~8.3 倍的加速。
from sentence_transformers import CrossEncoder
model = CrossEncoder(
"cross-encoder/ettin-reranker-32m-v1",
model_kwargs={"dtype": "bfloat16", "attn_implementation": "flash_attention_2"},
)
下面是一个完整示例:使用快速嵌入模型进行检索,再使用重排序模型完成最终排序。
from sentence_transformers import SentenceTransformer, CrossEncoder
# Fast retrieval with a static embedder (sub-millisecond on CPU per query)
embedder = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1")
reranker = CrossEncoder("cross-encoder/ettin-reranker-68m-v1")
corpus = [
"Apple Inc. was founded in Cupertino, California in 1976 by Steve Jobs, Steve Wozniak, and Ronald Wayne.",
"The Fuji apple is an apple cultivar developed in the late 1930s.",
"Steve Jobs introduced the iPhone in 2007 at Macworld.",
"Macintosh computers were sold by Apple from 1984 onward.",
# ... thousands or millions more in production
]
query = "Where was Apple founded?"
# Step 1: encode + retrieve top-100
query_emb = embedder.encode_query(query, convert_to_tensor=True)
corpus_emb = embedder.encode_document(corpus, convert_to_tensor=True)
scores = embedder.similarity(query_emb, corpus_emb)[0]
top_k_idx = scores.topk(min(100, len(corpus))).indices.tolist()
# Step 2: rerank
top_k_docs = [corpus[i] for i in top_k_idx]
ranked = reranker.rank(query, top_k_docs, top_k=5, return_documents=True)
for r in ranked:
print(f"({r['score']:.2f}): {r['text']}")
# (11.63): Apple Inc. was founded in Cupertino, California in 1976 by Steve Jobs, Steve Wozniak, and Ronald Wayne.
# (4.71): Steve Jobs introduced the iPhone in 2007 at Macworld.
# (1.96): The Fuji apple is an apple cultivar developed in the late 1930s.
# (1.49): Macintosh computers were sold by Apple from 1984 onward.
大多数现代搜索系统都采用这种结构。检索器决定哪些内容能够进入漏斗,重排序模型则决定最终胜出者。
六款重排序模型采用相同的架构,差异仅在于骨干网络的规模。它们的骨干网络分别来自约翰斯·霍普金斯大学 Ettin 系列中的六款编码器。这些模型采用 ModernBERT 风格的架构,具备无填充注意力、RoPE 位置编码和 GeGLU,并使用 2T 个采用开放许可证的 token 进行预训练,最多支持 8192 个 token 的上下文。
在每个编码器之上,重排序模型使用了一个由 4 个模块组成的分类头。该分类头复刻了 ModernBertForSequenceClassification 的结构,但使用 Sentence Transformers 的模块化组件构建。底层 Transformer 使用普通的 AutoModel,而不是 AutoModelForSequenceClassification,因此我们可以在使用 Flash Attention 2 时,对长度可变的输入进行序列去填充。对于中等文档序列长度,根据模型规模的不同,与 fp32+SDPA 相比,这种方式可实现 1.7~8.3 倍的加速(完整基准测试请参阅“速度”部分):
1. Transformer(FA2)
2. Pooling(cls)
3. Dense(H, H, bias=False, GELU)
4. LayerNorm(H)
5. Dense(H, 1, scores)
在我的消融实验中,CLS 池化的表现优于平均池化,这一点有些出人意料。ModernBERT 每三层才使用一次全局注意力,其余三分之二的层都采用局部窗口注意力,无法从较远的位置访问 CLS。实验结果表明,仅凭少数几个全局注意力层就足以传递充分的信号,使 CLS 成为更好的池化选择。
六款模型均采用 Apache 2.0 许可证发布,与 Ettin 编码器保持一致。
我使用 MTEB 的两阶段重排序流程,在完整的 MTEB(eng, v2) Retrieval 基准测试(10 项任务,对 top-100 结果进行重排序)上评测了每款已发布的模型,并将各重排序模型分别与六款横跨不同速度和质量水平的嵌入模型搭配:
下方每张图中的虚线代表仅使用检索器时的结果,也是需要超越的基准数字。任何低于该虚线的结果,都意味着重排序模型实际上降低了流水线的平均表现:
按照降序排列的六种嵌入模型搭配方案的平均 NDCG@10。我们的六款模型以粗体显示,教师模型 mixedbread-ai/mxbai-rerank-large-v2 以下划线显示。
† 上限设为 max_seq_length=8192(基于 4B Qwen3 的重排序模型在原生上下文长度下无法装入单张 H100 80GB)。使用原生上下文长度进行评估的结果可能会更高。
NanoBEIR 是 BEIR 的一个快速子集,包含 13 个数据集,每个数据集使用 50 个查询,最多检索 5000 个文档。训练期间,metric_for_best_model 被设置为 NanoBEIR(参见“评估”部分),我的实验方向也是以它为依据。
这次发布的最小模型 17M,参数量约为 33M ms-marco-MiniLM-L12-v2 的一半,但在 MTEB 上的 NDCG@10 高出 0.051(0.5576 对 0.5066),在 NanoBEIR 上高出 0.038(0.6746 对 0.6369)。32M 模型则以 0.025 的优势在 MTEB 上击败了 568M 的 BAAI/bge-reranker-v2-m3(0.5779 对 0.5526),两者参数量相差 17 倍。如果你一直在检索后重排序流程中默认使用某个旧版 MiniLM 重排序模型,那么换成我们的 17M(或 32M)模型,是一种风险很低的直接替换方案,并且能在两个基准测试上获得明显的质量提升。
继续往表格上方看,我们的 150M 模型是我测试过的 600M 参数以下模型中,在 MTEB 上表现最强的重排序模型。它以 0.005 的优势略微超过近期发布的 Qwen/Qwen3-Reranker-0.6B(596M)(0.5994 对 0.5940),并领先所有 BAAI bge-reranker 变体 0.03 至 0.05。68M 模型也值得一提:它取得了 0.5915 的成绩,与 Qwen3-Reranker-0.6B 的 0.5940 几乎完全相当,却只使用了后者九分之一的参数。
在已发布模型的最高规格端,我们的 1B 模型表现得与其教师模型非常接近。在 MTEB 上,它与 1.54B 的 mxbai-rerank-large-v2 仅相差 0.0001(0.6114 对 0.6115);在 NanoBEIR 上也只相差 0.008,尽管其蒸馏来源模型的规模比自身大 54%。蒸馏实际上消除了与教师模型之间的差距,这正是我在准备此次发布时希望看到的结果。
对比中整体表现最强的重排序模型是 Qwen/Qwen3-Reranker-4B,MTEB 得分为 0.6367,比我们的 1B 模型高出 0.025。若想使用当前方案缩小这一差距,可能需要从更强的教师模型进行蒸馏(我们的教师模型本身就弱于 Qwen3-Reranker-4B)。对于大多数检索后重排序工作负载而言,我们的 1B 模型仅使用四分之一的参数(参见“速度”部分),是实用得多的选择。
质量指标只体现了重排序模型一半的重要性。另一半在于,它的延迟能否满足从完成检索到向用户展示结果之间的时间预算。下面介绍一下我的测量结果。
我在单张 NVIDIA H100 80GB 上,将此次发布的全部六个模型与十三个公开重排序模型进行了基准测试,其中包括参数规模约为 1B 的强基线模型。查询和文档来自 sentence-transformers/natural-questions,并保留其自然的文档长度分布:大多数 NQ 答案较短,也有一些较长。文档在 max_length=512 处截断,以免让旧模型处于不公平的劣势。每个模型都使用其支持的最佳注意力实现:支持的架构(BERT、XLM-RoBERTa、ModernBERT、Qwen2)使用 Flash Attention 2,不支持的架构使用 SDPA,而 DeBERTa-v2 使用 eager(目前 transformers 中的 DeBERTa-v2 既不支持 FA2,也不支持 SDPA)。
对于每个模型,自动批量大小搜索都从 batch size 8 开始,并不断翻倍,直到 GPU 内存不足。对于每种批量大小,我都会执行三次计时测试并取吞吐量中位数,避免某次偶然的异常运行影响结果。最终报告的是取得最佳结果的批量大小所对应的吞吐量。
表 1:吞吐量,单位为每秒处理的文档对数,所有模型均使用 bfloat16。我们的六个重排序模型以粗体标出。
我们的 17M 模型是整个对比中速度最快的重排序模型,吞吐量达到每秒 7517 对。这几乎是 ms-marco-MiniLM-L6-v2(3817)的两倍,甚至比规模更小的 ms-marco-MiniLM-L4-v2(4029)还要快。而且,正如你在前面的 MTEB 表格中看到的那样,我们的 17M 模型也比所有 MiniLM 变体更准确。如果你目前正在运行 MiniLM 交叉编码器,只需修改一行代码换成我们的 17M 模型,就能同时改善延迟和搜索质量。
我们的 150M 模型带来了一个更加有趣的对比,因为有两个参数量同为 150M 的直接架构对手:Alibaba-NLP/gte-reranker-modernbert-base 和 ibm-granite/granite-embedding-reranker-english-r2。两者都基于相同的 ModernBERT-base 主干网络构建。我们的 150M 模型吞吐量为每秒 3237 对,而另外两个模型分别为 1418 和 1404,速度相差 2.3 倍。
三个 150M 模型都使用 Flash Attention 2,但另外两个模型通过 AutoModelForSequenceClassification 加载,因此输入仍然保留填充。也就是说,注意力计算本身运行的是 FA2 内核,但模型的其余部分依旧会对毫无贡献的填充 token 执行稠密计算。我们的模块化 Transformer 模块(参见上文“架构细节”)会让去除填充后的输入贯穿整个模型,因此每一层都只在真实 token 上进行计算。这正是只能获得部分 FA2 收益与获得全部 FA2 收益之间的区别。
在表格底部,我们的 1B 模型达到每秒 928 对,比 1.54B 的教师模型 mxbai-rerank-large-v2(每秒 387 对)快 2.4 倍,同时 MTEB 得分与其仅相差 0.0001。教师模型基于 Qwen2,并且每个文档对都存在提示模板开销,因此蒸馏后的学生模型继承了教师模型的校准能力和判断能力,却省去了所有运行时负担。老实说,这是整个发布中最让我满意的单项数据。
有一点比较遗憾:基于 DeBERTa-v2 的 mxbai-rerank-{xsmall,base,large}-v1 系列比表中其他模型慢得多,因为 transformers 中的 DeBERTa-v2 目前既不支持 Flash Attention 2,也不支持 SDPA。70M 的 mxbai-rerank-xsmall-v1 吞吐量为每秒 2636 对,在参数量几乎相同的情况下,只有我们 68M 模型吞吐量的一半左右。这些模型本身完全没有问题,只是无法使用现代注意力内核。
如果你不是在数据中心 GPU 上部署,而是在消费级显卡上自行托管,下面是在 RTX 3090 上进行的同类吞吐量测试。基准测试设置与表 1 相同:使用 bfloat16、每个模型支持的最佳注意力实现,并在能够装入显存的最大批量大小下进行三次测试,取吞吐量中位数。
我们的 17M 模型仍然是表中最快的模型,吞吐量达到每秒 9008 对,实际上还高于它在 H100 上的数据。这说明在如此小的模型规模下,原始计算能力并不是瓶颈,H100 额外的算力并不能转化为性能优势。表格中段的排名略有变化:MiniLM 重排序模型超过了我们的 32M 和 68M 模型,1B 模型也落后于 mxbai-rerank-base-v2(每秒 189 对与 221 对)。我们的 150M 模型仍明显领先另外两个基于 ModernBERT 的 150M 模型;替代教师模型的优势也依然成立,我们的 1B 模型吞吐量是 1.5B mxbai-rerank-large-v2 的 2.7 倍(每秒 189 对与 69 对)。
在 CPU 上,我们无法利用 bf16、Flash Attention 2 或去填充,因此延迟情况简单得多:参数量越大,模型越慢。17M 模型明显快于 ms-marco-MiniLM-L6-v2(每秒 267.4 对与 143.9 对),甚至快于规模更小的 ms-marco-MiniLM-L4-v2(每秒 206.2 对)。正如预期,在无法应用去填充优化后,我们的 150M 模型与另外两个 150M 模型处于同一水平(每秒 14.0 对,而它们分别为 14.5 对和 14.7 对)。如果你的工作负载受 CPU 限制,我们的 17M 和 32M 模型是更实用的选择。
为了说明速度提升来自哪里,下一张表使用相同的基准测试配置,对我们的六个模型分别测试了 fp32+SDPA、bf16+SDPA 和 bf16+FA2。FA2 列又分为两种情况:一种仍然使用填充后的输入(封装模型看到的输入),另一种使用去除填充后的输入(我们的模块化 Transformer 实际使用的输入)。最右侧一列是启用 FA2 时我们的模型默认使用的配置。
表 2:六种已发布模型规格在 max_length=512、使用自然长度 NQ 文档时的精度和注意力消融实验。每个单元格第一行显示每秒处理的文档对数,并在括号中标出相对于 fp32+SDPA 的倍数;第二行显示 GPU 峰值内存。最右侧一列(粗体)是启用 FA2 时我们的模型默认使用的配置。
与 fp32+SDPA 基线相比,使用不带填充的 bf16+FA2 所获得的总体加速,会随着模型规模增大而急剧上升:17M 模型为 1.71 倍,1B 模型则达到 8.26 倍。其中大部分增幅仅来自 bf16:从 fp32+SDPA 切换到 bf16+SDPA 时,17M 模型只获得 1.03 倍加速,而 1B 模型则获得了整整 5.60 倍加速,这也得益于内存开销降低后可以使用更大的批量大小。简而言之,bfloat16 是整体加速中贡献最大的单一因素。
出乎意料的是,在输入仍然包含填充的情况下启用 FA2,在此次发布的所有模型规格上实际上都比 bf16+SDPA 更慢。FA2 内核更适合去除填充后的格式;如果向其传入填充后的输入,你既要承担格式转换的管理开销,又仍然需要在填充 token 上花费计算资源。因此,bf16+FA2 w. padding 这一列大致对应这样的情况:你只是在 model_kwargs 中将 sdpa 替换为 flash_attention_2,却没有对模型加载器进行任何其他改动。表 1 中的 gte-reranker-modernbert-base 和 granite-embedding-reranker-english-r2 就属于这种情况。
最后,从bf16+FA2(带padding)升级到bf16+FA2(不带padding)可获得1.78倍(1B模型)到2.45倍(68M模型)的额外吞吐量提升,同时也能显著降低峰值内存占用,允许更大的批次大小。
因此我的建议很简单:同时启用bf16和FA2。六个Ettin重排器默认将使用不带padding的输入,因为这正是架构详情部分的模块化Transformer模块的配置方式。完整代码片段与上面的使用部分相同:
from sentence_transformers import CrossEncoder
model = CrossEncoder(
"cross-encoder/ettin-reranker-150m-v1",
model_kwargs={
"dtype": "bfloat16",
"attn_implementation": "flash_attention_2", # See tip below
},
)
使用 pip install kernels 安装FA2。它为多种GPU架构、CUDA版本和操作系统提供了预构建内核。
对其他CrossEncoder的一个注意事项:完整的加速只对使用类似Ettin重排器这样的模块化Transformer构建的模型可用。将相同的两个标志应用于通过AutoModelForSequenceClassification加载的CrossEncoder,你会落在表2中较慢的bf16+FA2(带padding)列中。
下面的训练脚本最初来自新的train-sentence-transformers智能体技能的输出,随Sentence Transformers v5.5.0一同发布。如果你使用AI编码智能体(Claude Code、Codex、Cursor、Gemini CLI等),你可以安装该技能并让它在你的数据上微调SentenceTransformer、CrossEncoder或SparseEncoder模型。该技能提供了版本感知的指导,包括基础模型选择、损失函数和评估器选择、硬负样本挖掘、蒸馏、LoRA、Matryoshka、多语言训练和静态嵌入,以及每种模型类型的模板脚本。
hf skills add train-sentence-transformers --claude # symlinks into .claude/skills/
hf skills add train-sentence-transformers --global # under ~/.agents/skills/
一个类似"在我的数据集中的(query, document)对上微调交叉编码器重排器,挖掘硬负样本,并推送到我的Hub仓库"的提示会生成一个可运行的脚本,你可以继续迭代。这就是我开始处理下面配方的方式。
所有六个重排器都使用相同的单阶段配方进行训练。只有学习率和每设备批次大小因模型大小而异。完整的训练脚本约150行,使用一个已发布的数据集。
该配方在单次跨模型大小的扫描后收敛。每个大小的学习率通过在最终训练数据的约15%子集上进行小规模网格搜索进行调整,得到的学习率干净利落地转移到全数据运行中而无需重新调整。除了学习率外,不需要进行每尺寸的调整。
大多数已发布的重排器配方在人工标注的相关性三元组(一个query、一个正文档和可选的硬负样本)上进行训练,使用对比损失、逐点损失、成对损失或列表损失,例如MultipleNegativesRankingLoss、BinaryCrossEntropyLoss、RankNetLoss或LambdaLoss。例如,参见我早先的"使用Sentence Transformers训练和微调重排器模型"博文。
但这种方法有一些实际和理论上的缺陷。首先,正样本需要人工标注,这既昂贵又难以跨许多领域扩展。其次,模型只看到某人处理过的(query, document)对的小子集的标签。特别是在硬负样本挖掘之后,你最终会得到大量的假负样本,例如"硬负样本,艰难的课程"中所示的那样。第三,这种标注的二元性质与现实不符,现实中某些文档只是比其他文档更相关。
我在这里采取了不同的方式:来自现有强大的教师重排器的逐点MSE蒸馏。该设置足够简单,可以用三行描述:
Teacher: mixedbread-ai/mxbai-rerank-large-v2 (1.54B parameters)
Loss: MSELoss on the raw teacher logits (range ~[−12, 22]),即无需重新缩放
Training data: ~143M (query, document, teacher_score) 三元组
我已将训练数据作为单个Hugging Face数据集发布,cross-encoder/ettin-reranker-v1-data,由两个来源组合而成。每个来源都保留为自己的分割,以便来源清晰透明:
LightOn预训练数据(lightonai/embeddings-pre-training,未策划):32个分割,覆盖广泛领域文本相似度信号(MTP、FW-EDU、Reddit、PAQ、S2ORC、Amazon、Wikipedia、MS MARCO等)。我限制了某些分割的样本数量,总共产生约110M个(query, document, similarity)三元组。
来自lightonai/embeddings-fine-tuning的重新评分检索数据:7个分割(msmarco、hotpotqa、trivia、nq、squadv2、fiqa、fever)。源数据集每个query最多有2048个候选文档(最初使用Alibaba-NLP/gte-modernbert-base评分),我使用mixedbread-ai/mxbai-rerank-large-v2进行了重新评分并上传为cross-encoder/lightonai-embeddings-fine-tuning-reranked-v1。该数据集使用Jang et al.分位数锚点配方将每个query的2048个候选降采样至256个(所有正样本+前16个硬样本+约239个分位数锚点分层)。对于训练,我从每个query的那256个中选择64个:32个来自评分排序的头部(正样本加最难的负样本),32个中等难度的负样本从教师排名中较后的波段中采样。详见数据集卡获取确切的排名位置。
总计:约143M个(query, document, score)三元组,加上一个保留的5K行评估分割(来自quora的尾部),驱动训练中的评估损失。
大多数超参数在模型大小间保持恒定:
CrossEncoderTrainingArguments(
num_train_epochs=1, # I chose more data over more epochs
per_device_train_batch_size=..., # global_batch_size // world_size (see table below)
gradient_accumulation_steps=1,
learning_rate=..., # per-size, see table
warmup_ratio=0.03, # ~3% linear warmup, then linear decay (default)
bf16=True, # FA2 + bf16 throughout
eval_strategy="steps",
eval_steps=0.05, # NanoBEIR