分布式训练¶

Sentence Transformers 实现了两种形式的分布式训练：数据并行（DP）和分布式数据并行（DDP）。阅读 Hugging Face 上的数据并行性文档以获取有关这些策略的更多详细信息。其中一些关键区别包括：

DDP 通常比 DP 更快，因为它需要传输的数据更少。
使用DP时，GPU 0 承担了大部分工作，而使用DDP时，工作则更均匀地分布在所有GPU上。
DDP 允许在多台机器上进行训练，而 DP 仅限于单台机器。

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

Via torchrun

torchrun 文档

torchrun --nproc_per_node=4 train_script.py

Via 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 实例，即 4x 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 仅在大模型中才有意义。如果你想在 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 文档以获取更多详情。