一、前言：從推理自動伸縮到訓練水平擴展

在上一篇我們聊到 LLM 推理的自動伸縮策略，解決了「模型在服務階段的彈性伸縮」問題。
然而，模型要能被穩定部署之前，還得先解決另一個痛點——如何有效地訓練它。

訓練大型模型（如 Llama、BERT、Stable Diffusion）不只是「多用幾張 GPU」這麼簡單。
單機的 GPU 訓練常會遇到三個限制：

1. 顯示記憶體不足：模型或 batch 太大直接 OOM。
2. 訓練時間過長：幾十億參數的模型要跑數週。
3. 資源利用不均：GPU 閒置或 CPU 成瓶頸。

因此，「分散式訓練（Distributed Training）」成了進階 AI 工程的必經之路。
這篇文章會示範如何在 Kubernetes + Kubeflow 環境中部署分散式訓練任務，
並結合 PyTorch Distributed 與 DeepSpeed 來微調 BERT 模型。

二、分散式訓練核心概念快速導覽

分散式訓練有三種常見模式：

| 模式 | 說明 | 案例 |
|------|------|------|
| Data Parallel | 每個節點訓練不同資料 batch，最後平均梯度。 | PyTorch DDP |
| Model Parallel | 模型切成多份分布到不同 GPU。 | Megatron-LM |
| Pipeline Parallel | 模型分層串接，前後傳遞交錯執行。 | DeepSpeed |

在雲端集群上，我們主要用 AllReduce 架構（例如 NCCL backend）來同步梯度，
這讓 GPU 之間的通訊延遲最小化，也最適合透過 Kubernetes 進行動態調度。

三、Kubeflow Training Operator：讓訓練任務變成 Kubernetes 原生工作負載

Kubeflow 提供多種訓練資源控制器（CRD）：

• PyTorchJob：分散式 PyTorch 任務
• TFJob：TensorFlow 訓練
• MPIJob：高性能計算
• XGBoostJob：分散式樹模型訓練

這些 Operator 會自動幫你：

• 建立 Master / Worker Pod
• 設定環境變數（MASTER_ADDR, RANK, WORLD_SIZE）
• 管理訓練生命週期（失敗重啟、完成清理）

你只需要撰寫一份 YAML，Kubeflow 就能幫你把分散式訓練變成雲原生工作負載。

四、實作：在 Kubeflow 上進行 BERT 微調

🎯 目標

在 Kubeflow 環境中使用兩個 GPU 節點（各一張 GPU）進行 BERT 的情感分類微調。

🧱 專案結構

bert-finetune/
├── train.py
├── requirements.txt
├── ds_config.json
└── pytorchjob.yaml

🧩 train.py（訓練腳本）

這裡使用 HuggingFace 的 transformers 與 datasets，並啟用 DeepSpeed：

from transformers import BertForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset
import torch

def main():
    dataset = load_dataset("imdb")
    model = BertForSequenceClassification.from_pretrained("bert-base-uncased")

    training_args = TrainingArguments(
        output_dir="/mnt/output",
        per_device_train_batch_size=8,
        num_train_epochs=1,
        evaluation_strategy="epoch",
        deepspeed="/mnt/config/ds_config.json",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
    )
    trainer.train()

if __name__ == "__main__":
    torch.distributed.init_process_group(backend="nccl")
    main()

⚙️ DeepSpeed 設定（ds_config.json）

{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2
  }
}

📦 requirements.txt

transformers==4.44.0
datasets==2.20.0
deepspeed==0.15.1
torch==2.2.0

🧾 Kubeflow PyTorchJob（pytorchjob.yaml）

這份 YAML 定義了 1 個 Master 與 1 個 Worker。

apiVersion: kubeflow.org/v1
kind: PyTorchJob
metadata:
  name: bert-finetune
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: your-registry/bert-finetune:latest
              command: ["python", "train.py"]
              volumeMounts:
                - mountPath: /mnt
                  name: shared
              resources:
                limits:
                  nvidia.com/gpu: 1
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: your-registry/bert-finetune:latest
              command: ["python", "train.py"]
              volumeMounts:
                - mountPath: /mnt
                  name: shared
              resources:
                limits:
                  nvidia.com/gpu: 1
  volumes:
    - name: shared
      persistentVolumeClaim:
        claimName: data-pvc

💡 你可以在 GKE 或 Kind + GPU 節點環境中直接 apply：

kubectl apply -f pytorchjob.yaml
kubectl get pytorchjobs

📂 Persistent Volume（資料與模型保存）

Kubeflow 會幫每個容器掛載 /mnt 下的 PVC，
用來保存：

• Dataset（例如 /mnt/data/imdb/）
• Checkpoint（例如 /mnt/output/checkpoint.pt）
• DeepSpeed 配置（例如 /mnt/config/ds_config.json）

這樣訓練中斷也不會遺失進度。

🧪 驗證訓練結果

訓練完成後，你可以檢查輸出：

kubectl logs bert-finetune-master-0

在 /mnt/output 會產生 trainer_state.json、pytorch_model.bin 等檔案，
可直接上傳到 HuggingFace Hub 或部署到推理服務中。

五、進階：StatefulSet 與多節點狀態管理

雖然 PyTorchJob 已能自動管理訓練流程，但若你需要更細緻控制（如固定 rank 或復原順序），
可以改用 StatefulSet：

• 每個 Pod 會有固定名稱（trainer-0, trainer-1, …）
• 適合需要長期穩定節點關係的訓練任務
• 可搭配 Headless Service 做節點間通訊

此架構常見於 Parameter Server 模式，或有明確主從依賴的任務。

六、監控與資源自動化

分散式訓練是高成本任務，因此需要良好的監控與彈性：

• Prometheus + Grafana：監控 GPU 使用率、Pod 狀態、I/O
• KEDA ScaledObject：根據任務佇列數量自動擴縮
• 訓練完成自動觸發部署流程：整合 Argo Workflow 或 Tekton Pipeline

這樣，你的模型訓練就能像 CI/CD 一樣自動化與可觀測。

七、結語

「推理自動伸縮」解決了如何讓模型穩定提供服務；
「分散式訓練與微調」則解決了如何讓模型高效誕生與進化。

透過 Kubeflow Training Operator、PyTorch Distributed、DeepSpeed、PVC 管理，
我們得以在 Kubernetes 內部建構出一套可伸縮、可觀測、可重現的 AI 訓練系統。

下一篇，要來自我回顧一下鐵人賽