笔者最近在做一些mllm相关的工作,书写此文对mllm框架进行梳理总结,定有不少纰漏,请读者立即指出,谢谢。mllm目前在做一些其他工作,这篇文章的书写时间为发布时间。在mllm的其他工作合并进主仓库后,本文还会进一步的跟进。读者请注意本文的时效性。


1. 简介

mllm是一款适用于移动设备和边缘设备的快速、轻量的多模态LLM推理引擎。

  • 完全的C/C++实现,无第三方依赖
  • 针对fuyu-8B等多模态LLM进行了优化
  • 支持ARM NEON和X86 AVX2向量指令
  • 支持4 bits和6 bits整数量化

本文将更多的以工程的视角来解析mllm框架,在行文过程中,本文会将mllm与其他框架的设计方法做对比。接下来,本文将会用项目组织结构、框架执行流程、自定义Op/Layer、Tokenizer和如何支持新模型五个章节来详细描述mllm框架的各项特性和总体结构。读者可以把该文章做mllm的使用文档。在最后,本文将会指出mllm的不足之处和可以尝试跟进的工作


在开始正式解析mllm之前,读者可以先clone下mllm的代码库,以便于跟进分析流程。mllm不依赖于git submodule,项目配置起来很方便,目前mllm可以在linux上使用Clang/GCC编译器进行编译。目前mllm支持的目标设备体系结构是X86和Arm。

git clone https://github.com/UbiquitousLearning/mllm

mllm团队将所有LLM相关的vocab文件都放在了git仓库中(这个其实可以移动到HuggingFace的仓库上),LLM量化后的模型文件都存储在HuggingFace上,读者可以在https://huggingface.co/mllmTeam上找到mllm提供的模型文件。

2. 框架执行流程

2.1 以两层Linear层运行为例

首先,考虑下面的代码,定义了两个Linear Layers,并且输入$X$通过两个Linear Layers来得到输出:

class TwoLinear final : public Module {
public:
    TwoLinear() = default;
    TwoLinear() {
        linear1 = Linear(in_f, out_f, /*bias*/true, "linear1");
        linear2 = Linear(out_f, out_f, /*bias*/true, "linear2");
    }

    std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
        x = inputs[0];
        x = linear1(x);
        x = linear2(x);
        return x;
    }

private:
    Layer linear1;
    Layer linear2;
}

TwoLinear tl;

2.1.1 加载参数

读者可以使用 tl.load(path)来加载参数。那么mllm是如何实现参数加载的呢?在load函数中,mllm会创建一个ParamLoader,这个ParamLoader是Static的,在全局可以访问。然后mllm会设置另一个全局参数doLoad为True,进而进入推理流程operator()(tmps, tmpt);。在推理流程中,要是执行层发现doLoad为True,那么就执行每个算子内定义好的load指令,而不是执行每个算子的原本逻辑。 load的执行在Layer.hpp文件的INIT_OP()中。

2.1.2Module的Operator()是如何调用Forward函数的?

对于常见的INPUT_TENSOR类型的Tensor,mllm首先会设置这个Tensor的类型为TENSOR_STATIC_INIT,进行一遍Forward推理;第一遍Forward推理完毕以后再把Tensor的类型设为TENSOR_STATIC_READY,然后进行第二遍Forward推理。

if (inputs[0].ttype() == TensorType::INPUT_TENSOR) {
    for (auto &input : inputs) {
        input.setTtype(TensorType::NORMAL_TENSOR);
        input.status() = TENSOR_STATIC_INIT;
        if(input.batch() == 0){
            Tensor::gph_[input.name()] = input;
        }
    }
    tensor_status = TENSOR_STATIC_INIT;

    Forward(inputs, anyArgs);
    for (auto &input : inputs) {
        input.status() = TENSOR_STATIC_READY;
    }
    tensor_status = TENSOR_STATIC_READY;

    return Forward(inputs, anyArgs);
}

第一次Forward推理的目的是调用Op定义的Reshape和SetUp函数,Reshape函数会推理出这一次模型推理的过程中每个Tensor的形状大小。SetUp函数会对Op需要输出的Tensor做内存的申请。 第二次Forward推理才是真正的计算。

2.1.3 Linear层的执行

每个Layer在实现的时候都会重载operator(),比如linear layer的operator()函数如下:

Tensor &operator()(Tensor &input) {
    return _1I1O_OP(input);
}

其中,_1I1O_OP表示的意思是,这是需要使用1个输入和1个输出的函数来处理这个算子。mllm还提供了许多类似于_1I1O_OP的函数来处理不同的算子。

2.2 总结

大体来说,mllm使用了类似于状态机的参数来设置了当前推理过程的运行状态。每一次都是通过Forward函数来进行全模型的遍历,在Op的执行过程中,用这些设定的参数来区分每次Op需要表现的行为。

3. 如何编写Op与自定义Layer

3.1 新增对应Backend的Op文件

mllm提供了src/backends/new_op.py实用工具来帮助创建Op Class。该文件会帮助读者创建下述基本函数:

ErrorCode reshape(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;
ErrorCode execute(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;
ErrorCode load(AbstructLoader &loader) override;
ErrorCode free(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;
ErrorCode setUp(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;

3.2 Op参数自定义

比如对于CPU上的LinearOp,需要in_featuresout_featureshas_bias三个参数。那么可以在3.1自动生成的class中加入:

class CPULinear final : public Op {
...
private:
    int in_features_;
    int out_features_;
    bool support_bias_;
    int thread_count = 4;
    Tensor weight_;
    Tensor bias_;
};

在CPULinearCreator中加入:

class CPULinearCreator : public CPUBackend::Creator {
public:
    virtual Op *create(OpParam op_param, Backend *bn, string name, int threadCount) const {
        int in_features = op_param["in_features"];
        int out_features = op_param["out_features"];
        int bias = op_param["bias"];
        return new CPULinear(bn, name, in_features, out_features, (bool)bias, threadCount);
    }
};

请注意,OpParam是一个string-float map。

3.3 重载函数

读者需要自行实现reshape,execute,load,free函数,视情况重载setUp函数。 以Linear Op为例,reshape函数就会通过in_features_变量来检查输入的Tensor的维度是否正确,然后对output Tensor做outputs[0]->reshape(inputs[0]->batch(), inputs[0]->head(), inputs[0]->sequence(), out_features_)

在load函数中,实现Weight和Bias的加载。

在execute函数中,具体实现矩阵乘法等计算操作。

在free函数中释放Weight和Bias。

3.4 Op是如何被注册和创建的?

在定义完成Op后,读者还需要把该Op注册到相应的Backend中,以及将Op抽象成Layer。

3.4.1 在Backend中注册Op

以CPU Backend为例,读者需要再CPUBackend文件中加入addCreator(LINEAR, (CPUBackend::Creator *)(new CPULinearCreator()));

如果这是一个新的算子,读者还需要在OpDefined文件中加入新Op的Enum项。

3.4.2 在Layer.hpp中加入对应的Op Layer

如Linear Layer:

class Linear final : public Layer {
public:
    explicit Linear(int in_features, int out_features, bool bias, std::string name) {
        param_["in_features"] = in_features;
        param_["out_features"] = out_features;
        param_["bias"] = (float)bias;
        init(std::move(name), OpType::LINEAR);
    }
    Tensor &operator()(Tensor &input) {
        return _1I1O_OP(input);
    }
};

其中,在构造函数中的**init()**函数并没有创建这个Linear算子。它只是负责给这个Linear指派了Backend。 真正的算子创建还是在INIT_OP()函数中。在这个函数中,它会通过backend_->opCreate(param_, name_);来创建算子。

4. Tokenizer

mllm提供了基础的Tokenizer支持,目前支持BPE和Unigram两种分词算法。

5. 如何对新模型进行支持

在mllm中,对模型组件(model、Tokenizer、Configuration)的定义和HuggingFace Transformer库中的定义方法基本一致。以支持QWen0.5B模型为例,需要编写三个文件:

configuration_qwen.cpp
modeling_qwen.cpp
tokenization_qwen.cpp

其中configuration_qwen.cpp定义了Qwen LLM的各类参数,如Head数量,hidden dim等。modeling_qwen.cpp定义了Qwen LLM网络。tokenization_qwen.cpp包含了将句子转化为Token的预处理行为。

5.1 生成mllm支持的vocab和模型参数

5.1.1 模型转换

使用mllm提供的Converter实用工具来进行转换:

cd tools/convertor
pip install -r ./requirements.txt

# for one file pytorch model
python convert.py --input_model=model.pth --output_model=model.mllm --type=torch

# for multi-file pytorch model
python convert.py --input_model=pytorch_model.bin.index.json --output_model=model.mllm --type=torch

# for one file safetensor model
python convert.py --input_model=model.bin --output_model=model.mllm --type=safetensor

# for multi-file safetensor model
python convert.py --input_model=model.safetensors.index.json --output_model=model.mllm --type=safetensor

5.1.2 Vocab转换

使用mllm提供的Converter实用工具来进行转换:

cd tools/convertor
python vocab.py --input_file=tokenizer.json --output_file=vocab.mllm --type=Unigram

5.1.3 量化

mllm提供了量化工具,该工具支持4 bits和6 bits整数量化,你可以使用下述指令来对模型参数进行量化

cd bin
./quantize model.mllm model_q4_0.mllm Q4_K

5.2 Configuration

设置文件里面主要实现两个类,一个是QWenNameConfig,一个是QWenConfig,其中QWenNameConfig包含QWenConfig。在一个mllm模型参数文件中,模型参数是以key-value对的形式统一起来的。QWenNameConfig的目的就是给出每个参数的名称,以便于mllm框架索引到正确的模型参数。

class QWenNameConfig : public TransformerNameConfig {
public:
    /**
     * @brief QWen2 following the hugging face naming method
     *
     * @param type RoPEType
     */
    void init(RoPEType type = RoPEType::HFHUBROPE) {
        switch (type) {
        case RoPEType::HFHUBROPE: {
            blk_name = "model.layers.";
            _attn_base_name = "self_attn.";
            _ffn_base_name = "mlp.";
            _q_proj_name = "q_proj";
            _k_proj_name = "k_proj";
            _v_proj_name = "v_proj";
            _o_proj_name = "o_proj";
            _gate_proj_name = "gate_proj";
            _up_proj_name = "up_proj";
            _down_proj_name = "down_proj";
            _attn_norm_name = "input_layernorm";
            _ffn_norm_name = "post_attention_layernorm";
            token_embd_name = "model.embed_tokens";
            post_norm_name = "model.norm";
            lm_head_name = "lm_head";
            break;
        }
        ...
        }
    }
    std::string blk_name;
    std::string token_embd_name;
    std::string post_norm_name;
    std::string lm_head_name;
    std::string _gate_proj_name;
};

QWenConfig中则主要定义各层的超参数,如rope的theta值、中间层维度大小等,如下面的代码所示:

struct QWenConfig {
    explicit QWenConfig(int token_limit, string billions = "0.5B", RoPEType type = RoPEType::HFHUBROPE) :
        cache_limit(token_limit) {
        ...
    };

    float attention_dropout = 0.0;
    int bos_token_id = 151643;
    int eos_token_id = 151643;
    std::string hidden_act = "silu";
    int hidden_size = 1024;
    float initializer_range = 0.02;
    int intermediate_size = 2816;
    int max_position_embeddings = 32768;
    int max_window_layers = 21;
    std::string model_type = "qwen2";
    int num_attention_heads = 16;
    int num_hidden_layers = 24;
    int num_key_value_heads = 16;
    double rms_norm_eps = 1e-6;
    float rope_theta = 1000000.0;
    int sliding_window = 32768;
    int vocab_size = 151936;
    bool tie_embedding_words = false;

    int cache_limit;
    RoPEType RoPE_type = RoPEType::HFHUBROPE;
    QWenNameConfig names_config;
};

5.3 Tokenization

Tokenization是一个非常客制化的步骤,每个LLM的Tokenization方法都不尽相同。以QWen为例子,QWen使用了BBPE方法,那么读者在支持QWen模型的时候,就要给出实现了BBPE的Tokenizer。mllm内部已经实现一个BPE算法,读者可以复用该实现来实现自己的Tokenizer。

5.4 Modeling

使用mllm框架提供的算子来实现模型是非常简单和便利的,熟悉Pytorch的读者可以快速的上手mllm。本文在这里默认读者对llama/qwen/mistral等常见LLM的模型有着基本的了解。在下文中,本文以Attention模块为例来演示如何使用mllm来搭建模型。 首先,所有的class需要继承Module父类。Module父类提供了Forward函数,读者需要重载该函数来实现相应的计算流程。

class QWenAttention final : public Module ...

5.4.1 创建该Module需要使用的Layers

class QWenAttention final : public Module {
public
    QWenAttention() = default;
    QWenAttention(const QWenConfig &config, const QWenNameConfig &names, const string &base_name) {
        hidden_size = config.hidden_size;
        num_heads = config.num_attention_heads;
        head_dim = config.hidden_size / num_heads;
        num_key_value_heads = config.num_key_value_heads;
        num_key_value_groups = num_heads / num_key_value_heads;

        // init layers
        q_proj = Linear(hidden_size, num_heads * head_dim, true, base_name + names._q_proj_name);
        k_proj = Linear(hidden_size, num_key_value_heads * head_dim, true, base_name + names._k_proj_name);
        v_proj = Linear(hidden_size, num_key_value_heads * head_dim, true, base_name + names._v_proj_name);
        o_proj = Linear(num_heads * head_dim, hidden_size, false, base_name + names._o_proj_name);
        q_rope = RoPE(config.RoPE_type, config.rope_theta, config.max_position_embeddings, base_name + "q_rope");
        k_rope = RoPE(config.RoPE_type, config.rope_theta, config.max_position_embeddings, base_name + "k_rope");
        k_cache = KVCache(num_key_value_groups, config.cache_limit, base_name + "k_cache");
        v_cache = KVCache(num_key_value_groups, config.cache_limit, base_name + "v_cache");
        mask = Causalmask(base_name + "mask");
        softmax = Softmax(DIMENSION, base_name + "softmax");
    }

private:
    int hidden_size;
    int num_heads;
    int head_dim;
    int num_key_value_heads;
    int num_key_value_groups;
    Layer q_proj;
    Layer k_proj;
    Layer v_proj;
    Layer o_proj;
    Layer q_rope;
    Layer k_rope;
    Layer k_cache;
    Layer v_cache;
    Layer mask;
    Layer softmax;
}

细心的读者可能已经发现了,在QWenAttention的构造函数中,创建每个Layer的时候都在最后一个参数上传递了Layer名称(std::string type),这是因为mllm依赖于Layer的名称来寻找该Layer所需要的参数。

5.4.2 重载Forward前向推理函数

创建完了所有我们需要的Layers以后,就可以编写Forward函数来定义Attention模块的计算流程,Forward函数接收一个Tensor Array和一个std::any Array,返回Tensor Array:

std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
    auto query_states = q_proj(inputs[0]);
    auto key_states = k_proj(inputs[1]);
    auto value_states = v_proj(inputs[2]);

    // [batch, heads, sequence, dims]
    query_states = query_states.view(-1, num_heads, -1, head_dim);
    key_states = key_states.view(-1, num_key_value_heads, -1, head_dim);
    value_states = value_states.view(-1, num_key_value_heads, -1, head_dim);

    // embedding
    query_states = q_rope(query_states);
    key_states = k_rope(key_states);

    // kv cache
    key_states = k_cache(key_states);
    value_states = v_cache(value_states);

    // attention weight
    auto atten_weight = Tensor::mm(query_states, key_states.transpose(Chl::SEQUENCE, Chl::DIMENSION)) / std::sqrt(head_dim);
    atten_weight = mask(atten_weight);
    atten_weight = softmax(atten_weight);

    // attention output
    auto atten_output = Tensor::mm(atten_weight, value_states);
    atten_output = atten_output.view(-1, 1, -1, head_dim * num_heads);
    atten_output = o_proj(atten_output);
    return {atten_output};
}

5.5 运行

完整的Qwen模型定义代码可以在附录1中找到。读者可以像Torch一样调用定义好的模型:首先,创建模型:

QWenConfig config(tokens_limit, "0.5B", RoPEType::HFHUBROPE);
auto model = QWenForCausalLM(config);
model.load(model_path);

moduleclass重载了()operator,读者可以使用model({input_tensor})来进行推理。

6. mllm框架的不足

这里写的有点mean,本人专业知识浅薄,在学术上是依托答辩,对mllm的理解更是不到位,大家轻喷。

6.1 Benchmark

  1. 缺少算子的Benchmark

本文认为,mllm在实现的时候极力的避免使用第三方的库,因为mllm需要迁移到移动设备上,一些三方库可能不能正常工作。但是手工实现的Kernel还是需要一个Benchmark来和目标平台上提供的算子库来进行性能比较的。就mllm目前提供的MatMul Kernel来看,似乎缺少Pack优化和/micro Kernel的优化?

  1. 缺少prefill/decode的Benchmark

mllm的issues中也有人提到过这个问题。作为具有LLM推理能力的引擎,应当测一下这两个基本能力。

6.2 对于移动端LLM推理的特定优化

  1. KV Cache量化

IIRC,在OPPO的Transformer-Lite[2]中,用到了KV Cache量化的小技巧。这对移动设备有限的内存来说可能会更加友好,当然还需要考量量化带来的CPU负载问题。

  1. 动态形状推理/内存复用/KV Cache搬移优化

目前mllm是没有做内存复用的,可以考虑使用符号推理方法来做动态形状的支持进而便于求解下一轮的内存使用情况。或许可以考虑一下PageAttention[3]的Tensor管理方法或者[2]中的KV Cache规划方法来进一步减少内存的搬移。

  1. 异构算力

可以考虑把形状推理(CPU)和计算(GPU/NPU)并行执行起来。或者是6.2.4中提到的内容与计算并行起来。

  1. 对模型参数的Lazy Fetch和Pre Fetch

目前,mllm会把参数一次性的读入内存?考虑到移动设备的内存有限,可以在合适的时机提前从外存上预取而不是全数载入。

6.3 易用性

  1. 模型结构需要手动编写且无法保存

目前,mllm的模型结构还是需要在C++文件中进行显示的手动定义。或许可以考虑创建自己的计算图和算子描述方式,使用flatbuffers来存储计算图。

  1. 如果要很好的使用所有的算力,可能还是需要完善的计算图机制,这样便于优化分析。
  2. 尝试引入三方易用的库如icu等来弥补C++ utf-8处理能力的不足。

Ref:

[1] mllm, https://github.com/UbiquitousLearning/mllm

[2] transformer-lite, https://arxiv.org/abs/2403.20041

[3] PageAttention, https://arxiv.org/abs/2309.06180

A1. Qwen模型定义

#ifndef MODELING_QWEN_HPP
#define MODELING_QWEN_HPP

#include "Backend.hpp"
#include "Layer.hpp"
#include "Module.hpp"
#include "Tensor.hpp"
#include "configuration_qwen.hpp"
#include <cmath>
using namespace mllm;

// Copied from GemmaMLP with Gemma->Qwen and using silu
class QWenMLP final : public Module {
public:
    QWenMLP() = default;
    QWenMLP(int hidden_size, int intermediate_size, const QWenNameConfig &names, const std::string &base_name) {
        gate_proj = Linear(hidden_size, intermediate_size, false, base_name + names._gate_proj_name);
        silu = SiLU(base_name + "act");
        up_proj = Linear(hidden_size, intermediate_size, false, base_name + names._up_proj_name);
        down_proj = Linear(intermediate_size, hidden_size, false, base_name + names._down_proj_name);
    }

    std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
        auto x = gate_proj(inputs[0]);
        x = silu(x);
        auto y = up_proj(inputs[0]);
        x = x * y;
        x = down_proj(x);
        return {x};
    }

private:
    Layer gate_proj;
    Layer up_proj;
    Layer down_proj;

    Layer silu;
};

// Copied from GemmaAttention with Gemma->Qwen and using SWA
class QWenAttention final : public Module {
public:
    QWenAttention() = default;
    QWenAttention(const QWenConfig &config, const QWenNameConfig &names, const string &base_name) {
        hidden_size = config.hidden_size;
        num_heads = config.num_attention_heads;
        head_dim = config.hidden_size / num_heads;
        num_key_value_heads = config.num_key_value_heads;
        num_key_value_groups = num_heads / num_key_value_heads;

        // init layers
        q_proj = Linear(hidden_size, num_heads * head_dim, true, base_name + names._q_proj_name);
        k_proj = Linear(hidden_size, num_key_value_heads * head_dim, true, base_name + names._k_proj_name);
        v_proj = Linear(hidden_size, num_key_value_heads * head_dim, true, base_name + names._v_proj_name);
        o_proj = Linear(num_heads * head_dim, hidden_size, false, base_name + names._o_proj_name);
        q_rope = RoPE(config.RoPE_type, config.rope_theta, config.max_position_embeddings, base_name + "q_rope");
        k_rope = RoPE(config.RoPE_type, config.rope_theta, config.max_position_embeddings, base_name + "k_rope");
        k_cache = KVCache(num_key_value_groups, config.cache_limit, base_name + "k_cache");
        v_cache = KVCache(num_key_value_groups, config.cache_limit, base_name + "v_cache");
        // mask = SlidingWindowMask(config.sliding_window, base_name + "mask");
        mask = Causalmask(base_name + "mask");
        softmax = Softmax(DIMENSION, base_name + "softmax");
    }

    std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
        auto query_states = q_proj(inputs[0]);
        auto key_states = k_proj(inputs[1]);
        auto value_states = v_proj(inputs[2]);

        // [batch, heads, sequence, dims]
        query_states = query_states.view(-1, num_heads, -1, head_dim);
        key_states = key_states.view(-1, num_key_value_heads, -1, head_dim);
        value_states = value_states.view(-1, num_key_value_heads, -1, head_dim);

        // embedding
        query_states = q_rope(query_states);
        key_states = k_rope(key_states);

        // kv cache
        key_states = k_cache(key_states);
        value_states = v_cache(value_states);

        // attention weight
        auto atten_weight = Tensor::mm(query_states, key_states.transpose(Chl::SEQUENCE, Chl::DIMENSION)) / std::sqrt(head_dim);
        atten_weight = mask(atten_weight);
        atten_weight = softmax(atten_weight);

        // attention output
        auto atten_output = Tensor::mm(atten_weight, value_states);
        atten_output = atten_output.view(-1, 1, -1, head_dim * num_heads);
        atten_output = o_proj(atten_output);
        return {atten_output};
    }

private:
    int hidden_size;
    int num_heads;
    int head_dim;
    int num_key_value_heads;
    int num_key_value_groups;
    Layer q_proj;
    Layer k_proj;
    Layer v_proj;
    Layer o_proj;
    Layer q_rope;
    Layer k_rope;
    Layer k_cache;
    Layer v_cache;
    Layer mask;
    Layer softmax;
};

// Copied from GemmaDecoder with Gemma->Qwen and set RmsNorm(without add_unit_offset)
class QWenDecoder final : public Module {
public:
    QWenDecoder() = default;
    QWenDecoder(const QWenConfig &config, const QWenNameConfig &names, const string &base_name) {
        self_atten = QWenAttention(config, names, base_name + names._attn_base_name);
        mlp = QWenMLP(config.hidden_size, config.intermediate_size, names, base_name + names._ffn_base_name);
        input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps, base_name + names._attn_norm_name);
        post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps, base_name + names._ffn_norm_name);
    }

    std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
        auto x = input_layernorm(inputs[0]);
        x = self_atten({x, x, x})[0];
        auto tmp = x + inputs[0];
        x = post_attention_layernorm(tmp);
        x = mlp({x})[0];
        x = x + tmp;
        return {x};
    }

private:
    QWenAttention self_atten;
    QWenMLP mlp;
    Layer input_layernorm;
    Layer post_attention_layernorm;
};

// Copied from GemmaModel with Gemma->Qwen and set RmsNorm(without add_unit_offset)
class QWenModel final : public Module {
public:
    QWenModel() = default;
    QWenModel(const QWenConfig &config, const QWenNameConfig &names, const string &base_name) {
        blocks = List<QWenDecoder>(config.num_hidden_layers, config, names, base_name);
        norm = RMSNorm(config.hidden_size, config.rms_norm_eps, names.post_norm_name);
    }

    std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
        auto x = inputs[0];
        for (auto &block : blocks) {
            x = block({x})[0];
        }
        x = norm(x);
        return {x};
    }

private:
    std::vector<QWenDecoder> blocks;
    Layer norm;
};

class QWenForCausalLM final : public Module {
public:
    QWenForCausalLM(QWenConfig &config) {
        auto names = config.names_config;
        hidden_size = config.hidden_size;
        tie_embedding_words = config.tie_embedding_words;
        embedding = Embedding(config.vocab_size, config.hidden_size, names.token_embd_name);
        model = QWenModel(config, names, names.blk_name);

        // FIXME Qwen-0.5 use tied embedding
        // Others use nn.Linear()
        if (tie_embedding_words) {
            lm_head = Parameter(1, config.vocab_size, 1, config.hidden_size, names.token_embd_name + ".weight");
        }
    }

    std::vector<Tensor> Forward(std::vector<Tensor> inputs, std::vector<std::any> args) override {
        auto x = embedding(inputs[0]);

        // go through model
        auto outputs = model({x})[0];
        if (tie_embedding_words) {
            outputs = Tensor::mm(outputs, lm_head().transpose(Chl::SEQUENCE, Chl::DIMENSION));
        }
        return {outputs};
    }

private:
    int hidden_size;
    bool tie_embedding_words;
    Layer embedding;
    Parameter lm_head;
    QWenModel model;
};

#endif //! MODELING_QWEN_HPP