sentence_transformers

2025年6月6日

16:18

    分布式训练

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

    1. DDP 通常比 DP 更快，因为它需要通信的数据更少。
    2. 使用 DP 时，GPU 0 完成大部分工作，而使用 DDP 时，工作更均匀地分布在所有 GPU 上。
    3. DDP 允许跨多台机器进行训练，而 DP 仅限于单台机器。

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

    通过 torchrun

    torchrun --nproc_per_node=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()

     

 

已使用 OneNote 创建。