使用训练运算符进行分布式训练

此页面展示了训练操作符可以使用的不同分布式策略。

PyTorch的分布式训练

该图展示了训练操作符如何为ring all-reduce algorithm创建PyTorch工作者。

您负责使用原生 PyTorch分布式API 编写训练代码，并使用训练操作符Python SDK创建具有所需数量的工作者和GPU的PyTorchJob。 然后，训练操作符创建具有适当环境变量的Kubernetes pods，以便 torchrun CLI启动分布式 PyTorch训练任务。

在环形全归约算法的最后，所有工作节点中的梯度被同步（g1, g2, g3, g4），模型被训练。

您可以在训练代码中定义 PyTorch 支持的各种分布式策略（例如 PyTorch FSDP），训练操作符将为 torchrun 设置适当的环境变量。

分布式训练用于TensorFlow

该图显示了训练操作符如何为PS 分布式训练创建 TensorFlow 参数服务器 (PS) 和工人。

您负责使用原生 TensorFlow 分布式 API 编写训练代码，并使用训练操作符 Python SDK 创建具有所需数量的 PS、工作节点和 GPU 的 TFJob。然后，训练操作符创建具有适当环境变量的 Kubernetes pod，以便 TF_CONFIG 启动分布式 TensorFlow 训练任务。

参数服务器为每个工作节点分割训练数据，并根据每个工作节点产生的梯度平均模型权重。

您可以在训练代码中定义各种TensorFlow支持的分布式策略，训练操作符将为TF_CONFIG设置适当的环境变量。