K8s + Megatron-LM + vLLM：打通超大模型训练全链路实战-平芜编程栈

1. 背景与目标

随着大模型参数规模突破千亿甚至万亿级别，单机训练已无法满足需求。本文以 Kubernetes（K8s）为底座，结合 NVIDIA Megatron-LM 完成大规模分布式预训练，再通过 vLLM 实现高效推理部署，打通从训练到推理的完整链路。全文包含可运行的代码与配置示例，帮助读者在真实集群中落地。

整体架构分为三层：K8s 负责资源调度与弹性伸缩，Megatron-LM 负责张量并行、流水线并行与数据并行训练，vLLM 负责将训练产出的模型权重高效加载并对外提供推理服务。

2. 环境准备与集群规划

本文示例基于以下环境：Kubernetes 1.28+、NVIDIA GPU Operator、Helm 3、PyTorch 2.1、CUDA 12.2。建议使用 A100 或 H800 等大显存 GPU 节点，每个节点至少 8 卡。

首先确认集群 GPU 资源可用：

kubectl get nodes -l gpu=available -o wide kubectl describe node gpu-node-01 | grep -A 5 "Allocated resources"

安装 NVIDIA Device Plugin 与 GPU Feature Discovery，确保 Pod 能申请到 GPU 资源：

helm repo add nvdp https://nvidia.github.io/k8s-device-plugin helm repo update helm install nvidia-device-plugin nvdp/nvidia-device-plugin \ --namespace kube-system \ --set gfd.enabled=true

3. 构建训练镜像

训练镜像需要包含 Megatron-LM 源码、PyTorch、NCCL 以及必要的依赖。以下 Dockerfile 基于官方 NGC PyTorch 镜像构建：

FROM nvcr.io/nvidia/pytorch:23.12-py3 WORKDIR /workspace 安装 Megatron-LM RUN git clone https://github.com/NVIDIA/Megatron-LM.git && cd Megatron-LM && pip install -e . 安装 vLLM（用于后续推理验证） RUN pip install vllm==0.4.2 安装常用工具 RUN pip install tensorboard datasets accelerate COPY entrypoint.sh /workspace/entrypoint.sh RUN chmod +x /workspace/entrypoint.sh ENTRYPOINT ["/workspace/entrypoint.sh"]

构建并推送到镜像仓库：

docker build -t registry.example.com/llm/megatron-train:latest . docker push registry.example.com/llm/megatron-train:latest

4. 编写训练入口脚本

训练入口脚本负责启动 Megatron-LM 的预训练流程。以下脚本通过环境变量接收分布式训练参数：

#!/bin/bash set -e GPUS_PER_NODE=${GPUS_PER_NODE:-8} NNODES=${NNODES:-1} NODE_RANK=${NODE_RANK:-0} MASTER_ADDR=${MASTER_ADDR:-localhost} MASTER_PORT=${MASTER_PORT:-6000} DISTRIBUTED_ARGS=" --nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT " TRAIN_ARGS=" --tensor-model-parallel-size 4 --pipeline-model-parallel-size 2 --num-layers 24 --hidden-size 2048 --num-attention-heads 32 --seq-length 2048 --max-position-embeddings 2048 --micro-batch-size 4 --global-batch-size 64 --train-iters 5000 --lr 1.5e-4 --lr-decay-style cosine --min-lr 1.0e-5 --weight-decay 0.1 --clip-grad 1.0 --fp16 --data-path /workspace/data/wikipedia_text --tokenizer-type GPT2BPETokenizer --vocab-file /workspace/data/gpt2-vocab.json --merge-file /workspace/data/gpt2-merges.txt --save /workspace/checkpoints --load /workspace/checkpoints --log-interval 10 --save-interval 1000 --eval-interval 1000 --eval-iters 10 " torchrun $DISTRIBUTED_ARGS /workspace/Megatron-LM/pretrain_gpt.py $TRAIN_ARGS

5. 编写 K8s 训练任务清单

使用原生 K8s StatefulSet 或 Job 编排多节点训练任务。这里以 StatefulSet 为例，保证每个 Pod 有稳定的网络标识，便于 NCCL 通信：

apiVersion: apps/v1 kind: StatefulSet metadata: name: megatron-gpt namespace: llm spec: serviceName: megatron-gpt replicas: 2 selector: matchLabels: app: megatron-gpt template: metadata: labels: app: megatron-gpt spec: affinity: podAntiAffinity: requiredDuringSchedulingIgnoredDuringExecution: - labelSelector: matchExpressions: - key: app operator: In values: - megatron-gpt topologyKey: kubernetes.io/hostname containers: - name: trainer image: registry.example.com/llm/megatron-train:latest command: ["/workspace/entrypoint.sh"] env: - name: GPUS_PER_NODE value: "8" - name: NNODES value: "2" - name: NODE_RANK valueFrom: fieldRef: fieldPath: metadata.labels['statefulset.kubernetes.io/pod-name'] - name: MASTER_ADDR value: "megatron-gpt-0.megatron-gpt.llm.svc.cluster.local" - name: MASTER_PORT value: "6000" resources: limits: nvidia.com/gpu: 8 volumeMounts: - name: data mountPath: /workspace/data - name: checkpoints mountPath: /workspace/checkpoints volumes: - name: data persistentVolumeClaim: claimName: llm-data-pvc - name: checkpoints persistentVolumeClaim: claimName: llm-checkpoints-pvc

注意：NODE_RANK 需要从 Pod 名称中解析出序号，可通过 initContainer 或自定义 entrypoint 处理。这里给出一个更稳妥的解析方式：

#!/bin/bash # 在 entrypoint.sh 中解析节点序号 POD_NAME=$(hostname) NODE_RANK=$(echo $POD_NAME | awk -F'-' '{print $NF}') export NODE_RANK=$NODE_RANK echo "Node rank: $NODE_RANK"

6. 创建 Headless Service 与存储

StatefulSet 需要配套 Headless Service 提供稳定的 DNS 解析：

apiVersion: v1 kind: Service metadata: name: megatron-gpt namespace: llm spec: clusterIP: None selector: app: megatron-gpt ports: - name: master port: 6000 targetPort: 6000

创建共享存储 PVC，用于存放训练数据和模型检查点。这里以 NFS 为例：

apiVersion: v1 kind: PersistentVolumeClaim metadata: name: llm-data-pvc namespace: llm spec: accessModes: - ReadWriteMany storageClassName: nfs-csi resources: requests: storage: 2Ti --- apiVersion: v1 kind: PersistentVolumeClaim metadata: name: llm-checkpoints-pvc namespace: llm spec: accessModes: - ReadWriteMany storageClassName: nfs-csi resources: requests: storage: 5Ti

7. 提交训练任务并监控

应用清单并观察训练状态：

kubectl apply -f megatron-gpt.yaml kubectl get pods -n llm -l app=megatron-gpt -w

查看训练日志：

kubectl logs -n llm megatron-gpt-0 -f | tail -50

训练过程中可通过 TensorBoard 监控 loss 曲线。在本地端口转发后访问：

kubectl port-forward -n llm megatron-gpt-0 6006:6006

8. 训练完成后的权重导出

Megatron-LM 训练产出的检查点格式与 HuggingFace 不兼容，需要先转换为 HuggingFace 格式，才能被 vLLM 加载。Megatron-LM 仓库提供了转换脚本：

python /workspace/Megatron-LM/tools/checkpoint_util.py \ --model-type GPT \ --loader megatron \ --saver huggingface \ --load-dir /workspace/checkpoints \ --save-dir /workspace/hf-checkpoints \ --target-tensor-parallel-size 1 \ --target-pipeline-parallel-size 1

转换完成后，将 HuggingFace 格式的权重上传到对象存储或直接放入推理服务可访问的 PVC：

kubectl exec -n llm megatron-gpt-0 -- \ tar czf /workspace/hf-checkpoints.tar.gz -C /workspace hf-checkpoints kubectl cp llm/megatron-gpt-0:/workspace/hf-checkpoints.tar.gz ./hf-checkpoints.tar.gz

9. 使用 vLLM 部署推理服务

vLLM 支持直接加载 HuggingFace 格式的模型权重，并提供 OpenAI 兼容的推理接口。以下 Dockerfile 构建推理镜像：

FROM nvcr.io/nvidia/pytorch:23.12-py3 RUN pip install vllm==0.4.2 WORKDIR /workspace COPY serve.py /workspace/serve.py EXPOSE 8000 CMD ["python", "/workspace/serve.py"]

编写 vLLM 推理服务入口脚本：

from vllm import LLM, SamplingParams llm = LLM( model="/workspace/hf-checkpoints", tensor_parallel_size=4, dtype="float16", max_model_len=2048, ) sampling_params = SamplingParams( temperature=0.7, top_p=0.9, max_tokens=512, ) from fastapi import FastAPI from pydantic import BaseModel app = FastAPI() class GenerateRequest(BaseModel): prompt: str max_tokens: int = 512 temperature: float = 0.7 @app.post("/v1/completions") async def generate(req: GenerateRequest): outputs = llm.generate([req.prompt], SamplingParams( temperature=req.temperature, max_tokens=req.max_tokens, )) text = outputs[0].outputs[0].text return {"choices": [{"text": text}]} if name == "main": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)

10. 编写 vLLM 推理 Deployment 与 Service

推理服务同样部署在 K8s 中，通过 Service 对外暴露接口：

apiVersion: apps/v1 kind: Deployment metadata: name: vllm-server namespace: llm spec: replicas: 1 selector: matchLabels: app: vllm-server template: metadata: labels: app: vllm-server spec: containers: - name: vllm image: registry.example.com/llm/vllm-serve:latest ports: - containerPort: 8000 resources: limits: nvidia.com/gpu: 4 volumeMounts: - name: hf-checkpoints mountPath: /workspace/hf-checkpoints volumes: - name: hf-checkpoints persistentVolumeClaim: claimName: llm-checkpoints-pvc --- apiVersion: v1 kind: Service metadata: name: vllm-server namespace: llm spec: selector: app: vllm-server ports: - port: 8000 targetPort: 8000 type: ClusterIP

11. 验证推理服务

部署完成后，通过端口转发或 Ingress 访问推理接口：

kubectl apply -f vllm-deployment.yaml kubectl get pods -n llm -l app=vllm-server -w kubectl port-forward -n llm svc/vllm-server 8000:8000

使用 curl 验证推理效果：

curl -X POST http://localhost:8000/v1/completions \ -H "Content-Type: application/json" \ -d '{ "prompt": "Kubernetes 是", "max_tokens": 100, "temperature": 0.7 }'

返回结果示例：

{ "choices": [ { "text": "一个开源的容器编排平台，用于自动化部署、扩展和管理容器化应用。它提供了服务发现、负载均衡、自动伸缩等能力……" } ] }

12. 常见问题与调优建议

训练过程中常见的几个问题及解决方案：

NCCL 通信超时：检查 Pod 间网络连通性，确认 Headless Service 正常解析，必要时增大NCCL_TIMEOUT环境变量。
GPU 显存不足：降低micro-batch-size或seq-length，或增大张量并行度。
检查点保存失败：确认 PVC 有足够空间，且所有 Pod 都能读写共享存储。
vLLM 加载慢：首次加载需要将权重从磁盘读入显存，建议使用高性能存储或提前预热。

调优建议：训练阶段优先保证global-batch-size与数据规模匹配；推理阶段根据并发量调整tensor_parallel_size与副本数，必要时配合 HPA 自动伸缩。

13. 总结

本文从环境准备、镜像构建、分布式训练、权重导出到 vLLM 推理部署，完整演示了 K8s + Megatron-LM + vLLM 的超大模型训练全链路。核心要点包括：使用 StatefulSet 保证分布式训练节点稳定通信，通过共享存储管理数据集与检查点，训练完成后将 Megatron 权重转换为 HuggingFace 格式供 vLLM 加载。读者可根据实际模型规模调整并行策略与资源配额，将这套链路应用到自己的业务场景中。