train-sentence-transformers - hardware

硬件指南

训练嵌入模型通常是内存受限而不是计算受限。

如果遇到 OOM

按此顺序尝试：

减小per_device_train_batch_size。提高gradient_accumulation_steps以保持回归损失的有效批次大小。（对于 MNRL，通过梯度累积得到的有效批次不等价——见第 3 点。）
启用gradient_checkpointing=True。慢约 30%，激活内存减少约 40%。与Cached*损失不兼容。
切换到Cached*损失：
- CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)— 以小批量前向传播，在整个批次上累积对比损失。可以在 24GB GPU 上模拟 1024+ 的批次大小。
- CachedSpladeLoss(model, loss=..., mini_batch_size=16)— 对稀疏使用同样的技巧。
- CachedGISTEmbedLoss(model, guide_model, mini_batch_size=32)— GIST 变体。
对 >1B 的解码器模型启用 PEFT / LoRA。LoraConfig(r=64, lora_alpha=128, task_type="FEATURE_EXTRACTION")。参见../scripts/train_sentence_transformer_with_lora_example.py（docstring 涵盖何时使用、超参数、QLoRA、共享）。
转向多 GPU。见下文。
缩短序列。如果截断到 128 对你的任务已经足够，请在 transformer 模块上设置max_seq_length。

多 GPU

sentence-transformers在底层使用accelerate。分布式训练无需修改代码即可工作。

数据并行（DDP）

启动：

accelerate launch train.py# 或显式：accelerate launch--multi_gpu--num_processes=4train.py

per_device_train_batch_size保持为每 GPU 的值。有效批次大小线性扩展。MNRL 的批内负样本保持为每设备，而非全局，除非你向支持它的损失函数传入gather_across_devices=True（MultipleNegativesRankingLoss、CachedMultipleNegativesRankingLoss、对称变体、GISTEmbedLoss、CachedGISTEmbedLoss、SparseMultipleNegativesRankingLoss）。

FSDP / DeepSpeed

对于 >3B 的模型，使用accelerate config启用 FSDP 或 DeepSpeed ZeRO。两者都受支持——sentence-transformers不需要任何代码更改，只需要启动配置。

accelerate config# 交互式；选择 FSDP 或 DeepSpeedaccelerate launch train.py

使用 FSDP 全分片：7B 模型可以在 4×24GB GPU 上训练，而其中任何单个 GPU 都会 OOM。

FSDP 注意事项（来自分布式训练文档）：

截至编写时，评估器无法在 FSDP 下运行—— 评估钩子调用model.encode()，而 FSDP 包装的模块在训练中途无法处理它。改为计划在训练后以单 GPU 加载最终检查点进行评估，或者如果你需要训练中评估，使用 DDP 训练。
必须指定层包装，例如fsdp_config={"transformer_layer_cls_to_wrap": "BertLayer"}（为你的模型替换正确的层类：BertLayer、LlamaDecoderLayer、Qwen2DecoderLayer等）。不这样做，FSDP 分片可能静默出错。
对适合单 GPU 的模型比 DDP 慢—— 只有在你确实需要内存节省时才使用 FSDP。

DeepSpeed ZeRO-2/3 是一个替代方案，有自己的配置；在accelerate config层面工作方式相同。

对比损失的有效批次大小

对于MultipleNegativesRankingLoss及其变体，批次大小是质量旋钮，而不仅仅是速度旋钮。更大的批次 = 更多的批内负样本 = 更丰富的梯度。

每个锚点的有效批内负样本池：

设置	每个锚点的批内负样本数
单 GPU，批次 64	63
4× DDP，每设备批次 64	默认仅本地 63；使用`MultipleNegativesRankingLoss(model, gather_across_devices=True)`为 255
单 GPU，CachedMNRL，mini_batch 32，批次 256	255
4× DDP，CachedMNRL，每设备 256	本地 255；使用`gather_across_devices=True`为 1023

对于大型语料库（检索），朝 512+ 有效负样本推进。对于小而干净的数据集（STS），64 就足够了。

按 GPU 的精度选择

GPU 代	推荐
T4、V100、GTX 1xxx、RTX 2xxx	`fp16=True`
RTX 3xxx、A10G、A100、L4	`bf16=True`
RTX 4xxx、H100、B200	`bf16=True`（或通过特定内核在 H100 上用 fp8——不是默认）
Apple M 系列 / ROCm	MPS/ROCm 支持不稳定；`fp16`或`fp32`最可靠

bf16 数值上更稳定，在可用时几乎总是首选。

Hugging Face Jobs 类型指南

Hugging Face Jobs 需要 Pro/Team/Enterprise 计划。价格是近似值，可能变动——参见 Jobs 定价页面。

类型	内存	典型用途	估计 $/小时
`cpu-basic`	~2 GB	数据准备、验证、困难负样本挖掘（小）	<$0.10
`cpu-upgrade`	~4 GB	同上，稍大	$0.10
`t4-small`	16 GB	演示、小批次的 MiniLM/DistilBERT	~$0.75
`t4-medium`	16 GB	较大批次的 MiniLM / DistilBERT	~$1.50
`l4x1`	24 GB	BERT-base、MPNet、ModernBERT-base	~$2.50
`a10g-small`	24 GB	BERT-base 到 BERT-large	~$3.50
`a10g-large`	48 GB	ModernBERT-large、Qwen3-0.6B	~$5.00
`a10g-largex2`	96 GB（2× 48GB）	中型多 GPU	~$10
`a100-large`	80 GB	大模型或大对比批次	~$10–12
`h100`	80 GB	最大的单 GPU	~$12
`h100x8`	640 GB	LLM 规模分布式	~$96