如何扩展Keras GPT2模型？混合专家（MoE）实现示例

最终更新 #

已成功实现功能，下方代码可正常运行，欢迎提出反馈或优化建议。

一些思考点：

• 由于当前全连接层已支持3D张量，门控网络无需使用TimeDistributed层，但不确定针对基础GPT2模型，搭配2、4个等不同数量专家时，门控网络的规模应设置为多大。
• 当前实现不会返回每个查询的选择结果，可能是该特性在最初实现时并未纳入。
• 之前遇到的诸多问题，除了代码bug外，还与内存不足有关。

更新2 #

在Colab中运行时出现报错：

Received a label value of 50256 which is outside the valid range of [0, 768).  Label values: 31373 11 703 389 345 30 50256 0 0 0 0... 

一开始困惑为何损失函数期望标签值在[0,768]范围内，而token ID的取值范围是0到词汇表最大值。后来发现问题根源：在GPT层中错误调用了backbone模型而非GPT2CausalLM，前者有其他用途。

更新1 #

核心问题：如何链式扩展Keras GPT模型，例如构建像MoE这样的更大模型？

可运行代码 #

import tensorflow as tf
import keras_nlp


def create_gating_network(sequence_length, num_experts, feature_dim=768):
    inputs = tf.keras.layers.Input(shape=(sequence_length, feature_dim))
    x = tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(64, activation="relu"))(
        inputs
    )
    outputs = tf.keras.layers.TimeDistributed(
        tf.keras.layers.Dense(num_experts, activation="softmax")
    )(x)
    gating_model = tf.keras.models.Model(inputs=inputs, outputs=outputs)
    return gating_model


def moe_function(args):
    expert_outputs, gating_coefficients = args
    weighted_experts = expert_outputs * gating_coefficients
    intermediate_sum = tf.reduce_sum(weighted_experts, axis=2)
    weighted_sum = tf.reduce_sum(intermediate_sum, axis=2)
    return weighted_sum


class ExpertGPT2Layer(tf.keras.layers.Layer):
    def __init__(self, name="gpt2_base_en", sequence_length=128, **kwargs):
        super(ExpertGPT2Layer, self).__init__(**kwargs)
        self.sequence_length = sequence_length
        self.preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
            name, sequence_length=sequence_length
        )
        self.gpt2_model = keras_nlp.models.GPT2CausalLM.from_preset(
            name,
            preprocessor=self.preprocessor,
        )

    def call(self, inputs, training=False):
        preprocess = self.preprocessor(inputs)
        outputs = self.gpt2_model(preprocess[0], training=True)
        return outputs


class CustomGPT2Model(tf.keras.Model):
    def __init__(
        self,
        gating_network,
        name="gpt2_base_en",
        sequence_length=128,
        feature_dim=768,
        num_experts=4,
        **kwargs
    ):
        super(CustomGPT2Model, self).__init__(**kwargs)
        self.sequence_length = sequence_length
        self.feature_dim = feature_dim
        self.num_experts = num_experts
        self.tokenizer = keras_nlp.models.GPT2Tokenizer.from_preset(name)
        self.preprocessor = keras_nlp.models.GPT2CausalLMPreprocessor.from_preset(
            name, sequence_length=sequence_length
        )
        self.expert_layers = [
            ExpertGPT2Layer(sequence_length=sequence_length, name=name)
            for _ in range(num_experts)
        ]
        self.gating_network = gating_network

    def apply_expert(self, expert, inputs, training):
        result = expert(inputs, training=training)
        return result

    def build(self, input_shape):
        inputs = tf.keras.layers.Input(
            shape=input_shape, dtype=tf.string, name="text-input"
        )

        # 预处理器返回x, y, w
        x, labels, w = self.preprocessor(inputs)
        time_dim_token_ids = tf.expand_dims(x["token_ids"], axis=-1)
        replicated_token_ids = tf.tile(time_dim_token_ids, [1, 1, self.feature_dim])

        # 计算专家模型输出
        expert_outputs = [
            self.apply_expert(expert, inputs, training=True)
            for expert in self.expert_layers
        ]
        stacked_expert_outputs = tf.stack(expert_outputs, axis=1)

        # 计算门控系数
        gating_coefficients = self.gating_network(replicated_token_ids)
        expanded_gating_coefficients = tf.expand_dims(
            tf.expand_dims(gating_coefficients, axis=-1), axis=-1
        )

        moe_output = moe_function(
            [stacked_expert_outputs, expanded_gating_coefficients]
        )
        self.model = tf.keras.Model(inputs=inputs, outputs=[moe_output, labels])
        super(CustomGPT2Model, self).build(input_shape)

    def call(self, inputs, training=False):
        return self.model(inputs, training)

    @tf.function
    def train_step(self, data):
        x = data

        with tf.GradientTape() as tape:
            y_pred, y_true = self.model(x, training=True)
            loss = self.compiled_loss(y_true, y_pred, regularization_losses=self.losses)

        gradients = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))

        self.compiled_metrics.update_state(y_true, y_pred)
        return {m.name: m.result() for m in self.metrics}


def main():
    text = ["hello, how are you?", "I am good"]
    batch_size = 1
    num_experts = 2
    sequence_length = 64

    dataset = tf.data.Dataset.from_tensor_slices(text)
    dataset = dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)

    gating_network = create_gating_network(sequence_length, num_experts)
    moe_model = CustomGPT2Model(
        gating_network, sequence_length=sequence_length, num_experts=num_experts
    )

    moe_model.compile(
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        optimizer=tf.keras.optimizers.Adam(2e-5),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
    )
    moe_model.build(input_shape=(1,))
    moe_model.summary()
    moe_model.fit(dataset, epochs=3, verbose=1)


if __name__ == "__main__":
    main()

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