Keras 中的 Llama 3.2
這將是有史以來最短的部落格文章。
問題：Llama 3.2 兩週前釋出在 Hugging Face / Transformers 上。它什麼時候能在 Keras 中使用？
答案：它從第一天起就一直在工作 😀。沒有什麼可等待的。
是的，Keras Llama3 可以從任何標準 (即 safetensors) Hugging Face 檢查點載入，包括 3.2 檢查點。如果需要轉換，它會即時進行。試試這個
!pip install keras_hub
from keras_hub import models.Llama3CausalLM
model = Llama3CausalLM.from_preset("hf://meta-llama/Llama-3.2-1B-Instruct", dtype="bfloat16")
model.generate("Hi there!")
這是一個 Colab 來試試看。盡情享受！🤗
好的，好的，有人告訴我，如果我想釋出一篇部落格文章，我必須填滿內容。以下是關於 Keras 的一些額外資訊。
Keras 是多後端框架
Keras 是經過時間考驗的 JAX、PyTorch 和 TensorFlow 建模庫。你可能在 演示 Colab 中注意到這行程式碼
import os
os.environ["KERAS_BACKEND"] = "jax" # or "torch", or "tensorflow"
它必須出現在 import keras 之前，並控制模型是在 JAX、PyTorch 還是 TensorFlow 上執行。這對於在 JAX 上使用 XLA 編譯來嘗試你喜歡的模型非常方便 🚀。
什麼是 keras-hub？
Keras 是一個建模庫，而 keras-hub 是其預訓練模型的集合。它以前被稱為 KerasNLP 和 KerasCV。正在進行重新命名。它包含所有流行的預訓練模型（Llama3、Gemma、StableDiffusion、Segment Anything 等），以及它們在 Keras 中的標準實現。
Keras 中的 LLM “開箱即用”
我的意思是，“內建分詞器”。model.generate() 直接在字串上工作
model.generate("Hi there!")
> "Hi there! I'm looking for information on how to ...
訓練也是如此。你可以直接在一組字串上訓練
model.fit(strings) # list or dataset of input strings
與 LLM 聊天
流行 LLM 的指令調整變體可用於輪流對話。在這裡，Llama-3.2-1B-Instruct 理解以下對話標籤（參見 meta docs）。
<|start_header_id|>system<|end_header_id|>You are a helpful assistant<|eot_id|>\n
\n
<|start_header_id|>user<|end_header_id|>Hello_<|eot_id|>\n
\n
<|start_header_id|>assistant<|end_header_id|>\n
\n
對話一旦以這種方式格式化，就可以直接輸入到 model.generate() 中。
為了方便起見，演示 Colab 實現了一個名為 ChatState 的輔助類，它自動完成必要的字串連線。
底層訪問：分詞器、骨幹網路
如果你不喜歡“內建電池”，想訪問底層分詞器和模型，它們很容易訪問
# tokenizer
model.preprocessor.tokenizer
# the model itself
model.backbone
# You can even load them separately from the same preset
backbone = keras_hub.models.Llama3CausalLM.from_preset("hf://meta-llama/Llama-3.2-1B-Instruct", dtype="float16")
tokenizer = keras_hub.models.Llama3Tokenizer.from_preset("hf://meta-llama/Llama-3.2-1B-Instruct")
等等，分詞器、預處理器？我有點困惑
分詞器只將文字轉換為整數向量。“Hello”在這裡轉換為一個標記
tokenizer("Hello")
> Array([9906], dtype=int32)
預處理器是一個包羅永珍的概念，用於執行模型所需的所有資料轉換。例如，對於涉及影像的任務，這可以是影像大小調整或增強，或者對於文字模型，可以是文字分詞。對於 CausalLM 任務，預處理器負責另外三個細節
- 新增模型預期的文字開始和文字結束標記
- 填充標記序列並生成掩碼
- 為訓練和微調生成“預期輸出”。對於 CausalLM 任務，這是輸入字串偏移一位的結果。
tokens = model.preprocessor("Hello")
tokens[0] # 128000 and 128009 are the start and end text tokens
> {'token_ids': Array([128000, 9906, 128009, 0, 0, 0], dtype=int32), 'padding_mask': Array([True, True, True, False, False, False], dtype=bool)}
tokens[1] # input sequence shifted by one
> [9906, 128009, 0, 0, 0, 0]
# feeding the model manually
model.backbone(model.preprocessor(["Hello", "Hi!"])[0]) # raw logits as output
> [[[ 0.9805 0.1664 0.625 ... -0.834 -0.264 0.05203]
...]]
# More typically you would use Keras built-in functions model.generate, model.fit, model.predict, model.evaluate
Keras 有內建的訓練器
只需在你的訓練資料集上呼叫 model.fit(ds)。這個訓練器與 Keras 中各種可用的分散式訓練、混合精度、量化或LoRA/QLoRA 選項相容。它也是完全可選的。如果你願意，可以編寫自定義訓練迴圈。
有關完整示例，請參見演示 Colab，我們將在其中微調 Llama 3.2，使其像海盜一樣說話
你可以上傳到 Hub
微調模型完成後，你可以直接從 Keras 上傳它
model.save_to_preset("./pirate-llama")
# Use your own repo here
keras_hub.upload_preset(
uri = "hf://martin-gorner/llama-3.2-1B-pirate-instruct",
preset = "./pirate-llama")
上傳的模型可以在這裡看到。
用於推理或訓練的分散式模型並行
本節演示 Colab：Llama 3.1 Keras 模型並行
你們中的一些人可能會想，既然已經可以使用 Hugging Face 上的 Transformers 來處理 LLM，為什麼還要使用 Keras 呢？答案是：即使你不關心 Keras 作為建模框架的靈活性和可用性（你應該關心！），Keras 也是你透過 JAX 及其強大的 XLA 編譯器實現高階模型並行的最快途徑。
讓我們選擇一個 8B 引數模型來演示：meta-llama/Llama-3.1-8B-Instruct（演示 Colab 在這裡）。在不量化的情況下，這個模型對於任何單個加速器來說都太大了。使用 Keras，你可以將其分片載入到多個加速器上，無論是 GPU 還是 TPU。如果你不確定“正確”的權重分片，大多數模型都提供了合理的預設值。在這裡，呼叫 keras_hub.models.Llama3Backbone.get_layout_map(device_mesh)
devices = keras.distribution.list_devices() # 8 TPU cores: let's do a 2x4 mesh
device_mesh = keras.distribution.DeviceMesh((2, 4), ["batch", "model"], devices)
layout_map = keras_hub.models.Llama3Backbone.get_layout_map(device_mesh) # defaults
distrib = keras.distribution.ModelParallel(layout_map=layout_map, batch_dim_name="batch")
keras.distribution.set_distribution(distrib)
# Now load the model, distributed across the accelerators
model = keras_hub.models.Llama3CausalLM.from_preset("hf://meta-llama/Llama-3.1-8B-Instruct")
如果你不相信模型提供的預設佈局對映，你可以定義自己的佈局對映。在這個執行在只有 8 個核心的“小型”TPU 設定上的例子中，以下佈局對映比預設的更快（每 epoch 54 秒，而不是 62 秒）
layout_map = keras.distribution.LayoutMap(device_mesh)
layout_map["token_embedding/embeddings"] = ("model", None)
layout_map["token_embedding/reverse_embeddings"] = ("model", None)
layout_map["self_attention.*(query|key|value).kernel"] = ("model", None, None)
layout_map["self_attention.*attention_output.kernel"] = ("model", None, None)
layout_map["feedforward_intermediate_dense.kernel"] = (None, "model")
layout_map["feedforward_gate_dense.kernel"] = (None, "model")
layout_map["feedforward_output_dense.kernel"] = ("model", None)
檢視這裡的演示 Colab，其中在 Google TPU v5e（可在Hugging Face Spaces 上的 JupyterLab 中獲取）上，在不到 8 分鐘內微調了更大的 8B Llama，使其說出海盜語。微調後的模型在這裡。如果你需要關於模型並行和 Keras 的簡短模型直譯器，我在這裡有介紹。
