針對7B模型的deepspeed_launcher.sh 啟動腳本

針對兩張 RTX 5090 訓練 7B 模型，這個啟動腳本會整合我們之前討論的所有技術（DeepSpeed ZeRO-2、Liger Kernel、Flash Attention 3），它可以一鍵啟動兩張顯卡進行分散式訓練。

1. 建立啟動腳本 deepspeed_launcher.sh

請將以下內容存為文件，並確保它與你的訓練腳本（例如 train.py）放在同一個目錄下。

</>Bash
#!/bin/bash

# --- 1. 硬體與環境設定 ---
export CUDA_VISIBLE_DEVICES=0,1        # 指定使用的 GPU 編號
export NCCL_P2P_DISABLE=0               # 5090 支援高效 P2P，確保開啟
export NCCL_IB_DISABLE=1                # 若非集群環境，禁用 InfiniBand 以避免錯誤
export PYTHONUNBUFFERED=1               # 確保 Log 即時輸出

# --- 2. 訓練參數設定 ---
MODEL_NAME_OR_PATH="yours/llama-7b-config" # 模型路徑或 HuggingFace ID
DATA_PATH="./train_data.bin"               # 我們之前建立的 mmap 二進制文件
OUTPUT_DIR="./output_7b_pretrain"          # 模型存檔路徑

# --- 3. 啟動命令 ---
# --num_gpus=2 代表使用兩張 5090
deepspeed --num_gpus=2 train.py \
    --deepspeed ds_config.json \
    --model_name_or_path $MODEL_NAME_OR_PATH \
    --data_path $DATA_PATH \
    --output_dir $OUTPUT_DIR \
    --max_steps 100000 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 8 \
    --learning_rate 2e-5 \
    --bf16 True \
    --tf32 True \
    --gradient_checkpointing True \
    --logging_steps 10 \
    --save_steps 500 \
    --lr_scheduler_type cosine

2. 配套的 ds_config.json (DeepSpeed 設定)

為了讓上面的腳本運作，你需要這個針對 32GB 顯存 優化的 ZeRO-2 配置文件：

</>JSON
{
  "fp16": { "enabled": false },
  "bf16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "gradient_clipping": 1.0,
  "steps_per_print": 10,
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "wall_clock_breakdown": false
}

3. 如何執行與權限設定

1. 賦予執行權限：

   </>Bash
   chmod +x deepspeed_launcher.sh

2. 啟動訓練：

   </>Bash
   ./deepspeed_launcher.sh

4. 針對 5090 的專業建議 (黑井架構特化)

• Micro Batch Size (MBS)： 雖然 5090 有 32GB 顯存，但預訓練時建議將 per_device_train_batch_size 設為 4 或 8，其餘靠 gradient_accumulation_steps 來補足。這樣可以留出足夠的顯存給 Flash Attention 3 處理長文本。
• 監控顯存： 訓練開始後，請另開視窗執行 watch -n 1 nvidia-smi。如果兩張卡的顯存佔用都在 28GB 左右，說明你的配置達到了最佳平衡。
• CPU 負載： 因為我們使用了 offload_optimizer（將優化器狀態放 RAM），請確保你的 CPU 核心數夠多，否則 CPU 計算梯度更新會變成瓶頸。

最後的技術檢查點

在啟動之前，請確認你的 train.py 裡面是否已經包含了這兩行，這是啟動 Liger Kernel 加速的關鍵：

</>Python
from liger_kernel.transformers import apply_liger_kernel_to_llama
apply_liger_kernel_to_llama()

準備開始第一次 Pre-train 測試。