深度学习分布式训练任务模板

需要在 cube-studio pipeline 上发起 TensorFlow / PyTorch / Paddle / DeepSpeed / Horovod / Megatron / ColossalAI / MindSpore / MindFormers / MXNet 等分布式训练时

05-任务模板 / 深度学习

深度学习分布式训练tensorflowtftfjobpytorchpytorchjobddpkubeflowtraining-operatordeep-learning-frameworkpaddlepaddlejobdeepspeedhorovodmpi

深度学习分布式训练任务模板

本篇覆盖三组深度学习分布式训练模板：

tf：Kubeflow TFJob 分布式（TensorFlow）。
pytorch：Kubeflow PyTorchJob 分布式。
deep-learning-framework：一个统一镜像内含多框架 launcher，子框架包括 paddle / deepspeed / horovod / megatron / colossalai / mindformers / mindspore / mpi / mxnet。

这些模板都遵循同一套思路：Launcher 容器只负责按平台下发的环境变量和启动参数组装一个分布式 CRD（TFJob / PyTorchJob / PaddleJob / MPIJob / MXJob / Volcano Job）并提交、用 stern 汇聚日志、轮询状态决定退出码，不包含你的训练逻辑。真正的训练代码跑在你指定的 worker 镜像（--image / --work_images）里。

各分布式模板均依赖对应的 Operator/CRD（TFJob/PyTorchJob/PaddleJob/MPIJob/MXJob 需 Kubeflow training-operator，Volcano Job 需 Volcano 调度器）。集群、GPU、Operator 的准备见 02-部署安装 段。

tf

源码：job-template/job/tf/launcher.py。创建 Kubeflow TFJob（kubeflow.org/v1），发起 TensorFlow 多 Worker 分布式训练。

镜像（launcher 镜像）：ccr.ccs.tencentyun.com/cube-studio/tf:20230801（build.sh:5）。
启动命令：python3 launcher.py。K8s 账号：kubeflow-pipeline。
角色与副本：仅 Worker 角色（无 PS/Chief），副本数 = --num_worker，适配 TF 2.x 的 tf.distribute 多 Worker。容器名固定为 tensorflow（TFJob 规范要求，见 launcher.py:211）。
策略：cleanPodPolicy: None（不自动删 Pod 便于排查）、backoffLimit = num_worker（launcher.py:282-283）。
结束：TFJob 状态为 Succeeded 退出码 0，否则 1。

启动参数（核对 launcher.py:336-339）：

参数	默认值	说明
`--image`	空	worker 镜像；为空时用环境变量 `KFJ_TASK_IMAGES`
`--working_dir`	`/mnt/`	容器内工作目录
`--command`	`python3 mnist.py`	实际执行命令（`bash -c` 执行）
`--num_worker`	`3`	Worker 副本数

TFJob 名称形如 tfjob-<pipeline-name>-xxxx。示例训练代码可参考 Kubeflow 官方 mnist_with_summaries。

pytorch

源码：job-template/job/pytorch/launcher.py（含分布式示例 demo.py）。创建 Kubeflow PyTorchJob（kubeflow.org/v1，plural: pytorchjobs）。

镜像（launcher 镜像）：ccr.ccs.tencentyun.com/cube-studio/pytorch:20250801（build.sh:5）。
角色与副本：Master + Worker 两个角色（launcher.py:332-333）；Worker 副本数 = --num_worker - 1（因为 Master 也算一个 worker，见 launcher.py:307）。容器名为 pytorch，PyTorchJob 名称形如 pytorchjob-<pipeline-name>-xxxx。
策略：cleanPodPolicy: None、backoffLimit = num_worker。
启动参数：与 tf 一致 —— --image / --working_dir(默认/mnt/) / --command(默认python3 mnist.py) / --num_worker(默认3)，见 launcher.py:389-392。

分布式代码写法（你的训练脚本需自己改造）

每个 Pod 注入如下环境变量（MASTER_PORT="23456"）：

NCCL_DEBUG=INFO
MASTER_PORT=23456
NCCL_SOCKET_IFNAME=eth0
MASTER_ADDR=pytorchjob-xxx-master-0
WORLD_SIZE=3       # worker 总数
RANK=0            # 当前 worker 序号；master=0, worker-0=1, worker-1=2 ...

要点：

每个进程 rank 唯一，master 必须 rank=0；WORLD_SIZE>1 时调 dist.init_process_group(backend=..., init_method=None)（系统自动识别 RANK/WORLD_SIZE）。
用 DistributedDataParallel 封装 model；用 DistributedSampler 配合 DataLoader，每 epoch train_sampler.set_epoch(epoch)。
启动可直接 python your_start.py，或 python -m torch.distributed.launch --nproc_per_node=<每worker卡数> --nnodes=$WORLD_SIZE --node_rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT train.py ...。

GPU 利用率优化（要点）

GPU 整卡占用时利用率低，核心原因是「CPU 操作慢阻塞 GPU 计算」（数据加载/预处理/频繁 IO/日志上报等）。常见优化：

数据加载：数据导入集群存储、合并小文件、提高 num_workers、配置 prefetch_factor、pin_memory=true、增大 batch_size；
数据预处理与训练拆成两个 task，或用 Nvidia DALI 在 GPU 上做预处理；
减少 checkpoint / tensorboard / 日志的保存频率。
多进程共享单卡：用 start.sh 或 start.py 拉起多个进程，每个进程加 --process_index xx --process_num xx，后台运行并 wait。

监控可在平台「监控」按钮查看资源使用率，或在容器内 watch nvidia-smi / gpustat。

deep-learning-framework

源码目录：job-template/job/deep-learning-framework/，内含 launcher-<框架>.py 与子目录。多个子框架共享同一个 launcher 镜像：

统一 launcher 镜像：ccr.ccs.tencentyun.com/cube-studio/deep-learning-framework:20260419（多架构 amd64+arm64，见 build.sh）。镜像内含 stern、kubernetes==28.1.0 等，没有默认 ENTRYPOINT，因此每个子框架靠任务模板的「启动命令」指定，如 python3 launcher-paddle.py、python3 launcher-deepspeed.py 等。

旧版各子框架 README 里标注的 paddle:20230801、megatron:20230801、mxnet:20221010、mindspore:20230801、mindformers:20250101、deepspeed:20230801 等，是历史上各自独立的 launcher 镜像；当前已统一为上面的 deep-learning-framework:20260419。worker 镜像（--image/--work_images）则各框架不同，见各小节。
共用机制（Volcano 类 launcher：deepspeed/megatron/colossalai/mindspore；Kubeflow CRD 类：paddle/horovod/mpi/mxnet/mindformers）：
- 从环境变量读取资源与调度：KFJ_TASK_RESOURCE_CPU/MEMORY/GPU、GPU_RESOURCE_NAME、KFJ_TASK_NODE_SELECTOR、KFJ_TASK_VOLUME_MOUNT（解析出 volumes/volumeMounts）、KFJ_TASK_RESOURCE_RDMA/RDMA_RESOURCE_NAME、KFJ_TASK_IMAGES、HUBSECRET、IMAGE_PULL_POLICY(默认IfNotPresent) 等；
- 容器内注入 NCCL_DEBUG=INFO、GPU_NUM；无 GPU 时设 NVIDIA_VISIBLE_DEVICES=none；有 RDMA 时申请 RDMA 资源并加 IPC_LOCK capability；
- 节点亲和按 KFJ_TASK_NODE_SELECTOR（如 cpu=true,train=true），Pod 反亲和尽量把同 Job 的 Worker 分散到不同节点；
- 用 stern 实时汇聚日志、后台线程轮询状态，结束按状态返回退出码（成功 0 / 失败 1）。

说明：任务家族清单里的 paddle 不是独立顶层目录，实际是 deep-learning-framework/README-paddle.md + launcher-paddle.py（PaddleJob CRD），在本文档以 paddle 子节呈现。

paddle

launcher-paddle.py 创建 Kubeflow PaddleJob（kubeflow.org/v1）。

角色：Worker（副本 = --num_worker）+ 可选 Master/PS（--num_ps > 0 时，副本 = --num_ps，与 Worker 同镜像/目录/命令）。
弹性策略：elasticPolicy.minReplicas = maxReplicas = num_worker，maxRestarts = 0（固定副本、不自动重启，见 launcher-paddle.py:299-302）。
启动参数（launcher-paddle.py:370-374）：--image、--working_dir(默认/mnt/)、--command(默认python3 mnist.py)、--num_ps(默认0)、--num_worker(默认3)。
worker 镜像：可用官方 registry.baidubce.com/paddlepaddle/paddle:2.4.0rc0-gpu-cuda11.2-cudnn8.1-trt8.0，启动命令如 python -m paddle.distributed.launch train.py。PaddleJob 名形如 paddlejob-<pipeline-name>-xxxx，状态 Succeeded/Completed 视为成功。

deepspeed

launcher-deepspeed.py 创建 Volcano Job（batch.volcano.sh/v1alpha1，schedulerName: volcano），一组相同规格的 Worker Pod，名称形如 {job}-worker-0、{job}-worker-1…。

启动参数（launcher-deepspeed.py:382-385）：--working_dir(默认/mnt/)、--command(默认python3 mnist.py)、--num_worker(默认3)、--image(空则用 KFJ_TASK_IMAGES)。
launcher 跟踪 master Pod（{job}-worker-0）日志并等待结束；非 Succeeded 则非 0 退出。
多机时 worker 可读 VC_WORKER_HOSTS、VC_TASK_INDEX 等做节点发现；若用基于 SSH 的 DeepSpeed launcher，镜像内需装 openssh-server 并在 --command 中先起 sshd 再起 deepspeed，hostfile 自行生成。

horovod

launcher-horovod.py 创建 Kubeflow MPIJob（kubeflow.org/v1），1 个 Launcher + N 个 Worker，Launcher 用 mpirun -np <num_worker> 在 --working_dir 下执行 --command。

启动参数（launcher-horovod.py:416-419）：--num_worker(默认2)、--command、--work_images(默认 ccr.ccs.tencentyun.com/cube-studio/horovod:20210401)、--working_dir(默认/mnt/admin)。
- 注意 horovod 用的是 --work_images（不是 --image）。
worker 镜像里需装好深度学习框架 + Horovod/MPI；只要在镜像里 --command 能正常跑单机训练，多机多卡封装由模板和 Kubeflow 接管。

megatron

launcher-megatron.py 创建 Volcano Job（batch.volcano.sh/v1alpha1）。

启动参数（核对 launcher-megatron.py:379-382）：--working_dir(默认/mnt/)、--command(默认python3 mnist.py)、--cluster(默认worker=3)、--image。
⚠️ 注意：megatron 使用 --cluster 而非 --num_worker（与旧 README 里写的 --num_worker JSON 不符，已按代码 launcher-megatron.py:381 改正）。--cluster 是按行的角色规格，每行 角色=数量 或 角色:数量（用 : 或 = 分隔，见 launcher-megatron.py:284-298），可定义多个角色，如：
```
master=1
worker=3
```
每个角色生成一个 Volcano task，容器内注入 VC_TASK_ROLE。
worker 侧可用 VC_WORKER_NUM（总 worker 数）、VC_TASK_INDEX（当前序号）做分工。

colossalai

launcher-colossalai.py 创建 Volcano Job（batch.volcano.sh/v1alpha1，schedulerName: volcano，minAvailable = num_workers 保证一次性调齐）。Job 名形如 colossalai-<pipeline_name>-<4hex>。

启动参数（核对 launcher-colossalai.py:375-378）：--working_dir(默认/mnt/admin/pipeline/example/colossalai/)、--command(默认bash start.sh)、--num_worker(默认2)、--work_images(默认 hpcaitech/colossalai:0.3.0)。

⚠️ 注意：colossalai 镜像参数是 --work_images，不是 --image（旧 README 表格里写的是 --image，已按代码 launcher-colossalai.py:378 改正）。
监控：后台线程每 60s 检查 master Pod（<job>-worker-0）状态，stern 跟踪日志（每轮约 10s，约 1 小时重启一次），master 为 Succeeded 则 exit(0) 否则 exit(1)。

mindspore

launcher-mindspore.py 创建 Volcano Job（batch.volcano.sh/v1alpha1，schedulerName: volcano，queue: default），Job 名形如 msjob-<pipeline-name>-<short-uuid>。

启动参数（launcher-mindspore.py:376-379）：--working_dir(默认/mnt/)、--command(默认python3 mnist.py)、--num_worker(默认3)、--image(空则用 KFJ_TASK_IMAGES)。
状态为 Completed/Failed(或 aborted/terminated) 时退出；Failed 则非 0 退出。worker 数需与 MindSpore 的 device_num/rank_size 等配置一致。

mindformers

launcher-mindformers.py 用于 MindFormers 分布式训练。

⚠️ 重要更正：旧 README-mindformers.md 称其「launcher 位于 job/mindspore/launcher.py，创建 Volcano Job，作业名 msjob-」，但当前实际代码 launcher-mindformers.py 创建的是 Kubeflow PyTorchJob（kubeflow.org/v1，plural: pytorchjobs，见 launcher-mindformers.py:48-49），作业名前缀为 mindjob-（launcher-mindformers.py:77），角色为 Master + Worker（:333-334），并非 Volcano Job。本节按代码为准。

启动参数（launcher-mindformers.py:390-393）：--working_dir(默认/mnt/)、--command(默认python3 mnist.py)、--num_worker(默认3)、--image。
worker 命令会被追加 <WORKER_NUM> <LOCAL_WORKER> $MASTER_ADDR $MASTER_PORT $RANK 等参数（见 launcher-mindformers.py:255）。worker 数需与 MindFormers 的 device_num/rank_size 一致。

mpi

launcher-mpi.py 创建 Kubeflow MPIJob（kubeflow.org/v1，plural: mpijobs），实现 Horovod/MPI 等分布式。

角色：Launcher 1 个（通过 mpirun/horovodrun 启动训练）+ Worker --num_worker 个（容器初始长时间 sleep，由 Launcher 经 kubectl exec 下发训练命令）。
启动参数（launcher-mpi.py:402-405）：--num_worker(默认2)、--command(默认python /examples/tensorflow2_mnist.py)、--work_images(默认 horovod/horovod:0.20.0-tf2.3.0-torch1.6.0-mxnet1.5.0-py3.7-cpu)、--working_dir(默认/mnt/admin)。同样用 --work_images。
Launcher Pod 起来后最长等 1 小时拉镜像，再实时跟踪 Launcher 容器日志；结束读 MPIJob conditions，最近一次 type == "Succeeded" 退出码 0，否则 1。命令里不要写死 Node IP/Pod 名，集群拓扑由 Launcher 发现注入。

mxnet

launcher-mxnet.py 创建 Kubeflow MXJob（kubeflow.org/v1），支持 Scheduler / Server / Worker 三类角色。

launcher 镜像（历史）：ccr.ccs.tencentyun.com/cube-studio/mxnet:20221010（现统一为 deep-learning-framework:20260419）。
启动参数（launcher-mxnet.py:391-395）：--working_dir(默认/mnt/)、--command(默认python3 mnist.py)、--num_ps(默认0，对应 Server 角色 replicas)、--num_worker(默认3，对应 Worker replicas)、--image。
所有角色用同一镜像与命令（bash -c <--command>）。MXJob 名形如 mxjob-<pipeline-name>-xxxx，状态 Succeeded 视为成功。

各模板的 CRD 字段、资源组装与默认值以对应 launcher-*.py 源码为准，本篇仅供导航与检索。

最后更新 2026-06-30完整文档以官方仓库为准：GitHub Wiki