K8s+Megatron-LM+vLLM打通超大模型训练全链路

摘要

Megatron-LM 是 NVIDIA 开源的大模型训练库，核心是张量并行（TP）、流水线并行（PP）等模型并行技术------和前两篇讲的 Ray、DeepSpeed 不同，它解决的不是"怎么调度"或"怎么省显存"，而是"单个模型本身大到一张卡、甚至一台机器都装不下"这个更底层的问题。这篇文章先讲清楚 Megatron-LM 怎么装、怎么用，再把它和上一篇 DeepSpeed 的 ZeRO 方案做正面对比，最后同样给出一条能在 K8s 上跑通的企业级链路：Kubeflow Trainer 编排张量并行预训练 → Megatron Bridge 转换 checkpoint → vLLM 部署上线，命令和 API 都对照 Megatron-LM、Kubeflow Trainer、Megatron Bridge 官方文档核实过。

背景与问题

Megatron-LM 解决的是"模型本身切不开就装不下"这个问题。 上一篇讲 DeepSpeed 的 ZeRO 时提到过，数据并行训练里每张卡都要存一份完整的模型参数、梯度和优化器状态------ZeRO 的做法是把这些状态在逻辑上"去冗余"，按 GPU 数切分存储，但模型的计算路径本身没变，每张卡跑的还是完整的前向/反向。Megatron-LM 走的是另一条路：把模型的计算本身切开分布到多张卡上------张量并行把一层内部的矩阵乘法按行/列切给不同 GPU 各算一部分，流水线并行把不同的 Transformer 层分给不同 GPU 顺序执行。两者不是谁比谁更先进，而是切的对象不一样：ZeRO 切"状态"，Megatron 切"计算"。

围绕它有两个容易搞混的点：

第一个误区：把 Megatron-LM 当成一个具体的模型架构。 名字里带"Megatron"容易让人以为它是某个特定模型（类似 GPT、LLaMA），但它实际上是一套训练系统：底层的 Megatron-Core 提供并行感知的可组合组件（GPTModel、TransformerConfig 等），上层的 Megatron-LM 仓库是 Megatron-Core 加一套参考训练脚本的完整实现，配套的 Megatron Bridge 负责和 HuggingFace 生态做双向 checkpoint 转换。你完全可以用 Megatron-Core 训练 Qwen、LLaMA 或任何 Transformer 架构，"Megatron"指的是并行训练的方法论，不是某个具体模型。

第二个误区：认为 Megatron-LM 和 DeepSpeed 是互斥的二选一。 二者历史上其实被组合使用过------训练 5300 亿参数的 Megatron-Turing NLG 时，就是用 Megatron-LM 的张量并行在节点内切分模型，用 DeepSpeed 的流水线并行和 ZeRO-1 在跨节点场景做进一步优化，这也是"Megatron-DeepSpeed"这个组合仓库存在的原因。但组合方案会同时带上两边的复杂度，本文和上一篇一样只走一条干净的单一路径，方便直接对比两种编排哲学的边界，而不是教怎么把两者揉在一起。

核心思路与优势

四层职责定位（对比前两篇）

层级	组件	具体职责
物理资源调度	Kubernetes	把 Pod 调度到具体节点，管理 GPU 配额、网络、存储卷
训练任务编排	Kubeflow Trainer（torch-distributed 运行时）	Megatron-Core 原生用 torchrun 做分布式启动器，不需要专门的 Megatron 运行时，直接复用已有的 torch-distributed ClusterTrainingRuntime
分布式训练优化	Megatron-Core	张量并行（TP）、流水线并行（PP，含 interleaved 调度减少 pipeline bubble）、序列并行（SP）、上下文并行（CP，长序列场景）
推理引擎	vLLM	PagedAttention + continuous batching 做高吞吐推理，提供 OpenAI 兼容 API

三篇文章走到这里，编排层出现了三种不同的选择：Ray 是应用层调度器，负责把 Actor 动态分配到资源上；DeepSpeed 走 MPI 范式，靠 Kubeflow Trainer 的 deepspeed-distributed 运行时生成 hostfile 拉起 mpirun；Megatron-Core 则更"轻"------它本身就是标准的 PyTorch 分布式训练（torchrun + torch.distributed），不需要 MPI，也不需要专门的运行时适配层，torch-distributed 这个通用运行时就够用。训练优化层三条路径也各不相同：Ray 侧的优化落在 FSDP/verl 这类训练框架里，DeepSpeed 靠 ZeRO 切分状态，Megatron-Core 靠切分计算本身。三条路径殊途同归，最终都收敛到 vLLM 做推理服务化。

Megatron-LM 与 DeepSpeed 正面对比

维度	Megatron-LM（张量+流水线并行）	DeepSpeed（ZeRO 显存切分）
核心思路	把模型计算本身切开分布到多卡	模型计算不变，只把优化器状态/梯度/参数这些冗余存储按 rank 切分
对模型代码的侵入性	需要模型按 row/column-parallel 方式改造（GPTModel 这类组件已经内置了这套改造）	对已有模型代码几乎无侵入，deepspeed.initialize 包一层即可
依赖体积	较重：需要编译 C++ 数据集 helper，生产配置通常还要装 Transformer Engine	较轻：纯 Python 依赖为主，安装门槛低
相同 batch size 下的显存占用	更低（张量并行天然减少单卡显存压力）	更高
大 batch size 下的吞吐	相对较弱	通常更好
长序列支持	需要额外开启上下文并行（CP）	Ulysses 序列并行在支持场景下普遍更优，能跑更长序列
checkpoint 可移植性	和训练时的 TP/PP/DP 具体切分方式绑定，换并行度需要 reshape 或用 Megatron Bridge 转换	按 rank 数切分，可直接合并成单一 fp32 state_dict（见上一篇）
典型适用场景	单模型规模巨大、节点内有 NVLink/InfiniBand 高速互联支撑张量并行通信	已有模型代码，想低成本加上显存优化，或需要灵活切换并行度

结论不是"哪个更好"，而是两种方案在解决不同层面的约束：如果模型大到即使把 ZeRO stage 3 的 offload 都用上依然装不下，或者你有充足的节点内高速互联可以承受张量并行的通信开销，Megatron-LM 的模型并行是更合适的答案；如果模型规模在 ZeRO 能应付的范围内，且想尽量少改动已有训练代码，DeepSpeed 依然是更省事的选择。

安装配置

本地/开发环境安装 Megatron-Core：

# 轻量安装：不含 Transformer Engine 和 ModelOpt，够用于验证 TP/PP 逻辑
uv pip install --no-build-isolation "megatron-core[training,lts]"

# 完整安装：含 Transformer Engine，生产训练推荐
git clone https://github.com/NVIDIA/Megatron-LM.git
cd Megatron-LM
uv pip install --group build
uv pip install --no-build-isolation -e ".[training,dev]"

完整安装依赖匹配版本的 CUDA/cuDNN/NCCL 和 Transformer Engine，从源码编译内存占用较大，如果编译过程 OOM，用 MAX_JOBS=4 这类环境变量限制并行编译任务数。生产环境不建议每次都从源码编译------NVIDIA 官方发布的 NGC PyTorch 容器已经预装好这套依赖，直接基于它构建训练镜像通常更省事。

安装 Megatron Bridge（用于和 HuggingFace 生态互转 checkpoint，训练前导入/训练后导出都要用到）：

pip install megatron-bridge

面向人群

• 招聘要求里看到张量并行/流水线并行/Megatron-LM 关键词、想补一个能讲清楚完整实践项目的工程师
• 正在评估该用 Megatron-LM 还是 DeepSpeed 做模型并行策略、需要一份对比依据的平台工程师
• 想理解"为什么百亿参数以上的模型训练绕不开模型并行"这个问题本质、而不只是记住 TP/PP 这两个名词的技术决策者

实践步骤：从张量并行预训练到推理上线

以下三个阶段接着覆盖 Kubeflow Trainer 配置→提交张量并行训练任务→checkpoint 转换→vLLM 上线，命令和 API 用法对照 Kubeflow Trainer Megatron 指南、Megatron-LM、Megatron Bridge 官方文档核实。和前两篇一样，受限于篇幅，训练阶段用的是能验证张量并行逻辑的最小示例（小模型 + mock 数据），而非真实的千卡预训练规模，但架构和命令是生产可直接复用的，把模型定义和 num_nodes/resources_per_node 换大就是生产配置。

阶段一：安装 Kubeflow Trainer + 配置 Megatron 专属依赖

Kubeflow Trainer 的安装步骤和上一篇 DeepSpeed 文章完全一样（同一个 Trainer 装一次，各种运行时可以共存），这里不重复贴命令。区别在于 Megatron-Core 不需要像 DeepSpeed 那样额外装一个专门的 ClusterTrainingRuntime------它原生用 torchrun 启动，直接复用内置的 torch-distributed 运行时：

kubectl get clustertrainingruntimes | grep torch-distributed
# 应看到 torch-distributed，Kubeflow Trainer 默认自带，无需额外安装

但有一个 Megatron 特有的坑：K8s Pod 默认 /dev/shm 只有 64MB。 Megatron-Core 会为每一个并行维度（TP、DP......）各创建一个独立的 NCCL 通信组，通信组一多，默认的 64MB 共享内存很容易不够用而报错。需要把 /dev/shm 挂载成内存盘扩容，可以用 runtimePatches 只对单次训练任务生效，不影响集群里跑其他运行时的任务：

runtimePatches:
  - manager: trainer.kubeflow.org/kubeflow-sdk
    trainingRuntimeSpec:
      template:
        spec:
          replicatedJobs:
            - name: node
              template:
                spec:
                  template:
                    spec:
                      volumes:
                        - name: dshm
                          emptyDir:
                            medium: Memory
                      containers:
                        - name: node
                          volumeMounts:
                            - name: dshm
                              mountPath: /dev/shm

另外，Megatron 的数据集加载会先编译一段 C++ helper（compile_helpers），这段编译需要 make/g++，而官方默认的 pytorch/pytorch:*-runtime 系列镜像不带 C/C++ 工具链。官方示例的做法不是重新构建镜像，而是直接在训练函数里运行时安装：apt-get install -y make g++，再从 Megatron-LM 仓库拉取 Makefile/helpers.cpp 源文件放到本地 megatron.core.datasets 目录下，最后只让 rank 0 跑一次 compile_helpers()、其余 rank 用 torch.distributed.barrier() 等它编译完------这样就不用为了一段 C++ helper 单独维护一个自定义镜像。

阶段二：提交张量并行训练任务（Kubeflow Trainer SDK）

下面这段训练函数改编自 Kubeflow Trainer 官方 Megatron TP 示例 notebook，该 notebook 本身又是基于 Megatron-LM 仓库的 run_simple_mcore_train_loop.py 改的：

from kubeflow.trainer import TrainerClient, CustomTrainer

def train_megatron_gpt_tp():
    import os
    import subprocess
    import urllib.request
    from functools import partial

    import torch
    from torch.optim import Adam
    from torch.utils.data import DataLoader

    from megatron.core import parallel_state
    from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
    from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
    from megatron.core.transformer.transformer_config import TransformerConfig
    from megatron.core.models.gpt.gpt_model import GPTModel
    from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
    from megatron.core.datasets.utils import compile_helpers
    from megatron.core.datasets.blended_megatron_dataset_builder import BlendedMegatronDatasetBuilder
    from megatron.core.datasets.gpt_dataset import GPTDatasetConfig, MockGPTDataset
    from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
    from megatron.core.distributed.finalize_model_grads import finalize_model_grads
    from megatron.core.tokenizers import MegatronTokenizer

    seq_length = 64

    # 第一步：初始化 torch.distributed 和 Megatron 的模型并行分组
    parallel_state.destroy_model_parallel()
    rank, world_size, local_rank = (int(os.environ[k]) for k in ("RANK", "WORLD_SIZE", "LOCAL_RANK"))
    torch.cuda.set_device(local_rank)
    torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size)

    tp_size = int(os.environ["TP_SIZE"])  # TP_SIZE 就是要把每层切成几份，由外面 env 传入
    parallel_state.initialize_model_parallel(tp_size, pipeline_model_parallel_size=1)
    model_parallel_cuda_manual_seed(123)

    # 第二步：搭一个小规模 GPT，验证 TP 逻辑而非真实训练
    transformer_config = TransformerConfig(
        num_layers=2, hidden_size=12, num_attention_heads=4,
        use_cpu_initialization=True, pipeline_dtype=torch.float32,
    )
    gpt_model = GPTModel(
        config=transformer_config, transformer_layer_spec=get_gpt_layer_local_spec(),
        vocab_size=100, max_sequence_length=seq_length,
    ).to("cuda")

    ddp_config = DistributedDataParallelConfig(
        grad_reduce_in_fp32=False, overlap_grad_reduce=False, use_distributed_optimizer=False,
    )
    gpt_model = DistributedDataParallel(config=transformer_config, ddp_config=ddp_config, module=gpt_model)
    optim = Adam(gpt_model.parameters())

    # 第三步：准备 mock 数据集------数据集加载依赖一段 C++ helper，运行时现装现编译，
    # 不需要为此单独维护一个自定义镜像
    subprocess.run(["apt-get", "update", "-qq"], capture_output=True)
    subprocess.run(["apt-get", "install", "-y", "-qq", "make", "g++"], capture_output=True)
    datasets_dir = os.path.dirname(compile_helpers.__code__.co_filename)
    base_url = "https://raw.githubusercontent.com/NVIDIA/Megatron-LM/main/megatron/core/datasets/"
    for f in ["Makefile", "helpers.cpp"]:
        urllib.request.urlretrieve(base_url + f, os.path.join(datasets_dir, f))
    if torch.distributed.get_rank() == 0:
        compile_helpers()
    torch.distributed.barrier()

    dataset_config = GPTDatasetConfig(
        random_seed=0, sequence_length=seq_length,
        reset_position_ids=False, reset_attention_mask=False, eod_mask_loss=False,
        tokenizer=MegatronTokenizer.from_pretrained(metadata_path={"library": "null-text"}, vocab_size=seq_length),
        mid_level_dataset_surplus=0.005,
    )
    datasets = BlendedMegatronDatasetBuilder(MockGPTDataset, [1000, None, None], lambda: True, dataset_config).build()
    train_iterator = iter(DataLoader(datasets[0], batch_size=8, shuffle=True))

    # 第四步：定义 forward_step_func------Megatron 的流水线调度器要求这个签名
    def forward_step_func(data_iterator, model):
        def loss_func(loss_mask, output_tensor):
            losses = output_tensor.float()
            loss_mask = loss_mask.view(-1).float()
            loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
            return loss, {"lm loss": loss}

        data = next(data_iterator)
        tokens, attention_mask, position_ids, labels, loss_mask = (
            data[k].to("cuda") for k in ("tokens", "attention_mask", "position_ids", "labels", "loss_mask")
        )
        output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
        return output_tensor, partial(loss_func, loss_mask)

    # 第五步：跑 5 轮训练循环
    forward_backward_func = get_forward_backward_func()
    for iteration in range(5):
        optim.zero_grad()
        losses_reduced = forward_backward_func(
            forward_step_func=forward_step_func, data_iterator=train_iterator, model=gpt_model,
            num_microbatches=1, seq_length=seq_length, micro_batch_size=8,
            decoder_seq_length=seq_length, forward_only=False,
        )
        finalize_model_grads([gpt_model])  # 同步 TP 和 DP 分组间的梯度
        optim.step()
        print(f"iteration {iteration}: losses_reduced={losses_reduced}")

    torch.distributed.destroy_process_group()

num_nodes, num_gpu, tensor_model_parallel_size = 2, 1, 2

job_name = TrainerClient().train(
    runtime="torch-distributed",
    trainer=CustomTrainer(
        func=train_megatron_gpt_tp,
        num_nodes=num_nodes,
        resources_per_node={"memory": "16Gi", "gpu": num_gpu},
        packages_to_install=["megatron-core", "pybind11"],
        env={"TP_SIZE": str(tensor_model_parallel_size)},
    ),
)
client = TrainerClient()
client.wait_for_job_status(name=job_name, status={"Running"})
for c in client.get_job(name=job_name).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}")

RANK/WORLD_SIZE/LOCAL_RANK 由 torch-distributed 运行时自动注入到每个 Pod，TP_SIZE 则是通过 CustomTrainer 的 env 参数传进去的自定义变量，必须等于总 GPU 数（num_nodes * num_gpu，这里是 2×1=2）------张量并行是把每一层切成 TP 份分别放到不同 GPU 上算，卡数和切分份数必须对齐。所有 import 和依赖安装都写在训练函数内部 不是随手的代码风格，而是硬性要求------CustomTrainer 会把这个函数序列化后分发到每个节点单独执行，写在函数外的代码在远端进程里根本不会被执行到。对比上一篇 DeepSpeed 版本：那边训练函数里调的是 deepspeed.initialize，这里调的是标准 PyTorch 的 torch.distributed.init_process_group 加 Megatron 自己的 parallel_state.initialize_model_parallel，没有 MPI、没有 hostfile，两条编排路径在这一步的差异清晰可见。

阶段三：Checkpoint 转换（Megatron Bridge）

Megatron-Core 训练产出的原生 checkpoint 是按训练时的 TP/PP/DP 具体切分方式存储的分片，既不能直接给 vLLM 加载，也不能像上一篇 DeepSpeed 的 ZeRO checkpoint 那样简单地做"物理去冗余合并"------这里的分片本来就是模型计算的一部分，合并的本质是按 TP/PP 的切分顺序把权重重新拼接排列成标准的 dense 权重。用 Megatron Bridge 做这一步转换：

from megatron.bridge import AutoBridge

bridge = AutoBridge.from_hf_pretrained("Qwen/Qwen2.5-7B")
bridge.export_ckpt(
    "/mnt/shared/megatron_checkpoints/qwen25_7b",
    "/mnt/shared/hf_exports/qwen25_7b",
)

export_ckpt 整个转换过程可以在 CPU 上单进程完成，不需要像训练那样拉起多进程------这一点和上一篇的 checkpoint 合并 Job 思路一致：另起一个只申请 CPU/内存、不占 GPU 的一次性 Job，挂载和训练 Pod 相同的共享 PVC 跑这段转换代码，跑完自动退出：

apiVersion: batch/v1
kind: Job
metadata:
  name: megatron-checkpoint-export
spec:
  template:
    spec:
      restartPolicy: Never
      containers:
        - name: export
          image: <和训练用的同一个镜像，已装好 megatron-bridge>
          command: ["python", "/scripts/export_checkpoint.py"]
          resources:
            requests: { cpu: "4", memory: "32Gi" }
            limits: { cpu: "8", memory: "48Gi" }
          volumeMounts:
            - { name: shared, mountPath: /mnt/shared }
      volumes:
        - name: shared
          persistentVolumeClaim: { claimName: shared-checkpoint-pvc }

转换完的 /mnt/shared/hf_exports/qwen25_7b 是一个标准 HuggingFace 格式目录，和上一篇 DeepSpeed 合并出来的 merged-model 目录形态完全一致------两条不同的训练路径，在"交给 vLLM 之前"这一步殊途同归。

阶段四：用 vLLM 部署到 K8s

这一步和上一篇 DeepSpeed 文章的阶段四几乎一样，只是模型路径换成了 Megatron Bridge 转换出来的目录：

apiVersion: apps/v1
kind: Deployment
metadata:
  name: vllm-megatron-qwen-server
spec:
  replicas: 2
  selector:
    matchLabels: { app: vllm-megatron-qwen }
  template:
    metadata:
      labels: { app: vllm-megatron-qwen }
    spec:
      containers:
        - name: vllm
          image: vllm/vllm-openai:v0.27.1   # 部署前去 https://github.com/vllm-project/vllm/releases 确认是否有更新的稳定版
          args:
            - --model=/models/qwen25_7b
            - --served-model-name=qwen2.5-7b-megatron
            - --tensor-parallel-size=1
            - --gpu-memory-utilization=0.9
            - --max-model-len=8192
          resources:
            limits: { nvidia.com/gpu: 1 }
          readinessProbe:
            httpGet: { path: /health, port: 8000 }
            initialDelaySeconds: 30
            periodSeconds: 10
          volumeMounts:
            - { name: model-storage, mountPath: /models, readOnly: true }
      volumes:
        - name: model-storage
          persistentVolumeClaim: { claimName: shared-checkpoint-pvc }
---
apiVersion: v1
kind: Service
metadata:
  name: vllm-megatron-qwen-svc
spec:
  selector: { app: vllm-megatron-qwen }
  ports:
    - { port: 8000, targetPort: 8000 }

这里有个容易搞混的地方：--tensor-parallel-size 这个 vLLM 参数和阶段二训练时的 TP_SIZE=2 没有任何绑定关系。 训练时的张量并行是为了把模型塞进有限的训练显存里，切分方式已经在阶段三被 Megatron Bridge 转换成标准 dense 权重时"抹平"了；推理时用几张卡做张量并行，完全是根据推理阶段自己的显存和吞吐需求重新决定的独立参数------这也是为什么阶段三的 checkpoint 转换是必要的一步，而不是可以跳过的形式主义。部署后验证：

kubectl port-forward svc/vllm-megatron-qwen-svc 8000:8000
curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "qwen2.5-7b-megatron", "messages": [{"role": "user", "content": "你好"}]}'

到这里，三篇文章的链路在同一个终点会合：Ray 那条路径靠动态调度让训练推理反复切换，DeepSpeed 那条路径靠 ZeRO 切分显存状态、MPI 编排多机训练，Megatron-LM 这条路径靠张量/流水线并行直接切分模型计算本身------三种截然不同的分布式训练哲学，最后都在 vLLM 这一层收敛成同一种推理服务化方式。

小结

Megatron-LM 在这三篇系列里补上的是"模型本身太大切不开就装不下"这一类问题的答案：K8s 管物理资源不变，torch-distributed 运行时接手编排（不需要像 DeepSpeed 那样引入专门的运行时，因为 Megatron-Core 原生就是标准 PyTorch 分布式），Megatron-Core 负责把模型计算切分到多卡，Megatron Bridge 负责把训练产出的切分 checkpoint 转换回标准格式，vLLM 依旧负责推理服务化。三篇文章走完一遍：训练推理要反复切换选 Ray，已有模型代码想低成本上显存优化选 DeepSpeed，模型规模大到必须做模型并行、且有高速互联支撑通信开销选 Megatron-LM------三条路径没有绝对的优劣，只是解决的约束层面不同，选型时先看清楚自己卡在哪一层，再决定往下走哪条编排路径。