分布式训练

Sentence Transformers 实现了两种形式的分布式训练：数据并行 (DP) 和分布式数据并行 (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()

注意

当使用评估器时，评估器仅在第一个设备上运行，这与在所有设备之间共享的训练和评估数据集不同。

比较

下表显示了在给定特定硬件设置的情况下，DDP 相对于 DP 和无并行化的加速效果。

硬件：一个 p3.8xlarge AWS 实例，即 4 个 V100 GPU
正在训练的模型：microsoft/mpnet-base (133M 参数)
最大序列长度：384（遵循 all-mpnet-base-v2）
训练数据集：MultiNLI、SNLI 和 STSB（注意：这些数据集包含短文本）
损失：MultiNLI 和 SNLI 使用 SoftmaxLoss，STSB 使用 CosineSimilarityLoss
每个设备的批次大小：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

完全分片数据并行 (FSDP) 是另一种分布式训练策略，Sentence Transformers 尚未完全支持。它是 DDP 的更高级版本，特别适用于非常大的模型。请注意，在之前的比较中，FSDP 达到了每秒 5782 个样本（2.122 倍加速），即比 DDP 更差。FSDP 仅对非常大的模型有意义。如果您想将 FSDP 与 Sentence Transformers 一起使用，则必须注意以下限制

您不能将 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 文档以获取更多详细信息。