使用 Training Operator 进行分布式训练
旧版本
本页面介绍的是 Kubeflow Training Operator V1，最新信息请查看 Kubeflow Trainer V2 文档。
本页面展示了 Training Operator 可以使用的不同分布式策略。
PyTorch 分布式训练
此图展示了 Training Operator 如何为 ring all-reduce 算法创建 PyTorch worker。
您需要负责使用原生的 PyTorch Distributed API 编写训练代码，并使用 Training Operator Python SDK 创建具有所需 worker 数量和 GPU 数量的 PyTorchJob。然后，Training Operator 会创建具有适当环境变量的 Kubernetes pod，以便 torchrun CLI 启动分布式 PyTorch 训练作业。
在 ring all-reduce 算法结束时，每个 worker (g1, g2, g3, g4) 中的梯度会同步，并且模型得到训练。
您可以在训练代码中定义 PyTorch 支持的各种分布式策略（例如 PyTorch FSDP），Training Operator 将为 torchrun 设置适当的环境变量。
TensorFlow 分布式训练
此图展示了 Training Operator 如何为 PS 分布式训练创建 TensorFlow 参数服务器 (PS) 和 worker。
您需要负责使用原生的 TensorFlow Distributed API 编写训练代码，并使用 Training Operator Python SDK 创建具有所需 PS、worker 数量和 GPU 数量的 TFJob。然后，Training Operator 会创建具有适当环境变量的 Kubernetes pod，以便 TF_CONFIG 启动分布式 TensorFlow 训练作业。
参数服务器为每个 worker 分割训练数据，并根据每个 worker 生成的梯度平均模型权重。
您可以在训练代码中定义 TensorFlow 支持的各种分布式策略，Training Operator 将为 TF_CONFIG 设置适当的环境变量。