分布式训练

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()

注意

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

比较

下表显示了在给定特定硬件设置下，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
每个设备的批次大小：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，您必须注意以下限制：

您不能在 FSDP 中使用 evaluator 功能。
您必须使用 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 文档。