Tutorial_for_Llama3_Compute 运行qnn_model_prepare.ipynb，导出的onnx效果很差，并没有等效替换的效果

[chenfeng0232]

你好，请问一下，我在使用Tutorial_for_Llama3_Compute ，针对Qwen模型，安装官方教程进行model_prepare是，困惑度是正常的，但是当加载等效替换导出的onnx模型时，效果很差。请问是为什么呢。onnx推理代码如下，输入预处理函数，都是调用高通自定义的函数。：

class LLMForwardPassManager:

    def __init__(self, cfg, model_path, tokenizer, separate_tuple_input_output, num_tokens):

        self.tokenizer = tokenizer

        self.session = ort.InferenceSession(model_path)

        self.device = 'cpu'  # ONNX typically uses CPU unless GPU is configured

        self.num_heads = getattr(cfg, 'num_attention_heads', 14)

        self.num_kv_heads = getattr(cfg, 'num_key_value_heads', 2)

        self.num_layers = getattr(cfg, 'num_hidden_layers', 24)

        self.embed_dim = getattr(cfg, 'hidden_size', 896)

        self.rope_theta = getattr(cfg, "rope_theta", 10000.0)

        self.max_tokens = getattr(cfg, 'max_tokens', 4096)

        self.num_tokens = num_tokens  # Fixed at 2073

        self.max_kv_size = self.max_tokens - self.num_tokens  # 2023

        self.use_position_embedding_input = True

        self.use_combined_mask_input = True

        self.transposed_key_cache = True

        self.mask_neg = getattr(cfg, 'mask_neg', -50)

        self.use_input_embeddings = False

        self.separate_tuple_input_output = separate_tuple_input_output

        self.dtype = np.float32

    def replace_model(self, new_model_path):

        self.session = ort.InferenceSession(new_model_path)

    @contextlib.contextmanager

    def place_on_device(self, device):

        original_device = self.device

        try:

            self.to(device)

            yield

        finally:

            self.to(original_device)

    def to(self, device):

        self.device = device  # Note: ONNX Runtime device handling may require additional setup for GPU

    def _tokenize_text(self, text, max_length=None):

        if self.tokenizer is None:

            raise ValueError("No tokenizer registered with forward pass manager.")

        encoded = self.tokenizer(text, return_tensors="np", padding="max_length" if max_length else False,

                                 truncation=True, max_length=max_length or self.max_tokens)

        return encoded

    def _update_kv_cache(self, prev_key_value, new_key_value, max_cache_size, is_concatenated=False):

        def _concat(a, b, dim):

            if isinstance(a, tuple):

                assert len(a) == len(b), 'Unexpected key/value pair'

                return tuple(_concat(ai, bi, dim) for ai, bi in zip(a, b))

            return np.concatenate((a, b), axis=dim)

        def _do_concat(a, b, key_dim, value_dim):

            return tuple((_concat(ak, bk, key_dim), _concat(av, bv, value_dim)) for (ak, av), (bk, bv) in zip(a, b))

        def _shift(a, dim, shift_size):

            if isinstance(a, tuple):

                return tuple(_shift(ai, dim, shift_size) for ai in a)

            assert dim in (2, 3), 'Unexpected shift axis'

            return a[:, :, shift_size:, :] if dim == 2 else a[:, :, :, shift_size:]

        def _do_shift(a, key_dim, value_dim, shift_size):

            return tuple((_shift(k, key_dim, shift_size), _shift(v, value_dim, shift_size)) for k, v in a)

        value_dim = 2

        key_dim = 3 if self.transposed_key_cache else 2

        if prev_key_value is None or is_concatenated:

            next_key_value = new_key_value

        elif new_key_value is None:

            next_key_value = prev_key_value

        else:

            next_key_value = _do_concat(prev_key_value, new_key_value, key_dim, value_dim)

        shift_size = next_key_value[0][1].shape[-2] - max_cache_size

        if shift_size > 0:

            next_key_value = _do_shift(next_key_value, key_dim, value_dim, shift_size)

        return next_key_value

    def validate_inputs(self, input_text=None, input_ids=None, input_embeddings=None, past_key_values=None):

        input_count = sum(1 for x in (input_text, input_ids, input_embeddings) if x is not None)

        if input_count != 1:

            print("Incorrect number of arguments: one of (input_text, input_ids, input_embeddings) expected.")

            return False

        if past_key_values is not None and past_key_values[0][1].shape[-2] > self.max_kv_size:

            print(f"Provided past_key_values are too long: {past_key_values[0][1].shape[-2]} > {self.max_kv_size}")

            return False

        return True

    def validate_input_lengths(self, input_length, mask_length, attn_length):

        if not (1 <= input_length <= self.num_tokens):

            print(f"input_length({input_length}) must be between 1 and num_tokens({self.num_tokens}).")

            return False

        if not (input_length <= mask_length <= attn_length):

            print(f"mask_length({mask_length}) must satisfy input_length <= mask_length <= attn_length({attn_length}).")

            return False

        return True

    def prepare_inputs(self, input_text=None, input_ids=None, input_embeddings=None, attention_mask=None,

                       past_key_values=None, **kwargs):

        assert self.validate_inputs(input_text, input_ids, input_embeddings, past_key_values)

        kvcache_info_bundle = {}

        if input_text is not None:

            encoded = self._tokenize_text(input_text, max_length=self.num_tokens)

            input_ids = encoded['input_ids']

            attention_mask = encoded['attention_mask']

            kvcache_info_bundle["input_length"] = input_ids.shape[1]

        else:

            kvcache_info_bundle["input_length"] = input_ids.shape[1]

        input = input_ids.astype(np.int64)

        batch_size = input.shape[0]

        input_length = input.shape[1]

        kv_length = past_key_values[0][1].shape[-2] if past_key_values is not None else 0

        attn_length = min(input_length + kv_length, self.max_tokens)

        if attention_mask is None:

            attention_mask = np.ones((batch_size, input_length), dtype=np.int64)

        attention_mask = attention_mask.astype(np.int64)

        mask_length = attention_mask.shape[1]

        assert self.validate_input_lengths(input_length, mask_length, attn_length)

        if input_length < self.num_tokens:

            shape = (batch_size, self.num_tokens - input_length)

            input_extensions = np.full(shape, self.tokenizer.pad_token_id, dtype=np.int64)

            input = np.concatenate((input_extensions, input), axis=1)

            attention_mask_extension = np.zeros((batch_size, self.num_tokens - input_length), dtype=np.int64)

            attention_mask = np.concatenate((attention_mask_extension, attention_mask), axis=1)

        desired_length = self.max_tokens

        if mask_length < desired_length:

            attention_mask_extension = np.zeros((batch_size, desired_length - mask_length), dtype=np.int64)

            attention_mask = np.concatenate((attention_mask_extension, attention_mask), axis=1)

        elif mask_length > desired_length:

            attention_mask = attention_mask[:, -desired_length:]

        if past_key_values is None:

            past_key_values = get_padded_kv_values(past_size=self.max_kv_size, num_layers=self.num_layers,

                                                   hidden_size=self.embed_dim, num_attention_heads=self.num_heads,

                                                   num_kv_heads=self.num_kv_heads, device=self.device, dtype=self.dtype)

        else:

            past_key_values = self._update_kv_cache(None, past_key_values, self.max_kv_size)

 

        position_ids = np.cumsum(attention_mask, axis=1) - 1

        position_ids = np.clip(position_ids, 0, self.max_tokens - 1)

        position_ids = position_ids[:, -self.num_tokens:]

        position_ids_cos, position_ids_sin = get_position_embeddings_from_position_ids(position_ids,

                                                                                      head_dim=self.embed_dim // self.num_heads,

                                                                                      max_length=self.max_tokens,

                                                                                      rope_theta=self.rope_theta,

                                                                                      device=self.device,

                                                                                      dtype=self.dtype)

        attention_mask = prepare_combined_attention_mask(attention_mask, (batch_size, self.num_tokens), self.max_kv_size,

                                                         device=self.device, mask_neg=self.mask_neg, dtype=self.dtype)

        inputs = {

            'input_ids': input.reshape(1, self.num_tokens),

            'attention_mask': attention_mask.reshape(1, 1, self.num_tokens, self.max_tokens),

            'position_ids_cos': position_ids_cos.reshape(1, 1, self.num_tokens, 32),

            'position_ids_sin': position_ids_sin.reshape(1, 1, self.num_tokens, 32)

        }

        if self.separate_tuple_input_output:

            input_names = [input.name for input in self.session.get_inputs()]

            flattened_key_values = list(flatten_tensors(past_key_values))

            for key, value in zip(input_names[4:], flattened_key_values):

                inputs[key] = value

        return inputs, kvcache_info_bundle

    def prepare_outputs(self, outputs, prepared_inputs, kvcache_info_bundle):

        lm_logits = outputs[0]

        def _get_past_kv_from_outputs(outputs):

            if self.separate_tuple_input_output:

                return tuple((outputs[(2 * i) + 1], outputs[(2 * i) + 2]) for i in range(self.num_layers))

            else:

                return outputs[1:]

        new_past_key_values = _get_past_kv_from_outputs(outputs)

        new_past_key_values = self._update_kv_cache(None, new_key_value=new_past_key_values, max_cache_size=self.num_tokens)

        old_past_key_values = tuple((prepared_inputs[f"past_key_{i}_in"], prepared_inputs[f"past_value_{i}_in"]) for i in range(self.num_layers))

        past_key_values = self._update_kv_cache(old_past_key_values, new_past_key_values, self.max_kv_size)

        return {'lm_logits': lm_logits, 'past_key_values': past_key_values}

    def __call__(self, *args, **kwargs):

        prepared_inputs, kvcache_info_bundle = self.prepare_inputs(*args, **kwargs)

        outputs = self.session.run(None, prepared_inputs)

        prepared_outputs = self.prepare_outputs(outputs, prepared_inputs, kvcache_info_bundle)

        return prepared_outputs

    def generate(self, input_text, max_new_tokens=3):

        encoded = self._tokenize_text(input_text)

        input_ids = encoded['input_ids']

        attention_mask = encoded['attention_mask']

        past_key_values = get_padded_kv_values(past_size=self.max_kv_size, num_layers=self.num_layers,

                                               hidden_size=self.embed_dim, num_attention_heads=self.num_heads,

                                               num_kv_heads=self.num_kv_heads, device=self.device, dtype=self.dtype)

        generated_ids = input_ids[0].tolist()

        outputs = self(input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values)

        past_key_values = outputs['past_key_values']

        for _ in range(max_new_tokens):

            next_token_logits = outputs['lm_logits'][:, -1, :]

            next_token = np.argmax(next_token_logits, axis=-1).item()

            generated_ids.append(next_token)

            if next_token == self.tokenizer.eos_token_id:

                break

            input_ids = np.array([[next_token]], dtype=np.int64)

            attention_mask = np.ones((1, 1), dtype=np.int64)

            outputs = self(input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values)

            past_key_values = outputs['past_key_values']

        return self.tokenizer.decode(generated_ids, skip_special_tokens=True)

if __name__ == "__main__":

    class Config:

        num_attention_heads = 14

        num_key_value_heads = 2

        num_hidden_layers = 24

        hidden_size = 896

        rope_theta = 10000.0

        max_tokens = 4096

    cfg = Config()

    tokenizer = AutoTokenizer.from_pretrained('/home/Qwen2.5-0.5B-Instruct-old/huggface')

    model_path = '/home/Qwen2.5-0.5B-Instruct-old/onnx/qwen2.onnx'

    fpm = LLMForwardPassManager(cfg, model_path, tokenizer, separate_tuple_input_output=True, num_tokens=2073)

    prompt = "你能够做什么"

    generated_text = fpm.generate(prompt)

    print("Generated text:", generated_text)

[weixin_38498942]

在使用 Tutorial_for_Llama3_Compute 针对 Qwen 模型，按官方教程进行 model_prepare 时困惑度正常，但加载等效替换导出的 ONNX 模型后效果变差，可能由以下几方面原因导致：

模型转换问题

1. 算子支持问题
   • ONNX 对某些深度学习框架中的算子支持可能不完全，Qwen 模型可能使用了一些独特的自定义算子或复杂的组合算子。在将 PyTorch 模型转换为 ONNX 模型时，这些算子可能无法被正确转换或表示，从而导致模型在推理时出现错误或性能下降。
   • 例如，某些特殊的激活函数、归一化操作等可能没有合适的 ONNX 对应算子，或者在转换过程中被近似处理，影响了模型的准确性。
2. 转换参数设置问题
   • 在使用工具（如 torch.onnx.export）将模型导出为 ONNX 时，参数设置不正确可能会导致模型结构或权重信息丢失。例如，没有正确指定输入输出的名称、形状、动态轴等，可能会使 ONNX 模型在推理时无法正确处理输入数据，进而影响性能。
   • 另外，不同的导出版本（如 PyTorch 和 ONNX 的版本兼容性）也可能会导致转换问题。

环境差异问题

1. 推理引擎差异
   • 不同的 ONNX 推理引擎（如 ONNX Runtime、TensorRT 等）对模型的支持和优化程度不同。在加载 ONNX 模型进行推理时，选择的推理引擎可能没有对 Qwen 模型进行有效的优化，导致推理结果不准确。
   • 例如，某些推理引擎可能对某些层或算子的实现与原始模型不同，或者在计算精度、内存管理等方面存在差异，从而影响模型的性能。
2. 计算环境差异
   • 训练环境和推理环境的硬件配置、软件版本等可能存在差异。例如，训练时使用的 GPU 型号和驱动版本与推理时不同，可能会导致计算结果的细微差异，进而影响模型的性能。
   • 另外，不同的操作系统、Python 版本、库依赖等也可能会对模型的推理结果产生影响。

数据处理问题

1. 输入数据预处理不一致
   • 在训练和推理过程中，输入数据的预处理步骤必须保持一致。如果在将模型转换为 ONNX 后，推理时的输入数据预处理方式与训练时不同，例如数据的归一化范围、数据类型、维度顺序等不一致，会导致模型无法正确处理输入数据，从而影响性能。
2. 输出数据后处理问题
   • 同样，输出数据的后处理步骤也需要保持一致。如果在推理时对 ONNX 模型的输出数据进行了错误的后处理，例如解码方式、阈值设置等与训练时不同，也会导致模型效果变差。

解决方案建议

1. 检查模型转换过程
   • 仔细检查模型转换代码，确保所有参数设置正确，特别是输入输出的名称、形状和动态轴。
   • 查看 ONNX 模型的结构，确认是否有算子转换失败或丢失的情况。可以使用工具（如 Netron）可视化 ONNX 模型，检查模型结构是否与原始模型一致。
2. 选择合适的推理引擎
   • 尝试不同的 ONNX 推理引擎，并对其进行性能测试和调优。例如，对于 GPU 推理，可以尝试使用 TensorRT 进行优化，提高推理速度和性能。
3. 确保数据处理一致
   • 严格按照训练时的预处理和后处理步骤对输入输出数据进行处理。可以编写详细的数据处理脚本，确保在训练和推理过程中使用相同的代码。
4. 更新环境和库版本
   • 确保训练和推理环境的硬件配置、软件版本等一致。更新相关的库和驱动到最新版本，以避免版本兼容性问题。