Skip to content

필사 모드: 多 GPU、多节点训练平台全面梳理 — 从框架生态地图到 Slurm、Kubeflow 实战指南

中文
0%
정확도 0%
💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

引言 — 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_ADDRWORLD_SIZERANK 环境变量,于是在 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 训练数据预处理篇 中继续。

参考资料

현재 단락 (1/105)

单张 GPU 的训练,框架会替你打理好一切。而一旦变成八张 GPU、四台节点,基础设施问题就开始了 — **进程该怎么起、彼此该怎么找到对方、失败了又该怎么重启**。本文覆盖这整片地形 — 有哪些库存...

작성 글자: 0원문 글자: 6,020작성 단락: 0/105