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

1. 背景与目标随着大模型参数规模突破千亿甚至万亿级别单机训练已无法满足需求。本文以 KubernetesK8s为底座结合 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 gpuavailable -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.enabledtrue3. 构建训练镜像训练镜像需要包含 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 vllm0.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:latest4. 编写训练入口脚本训练入口脚本负责启动 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_ARGS5. 编写 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_RANK6. 创建 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: 5Ti7. 提交训练任务并监控应用清单并观察训练状态kubectl apply -f megatron-gpt.yaml kubectl get pods -n llm -l appmegatron-gpt -w查看训练日志kubectl logs -n llm megatron-gpt-0 -f | tail -50训练过程中可通过 TensorBoard 监控 loss 曲线。在本地端口转发后访问kubectl port-forward -n llm megatron-gpt-0 6006:60068. 训练完成后的权重导出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 格式的权重上传到对象存储或直接放入推理服务可访问的 PVCkubectl 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.gz9. 使用 vLLM 部署推理服务vLLM 支持直接加载 HuggingFace 格式的模型权重并提供 OpenAI 兼容的推理接口。以下 Dockerfile 构建推理镜像FROM nvcr.io/nvidia/pytorch:23.12-py3 RUN pip install vllm0.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_size4, dtypefloat16, max_model_len2048, ) sampling_params SamplingParams( temperature0.7, top_p0.9, max_tokens512, ) 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( temperaturereq.temperature, max_tokensreq.max_tokens, )) text outputs[0].outputs[0].text return {choices: [{text: text}]} if name main: import uvicorn uvicorn.run(app, host0.0.0.0, port8000)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: ClusterIP11. 验证推理服务部署完成后通过端口转发或 Ingress 访问推理接口kubectl apply -f vllm-deployment.yaml kubectl get pods -n llm -l appvllm-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 加载。读者可根据实际模型规模调整并行策略与资源配额将这套链路应用到自己的业务场景中。