分布式训练

Sentence Transformers 实现了两种形式的分布式训练：数据并行（Data Parallel, DP）和分布式数据并行（Distributed Data Parallel, DDP）。请阅读 Hugging Face 上的数据并行文档，了解这些策略的更多细节。一些关键区别包括：

DDP 通常比 DP 更快，因为它需要通信的数据更少。
使用 DP 时，GPU 0 会承担大部分工作；而使用 DDP 时，工作会更均匀地分布到所有 GPU 上。
DDP 支持跨多台机器进行训练，而 DP 仅限于单台机器。

简而言之，通常推荐使用 DDP。您可以通过 torchrun 或 accelerate 运行您的常规训练脚本来使用 DDP。例如，如果您有一个名为 train_script.py 的脚本，您可以使用以下命令通过 DDP 运行它：

通过 torchrun

torchrun 文档

torchrun --nproc_per_node=4 train_script.py

通过 accelerate

accelerate 文档

accelerate launch --num_processes 4 train_script.py

注意

在进行分布式训练时，您必须将代码包装在 main 函数中，并使用 if __name__ == "__main__": 来调用它。这是因为每个进程都会运行整个脚本，所以您不希望多次运行相同的代码。以下是一个示例：

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainingArguments, SentenceTransformerTrainer
# Other imports here

def main():
    # Your training code here

if __name__ == "__main__":
    main()

注意

当使用评估器 (Evaluator) 时，评估器仅在第一个设备上运行，而训练和评估数据集则会在所有设备间共享。

比较

下表展示了在给定的硬件配置下，DDP 相对于 DP 和无并行化的速度提升。

硬件：一个 p3.8xlarge AWS 实例，即 4 块 V100 GPU
训练模型：microsoft/mpnet-base (1.33 亿参数)
最大序列长度：384 (参考 all-mpnet-base-v2)
训练数据集：MultiNLI、SNLI 和 STSB（注意：这些数据集的文本较短）
损失函数：MultiNLI 和 SNLI 使用 SoftmaxLoss，STSB 使用 CosineSimilarityLoss
每设备批处理大小 (Batch size)：32

策略	启动器	每秒样本数
无并行化	`CUDA_VISIBLE_DEVICES=0 python train_script.py`	2724
数据并行 (DP)	`python train_script.py` (当使用 `python` 启动脚本时，默认使用 DP)	3675 (1.349 倍加速)
分布式数据并行 (DDP)	`torchrun --nproc_per_node=4 train_script.py` 或 `accelerate launch --num_processes 4 train_script.py`	6980 (2.562 倍加速)

FSDP

完全分片数据并行（Fully Sharded Data Parallelism, FSDP）是另一种分布式训练策略，但 Sentence Transformers 并未完全支持它。它是 DDP 的一个更高级版本，对非常大的模型特别有用。请注意，在之前的比较中，FSDP 每秒处理 5782 个样本（2.122 倍加速），即比 DDP 差。FSDP 仅在模型非常大时才有意义。如果您想在 Sentence Transformers 中使用 FSDP，需要注意以下限制：

您不能将 evaluator 功能与 FSDP 一起使用。
您必须先使用 trainer.accelerator.state.fsdp_plugin.set_state_dict_type("FULL_STATE_DICT")，然后再用 trainer.save_model("output") 来保存训练好的模型。
您必须在 SentenceTransformerTrainingArguments 中使用 fsdp=["full_shard", "auto_wrap"] 和 fsdp_config={"transformer_layer_cls_to_wrap": "BertLayer"}，其中 BertLayer 是编码器中包含多头注意力和前馈层的重复层，例如 BertLayer 或 MPNetLayer。

请阅读 Accelerate 的 FSDP 文档以获取更多细节。