Keras NLP GPT-2微调数据集兼容性问题：GPU环境fit报错

问题背景 #

在Google Colab的Python 3.11环境中微调GPT-2模型，安装keras-nlp后环境版本如下：

!pip install keras_nlp
print(tf.__version__)
print(keras.__version__)
print(keras_nlp.__version__)
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

输出结果：

2.16.1
3.3.3
0.11.1
1

该环境可识别GPU，但执行model.fit()时触发报错；而使用旧版本环境（tf2.12.1、keras2.12.0、keras-nlp0.5.0）代码可运行，但无法识别GPU。

最小复现代码 #

import numpy as np
import keras_nlp
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_text as tf_text
from tensorflow import keras
from tensorflow.lite.python import interpreter
import time
from google.colab import files
from google.colab import runtime

gpt2_tokenizer = keras_nlp.models.GPT2Tokenizer.from_preset("gpt2_base_en")
gpt2_preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
    "gpt2_base_en",
    sequence_length=512,
    add_end_token=True,
)
gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset("gpt2_base_en", preprocessor=gpt2_preprocessor)

# 创建用于复现错误的训练数据
training_list_manual = ['This is a sentence that I like', 'I went to school today.', 'I have a bike, you can ride it if you like.']
tf_train_ds = tf.data.Dataset.from_tensor_slices(training_list_manual)
processed_ds = tf_train_ds.map(gpt2_preprocessor, tf.data.AUTOTUNE).batch(64).cache().prefetch(tf.data.AUTOTUNE)

# 尝试微调
gpt2_lm.include_preprocessing = False

num_epochs = 1

lr = tf.keras.optimizers.schedules.PolynomialDecay(
    5e-5,
    decay_steps=part_of_ds.cardinality() * num_epochs,
    end_learning_rate=0.0,
)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
gpt2_lm.compile(
    optimizer=keras.optimizers.Adam(lr),
    loss=loss,
    weighted_metrics=["accuracy"])

gpt2_lm.fit(processed_ds, epochs=num_epochs)

报错信息 #

ValueError: Exception encountered when calling GPT2CausalLMPreprocessor.call().
Unsupported input for `x`. `x` should be a string, a list of strings, or a list of tensors. If passing multiple segments which should packed together, please convert your inputs to a list of tensors. Received `x={'token_ids': <tf.Tensor 'args_1:0' shape=(None, None, 1024) dtype=int32>, 'padding_mask': <tf.Tensor 'args_0:0' shape=(None, None, 1024) dtype=bool>}`

Arguments received by GPT2CausalLMPreprocessor.call():
  • x={'token_ids': 'tf.Tensor(shape=(None, None, 1024), dtype=int32)', 'padding_mask': 'tf.Tensor(shape=(None, None, 1024), dtype=bool)'}
  • y=tf.Tensor(shape=(None, None, 1024), dtype=int32)
  • sample_weight=tf.Tensor(shape=(None, None, 1024), dtype=bool)
  • sequence_length=None

预处理后单条数据格式 #

执行代码：

for example in tf_train_ds.take(1):
    print(gpt2_preprocessor(example))

输出结果：

({'token_ids': <tf.Tensor: shape=(1, 1024), dtype=int32, numpy=array([[50256,  1212,   318, ...,     0,     0,     0]], dtype=int32)>, 'padding_mask': <tf.Tensor: shape=(1, 1024), dtype=bool, numpy=array([[ True,  True,  True, ..., False, False, False]])>}, <tf.Tensor: shape=(1, 1024), dtype=int32, numpy=array([[1212,  318,  257, ...,    0,    0,    0]], dtype=int32)>, <tf.Tensor: shape=(1, 1024), dtype=bool, numpy=array([[ True,  True,  True, ..., False, False, False]])>)

解决方案 #

问题核心是**include_preprocessing = False**的使用逻辑与数据集格式不匹配：当设置该参数为False时，模型期望接收预处理后的token_ids和padding_mask字典作为输入，但当前数据集返回的是(输入字典, 标签, 样本权重)的三元组格式，导致预处理器被错误调用。

方案1：调整数据集格式适配模型 #

修改数据集的输出格式，将三元组转换为模型期望的(输入字典, 标签)结构：

# 预处理并调整数据集格式
processed_ds = tf_train_ds.map(
    lambda x: gpt2_preprocessor(x), 
    tf.data.AUTOTUNE
).batch(64).cache().prefetch(tf.data.AUTOTUNE)

# 转换为模型期望的(inputs_dict, labels)格式
processed_ds = processed_ds.map(
    lambda inputs, labels, sample_weight: (inputs, labels)
)

同时修正学习率调度中的未定义变量错误：

lr = tf.keras.optimizers.schedules.PolynomialDecay(
    5e-5,
    decay_steps=processed_ds.cardinality() * num_epochs,
    end_learning_rate=0.0,
)

方案2：简化流程，让模型自动处理预处理 #

新版本Keras-NLP支持直接传入原始文本数据，无需手动预处理，去掉include_preprocessing=False即可：

# 移除include_preprocessing=False，模型将自动使用内置预处理器
# gpt2_lm.include_preprocessing = False

# 直接传入原始文本数据集
gpt2_lm.fit(tf_train_ds.batch(64), epochs=num_epochs)

这种方式更简洁，避免格式兼容问题，同时能正常利用GPU加速。

内容的提问来源于stack exchange，提问作者Zhubarb