- Authors
- Name
- Youngju Kim
- @fjvbn20031
- 引言 — GPU 一旦变多，问题就成了软件问题
- 第 1 部分 — AI 库与框架生态地图
- 第 2 部分 — 并行化策略：先问「装不下的是什么」
- 第 3 部分 — torchrun：所有方案共同的分母
- 第 4 部分 — Slurm 使用指南：HPC 之王
- 第 5 部分 — Kubeflow 使用指南：Kubernetes 阵营给出的答案
- 第 6 部分 — 选择标准：Slurm vs Kubeflow vs Ray
- 结语
- 参考资料
引言 — GPU 一旦变多，问题就成了软件问题
单张 GPU 的训练，框架会替你打理好一切。而一旦变成八张 GPU、四台节点，基础设施问题就开始了 — 进程该怎么起、彼此该怎么找到对方、失败了又该怎么重启。本文覆盖这整片地形 — 有哪些库存在(生态地图)、模型该怎么切分(并行化策略)、以及作业该往哪儿扔(Slurm vs Kubeflow vs Ray)。这是把AI 模型开发生命周期一文中「分布式训练栈」一节，扩展到实战级别的续篇。
第 1 部分 — AI 库与框架生态地图
先认识一下登场人物。按层划分，地形就清楚了。
层 代表产品 一句话总结
────────────────── ──────────────────────────── ─────────────────────────────
基础框架 PyTorch 2.x 事实上的标准。用 torch.compile 加速编译
JAX 函数式 + XLA。TPU 与 Google 系研究的中心
模型与数据中枢 HuggingFace Transformers 模型定义的通用语言
Datasets / Tokenizers 数据加载、分词的标准
微调层 PEFT (LoRA), TRL (SFT/DPO) 用较少资源做调优、对齐
Accelerate 让单机与分布式代码保持一致的抽象层
分布式训练引擎 PyTorch DDP/FSDP 框架内置的并行化
DeepSpeed (ZeRO) 内存优化的代名词
Megatron-LM 含 TP/PP，超大规模预训练的标准
编排 Slurm HPC 集群之王
Kubeflow / Ray Kubernetes、Python 原生阵营
服务(参考) vLLM / SGLang / TensorRT-LLM 训练结束后该看这边 — 参见服务篇
选择的直觉是这样的：研究、微调默认用 PyTorch + HF 技术栈，TPU 或大规模并行研究用 JAX，超大规模预训练属于 Megatron 系，而中间地带的所有地方，都能看到 DeepSpeed 和 FSDP 的身影。
第 2 部分 — 并行化策略：先问「装不下的是什么」
并行化的选择，取决于「什么装不进 GPU」。
- DDP(Data Parallel) — 模型装得下，但想训练得更快。把模型在每张 GPU 上复制一份，拆分批次，反向传播后用 all-reduce 同步梯度。这是最简单、扩展效率也最好的方案，只要条件允许，DDP 永远是第一选择。
- ZeRO / FSDP(Fully Sharded Data Parallel) — 模型(准确说是参数 + 梯度 + 优化器状态)装不下。把这三样分片(shard)存放到各张 GPU 上，只在需要的那一刻才收集起来使用。DeepSpeed ZeRO 的各个阶段(1：优化器，2：+梯度，3：+参数)和 PyTorch FSDP，是同一个想法的两种实现。这是 7B~70B 微调的主力。
- TP(Tensor Parallel) — 连一层都装不下，或者延迟很关键。把矩阵乘法本身拆分到多张 GPU 上。通信十分频繁，原则上只在节点内(NVLink)使用。
- PP(Pipeline Parallel) — 把各层分给不同的 GPU 组，像流水线一样依次流过。用微批次(micro-batch)减少流水线气泡。适合跨越节点边界的切分。
- 实战中的超大规模训练会把这些组合成 3D 并行(DP × TP × PP)，如果是 MoE 模型，还会再加上 EP(Expert Parallel)。
便于记忆的公式：装得下就用 DDP → 装不下就上 FSDP/ZeRO → 还是不行就再加 TP、PP。无论哪种组合，混合精度(bf16)和梯度检查点(gradient checkpointing)都是标配。
第 3 部分 — torchrun：所有方案共同的分母
不管是 Slurm 还是 Kubeflow，底层大多是 torchrun(PyTorch 的分布式启动器)在跑。概念只有三个 — 进程总数(world size)、自己的序号(rank)、以及集合的地点(rendezvous)。
# 单节点 8 GPU
torchrun --nproc_per_node=8 train.py
# 多节点(例如：2 节点 × 8GPU = world size 16)
# 所有节点执行相同命令，只有 node_rank 不同
torchrun \
--nnodes=2 --nproc_per_node=8 --node_rank=0 \
--rdzv_backend=c10d --rdzv_endpoint=10.0.0.1:29500 \
train.py
训练代码只需要在调用 torch.distributed.init_process_group() 之后，用 LOCAL_RANK 环境变量抓到自己那张 GPU 就够了。用 HF Accelerate 或 Lightning 的话，这部分接线工作由库替你完成。多节点性能有一个前提条件 — 节点之间必须是 RDMA(InfiniBand/RoCE)，而不是普通以太网，否则 all-reduce 会成为瓶颈；NCCL 则在此之上负责具体的通信。
第 4 部分 — Slurm 使用指南：HPC 之王
Slurm 是在数十年 HPC 历史中打磨出来的批处理调度器，是 GPU 集群训练事实上的标准。版本号采用 年.月 体系(24.11、25.05……)，由 SchedMD 负责维护。四个概念就够用了：节点(机器) — 分区(partition，节点的集合，即队列) — 作业(job，资源申请 + 脚本) — GRES(通用资源，比如 GPU)。
# 快速查看集群状态
sinfo # 分区、节点状态
squeue --me # 我的作业队列
scontrol show node node01 # 节点详情
# 提交作业的三种方式
sbatch job.sh # 批处理提交(标准做法)
srun --pty bash # 交互式 shell(用于调试)
salloc --gres=gpu:2 # 先占用资源，再手动执行
多节点训练的核心在于 sbatch 脚本。这是一个 2 节点 × 8GPU 的 torchrun 示例：
#!/bin/bash
#SBATCH --job-name=llm-train
#SBATCH --partition=gpu
#SBATCH --nodes=2 # 2 个节点
#SBATCH --ntasks-per-node=1 # 每节点 1 个 torchrun(进程由 torchrun 自己生成 8 个)
#SBATCH --gres=gpu:8 # 每节点 8 张 GPU
#SBATCH --cpus-per-task=64
#SBATCH --time=48:00:00
#SBATCH --output=logs/%x-%j.out # %x=作业名 %j=作业 ID
# 集合点地址 = 第一个节点
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=29500
# 在每个节点上运行 torchrun — 用 SLURM 变量自动完成接线
srun torchrun \
--nnodes=$SLURM_JOB_NUM_NODES \
--nproc_per_node=8 \
--node_rank=$SLURM_NODEID \
--rdzv_backend=c10d \
--rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
train.py --config config.yaml
解读方式：srun 会在分配给该作业的 每个节点上各运行一次 torchrun，SLURM_NODEID 提供节点序号，scontrol show hostnames 自动填入主节点地址。光是这一种模式，就占了 Slurm 多节点训练的 80%。剩下的 20% 是运维技巧 — 检查点要写得比 --time 限制更频繁(以应对抢占和超时)，用 --signal=SIGUSR1@120 在结束前 2 分钟收到信号、留下最后一个检查点，再用数组作业(--array)跑超参数扫描。
第 5 部分 — Kubeflow 使用指南：Kubernetes 阵营给出的答案
如果你的 GPU 集群本来就是 Kubernetes(也就是用 GPU Operator 搭建的那个集群)，那把训练作业也当作 Kubernetes 资源来提交，就是很自然的事。Kubeflow 正是这个阵营的综合平台。
组件地图 — 不需要全部都用上：
组件 角色 是否必须
────────────────── ───────────────────────────── ─────────────────
Training Operator 分布式训练作业 CRD(核心！) ★ 光靠这个就能起步
Pipelines ML 工作流 DAG 编排 做流水线化时需要
Katib 超参数调优、AutoML 需要时再用
Notebooks 集群内的 Jupyter 环境 便利功能
KServe 模型服务(已独立成单独项目) 服务阶段需要
核心是 Training Operator。只需一个 PyTorchJob CRD，就能声明出多节点训练：
apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
name: llm-train
spec:
nprocPerNode: "8" # 每节点(pod)进程数 = GPU 数量
pytorchReplicaSpecs:
Master:
replicas: 1
template:
spec:
containers:
- name: pytorch
image: my-registry/train:latest
command: ["torchrun", "train.py"]
resources:
limits:
nvidia.com/gpu: 8
Worker:
replicas: 3 # 主节点 1 + 工作节点 3 = 4 个节点
template:
spec:
containers:
- name: pytorch
image: my-registry/train:latest
command: ["torchrun", "train.py"]
resources:
limits:
nvidia.com/gpu: 8
Training Operator 会创建这些 pod，并 自动注入 MASTER_ADDR、WORLD_SIZE、RANK 环境变量，于是在 Slurm 上要手动完成的接线工作就消失了。作业状态用 kubectl get pytorchjobs 查看。顺带一提，下一代的 Kubeflow Trainer(v2) 正朝着把各框架专属的 CRD 统一成单一 TrainJob 的方向演进 — 如果是新接入，建议先在文档里确认 v2 的支持状态再开始。Gang Scheduling(即所有节点必须同时到位才能启动)则通过接入 Volcano、Kueue 这类调度器来解决。
第 6 部分 — 选择标准：Slurm vs Kubeflow vs Ray
情况 推荐
───────────────────────────────────────────── ─────────────────────────
有 HPC 传统的组织、裸机 GPU 农场 Slurm
已经全靠 Kubernetes 运营一切 Kubeflow (Training Operator)
想在 Python 代码里灵活整合分布式训练与服务 Ray (Train/Serve)
研究室小规模、2~4 个节点 用普通的 torchrun + pdsh/tmux 就够了
偏好云端托管服务 SageMaker/Vertex 等 (概念相同)
只留下一个本质的话 — 不管哪个平台，底层都是一样的：torchrun(或其等价物)负责启动进程，NCCL 负责 GPU 间通信，检查点负责扛住故障。平台之间的差异，不过是「谁来出借节点、谁来填好环境变量」而已。所以只要在一个平台上把多节点这件事真正搞懂了，剩下的就只是一个翻译问题。
结语
把 GPU 从一张变成多张、多节点时会遇到的决策依序排开：并行化策略(装得下用 DDP，装不下用 FSDP/ZeRO，还是不行就再加 TP、PP)→ 启动器(torchrun)→ 编排器(Slurm 或 Kubeflow)→ 运维(检查点、RDMA、Gang Scheduling)。等训练在这套技术栈上跑起来之后，下一个瓶颈就是数据了 — 这个话题会在 LLM 训练数据预处理篇 中继续。