TranslateGemma 4B 在 NVIDIA DGX Spark 上的部署教程

SETUP_TUTORIAL.md

TranslateGemma 4B 在 NVIDIA DGX Spark 上的部署教程

環境說明

• 硬體: NVIDIA DGX Spark

  • GPU: NVIDIA GB10 (Blackwell 架構, sm_121)
  • 記憶體: 128GB Unified Memory
  • CUDA: 12.8

• 目標: 部署 Google TranslateGemma-4B-IT 翻譯模型，支援 55 種語言互譯

踩坑記錄

坑 1: PyTorch CUDA 版本不支援 Blackwell

DGX Spark 的 GB10 是 Blackwell 架構 (sm_121)，一般的 PyTorch CUDA 版本不支援。

# 錯誤訊息
NVIDIA GB10 with CUDA capability sm_121 is not compatible with the current PyTorch installation.

解法: 需要安裝 PyTorch nightly cu128 版本

pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128

坑 2: Conda 環境的 PyTorch 是 CPU 版本

用 conda 安裝的 PyTorch 可能是 CPU-only 版本。

import torch
print(torch.cuda.is_available())  # False

解法: 用 pip 指定 CUDA wheel 安裝

坑 3: 直接跑 transformers 太慢

用 HuggingFace transformers 直接載入模型，第一次推理要 60 秒（CUDA kernel 編譯），之後每次也要 9-14 秒。

即使加上 4-bit quantization (bitsandbytes)，還是要 9 秒左右。

# bitsandbytes 4-bit 設定
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
)

坑 4: vLLM 的 TranslateGemma 輸入格式

TranslateGemma 不是一般的 chat model，OpenAI-compatible /v1/chat/completions API 不直接支援。

需要用 /v1/completions 搭配特殊的 prompt 格式。

坑 5: max_tokens 設太高會很慢

vLLM 的 stop token 對 TranslateGemma 不生效，會一直生成到 max_tokens 為止。

# 錯誤做法
max_tokens=2048  # 即使翻譯只需要 10 個 token，也會生成 2048 個 → 超級慢！

解法: 根據輸入長度動態計算 max_tokens

# 正確做法：根據輸入估算
estimated_tokens = min(max(len(text) * 2, 20), 500)

坑 6: vLLM 不支援 TranslateGemma 的圖片格式

TranslateGemma 支援圖片翻譯（OCR + 翻譯），但 vLLM 的 OpenAI-compatible API 不支援其特殊格式。

# TranslateGemma 圖片格式（vLLM 不支援）
messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "source_lang_code": "en",
                "target_lang_code": "zh-TW",
                "url": image_url,
            }
        ],
    }
]

解法: 另外啟動一個 transformers pipeline server 處理圖片翻譯（見下方「圖片翻譯」章節）

---

最終方案: NVIDIA vLLM Container

Step 1: 準備 HuggingFace Token

TranslateGemma 是 gated model，需要：

1. 去 https://huggingface.co/google/translategemma-4b-it 同意條款
2. 取得 HF Token: https://huggingface.co/settings/tokens

export HF_TOKEN="your_token_here"

Step 2: 清理 GPU 記憶體

確保沒有其他程式佔用 GPU 記憶體：

# 查看 GPU 使用狀況
nvidia-smi

# 如果有其他 Python 程式佔用，kill 掉
pkill -f "python.*model"

Step 3: 啟動 vLLM Container (FP8 量化版)

docker run -d \
  --name tg-vllm \
  --gpus all \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e HF_TOKEN=$HF_TOKEN \
  -p 8603:8000 \
  --ipc=host \
  nvcr.io/nvidia/vllm:25.11-py3 \
  vllm serve google/translategemma-4b-it \
  --quantization fp8 \
  --gpu-memory-utilization 0.7

參數說明:

• --gpus all: 使用 GPU
• -v ~/.cache/huggingface:/root/.cache/huggingface: 掛載 HF cache 避免重複下載
• -e HF_TOKEN: 傳入 HuggingFace token
• -p 8603:8000: 對外開放 8603 port
• --ipc=host: 共享記憶體，提升效能
• vllm serve: 使用 vLLM 的 serve 命令
• --quantization fp8: FP8 量化，模型從 8GB 降到 5GB，速度提升 3 倍！
• --gpu-memory-utilization 0.7: 使用 70% GPU 記憶體

Step 4: 等待模型載入

模型載入需要幾分鐘，可以用 logs 追蹤：

docker logs -f tg-vllm

看到以下訊息表示準備完成：

INFO:     Application startup complete.

完整啟動流程：

1. 下載模型權重 (~77 秒)
2. 載入 safetensors shards (~66 秒)
3. torch.compile 優化 (~19 秒)
4. 分配 KV cache (~72 GiB)
5. 捕捉 CUDA graphs (~4 秒)

Step 5: 測試翻譯 API

TranslateGemma 需要特殊的 prompt 格式：

curl -s http://localhost:8603/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "google/translategemma-4b-it",
    "prompt": "<bos><start_of_turn>user\n<translategemma><source_lang_code>en</source_lang_code><target_lang_code>zh-Hans</target_lang_code><text>Hello, how are you?</text></translategemma><end_of_turn>\n<start_of_turn>model\n",
    "max_tokens": 50
  }'

回應:

{
  "choices": [{
    "text": "你好，你好吗？\n"
  }],
  "usage": {
    "prompt_tokens": 59,
    "completion_tokens": 50
  }
}

---

效能測試結果

方案	翻譯速度	備註
transformers (CPU)	224 秒	不可用
transformers (CUDA, 首次)	60 秒	CUDA kernel 編譯
transformers (CUDA, warmup 後)	9-14 秒	bfloat16
transformers + bitsandbytes 4-bit	9 秒	NF4 量化
vLLM bfloat16 (本地)	~1 秒	模型 8.15 GiB
vLLM FP8 量化 (本地)	0.3-1 秒	模型 5.03 GiB，推薦！

---

支援的語言代碼

TranslateGemma 支援 55 種語言，常用的包括：

代碼	語言
en	English
zh-Hans	簡體中文
zh-Hant	繁體中文
ja	日本語
ko	한국어
es	Español
fr	Français
de	Deutsch
ru	Русский
ar	العربية

---

封裝成 Python Client

import requests

class TranslateGemmaClient:
    def __init__(self, base_url="http://localhost:8603"):
        self.base_url = base_url

    def translate(self, text: str, source_lang: str = "en", target_lang: str = "zh-Hans") -> str:
        prompt = f"""<bos><start_of_turn>user
<translategemma><source_lang_code>{source_lang}</source_lang_code><target_lang_code>{target_lang}</target_lang_code><text>{text}</text></translategemma><end_of_turn>
<start_of_turn>model
"""
        # 動態計算 max_tokens，避免生成過多無用 token
        estimated_tokens = min(max(len(text) * 2, 20), 500)

        response = requests.post(
            f"{self.base_url}/v1/completions",
            json={
                "model": "google/translategemma-4b-it",
                "prompt": prompt,
                "max_tokens": estimated_tokens,
            }
        )
        result = response.json()
        return result["choices"][0]["text"].strip()

# 使用範例
client = TranslateGemmaClient("http://spark-bbba.local:8603")
print(client.translate("Hello, how are you?", "en", "zh-Hans"))
# 輸出: 你好，你好吗？

---

Gradio 測試介面

提供一個 Gradio 網頁介面方便測試：

# gradio_app.py
import gradio as gr
import requests
import time

API_URL = "http://spark-bbba.local:8603"

LANGUAGES = {
    "English": "en",
    "簡體中文": "zh-Hans",
    "繁體中文": "zh-Hant",
    "日本語": "ja",
    "한국어": "ko",
}

def translate(text, source_lang, target_lang):
    if not text.strip():
        return "", ""

    src_code = LANGUAGES.get(source_lang, "en")
    tgt_code = LANGUAGES.get(target_lang, "zh-Hans")

    prompt = f"""<bos><start_of_turn>user
<translategemma><source_lang_code>{src_code}</source_lang_code><target_lang_code>{tgt_code}</target_lang_code><text>{text}</text></translategemma><end_of_turn>
<start_of_turn>model
"""
    # 動態計算 max_tokens
    estimated_tokens = min(max(len(text) * 2, 20), 500)

    start = time.time()
    response = requests.post(
        f"{API_URL}/v1/completions",
        json={
            "model": "google/translategemma-4b-it",
            "prompt": prompt,
            "max_tokens": estimated_tokens,
        }
    )
    elapsed = time.time() - start

    result = response.json()
    translation = result["choices"][0]["text"].strip()
    info = f"⏱ {elapsed:.2f}s"

    return translation, info

with gr.Blocks(title="TranslateGemma") as demo:
    gr.Markdown("# TranslateGemma 翻譯測試")
    with gr.Row():
        source_lang = gr.Dropdown(list(LANGUAGES.keys()), value="English", label="來源語言")
        target_lang = gr.Dropdown(list(LANGUAGES.keys()), value="繁體中文", label="目標語言")
    input_text = gr.Textbox(label="輸入", lines=5)
    output_text = gr.Textbox(label="翻譯結果", lines=5)
    info_text = gr.Textbox(label="資訊")
    btn = gr.Button("翻譯", variant="primary")
    btn.click(translate, [input_text, source_lang, target_lang], [output_text, info_text])

demo.launch(server_name="0.0.0.0", server_port=7861)

啟動：

python gradio_app.py
# 開啟 http://localhost:7861

---

圖片翻譯 (OCR + 翻譯)

TranslateGemma 支援從圖片中提取文字並翻譯。由於 vLLM 不支援其特殊格式，需要另外啟動一個 transformers server。

架構

┌─────────────────────────────────────────────────┐
│  Client (Gradio)                                │
│  ├── 文字翻譯 → vLLM (port 8603)               │
│  └── 圖片翻譯 → image_server (port 8604)       │
└─────────────────────────────────────────────────┘

image_server.py

#!/usr/bin/env python3
"""TranslateGemma 圖片翻譯 Server"""

import base64
import time
import io
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import uvicorn

app = FastAPI(title="TranslateGemma Image API")
pipe = None


class ImageTranslateRequest(BaseModel):
    image_base64: str
    source_lang: str = "en"
    target_lang: str = "zh-TW"


class ImageTranslateResponse(BaseModel):
    translation: str
    source_lang: str
    target_lang: str
    elapsed_seconds: float


def load_pipeline():
    global pipe
    from transformers import pipeline

    print("Loading TranslateGemma pipeline...")
    pipe = pipeline(
        "image-text-to-text",
        model="google/translategemma-4b-it",
        device="cuda" if torch.cuda.is_available() else "cpu",
        torch_dtype=torch.bfloat16,
    )
    print("Ready!")


@app.on_event("startup")
async def startup_event():
    load_pipeline()


@app.get("/health")
async def health():
    return {"status": "ok", "pipeline_loaded": pipe is not None}


@app.post("/translate_image", response_model=ImageTranslateResponse)
async def translate_image(request: ImageTranslateRequest):
    from PIL import Image

    start_time = time.time()

    # Decode image
    image_data = base64.b64decode(request.image_base64)
    image = Image.open(io.BytesIO(image_data))

    # Resize if needed
    max_size = 896
    if image.width > max_size or image.height > max_size:
        ratio = min(max_size / image.width, max_size / image.height)
        new_size = (int(image.width * ratio), int(image.height * ratio))
        image = image.resize(new_size, Image.Resampling.LANCZOS)

    if image.mode in ("RGBA", "P"):
        image = image.convert("RGB")

    # TranslateGemma 圖片格式
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "source_lang_code": request.source_lang,
                    "target_lang_code": request.target_lang,
                    "image": image,
                }
            ],
        }
    ]

    output = pipe(text=messages, max_new_tokens=200, generate_kwargs={"do_sample": False})
    translation = output[0]["generated_text"][-1]["content"]

    return ImageTranslateResponse(
        translation=translation,
        source_lang=request.source_lang,
        target_lang=request.target_lang,
        elapsed_seconds=round(time.time() - start_time, 3),
    )


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8604)

啟動圖片翻譯服務

# 在 Spark 上執行
cd ~/apps/translategemma
conda activate translategemma
python image_server.py

# 或背景執行
nohup python image_server.py > image_server.log 2>&1 &

測試圖片翻譯 API

# 將圖片轉成 base64
IMAGE_BASE64=$(base64 -i test_image.jpg)

# 呼叫 API
curl -X POST http://localhost:8604/translate_image \
  -H "Content-Type: application/json" \
  -d "{
    \"image_base64\": \"$IMAGE_BASE64\",
    \"source_lang\": \"en\",
    \"target_lang\": \"zh-TW\"
  }"

注意事項

• 圖片翻譯會自動 OCR 提取文字再翻譯
• 繁體中文輸出可能需要額外用 OpenCC 轉換（模型訓練資料以簡體為主）
• 建議圖片解析度不要太高（會自動縮到 896x896）

---

進階: 使用 bfloat16 原始精度 (可選)

如果不需要量化，可以用原始 bfloat16 精度（模型較大，速度較慢）：

docker run -d \
  --name tg-vllm \
  --gpus all \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e HF_TOKEN=$HF_TOKEN \
  -p 8603:8000 \
  --ipc=host \
  nvcr.io/nvidia/vllm:25.11-py3 \
  vllm serve google/translategemma-4b-it \
  --dtype bfloat16 \
  --gpu-memory-utilization 0.7

---

常用 Docker 指令

# 查看容器狀態
docker ps

# 查看 logs
docker logs -f tg-vllm

# 停止容器
docker stop tg-vllm

# 刪除容器
docker rm tg-vllm

# 重啟容器
docker restart tg-vllm

---

故障排除

問題: GPU 記憶體不足

torch.OutOfMemoryError: CUDA out of memory

解法: 降低 --gpu-memory-utilization 或清理其他 GPU 程式

問題: HuggingFace 認證失敗

Cannot access gated repo

解法: 確認 HF_TOKEN 正確且已同意模型條款

問題: 連接被拒絕

curl: (7) Failed to connect

解法: 檢查 container 是否啟動完成 (docker logs)

---

總結

在 DGX Spark 上部署 TranslateGemma 最佳方案：

文字翻譯

使用 NVIDIA vLLM Container + FP8 量化：

• 翻譯速度：0.3-1 秒
• 模型大小：5GB（FP8 量化）

圖片翻譯

使用 transformers pipeline server：

• 支援 OCR + 翻譯
• 需要額外啟動 image_server.py

最終架構

┌──────────────────────────────────────────────────────────────┐
│                     DGX Spark Server                         │
│  ┌─────────────────────────────────────────────────────────┐ │
│  │  Docker: tg-vllm (port 8603)                            │ │
│  │  └── vLLM + FP8 量化                                    │ │
│  │  └── 文字翻譯：0.3-1 秒                                 │ │
│  └─────────────────────────────────────────────────────────┘ │
│  ┌─────────────────────────────────────────────────────────┐ │
│  │  Python: image_server.py (port 8604)                    │ │
│  │  └── transformers pipeline                              │ │
│  │  └── 圖片翻譯 (OCR + 翻譯)                              │ │
│  └─────────────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────────────┘
                           ↑
                      HTTP Request
                           ↑
┌──────────────────────────────────────────────────────────────┐
│  Client (Mac/PC)                                             │
│  └── Gradio UI (port 7861)                                   │
│      ├── 📝 文字翻譯 Tab → port 8603                         │
│      └── 🖼️ 圖片翻譯 Tab → port 8604                         │
└──────────────────────────────────────────────────────────────┘

關鍵優化點

1. vLLM FP8 量化: --quantization fp8
2. 動態 max_tokens: min(max(len(text) * 2, 20), 500)
3. 圖片自動縮放: 896x896 以內
4. 分離服務: 文字用 vLLM（快），圖片用 transformers（支援特殊格式）

快速啟動指令

# 1. 啟動文字翻譯 (vLLM)
docker start tg-vllm  # 或用 docker run 指令

# 2. 啟動圖片翻譯 (transformers)
cd ~/apps/translategemma
nohup python image_server.py > image_server.log 2>&1 &

# 3. 本地啟動 Gradio
python gradio_app.py