分布式训练

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 实例，即 4x V100 GPU
正在训练的模型：microsoft/mpnet-base（1.33 亿参数）
最大序列长度：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 文档。