RDK-OE-LLM工具链量化Gemma4全流程(TODO)

作者:SkyXZ

CSDN:SkyXZ~-CSDN博客

博客园:SkyXZ - 博客园

        2026年4月2日,Google 正式开源了新一代原生多模态模型 Gemma 4。相比我之前写过的 SigLip、InternVL,或者单纯的纯文本 LLM 适配,Gemma4 的工程特点要复杂得多。它并不是一个“视觉编码器 + 文本模型”的简单拼装,而是把 Vision、Audio、Text 三条分支统一纳入了同一套多模态体系中。对端侧部署来说,这种统一架构既带来了更强的表达能力,也带来了更高的重构和量化门槛。

        本文聚焦其中最关键、也是目前最适合在 OE-LLM 工具链中独立拆解落地的两部分:Vision 分支LLM/Text 分支。Audio 分支本文先不展开。全文会基于我在 leap_llm/models/gemma4leap_llm/apis/model/gemma4.py 中的实际工程实现,完整记录 Gemma4-E2B 如何完成 Leap 重构、校准、编译与精度验证。Vision 部分会重点讲 patchify、2D RoPE、pooler 与 vision projector 的适配方式;LLM 部分则会重点分析 chunk 化 prefill、sliding/full attention、KV cache 共享、per-layer input 分支,以及文本量化里最容易掉点的 attention/cache 路径是如何一步步被修稳的。

PS:感觉LLM部分的量化对我一个搞机器人和CV方向的大四学生来说还是有点困难,LLM部分从4.5号开始做到现在精度从0.80仅提到了0.91离可用的0.99+还很遥远,后续只能有时间慢慢慢慢做了,目前对话测试部分板端的代码还没有完成,仅完成了Vision分支和LLM主干,由于假期结束后有些忙了,所以本文由AI代写,如有不对的地方望指正

一、环境配置(按照开发手册配置即可)

开发机配置(PC电脑)

# Step1:下载D-Robotics_LLM_{version}.tar.gz安装包并正确解压。
wget https://d-robotics-aitoolchain.oss-cn-beijing.aliyuncs.com/llm_s100/1.0.0/D-Robotics_LLM_S100_1.0.0_SDK.tar.gz
# Step2:安装Conda环境
wget -c https://mirrors.bfsu.edu.cn/github-release/conda-forge/miniforge/LatestRelease/Miniforge3-Linux-x86_64.sh
chmod 777 Miniforge3-Linux-x86_64.sh
sh Miniforge3-Linux-x86_64.sh
conda create -n oellm python=3.10
conda activate oellm
# (oellm) xxx@xxx:~$
# Step3:安装必要的python环境
# D-Robotics_LLM_{version}路径下
pip install -r ./oellm_build/requirements.txt
pip install ./oellm_build/hbdk4_compiler-{version}-cp310-cp310-manylinux_2_17_x86_64.whl
pip install ./oellm_build/hbdk4_runtime_aarch64_unknown_linux_gnu_nash-{version}-py3-none-any.whl
pip install ./oellm_build/leap_llm-{version}-py310-none-any.whl

# Step4:将leapllm的model链接进本地开发目录
pip3 show leap-llm
ln -s /path/you/oellm/lib/python3.10/site-packages/leap_llm .

权重下载

        Gemma4 的 HuggingFace 仓库名可能会随着发布版本、基础版/指令版区分而变化,所以这里我不强写死某一个 repo name。对 OE-LLM 来说,真正需要的是一个完整的 Gemma4-E2B 权重目录,里面至少包含:

  • config.json
  • model.safetensorspytorch_model.bin
  • tokenizer 相关文件

如果你是从 HuggingFace 拉取,可以按下面的方式下载:

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download <your_gemma4_repo> --local-dir gemma4-e2b

二、使用OE-LLM量化编译

Leap工具链分析(仅个人分析,不代表地瓜官方)

        我个人理解,oe-llm与传统的量化流程(如离线 PTQ 或 QAT)有明显区别。它更像是一套“模型重构 + 校准 + 编译导出”的工具链:开发者需要先基于leap_llm提供的模块,将原始 PyTorch 网络改写为 Leap 可识别的实现;随后使用校准数据运行浮点forward()收集量化统计信息;最后切换到build()路径,将模型导出为 Leap/HBDK 计算图并编译成可部署到板端的 HBM 模型。

        从源码来看,leap_llm/nn/modules中常用的量化相关组件主要包括以下几类。需要注意的是,并不是所有模型都只依赖FakeQuant模块,部分模型也会混用DynamicQuantRMSNormLayerNormSplit等其他封装组件。

量化相关组件 所在文件 功能描述
FakeQuantEmbedding embedding.py 量化嵌入层
FakeQuantLinear linear.py 量化线性层
FakeQuantRMSNorm rms_norm.py 量化 RMS 归一化
ConstFakeQuant const_fake_quant.py 常量假量化/定点化
FakeQuantAdd ops.py 量化加法
FakeQuantMul ops.py 量化乘法
FakeQuantRsqrt ops.py 量化平方根倒数
FakeQuantReduceMean ops.py 量化均值归约
FakeQuantPow ops.py 量化幂运算
FakeQuantMatmul matmul.py 量化矩阵乘法
DynamicQuantMatmul matmul.py 动态量化矩阵乘法
FakeQuantSoftmax activation.py 量化 Softmax
FakeQuantSwish activation.py 量化 Swish 激活
FakeQuantGELU activation.py 量化 GELU 激活
FakeQuantPatchEmbedding vision_embedding.py 量化视觉 Patch 嵌入

        整体流程可以概括为三步。第一步是完成模型重构。通常需要将原始的nn.Module改写为继承ModelModule的 Leap 版本,并将其中的关键算子替换为对应的量化组件。每个重构后的模块通常都需要同时实现build()forward()两套逻辑:其中build()负责使用leap.*算子描述最终要编译到 BPU 上的计算图,forward()则保留 PyTorch 浮点路径,用于校准和精度对齐。

# 原始模型
class StandardLLM(nn.Module):
    def __init__(self):
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(num_layers)])
        self.norm = nn.RMSNorm(hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size)

# OE-LLM重构后的模型(示意)
class QuantLLM(Model):
    def __init__(self):
        self.embedding = FakeQuantEmbedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([QuantDecoderLayer() for _ in range(num_layers)])
        self.norm = FakeQuantRMSNorm(hidden_size)
        self.lm_head = FakeQuantLinear(hidden_size, vocab_size)
    def build(self, inputs):
        pass
    def forward(self, inputs):
        pass

        第二步是校准。校准阶段的核心目标是让各个量化模块在浮点前向过程中收集统计信息,例如absmax、缩放因子或归一化相关的范围信息。因此,校准数据的预处理流程必须与目标模型的真实输入分布尽可能一致。对于 Gemma4 来说,Vision 与 Text 两部分的校准方式完全不同:前者是图像 resize + patchify,后者则是 tokenizer 切 chunk + mask 构造 + cache 滚动。

# 收集量化统计信息
model.set_compile_mode(False)
for batch in calibration_data:
    model.forward(batch)  # 量化模块在forward过程中收集统计信息

        最后一步是编译导出。此时模型会切换到build()路径,由leap_llm将计算图导出、转换并进一步编译为板端可执行的模型文件。Gemma4 Vision 当前走 float16 编译更合适,而 Gemma4 Text 则是导出 prefill / decode 两个 stage,并且在 build 输入中显式携带 full/sliding mask 和 KV cache。

# 生成量化编译模型
model.set_compile_mode(True)
model.compile(
    dtype=leap.float16,
    march="nash-e",
    output_model_path="model.hbm"
)

Gemma4结构分析与工具链适配

        Gemma4 在当前工程里被拆成两个可以独立校准、独立编译、独立验证的子图:

  • Vision 编码器,对应 Gemma4Vision / Gemma4VisionApi
  • Text 模型,对应 Gemma4Text / Gemma4TextApi

        这种拆法对 OE-LLM 很友好,因为视觉侧本质上是固定分辨率的 Transformer Encoder,文本侧则是标准的 chunked prefill + decode LLM。Audio 分支虽然理论上也能继续拆,但其输入、特征提取和 runtime 接口会更复杂,所以本文先不展开。

Vision Branch

        Gemma4 Vision 在当前实现里对应 Gemma4VisionConfigPatchEmbeddingGemma4VisionEncoderLayerVisionPoolerVisionProjector。从配置上看,默认输入分辨率是 768 x 768,patch size 是 16,因此最终 patch 数是 48 x 48 = 2304,隐藏维是 768,编码器层数是 16,最后再通过 projector 投影到文本分支的 1536 维空间。

(1) Patchify 与 PatchEmbedding

        Gemma4 Vision 的输入不是直接把图像送进 PatchEmbedding,而是先在 API 层完成 resize 和 patchify。也就是说,真正送进编译图的输入已经是 [1, num_patches, 3*patch_size^2] 这种 patch 序列,而不是原始的 [1, 3, H, W] 图像张量。这样做的好处是能把 patch 展开逻辑留在 PyTorch 前处理里,减少编译图中的动态图复杂度。

def _patchify_image(image_tensor, patch_size=16):
    _, C, H, W = image_tensor.shape
    h_patches = H // patch_size
    w_patches = W // patch_size
    patches = image_tensor.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size)
    patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous()
    patches = patches.view(1, h_patches * w_patches, C * patch_size * patch_size)
    return patches

        进入模型后,PatchEmbedding 的逻辑也比较直接:先把输入从 [0, 1] 映射到 [-1, 1],然后经过一个 FakeQuantLinear 做 patch projection,再叠加预先算好的二维位置编码。这里的位置编码并不是运行时临时查表,而是在 _init_constants() 里先根据固定分辨率离线预计算好,再作为常量 buffer 挂在模型里。

class PatchEmbedding(Module):
    def __init__(self, config):
        self.input_proj = FakeQuantLinear(3 * config.patch_size ** 2, config.hidden_size, bias=False)
        self.input_norm = ConstFakeQuant(8)

    def build(self, pixel_values):
        x = leap.mul(pixel_values, 2.0)
        x = leap.add(x, -1.0)
        x = self.input_norm(x)
        hidden_states = self.input_proj(x)
        hidden_states = leap.add(hidden_states, self.position_embeddings.data.to(torch.float32))
        return hidden_states
(2) Vision Attention

        Gemma4 Vision 的 attention 与传统 ViT 有相似之处,但也有两个很关键的不同点。第一,它使用的是 2D 多维 RoPE,不是简单的一维位置编码。第二,它对 Q/K/V 都做了 norm,其中 v_norm 还是不带可学习 scale 的版本,因此官方实现并不再使用传统的 1/sqrt(d) attention scaling,而是固定 scaling = 1.0

class Gemma4VisionAttention(Module):
    def __init__(self, hidden_size, num_attention_heads, head_dim, rms_norm_eps):
        self.q_proj = FakeQuantLinear(hidden_size, num_attention_heads * head_dim, bias=False)
        self.k_proj = FakeQuantLinear(hidden_size, num_attention_heads * head_dim, bias=False)
        self.v_proj = FakeQuantLinear(hidden_size, num_attention_heads * head_dim, bias=False, quant_bits=8)
        self.o_proj = FakeQuantLinear(num_attention_heads * head_dim, hidden_size, bias=False)

        self.q_norm = FakeQuantRMSNorm(head_dim, eps=rms_norm_eps)
        self.k_norm = FakeQuantRMSNorm(head_dim, eps=rms_norm_eps)
        self.v_norm = FakeQuantRMSNorm(head_dim, eps=rms_norm_eps, fuse_norm=True)

        self.qk = FakeQuantMatmul(8, 16)
        self.sv = FakeQuantMatmul(None, 8)
        self.softmax = FakeQuantSoftmax(quant_bits=16, quantized=True)

        在 Leap 中适配这类 Vision attention 时,核心不在于“重新设计结构”,而在于把原本 PyTorch 可以隐式推导的 reshapetransposematmul、RoPE 拆分过程都显式写出来。特别是 Gemma4 Vision 把 head_dim 拆成了 (x, y) 两个空间维度分别做旋转,这部分如果只按一维 LLM RoPE 去套,结果肯定不对。

(3) Vision Pooler 与 Projector

        Gemma4 Vision 不会直接把 2304 个 patch token 全部喂给后续文本分支,而是先通过 VisionPooler 做固定网格池化。当前实现里 pooling kernel 是 3,因此 48 x 48 的 patch 网格最终会被缩成 16 x 16 = 256 个视觉 token。随后再经过 VisionProjector,把视觉隐藏维从 768 投影到文本分支的 1536

class VisionPooler(Module):
    def build(self, hidden_states):
        hidden_states = leap.reshape(hidden_states, [16, 3, 16, 3, 768])
        pooled = self.pool(hidden_states, [1, 3])
        pooled = leap.reshape(pooled, [256, 768])
        pooled = leap.mul(pooled, self.root_hidden_size)
        return pooled

class VisionProjector(Module):
    def build(self, hidden_states):
        hidden_states = self.norm(hidden_states)
        return self.projection(hidden_states)

        也就是说,Vision HBM 最终输出的就是 [256, 1536] 的视觉 token 序列。这一点在后续和 Text 模型对接时非常重要,因为它已经天然对齐到 LLM 的 hidden size,不需要额外再做一次跨模块维度变换。

Text Model

        Gemma4 Text 是当前整个适配里最有工程量的一部分。它不只是一个普通的 decoder-only Transformer,而是同时叠加了多种对量化并不友好的机制:

  • 35 层 decoder
  • sliding_attentionfull_attention 交替出现
  • 20 层 KV shared layers
  • 单独的 per-layer input 分支
  • chunk 化 prefill 与 decode 双 stage 编译

        从 Gemma4TextConfig 能看到,模型默认 chunk_size = 256cache_len = 4096sliding_window = 512。层类型不是统一的 full attention,而是每隔若干层插入一个 full attention,其余大部分层都走 sliding attention。full attention 层用的是 global_head_dim = 512,sliding attention 层用的是 head_dim = 256

(1) Embedding 与 Per-layer Input 分支

        Gemma4 Text 的一个很有意思的设计,是除了标准的 token embedding 之外,还额外构造了一条 per-layer input 支路。这条支路一部分来自 token embedding 的扩展版本,另一部分来自当前 hidden states 经过线性投影和归一化后的结果,最后两者相加后缩放。它相当于给每一层都额外提供一份更细粒度的条件输入。

class Gemma4TextScaledEmbedding(Module):
    def __init__(self, num_embeddings, embedding_dim, embed_scale):
        self.weight_fake_quant = ConstFakeQuant(8)
        self.scale_mul = FakeQuantMul(quantized=False)

    def build(self, x):
        weight = self.weight_fake_quant(self.weight.data)
        embeds = leap.gather_nd(weight, x, 0)
        return self.scale_mul(embeds, self.embed_scale)
def _build_per_layer_inputs(self, inputs_embeds, token_ids):
    per_layer_inputs = self.embed_tokens_per_layer(token_ids)
    per_layer_projection = self.per_layer_model_projection(inputs_embeds)
    per_layer_projection = self.per_layer_projection_mul(
        per_layer_projection,
        self.per_layer_projection_scale,
    )
    per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
    per_layer_inputs = self.per_layer_add(per_layer_projection, per_layer_inputs)
    return self.per_layer_input_scale_mul(per_layer_inputs, self.per_layer_input_scale)

        在 decoder layer 内部,这条 per-layer input 会先过一个小门控分支,再通过 per_layer_projection 投回 hidden_size,最后和主分支残差相加。这也解释了为什么 Gemma4 的每层除了 attention 和 MLP 之外,还多了一套 per_layer_input_gate / per_layer_act / per_layer_projection / post_per_layer_input_norm

(2) Text Attention、Sliding/Full 与 KV Shared

        Gemma4 Text attention 的整体思路是:

  • Q/K/V 都先做线性映射
  • Q 和 K 做 RMSNorm,再做 RoPE
  • V 也会先做 norm
  • full attention 与 sliding attention 共享统一的 attention 实现,但用不同的 RoPE 常量和 mask
  • 部分后段层不再自己生成 KV,而是直接复用前面某个同类型层保存下来的 full-length cache
class Gemma4TextAttention(Module):
    def __init__(self, hidden_size, num_attention_heads, num_key_value_heads, head_dim, rms_norm_eps):
        self.q_proj = FakeQuantLinear(hidden_size, num_attention_heads * head_dim, bias=False)
        self.k_proj = FakeQuantLinear(hidden_size, num_key_value_heads * head_dim, bias=False)
        self.v_proj = FakeQuantLinear(hidden_size, num_key_value_heads * head_dim, bias=False, quant_bits=8)
        self.o_proj = FakeQuantLinear(num_attention_heads * head_dim, hidden_size, bias=False)

        self.q_norm = FakeQuantRMSNorm(head_dim, eps=rms_norm_eps)
        self.k_norm = FakeQuantRMSNorm(head_dim, eps=rms_norm_eps)
        self.v_norm = FakeQuantRMSNorm(head_dim, eps=rms_norm_eps, fuse_norm=True)

        self.qk = FakeQuantMatmul(8, 8)
        self.sv = FakeQuantMatmul(16, 8, 16)
        self.add_mask = FakeQuantAdd(quantized=True, quant_bits=16)
        self.softmax = FakeQuantSoftmax(quant_bits=16, quantized=True)

        Gemma4 的一个关键工程点在于 KV shared。当前实现里 num_hidden_layers = 35num_kv_shared_layers = 20,因此前 15 层是非共享 KV 层,后 20 层会按 layer type 复用前面最近一个同类型层保存下来的 full cache。这样做的好处是能显著降低后段层的 KV 存储和带宽开销,但对量化和 verifier 的实现会带来额外复杂度。

self.first_kv_shared_layer_idx = (
    config.num_hidden_layers - config.num_kv_shared_layers
)
self.non_shared_layer_indices = list(range(self.first_kv_shared_layer_idx))
self.shared_source_by_type = {
    layer_type: len(prev_layers) - 1 - prev_layers[::-1].index(layer_type)
    for layer_type in set(prev_layers)
}

        在 build/forward 主循环里,这种共享逻辑体现得很明显:非共享层接收独立 cache,输出新的 new_k/new_v;当某一层是共享源层时,额外保存 full_k/full_v;后续 shared layer 则直接从 shared_caches[layer.layer_type] 中取出 cache 使用。

(3) Chunk 化 Prefill、Mask 构造与量化稳定性

        Gemma4 Text 在 OE-LLM 里不是一次性吃完整段文本,而是按 chunk_size=256 做 prefill。对应的输入准备逻辑放在 Gemma4TextCalibrationDataPreparer 里:先 tokenizer,再右侧补 pad 到 chunk_size 的整数倍,再为每个 chunk 分别构造 position_idsfull_masksliding_mask

for chunk_idx in range(len(input_chunks)):
    total_seen = min((chunk_idx + 1) * self.chunk_size, valid_len)
    chunk_start = chunk_idx * self.chunk_size
    chunk_valid = min(valid_len - chunk_start, self.chunk_size)
    full_masks.append(self._build_mask(total_seen, chunk_start, chunk_valid, None))
    sliding_masks.append(
        self._build_mask(total_seen, chunk_start, chunk_valid, self.sliding_window)
    )

        这里有一个很容易踩坑的地方:最后一个 chunk 可能并不满 256,但 cache 更新时依然是整块滚动的。如果 mask 仍按“最后 valid_total 列右对齐”的思路去写,就会把 pad 区域的 KV 也放开,最后导致长文本精度明显掉点。为了解决这个问题,我在 _build_mask() 里显式引入了 chunk_valid,确保最后一个不满 chunk 时,被放开的有效列只对应真实 token 的 cache 区间,而不是把最后一整段 pad KV 也误认为有效上下文。

        除了 mask,Gemma4 Text 的另一大掉点来源是 softmax @ V-cache 这条 sv 路径。它的典型表现是:qkadd_mask 看起来已经比较接近浮点基线,但一到 sv.out 和后续 o_proj 就开始快速掉点。围绕这条路径,我当前主要做了三类稳定化处理:

  • value_statescache_vattn_weights 分别增加显式 fake-quant 约束
  • sv 明确固定为 FakeQuantMatmul(16, 8, 16),避免编译器在这一层做过于激进的自由下沉
  • 对后段敏感层开启 chunked sv,把 attn_weights @ cache_v1024 一块拆开做,再把 partial result 累加

        这类改动不是“为了代码好看”,而是 Gemma4 Text 当前能把长文本 verifier 精度从明显异常拉回到可接受范围的关键工程手段。

(4) 我们在LLM精度优化里尝试过的路线

        如果只看现在的结果,Gemma4 Text 好像已经比较稳定了,但实际排查过程并不是一次到位的。这里我把这次比较关键的尝试路线也记录下来,后面如果要继续适配别的复杂 LLM,基本可以复用同样的思路。

  1. 第一阶段是先把结构一比一重构跑通。这个阶段的目标不是追求高精度,而是确保 forward()build()、prefill/decode 双 stage、cache 输入输出这些接口都能闭环。此时模型可以跑通,但文本 verifier 的最终输出只有大约 0.63,说明图是通的,但数值路径明显不稳。
  2. 第二阶段是先排除“假问题”。例如一些 cache_k_fq 条目虽然 cosine 为负,但本质上比较的是不同 layout 的张量,shape_torchshape_bc 本来就不一致,这类条目不能直接当成真实 bug。真正该重点看的,是 same-shape 的 softmax.out_quantsv_input_fqsv.out_fake_quanto_proj.out_fake_quant
  3. 第三阶段是修 partial chunk mask。最后一个 chunk 不满 256 时,如果还按旧逻辑构造 mask,就会把 pad 区域的 KV 一起放开。这个问题修掉之后,长文本的异常掉点先被消掉了一大截,最终输出均值从最初的 0.63 左右明显提升到 0.88 附近。
  4. 第四阶段是集中处理 V-cache -> sv 路径。我把 value_statescache_vattn_weights 都加上了更明确的 fake-quant 约束,同时把 sv 固定成 FakeQuantMatmul(16, 8, 16)。这一阶段的核心思路是:既然 qk 没有先崩,那问题就不该再泛泛地归咎于“attention 整体不行”,而应该把焦点放在 softmax @ V 这一条读 cache 的路径上。
  5. 第五阶段是对后段敏感层引入 chunked sv。最开始我只在 24/29/34 这几个 full-attention 层上试了按 4 x 1024 分块做 attn_weights @ cache_v,后面又进一步把启用范围扩到了后段 24~34 层。这个策略的作用不是改变数学结果,而是把一个超长的 sv kernel 拆成多个更稳的小块,从而减轻量化和编译后的数值漂移。
  6. 到目前为止,这条路线已经把 Gemma4 Text 的最终输出均值拉到了 0.909954,layer mean 拉到了 0.953301。虽然它还没有像 Vision 那样接近 1.0,但长文本已经不再出现最初那种系统性崩坏,说明主要矛盾已经被找对了。

模型注册

        Gemma4 在 OE-LLM 中的注册入口并不复杂,核心就是在 model_factory.py 里把 Vision 和 Text 分别注册成两个模型名:

@register_model("gemma4-e2b-vision", ["nash-e", "nash-m", "nash-p"])
def _build_gemma4_e2b_vision(args):
    from leap_llm.apis.model.gemma4 import Gemma4VisionApi
    return Gemma4VisionApi(...)

@register_model("gemma4-e2b-text", ["nash-e", "nash-m", "nash-p"])
def _build_gemma4_e2b_text(args):
    from leap_llm.apis.model.gemma4 import Gemma4TextApi
    return Gemma4TextApi(...)

        这一步完成之后,oellm_build.py 就能直接通过 --model_name gemma4-e2b-vision--model_name gemma4-e2b-text 来创建对应 API,后续流程就和 SigLip、Qwen、InternLM 这些模型保持统一了。

开始量化

(1) Vision 分支编译

        Vision 分支的 API 比较直接:加载图像校准集,跑浮点 forward() 收集统计信息,然后走 Gemma4Vision.compile() 导出 bc -> convert.bc -> hbo -> hbm。最终产物默认命名为:

  • gemma4-e2b_vit_ptq.hbm
python ./D-Robotics_LLM_S100_1.0.0_SDK/leap_llm/apis/oellm_build.py \
    --model_name gemma4-e2b-vision \
    --march nash-m \
    --input_model_path ./gemma4-e2b \
    --output_model_path ./output/gemma4_vision \
    --calib_image_path ./calibration_data/images \
    --device cuda:0 \
    --vit_core_num 1 \
    --verifier
(2) LLM 分支编译

        Text 分支的流程要复杂得多。它会先读取文本校准集,通过 Gemma4TextCalibrationDataPreparer 切 chunk,初始化空 cache,然后按 chunk 依次做 get_input_embeddings -> forward -> update_caches。校准完成后,再进入 stage="all" 编译模式,同时导出 prefill 和 decode 两个 stage。

        最终产物默认包括:

  • gemma4-e2b_lm_chunk_256_cache_4096_ptq.hbm
  • tok_embeddings.bin
python ./D-Robotics_LLM_S100_1.0.0_SDK/leap_llm/apis/oellm_build.py \
    --model_name gemma4-e2b-text \
    --march nash-m \
    --input_model_path ./gemma4-e2b \
    --output_model_path ./output/gemma4_text \
    --calib_text_path ./calibration_data/text \
    --chunk_size 256 \
    --cache_len 4096 \
    --device cuda:0 \
    --prefill_core_num 1 \
    --decode_core_num 1 \
    --verifier

        这里有两个参数必须注意:

  • chunk_size 必须和后续 runtime/prefill 配置对齐
  • cache_len 必须是 chunk_size 的整数倍,并且要和板端 KV cache 长度保持一致

精度验证

        如果在 oellm_build.py 里加了 --verifier,工具链会在编译后自动调用 verifier_cli.py。当然你也可以手动单独执行 verifier,这在调试掉点时会更方便。

(1) Vision 验证
python ./D-Robotics_LLM_S100_1.0.0_SDK/leap_llm/apis/verifier_cli.py \
    --model_name gemma4-e2b-vision \
    --model_dir ./gemma4-e2b \
    --hbm_vlm_model_path ./output/gemma4_vision/gemma4-e2b_vit_ptq.hbm \
    --input_image_path ./calibration_data/images

        结合我当前工程下的最新报告,Vision 分支的结果已经比较稳定:

  • layer output mean cosine: 0.957873
  • final output mean cosine: 0.999009
  • 两张测试图最终输出分别是:
    • image_0 = 0.999044
    • image_1 = 0.998975

        这个结果说明 Gemma4 Vision 的量化难点并不在 projector,而主要在 encoder 内部 attention / MLP 的数值保持。当前这版实现已经能够把最终视觉 token 输出做得非常接近 PyTorch 浮点基线。

(2) LLM 验证
python ./D-Robotics_LLM_S100_1.0.0_SDK/leap_llm/apis/verifier_cli.py \
    --model_name gemma4-e2b-text \
    --model_dir ./gemma4-e2b \
    --hbm_llm_model_path ./output/gemma4_text/gemma4-e2b_lm_chunk_256_cache_4096_ptq.hbm \
    --input_text_path ./calibration_data/text \
    --chunk_size 256 \
    --cache_len 4096

        Gemma4 Text 的 verifier 要比 Vision 更值得重点分析,因为它既有 chunk,又有 sliding/full attention,又有 shared cache,还涉及最后一个有效 token 的切片比较。根据我目前的最新实现与报告,Text 分支结果如下:

  • layer output mean cosine: 0.953301
  • final output mean cosine: 0.909954
  • final output min cosine: 0.724089
  • final output max cosine: 0.984893

        16 条文本样本中,当前最差和最好样本分别是:

Prompt Final Cosine
prompt_4 0.724089
prompt_10 0.791751
prompt_5 0.805644
prompt_11 0.826579
prompt_12 0.858014
prompt_8 0.984893

        这个结果比我最初直接适配时已经提升了很多。最开始掉点最严重的时候,文本分支的最终输出均值大约只有 0.63,而且长文本 attention 的 sv 路径会在后半段层数里明显失稳。后续沿着“先修 mask,再收敛 V-cache -> sv 路径,最后对敏感层做 chunked sv”这条路线逐步调整之后,最终输出均值已经被拉到了 0.909954,长文本也基本不再出现系统性崩坏。

三、RDK板端测试

性能测试

Vision

        Gemma4 Vision 量化后生成的是一个独立的视觉编码器 HBM,它的职责是把输入图像变成 [256, 1536] 的视觉 token 序列。因此在板端,它通常不会单独作为“可聊天”的模型运行,而是作为多模态系统的视觉前端被上层调用。对这一分支来说,板端验证的核心通常有两个:

  • HBM 是否能够稳定加载并完成 feature extraction
  • 输出 token 数和维度是否与上层文本分支预期一致

        如果你的目标是做完整 VLM,对话或图文问答时只需要保证视觉侧输出与文本侧 hidden size 对齐即可;当前实现里这一点已经通过 VisionProjector 固定到了 1536 维。

LLM

        Text 分支上板时,除了 *.hbm 文件本身,还需要把 tok_embeddings.bin 和 tokenizer 配置一起带上。原因是 Gemma4 当前的文本工程拆分方式并不是把所有 embedding 逻辑都包进 runtime,而是将 token embedding 权重额外导出,供运行时侧对齐使用。

        另外,板端 prefill / decode 的配置必须与编译参数保持一致,尤其是下面两项:

  • chunk_size = 256
  • cache_len = 4096

        如果这两个参数不一致,即使 HBM 能加载成功,也很容易出现 cache 排布或 mask 对不齐的问题,最终表现为对话质量异常或者长文本输出漂移。下面这两张图是我当前工程在板端的运行截图。

对话测试

TODO…

1 个赞