参数高效的GPT-2微调与LoRA

!pip install -q --upgrade keras-nlp
!pip install -q --upgrade keras  # 升级到Keras 3.

import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import keras_nlp
import keras
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_datasets as tfds
import time

keras.mixed_precision.set_global_policy("mixed_float16")

# 一般超参数
BATCH_SIZE = 32
NUM_BATCHES = 500
EPOCHS = 1  # 可以设置为更高的值以获得更好的结果
MAX_SEQUENCE_LENGTH = 128
MAX_GENERATION_LENGTH = 200

GPT2_PRESET = "gpt2_base_en"

# LoRA特定的超参数
RANK = 4
ALPHA = 32.0

reddit_ds = tfds.load("reddit_tifu", split="train", as_supervised=True)

for document, title in reddit_ds:
    print(document.numpy())
    print(title.numpy())
    break

train_ds = (
    reddit_ds.map(lambda document, _: document)
    .batch(BATCH_SIZE)
    .cache()
    .prefetch(tf.data.AUTOTUNE)
)
train_ds = train_ds.take(NUM_BATCHES)

class GPUMemoryCallback(keras.callbacks.Callback):
    def __init__(
        self,
        target_batches,
        print_stats=False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.target_batches = target_batches
        self.print_stats = print_stats

        self.memory_usage = []
        self.labels = []

    def _compute_memory_usage(self):
        memory_stats = tf.config.experimental.get_memory_info("GPU:0")
        # 转换字节为GB并存储在列表中。
        peak_usage = round(memory_stats["peak"] / (2**30), 3)
        self.memory_usage.append(peak_usage)

    def on_epoch_begin(self, epoch, logs=None):
        self._compute_memory_usage()
        self.labels.append(f"epoch {epoch} start")

    def on_train_batch_begin(self, batch, logs=None):
        if batch in self.target_batches:
            self._compute_memory_usage()
            self.labels.append(f"batch {batch}")

    def on_epoch_end(self, epoch, logs=None):
        self._compute_memory_usage()
        self.labels.append(f"epoch {epoch} end")

def generate_text(model, input_text, max_length=200):
    start = time.time()

    output = model.generate(input_text, max_length=max_length)
    print("\nOutput:")
    print(output)

    end = time.time()
    print(f"总耗时: {end - start:.2f}s")

def get_optimizer_and_loss():
    optimizer = keras.optimizers.AdamW(
        learning_rate=5e-5,
        weight_decay=0.01,
        epsilon=1e-6,
        global_clipnorm=1.0,  # 梯度裁剪。
    )
    # 将layernorm和偏差项排除在权重衰减之外。
    optimizer.exclude_from_weight_decay(var_names=["bias"])
    optimizer.exclude_from_weight_decay(var_names=["gamma"])
    optimizer.exclude_from_weight_decay(var_names=["beta"])

    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    return optimizer, loss

preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=MAX_SEQUENCE_LENGTH,
)
gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
    "gpt2_base_en", preprocessor=preprocessor
)

gpt2_lm.summary()

gpu_memory_callback = GPUMemoryCallback(
    target_batches=[5, 10, 25, 50, 100, 150, 200, 300, 400, 500],
    print_stats=True,
)

optimizer, loss = get_optimizer_and_loss()

gpt2_lm.compile(
    optimizer=optimizer,
    loss=loss,
    weighted_metrics=["accuracy"],
)

gpt2_lm.fit(train_ds, epochs=EPOCHS, callbacks=[gpu_memory_callback])
gpt2_lm_memory_usage = gpu_memory_callback.memory_usage

generate_text(gpt2_lm, "我喜欢篮球", max_length=MAX_GENERATION_LENGTH)
generate_text(gpt2_lm, "那家意大利餐厅是", max_length=MAX_GENERATION_LENGTH)

    kwargs.pop("name", None)

    super().__init__(name=name, trainable=trainable, **kwargs)

    self.rank = rank
    self.alpha = alpha

    self._scale = alpha / rank

    self._num_heads = original_layer_config["output_shape"][-2]
    self._hidden_dim = self._num_heads * original_layer_config["output_shape"][-1]

    # 层。

    # 原始密集层。
    self.original_layer = original_layer
    # 无论我们是在训练模型还是在推理模式下，
    # 这个层都应该被冻结。
    self.original_layer.trainable = False

    # LoRA 密集层。
    self.A = keras.layers.Dense(
        units=rank,
        use_bias=False,
        # 注意：原始论文提到使用
        # 正态分布进行初始化。然而，官方的 LoRA 实现
        # 使用“凯明/何初始化”。
        kernel_initializer=keras.initializers.VarianceScaling(
            scale=math.sqrt(5), mode="fan_in", distribution="uniform"
        ),
        trainable=trainable,
        name=f"lora_A",
    )
    # B 与原始层具有相同的 `equation` 和 `output_shape`。
    # `equation = abc,cde->abde`，其中 `a`：批大小，`b`：序列
    # 长度，`c`：`hidden_dim`，`d`：`num_heads`，
    # `e`：`hidden_dim//num_heads`。唯一的区别在于在层 `B` 中，
    # `c` 代表 `rank`。
    self.B = keras.layers.EinsumDense(
        equation=original_layer_config["equation"],
        output_shape=original_layer_config["output_shape"],
        kernel_initializer="zeros",
        trainable=trainable,
        name=f"lora_B",
    )

def call(self, inputs):
    original_output = self.original_layer(inputs)
    if self.trainable:
        # 如果我们正在微调模型，我们将把 LoRA 层的输出
        # 添加到原始层的输出中。
        lora_output = self.B(self.A(inputs)) * self._scale
        return original_output + lora_output

    # 如果我们处于推理模式，我们将“合并” LoRA 层的权重到
    # 原始层的权重中 - 更多内容将在文本生成
    # 部分中讨论！
    return original_output

### 注入 LoRA 层到模型中

我们现在将黑客攻击原始的 GPT-2 模型，并将 LoRA 层注入其中。让我们先做几件事：

- 删除之前的模型；
- 使用 [`tf.config.experimental.reset_memory_stats`](https://www.tensorflow.org/api_docs/python/tf/config/experimental/reset_memory_stats) 重置“峰值” GPU 内存使用；
- 加载一个新的 GPT-2 模型。


```python
del gpt2_lm
del optimizer
del loss

# 这将“重置”内存的“峰值”使用为“当前”内存使用。
tf.config.experimental.reset_memory_stats("GPU:0")

# 加载原始模型。
preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=128,
)
lora_model = keras_nlp.models.GPT2CausalLM.from_preset(
    "gpt2_base_en",
    preprocessor=preprocessor,
)

for layer_idx in range(lora_model.backbone.num_layers):
    # 更改查询密集层。
    decoder_layer = lora_model.backbone.get_layer(f"transformer_layer_{layer_idx}")
    self_attention_layer = decoder_layer._self_attention_layer
    # 允许对 Keras 层状态进行更改。
    self_attention_layer._tracker.locked = False

    # 更改查询密集层。
    self_attention_layer._query_dense = LoraLayer(
        self_attention_layer._query_dense,
        rank=RANK,
        alpha=ALPHA,
        trainable=True,
    )

    # 更改值密集层。
    self_attention_layer._value_dense = LoraLayer(
        self_attention_layer._value_dense,
        rank=RANK,
        alpha=ALPHA,
        trainable=True,
    )

lora_model(preprocessor(["LoRA is very useful for quick LLM finetuning"])[0])
pass

for layer in lora_model._flatten_layers():
    lst_of_sublayers = list(layer._flatten_layers())

    if len(lst_of_sublayers) == 1:  # "leaves of the model"
        if layer.name in ["lora_A", "lora_B"]:
            layer.trainable = True
        else:
            layer.trainable = False

lora_model.summary()  # 输出模型摘要

gpu_memory_callback = GPUMemoryCallback(
    target_batches=[5, 10, 25, 50, 100, 150, 200, 300, 400, 500],
    print_stats=True,
)

optimizer, loss = get_optimizer_and_loss()

lora_model.compile(
    optimizer=optimizer,
    loss=loss,
    weighted_metrics=["accuracy"],
)

lora_model.fit(
    train_ds,
    epochs=EPOCHS,
    callbacks=[gpu_memory_callback],
)
lora_model_memory_usage = gpu_memory_callback.memory_usage

plt.bar(
    ["GPT-2", "LoRA GPT-2"],
    [max(gpt2_lm_memory_usage), max(lora_model_memory_usage)],
    color=["red", "blue"],
)

plt.xlabel("时间")
plt.ylabel("GPU 内存使用 (GB)")

plt.title("GPU 内存使用比较")
plt.legend()
plt.show()

for layer_idx in range(lora_model.backbone.num_layers):
    self_attention_layer = lora_model.backbone.get_layer(
        f"transformer_layer_{layer_idx}"
    )._self_attention_layer

    # 合并查询密集层。
    query_lora_layer = self_attention_layer._query_dense

    A_weights = query_lora_layer.A.kernel  # (768, 1) (a, b)
    B_weights = query_lora_layer.B.kernel  # (1, 12, 64) (b, c, d)
    increment_weights = tf.einsum("ab,bcd->acd", A_weights, B_weights) * (ALPHA / RANK)
    query_lora_layer.original_layer.kernel.assign_add(increment_weights)

    # 合并值密集层。
    value_lora_layer = self_attention_layer._value_dense

    A_weights = value_lora_layer.A.kernel  # (768, 1) (a, b)
    B_weights = value_lora_layer.B.kernel  # (1, 12, 64) (b, c, d)
    increment_weights = tf.einsum("ab,bcd->acd", A_weights, B_weights) * (ALPHA / RANK)
    value_lora_layer.original_layer.kernel.assign_add(increment_weights)

    # 用更新的权重放回原层
    self_attention_layer._query_dense = query_lora_layer.original_layer
    self_attention_layer._value_dense = value_lora_layer.original_layer

# 在生成时不需要冻结权重，因为没有权重被更新。
generate_text(lora_model, "我喜欢篮球", max_length=MAX_GENERATION_LENGTH)
generate_text(
    lora_model, "那家意大利餐厅是", max_length=MAX_GENERATION_LENGTH
)