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 的简短模型解释器，我在这里有介绍。
