使用 MLX 微調 Meta 的 Llama 3.1 模型

Published: 約 12 分鐘

緣起

之前試過「提示工程(Prompt engineering)」來應用「大型語言模型(LLM)」，如 Llama 3, Mistral 等等。會開始玩大語言模型，主要是想建立一個對話式代理人(conversational agent)來協助使用者以口語方式描述部署的需求。 雖然結果還是會出現常見的幻覺(hallucination)問題，整體來說結果還是蠻有希望的。由於此代理人需要根據提供的 YAML 文件樣板來協助使用者產生部署的描述檔(manifests)，但是隨著樣板數量 的增加，就發現前後文(context)不夠長，以及處理時間變長的問題。試過 RAG(Retrieval Augumented Generation)， 但用在 YAML 樣板的效果不好。不管怎樣，趁著 Meta 出 Llama 3.1 新模型之際來試試微調的效果如何。

本文需要搭配GitHub 存儲庫 ⎘，其中山姆鍋也提供一個腳本以方便執行對應的步驟。

工具介紹

MLX-LM

山姆鍋目前公司都是以 Mac 系統作為工作機。加上，山姆鍋本人買了一台 Mac Studio M2 專門來玩大語言模型 。 所以，微調工具要以能夠在 Mac 系統上運行為主。Apple 自從 M1(Apple Silicon) 之後，就內建 GPU 以及神經網路專用的晶片。 另一方面，Apple Silicon 的統一記憶體模型，讓 CPU 以及 GPU 可以共用同一個記憶體空間，除了節省資料傳輸時間，也提供更多的記憶體空間來運行更大型的語言模型。

MLX 是 Apple 推出在自家 Apple Silicon 晶片上進行神經網路訓練與推體的框架，使用上跟 PyTorch 以及 Tenssorflow 類似。 但在本文撰寫時，MLX 能夠發揮 Apple Silicon 最大的效能。所以，山姆鍋選擇使用 MLX 來作為微調工具。 MLX 最大的缺點應該就是只能在 Mac 系統運行以及相關文件與範例不多，但這也是山姆鍋為何特別寫此文紀錄的原因！

本文使用的 mlx-lm 為 0.16.1。

Ollama

Ollama 是個讓多種語言模型可以使用者本機運行的工具，提供 OpenAI 相容的 API 服務，以方便整合現有的應用。Ollama 使用 llama.cpp 作為運行模型的引擎，除支援 CPU 之外，也可運用多種的硬體加速機制，如 GPU, NPU 或者 Apple 的 Metal等來進行推論(inference)。

本文使用的 Ollama 版本為 0.3.2。

Huggingace CLI

此命令列用來從 Hugging Face 網站下載模型。如果有使用 Homebrew, 可以使用下列指令安裝：

brew install huggingface-cli

環境準備

將配合本文的 GitHub 存儲庫複製到本地:

git clone --recurse-submodules https://github.com/sampot/mlx-llama-finetune.git

這個動作會聯同 llama.cpp 這個子模組一起複製。

安裝 Python 相依套件

./mlx-ft.sh install

除了安裝此存儲庫本身需要的相依套件外，也會安裝 llama.cpp 的。因為，需要用到 llama.cpp 的 Python 腳本來將模型轉換成 GGUF 格式。

撰寫本文時，Ollama 工具使用的 llama.cpp 與 llama.cpp 存庫庫版本不相同。遇到 llama.cpp 轉換的 GGUF 格式模型無法由 Ollama 載入的問題。所以，可能要使用舊版本的 llama.cpp。

如果需要使用特定版本的 llama.cpp, 使用下列指令(本文山姆鍋使用 b3418 這個版本):

cd llama.cpp
git checkout b3418

準備資料集

最終山姆鍋會使用自家的資料來進行微調，作為示範，本文使用 gretelai/synthetic_text_to_sql” ⎘這個開放資料集。

使用下列指令來下載以及轉換訓練用資料集：

./mlx-ft.sh data

這個指令背後的資料轉換實際由 prepare_data.py 這個腳本完成。

為了讓資料集可以被 MLX 使用，需要將資料集按照 MLX 期望的格式存放。MLX 支援下列三種格式：

  1. Completion：這種格式中，每筆資料包含 prompt 和 completion 這兩個鍵值，分別代表使用者的提示(prompt) 和預期模型的回覆。MLX 在內部會將這種格式的 JSON 物件轉換成 chat 形式資料。
{
  "prompt": "What is the capital of France?",
  "completion": "Paris."
}
  1. Chat：每個 JSON 物件包含一個 messages 鍵值，其內容是一個包含訊息的列表。MLX 會使用 Tokenizer 的 apply_chat_template 函式將聊天訊息轉換成訓練用的文本(text)資料。 以下是範例：
{
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful assistant."
    },
    {
      "role": "user",
      "content": "Hello."
    },
    {
      "role": "assistant",
      "content": "How can I assistant you today."
    }
  ]
}
  1. Text：這個格式要求每個 JSON 物件內有一個 text 鍵，其內容已經根據模型所需要的樣板格式安排。這個用在 Hugging Face 沒有尚未提供模型所需的 Chat Template 時自行根據模型要求預先格式化好訓練用文本。
{
  "text": "This is an example for the model."
}

這三種格式都需要在資料目錄有下列檔案:

  • train.jsonl：主要的訓練資料集。
  • valid.jsonl：用於訓練過程驗證結果的資料集。
  • test.jsonl：用於評估測試的資料集，訓練階段不使用。

訓練集檔案內容採用 JSONL 格式，也就是說：檔案內每一列都是完整的 JSON 文件，列跟列之間以換行(new line)字元分隔。 由於山姆鍋之後的模型是用於對話式代理人。所以，本文使用聊天訊息來作為訓練資料的樣板格式。

下載模型

要從 Hugging Face 下載 meta-llama/Meta-Llama-3.1-8B-Instruct這個模型，需要：

  • 有 Hugging Face 網站帳戶，並在模型頁面進行存取申請(i.e.同意使用條款)。
  • 在終端機使用 huggingface-cli login 進行登入。

本文使用到 huggingface_hub 這個套件來下載模型資料，需要完成上述步驟才能進行。

./mlx-ft.sh fetch

開始訓練

./mlx-ft.sh train

上面腳本主要是執行 mlx_lm.lora 這個指令。

mlx_lm.lora --config lora_config.yaml

MLX LORA 微調相關的參數都設定在 lora_config.yaml 這個檔案。

# The path to the local model directory or Hugging Face repo.
model: "meta-llama/Meta-Llama-3.1-8B-Instruct"
# Whether or not to train (boolean)
train: true

# Directory with {train, valid, test}.jsonl files
data: "data"

# The PRNG seed
seed: 0

# Number of layers to fine-tune
lora_layers: 16

# Minibatch size.
batch_size: 4

# Iterations to train for.
iters: 100

# Number of validation batches, -1 uses the entire validation set.
val_batches: 25

# Adam learning rate.
learning_rate: 1e-6

# Number of training steps between loss reporting.
steps_per_report: 10

# Number of training steps between validations.
steps_per_eval: 50

# Load path to resume training with the given adapter weights.
resume_adapter_file: null

# Save/load path for the trained adapter weights.
adapter_path: "adapters"

# Save the model every N iterations.
save_every: 20

# Evaluate on the test set after training
test: false
# Number of test set batches, -1 uses the entire test set.
test_batches: 100

# Maximum sequence length.
max_seq_length: 8192

# Use gradient checkpointing to reduce memory use.
grad_checkpoint: true

# LoRA parameters can only be specified in a config file
lora_parameters:
  # The layer keys to apply LoRA to.
  # These will be applied for the last lora_layers
  keys: ["self_attn.q_proj", "self_attn.v_proj"]
  #keys:
  #  [
  #    "mlp.gate_proj",
  #    "mlp.down_proj",
  #    "self_attn.q_proj",
  #    "mlp.up_proj",
  #    "self_attn.o_proj",
  #    "self_attn.v_proj",
  #    "self_attn.k_proj",
  #  ]
  rank: 16
  alpha: 32
  scale: 10.0
  dropout: 0.0

合併模型參數

./mlx-ft.sh fuse

背後執行下列指令：

  poetry run mlx_lm.fuse \
      --model ${local_model} \
      --save-path "${fused_model}" \
      --adapter-path ./adapters \
      --de-quantize

其中，local_model 是本文設定的模型名稱，fused_model 是要存放合併後模型的路徑名稱。這個指令同時會去除權重參數的量化， 因為後續轉換成 GGUF 格式的工具不支援量化後模型。

建立 Ollama 可用的模型

微調後的模型需要轉換成 GGUF 並載入 Ollama 才能使用。使用下列命令來建立 Ollama 的模型：

./mlx-ft.sh create

在建立 Ollama 模型之前，會執行 llama.cpp 中的 /convert_hf_to_gguf.py 的腳本來將合併(fused)的模型轉換成 GGUF 格式。

  poetry run python ./llama.cpp/convert_hf_to_gguf.py \
      --outfile ./models/model.gguf \
      --outtype q8_0 \
      --model-name ${model_name} \
      ${fused_model}

其中，--outtype 參數指定採用 Q8_0 的量化格式; ${fused_model} 指定已經合併的模型路徑。

運行微調後的模型

在終端機執行下列指令，會出現文字聊天介面，就可以開始測試了。

./mlx-ft.sh run

背後其實只是執行 ollama run <模型名稱>。

小結

本文的重點只是說明如何設定環境以及運用工具來進行模型微調，由於訓練資料的數量、訓練的時間以及 LORA 參數等諸多的因素，都可能影響微調後的效果。所以，真正的微調需要多次的嘗試與實驗才會接近期望的結果。

參考資料

郭信義 (Sam Kuo)

奔騰網路科技技術長，專長分散式系統、Web 應用與雲端服務架構、設計、開發、部署與維運。工作之餘，喜歡關注自由軟體的發展與應用，偶爾寫一下部落格文章。