SEM Decoder Model PTQ Quantization Results in Severe Output Collapse with Only 3 Unique Values

[tanjiong233]

Problem Description

I am experiencing a severe output collapse issue when performing PTQ quantization on the SEM (Spatial Enhanced Manipulation) decoder model using the Horizon Robotics toolchain. The quantized HBM
model produces only 3 unique values, while the original ONNX model generates 895 unique values.

Environment Information

• Hardware Platform: RDK S100
• Toolchain Version: Open Explorer (hb_compile)
• Python Environment: Python 3.10 + miniconda
• Model Architecture: SEM Robotwin (Encoder-Decoder separated structure)
• Affected Model: Decoder (encoder quantization works normally)

Reproduction Steps

1. Model Export

Following the official tutorial at https://forum.d-robotics.cc/t/topic/32657 for ONNX export:

  # Quantization-friendly modifications
  data["joint_relative_pos"] = data["joint_relative_pos"].to(torch.int8)
  timestep = timestep.to(torch.int16)
  # Changed float("-inf") to -15

  # Export ONNX
  python3 onnx_scripts/export_onnx.py \
      --config config_sem_robotwin.py \
      --model /path/to/model \
      --output_path /path/to/onnx \
      --num_joint 14 \
      --validate

2. Calibration Data Preparation

Prepared 100 calibration samples following the tutorial, preprocessed and passed through the encoder to generate decoder input features:

# Calibration data shape validation
 noisy_action:      (1, 64, 14, 8)   float32  range: [-3.005, 3.252]
 image_feature:     (1, 3, 400, 256) float32  range: [-4.295, 4.137]
 robot_feature:     (1, 14, 1, 256)  float32  range: [-3.518, 3.075]
 timestep:          (1,)              int16    value: [999]
 joint_relative_pos: (1, 14, 14)     int8     range: [0, 13]

3. Attempted Quantization Configurations

I have tried multiple quantization configurations, all of which failed:

Configuration 1: Int16 Quantization (Recommended by Official Tutorial)

  calibration_parameters:
    cal_data_dir: ./onnx_cup_42k/calibration_data_dir/...
    quant_config: {
      "model_config": {
        "all_node_type": "int16"
      },Duplicate of #
      "op_config": {
        "Resize": {"qtype": "float16"}
      }
    }
  compiler_parameters:
    optimize_level: O2
    compile_mode: latency

Result: Failed - 3 unique values

Issue: Input quantization completely saturated the int16 range [-32768, 32767], causing severe clipping

---

Configuration 2: Int8 Quantization with O0 Optimization

  quant_config: {
    "model_config": {
      "all_node_type": "int8"
    },
    "op_config": {
      "Resize": {"qtype": "float16"}
    }
  }
  compiler_parameters:
    optimize_level: O0  # Changed from O2 to O0 to avoid over-optimization
    compile_mode: latency

Result: Failed - 3 unique values

Improvement: Input quantization range normal [-128, 127], no clipping, but output still collapsed

---

Configuration 3: Int8 with Layerwise Search and Bias Correction (Optimal Configuration)

  quant_config: {
    "model_config": {
      "all_node_type": "int8",
      "activation": {
        "calibration_type": ["max", "kl"],
        "max_percentile": [0.99995, 0.99999, 1.0],
        "num_bin": [2048, 4096],
        "asymmetric": [false, true]
      },
      "weight": {
        "bias_correction": {
          "num_sample": 10,
          "metric": "cosine-similarity"
        }
      },
      "layerwise_search": {
        "metric": "cosine-similarity"
      }
    },
    "op_config": {
      "Resize": {"qtype": "float16"},
      "MatMul": {"qtype": "float16"},
      "Gemm": {"qtype": "float16"}
    }
  }
  compiler_parameters:
    optimize_level: O0
    compile_mode: bandwidth  # Changed to bandwidth mode, prioritizing accuracy

Result: Still failed - 3 unique values

---

Test Results Comparison

I created a detailed comparison test script to test both ONNX and HBM models using identical input data:

Test Output

  ======================================================================
  TESTING ONNX MODEL
  ======================================================================
  Output Analysis:
     Shape: (1, 64, 14, 8)
     Dtype: float32
     Range: [-0.971276, 1.467759]
     Unique position values: 895/896
     Status: GOOD DIVERSITY

  ======================================================================
  TESTING HBM MODEL (Int8 + Layerwise Search)
  ======================================================================
  Input Information:
     noisy_action: dtype=S8, scale=1.90911591e-02
     image_feature: dtype=S8, scale=3.66429798e-02
     robot_feature: dtype=S8, scale=3.07448395e-02

  Quantizing inputs:
     noisy_action (S8): range=[-128, 127]  (no clipping)
     image_feature (S8): range=[-117, 113] (no clipping)
     robot_feature (S8): range=[-114, 100] (no clipping)

  Running HBM inference...
  Output Analysis:
     Shape: (1, 64, 14, 8)
     Dtype: float32
     Range: [-6.037731, 1.848850]
     Unique position values: 3/896
     LOW DIVERSITY - All unique values:
     [-2.934769, -2.449742, -0.26831594]

  ======================================================================
  COMPARISON SUMMARY
  ======================================================================
  Output Diversity:
     ONNX unique values: 895/896
     HBM  unique values: 3/896
     Diversity loss: 892 values

  Numerical Difference:
     Mean absolute error: 1.534658
     Max absolute error: 6.541608
     Median absolute error: 0.916659

  Conclusion:
     HBM model has SEVERE output collapse

Eliminated Issues

1. Input data issue: Using identical calibration data, ONNX model outputs normally
2. Quantization scale issue: Int8 quantization scale is reasonable, no clipping
3. Compilation configuration issue: Tried O0/O1/O2, latency/bandwidth multiple combinations
4. Calibration method issue: Tried max/kl multiple calibration methods and parameter combinations
5. Output format issue: Output correctly configured as float32 (removed Dequantize via remove_node_type)
6. Compilation log: No warnings or errors, compilation successful

Debug Information

HBM Model Input/Output Information

  # hbm_runtime query results
  Model: decoder_opt_resize
  Inputs:
    - noisy_action:      [1, 64, 14, 8]   INT8
    - image_feature:     [1, 3, 400, 256] INT8
    - robot_feature:     [1, 14, 1, 256]  INT8
    - timestep:          [1]              INT16
    - joint_relative_pos: [1, 14, 14]     INT8

  Output:
    - pred_action:       [1, 64, 14, 8]   FLOAT32
      Quant type: NONE
      Scale: []

Compilation Log Key Information

  2025-10-23 21:54:10 hb_compile completes running
  remove_node_type: Quantize;Cast;Softmax
  pred_action output [1, 64, 14, 8] FLOAT32
  No warnings or errors

[jin-yiwei]

Hi @tanjiong233 , I don't quite understand the specific problem you're asking. Can you provide the output of the export onnx command, which can determine the conversion error of the onnx model? From the information you've provided, your problem may be in the dequantization and quantization of the decoder. Based on what you said, using the same input, you can check the similarity between onnx and hbm between the encoder and decoder to confirm which part has the problem. In addition, based on the configuration in the blog we provided, you should be able to run through all the processes and get an output with a similarity greater than 0.99 using the same input, which has been confirmed by other users.

[tanjiong233]

Hi @tanjiong233 , I don't quite understand the specific problem you're asking. Can you provide the output of the export onnx command, which can determine the conversion error of the onnx model? From the information you've provided, your problem may be in the dequantization and quantization of the decoder. Based on what you said, using the same input, you can check the similarity between onnx and hbm between the encoder and decoder to confirm which part has the problem. In addition, based on the configuration in the blog we provided, you should be able to run through all the processes and get an output with a similarity greater than 0.99 using the same input, which has been confirmed by other users.

Thank you for your response! I've exported the ONNX models following the tutorial's quantization-friendly modifications.The ONNX models were exported successfully and can perform inference normally. Here is the output from the onnx_export.sh command:
onnx_export.sh

#!/bin/bash
export PYTHONPATH="${PWD}:${PYTHONPATH}"
python3 onnx_scripts/export_onnx.py \
    --config config_sem_robotwin.py \
    --model workspace/checkpoint_42k \
    --output_path "./onnx_cup_42k" \
    --num_joint 14 \
    --validate

output

(sem) tanjiong@tanjiong-Legion-Y9000P-IAH7H:~/workspace/RoboOrchardLab/projects/sem/robotwin$ bash onnx_export.sh 
/home/tanjiong/workspace/RoboOrchardLab/projects/sem/robotwin/onnx_scripts/export_onnx.py:478: DeprecationWarning: Call to deprecated function (or staticmethod) basic_config. (Use `robo_orchard_core.utils.logging module` instead.) -- Deprecated since version 0.2.0.
  log_basic_config(
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/datasets/utils/experimental.py:37: UserWarning: 'register_feature' is experimental and might be subject to breaking changes in the future.
  warnings.warn(
10/23/2025 01:49:54 INFO base_lmdb_dataset.py:187 | dataset length: 40711, number of episode: 196
10/23/2025 01:49:54 WARNING robotwin_lmdb_dataset.py:127 | dataset T_base2world is not set, use default.
10/23/2025 01:49:54 WARNING robotwin_lmdb_dataset.py:127 | dataset T_base2world is not set, use default.
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:138: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  output_h = math.ceil(input_h / stride_h)
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:139: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  output_w = math.ceil(input_w / stride_w)
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:140: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  pad_h = max(
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:147: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  pad_w = max(
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:158: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if pad_h > 0 or pad_w > 0:
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:611: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert L == H * W, "input feature has wrong size"
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:707: TracerWarning: Converting a tensor to a Python integer might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  B = int(windows.shape[0] / (H * W / window_size / window_size))  # noqa: N806
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:686: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if pad_r > 0 or pad_b:
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/modules/swin_transformer.py:404: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert L == H * W, "input feature has wrong size"
/home/tanjiong/workspace/RoboOrchardLab/projects/sem/robotwin/onnx_scripts/export_onnx.py:128: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert depths.shape[1] == num_cams
/home/tanjiong/miniconda3/envs/sem/lib/python3.10/site-packages/robo_orchard_lab/models/sem_modules/layers.py:280: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if B != query_pos.shape[0]:
2025-10-23 01:49:59.164810378 [W:onnxruntime:, graph.cc:120 MergeShapeInfo] Error merging shape info for output. '/robot_encoder/Concat_output_0' source:{3} target:{4}. Falling back to lenient merge.
2025-10-23 01:49:59.181588700 [W:onnxruntime:, graph.cc:120 MergeShapeInfo] Error merging shape info for output. '/Concat_12_output_0' source:{3} target:{4}. Falling back to lenient merge.
2025-10-23 01:49:59.437847139 [W:onnxruntime:, graph.cc:120 MergeShapeInfo] Error merging shape info for output. '/Concat_1_output_0' source:{3} target:{4}. Falling back to lenient merge.
2025-10-23 01:49:59.438595681 [W:onnxruntime:, graph.cc:120 MergeShapeInfo] Error merging shape info for output. '/Concat_4_output_0' source:{3} target:{4}. Falling back to lenient merge.
10/23/2025 01:50:13 INFO export_onnx.py:334 | prepare calib data done
10/23/2025 01:50:13 INFO export_onnx.py:453 | Diff of image_feature: max_diff[0.0014872550964355469] mean[0.00010763126192614436]
10/23/2025 01:50:13 INFO export_onnx.py:458 | Diff of robot_feature: max_diff[1.430511474609375e-06] mean[2.2179271752520435e-07]
10/23/2025 01:50:13 INFO export_onnx.py:463 | Diff of decoder_pred: max_diff[0.0002473592758178711] mean[3.9341441151918843e-05]

export_onnx.py

# Project RoboOrchard
#
# Copyright (c) 2024-2025 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.

import argparse
import logging
import os
from functools import partial

import numpy as np
import torch
from misc import (
    DECODER_INPUT_NAMES,
    DECODER_OUTPUT_NAMES,
    ENCODER_INPUT_NAMES,
    ENCODER_OUTPUT_NAMES,
    get_pixels,
    unsqueeze_batch,
)
from onnx_model import OnnxSEMModel, OnnxSEMModelConfig
from torch import nn
from utils import load_config

from robo_orchard_lab.models.mixin import ModelMixin
from robo_orchard_lab.models.sem_modules.layers import AdaRMSNorm, rotate_half
from robo_orchard_lab.utils import log_basic_config

logger = logging.getLogger(__file__)


class DeployDataPreprocessor:
    def __init__(
        self,
        model,
        num_joint=14,
        pred_steps=64,
        do_unsqueeze_batch=True,
        get_noisy_action=True,
    ):
        self.preprocessor = model.data_preprocessor.eval()
        strides = model.cfg.data_preprocessor["batch_transforms"][0]["stride"]
        strides = tuple(strides[i] for i in model.decoder.feature_level)
        self.strides = strides
        self.pixels = None
        self.num_joint = num_joint
        self.pred_steps = pred_steps
        self.do_unsqueeze_batch = do_unsqueeze_batch
        self.get_noisy_action = get_noisy_action

    def __call__(self, data):
        data["projection_mat_inv"] = torch.linalg.inv(data["projection_mat"])
        if self.do_unsqueeze_batch:
            data = unsqueeze_batch(data)
        data = self.preprocessor(data)
        if self.pixels is None:
            self.pixels = get_pixels(data, self.strides)
        data["pixels"] = self.pixels.clone()
        if self.get_noisy_action:
            noise = torch.randn((self.pred_steps, self.num_joint)).to(
                data["imgs"]
            )
            noisy_action = data["kinematics"][0].joint_state_to_robot_state(
                noise, data.get("embodiedment_mat")[0]
            )
            data["noisy_action"] = noisy_action.unsqueeze(dim=0)
        return data


def rms_norm(layer, x):
    rms = x.pow(2).mean(dim=-1, keepdim=True)
    if layer.eps is not None:
        rms = rms + layer.eps
    rms = rms.sqrt()
    x = x / rms
    if layer.weight is not None:
        x = x * layer.weight
    return x


def ada_rms_norm(layer, x, c):
    x = rms_norm(layer, x)
    return layer.apply_ada_func(x, c)


class Encoder(nn.Module):
    def __init__(self, model):
        super().__init__()
        model.cuda().eval().requires_grad_(False)
        self.backbone = model.backbone
        self.neck = model.neck
        self.backbone_3d = model.backbone_3d
        self.neck_3d = model.neck_3d
        self.spatial_enhancer = model.spatial_enhancer
        self.robot_encoder = model.decoder.robot_encoder
        self.feature_level = model.decoder.feature_level
        for layer in self.robot_encoder.layers:
            if isinstance(layer, torch.nn.RMSNorm):
                layer.forward = partial(rms_norm, layer)

    def extract_feat(self, imgs, depths=None):
        if imgs.dim() == 5:
            bs, num_cams = imgs.shape[:2]
            imgs = imgs.flatten(end_dim=1)
        else:
            bs = imgs.shape[0]
            num_cams = 1

        feature_maps = self.backbone(imgs)
        if self.neck is not None:
            feature_maps = self.neck(feature_maps)
        feature_maps = [x.unflatten(0, (bs, num_cams)) for x in feature_maps]
        feature_maps = [feature_maps[i] for i in self.feature_level]
        feature_maps = self.flat_feature_maps(feature_maps)

        if self.backbone_3d is not None:
            assert depths.shape[1] == num_cams
            depths = depths.flatten(end_dim=1)
            feature_3d = self.backbone_3d(depths)
            if self.neck_3d is not None:
                feature_3d = self.neck_3d(feature_3d)
            feature_3d = [x.unflatten(0, (bs, num_cams)) for x in feature_3d]
            feature_3d = [feature_3d[i] for i in self.feature_level]
            feature_3d = self.flat_feature_maps(feature_3d)
        else:
            feature_3d = None
        return feature_maps, feature_3d

    def forward(
        self,
        imgs,
        depths,
        pixels,
        projection_mat_inv,
        hist_robot_state,
        joint_scale_shift,
        joint_relative_pos,
    ):
        feature_maps, feature_3d = self.extract_feat(imgs, depths)

        feature_maps = self.forward_spatial_enhancer(
            feature_maps, feature_3d, pixels, projection_mat_inv
        )

        hist_robot_state = torch.cat(
            [
                (
                    hist_robot_state[..., :1]
                    - joint_scale_shift[:, None, :, 1:2]
                )
                / joint_scale_shift[:, None, :, :1],
                hist_robot_state[..., 1:],
            ],
            dim=-1,
        )
        robot_feature = self.robot_encoder(
            hist_robot_state, joint_relative_pos
        )
        return feature_maps, robot_feature

    def flat_feature_maps(self, feature_maps):
        feature_maps = torch.cat(
            [x.flatten(start_dim=-2).transpose(-1, -2) for x in feature_maps],
            dim=-2,
        )
        return feature_maps

    def forward_spatial_enhancer(
        self, feature_2d, feature_3d, pixels, projection_mat_inv
    ):
        pts = self.get_pts(pixels, projection_mat_inv)

        se = self.spatial_enhancer
        if se.with_feature_3d:
            depth_prob_feat = se.pts_prob_pre_fc(feature_2d)
            depth_prob_feat = torch.cat([depth_prob_feat, feature_3d], dim=-1)
            depth_prob = se.pts_prob_fc(depth_prob_feat).softmax(dim=-1)
            feature_fused = [feature_2d, feature_3d]
        else:
            depth_prob = se.pts_prob_fc(feature_2d).softmax(dim=-1)
            feature_fused = [feature_2d]

        pts_feature = se.pts_fc(pts)
        pts_feature = (depth_prob.unsqueeze(dim=-1) * pts_feature).sum(dim=-2)
        feature_fused.append(pts_feature)
        feature_fused = torch.cat(feature_fused, dim=-1)
        feature_fused = se.fusion_fc(feature_fused) + feature_2d
        feature_fused = se.fusion_norm(feature_fused)
        return feature_fused

    def get_pts(self, pixels, projection_mat_inv):
        depths = torch.linspace(
            self.spatial_enhancer.min_depth,
            self.spatial_enhancer.max_depth,
            self.spatial_enhancer.num_depth,
        ).to(pixels)  # D
        depths = depths[None, :, None]  # 1 D 1
        pts = pixels * depths  # N 1 2 x 1 D 1 = N D 2
        depths = depths.tile(pixels.shape[0], 1, 1)  # N D 1
        pts = torch.cat([pts, depths, torch.ones_like(depths)], dim=-1)  # ND4
        pts = (
            pts
            @ projection_mat_inv.permute(0, 1, 3, 2).unsqueeze(dim=2)[..., :3]
        )  # ND4 x bc144
        return pts


class Decoder(nn.Module):
    def __init__(self, model, hist_steps=1, num_joint=14):
        super().__init__()
        model.cuda().eval().requires_grad_(False)
        self.decoder = model.decoder
        for op, layer in zip(
            self.decoder.operation_order, self.decoder.layers, strict=False
        ):
            if isinstance(layer, AdaRMSNorm):
                layer.forward = partial(ada_rms_norm, layer)
            elif isinstance(layer, torch.nn.RMSNorm):
                layer.forward = partial(rms_norm, layer)
            elif op == "temp_cross_attn":
                num_pred_chunk = (
                    self.decoder.pred_steps // self.decoder.chunk_size
                )
                num_hist_chunk = (
                    hist_steps // self.decoder.robot_encoder.chunk_size
                )
                query_pos = (
                    torch.arange(num_pred_chunk)[None].tile(num_joint, 1)
                    + num_hist_chunk
                )
                key_pos = torch.arange(num_pred_chunk + num_hist_chunk)[
                    None
                ].tile(num_joint, 1)
                layer.cos_query = layer.position_encoder.cos[
                    query_pos
                ].unsqueeze(1)
                layer.sin_query = layer.position_encoder.sin[
                    query_pos
                ].unsqueeze(1)
                layer.cos_key = layer.position_encoder.cos[key_pos].unsqueeze(
                    1
                )
                layer.sin_key = layer.position_encoder.sin[key_pos].unsqueeze(
                    1
                )

                def apply_position_encode(layer, q, query_pos, k, key_pos):
                    q = (q * layer.cos_query.to(q)) + (
                        rotate_half(q) * layer.sin_query.to(q)
                    )
                    k = (k * layer.cos_key.to(k)) + (
                        rotate_half(k) * layer.sin_key.to(k)
                    )
                    return q, k

                layer.apply_position_encode = partial(
                    apply_position_encode, layer
                )

        for layers in self.decoder.head.act_and_norm:
            for layer in layers:
                if isinstance(layer, torch.nn.RMSNorm):
                    layer.forward = partial(rms_norm, layer)

    def forward(
        self,
        noisy_action,
        image_feature,
        robot_feature,
        timestep,
        joint_relative_pos,
    ):
        image_feature = image_feature.flatten(1, 2)
        pred = self.decoder.forward_layers(
            noisy_action,
            image_feature,
            robot_feature=robot_feature,
            timesteps=timestep,
            joint_relative_pos=joint_relative_pos,
        )
        return pred


def prepare_calib_data(dataset, data_processor, encoder, sample_dataset=True):
    """准备PTQ量化的校准数据"""
    os.makedirs("calibration_data_dir", exist_ok=True)
    for e_name in ENCODER_INPUT_NAMES:
        os.makedirs(os.path.join("calibration_data_dir", e_name), exist_ok=True)
    for d_name in DECODER_INPUT_NAMES:
        os.makedirs(os.path.join("calibration_data_dir", d_name), exist_ok=True)

    if sample_dataset:
        new_dataset = []
        # 均匀采样数据集，这里取100份数据作为校准数据（可调整）
        num_samples = 100
        indices = np.arange(0, len(dataset), len(dataset) // num_samples)[:num_samples]
        for idx in indices:
            new_dataset.append(dataset[idx])
        dataset = new_dataset

    for i in range(len(dataset)):
        # 校准数据需经过前处理
        data = data_processor(dataset[i])
        # 对齐onnx导出时的输入信息
        data["joint_relative_pos"] = data["joint_relative_pos"].to(torch.int8)
        # 注意：不进行squeeze操作，保持batch维度，以匹配ONNX模型的实际输入形状

        for e_name in ENCODER_INPUT_NAMES:
            np.save(os.path.join("calibration_data_dir", e_name, f"{e_name}_{i}.npy"), data[e_name].cpu().numpy())

        image_feature, robot_feature = encoder(*[data[key] for key in ENCODER_INPUT_NAMES])
        data["image_feature"] = image_feature
        data["robot_feature"] = robot_feature
        data["timestep"] = torch.tensor([999], dtype=torch.int16)

        for d_name in DECODER_INPUT_NAMES:
            np.save(os.path.join("calibration_data_dir", d_name, f"{d_name}_{i}.npy"), data[d_name].cpu().numpy())

        # 不采样的话默认取前200份数据作为校准数据
        if not sample_dataset and i == 199:
            break

    logger.info("prepare calib data done")


def export_onnx(config, model, num_joint, output_path, validate=True):
    model = ModelMixin.load_model(model).eval().cuda().requires_grad_(False)

    config = load_config(config)
    base_config = config.config
    if base_config["multi_task"]:
        raise NotImplementedError(
            "Only single-task model export is currently supported"
        )
    build_dataset = config.build_dataset
    dataset = build_dataset(base_config)[0]

    data_processor = DeployDataPreprocessor(
        model,
        num_joint=num_joint,
        pred_steps=base_config["pred_steps"],
        do_unsqueeze_batch=True,
    )
    data = dataset[0]
    data = data_processor(data)

    encoder = Encoder(model)
    decoder = Decoder(model, base_config["hist_steps"], num_joint)

    # 量化友好修改：转换数据类型
    data["joint_relative_pos"] = data["joint_relative_pos"].to(torch.int8)

    image_feature, robot_feature = encoder(
        *[data[key] for key in ENCODER_INPUT_NAMES]
    )

    os.makedirs(output_path, exist_ok=True)
    opset_version = 19
    encoder_onnx = os.path.join(output_path, "encoder.onnx")
    decoder_onnx = os.path.join(output_path, "decoder.onnx")

    # export encoder onnx
    torch.onnx.export(
        encoder,
        tuple(data[key] for key in ENCODER_INPUT_NAMES),
        encoder_onnx,
        input_names=ENCODER_INPUT_NAMES,
        output_names=ENCODER_OUTPUT_NAMES,
        opset_version=opset_version,
    )

    # export decoder onnx
    timestep = torch.tensor([999]).to(data["imgs"])
    timestep = timestep.to(torch.int16)
    torch.onnx.export(
        decoder,
        (
            data["noisy_action"],
            image_feature,
            robot_feature,
            timestep,
            data["joint_relative_pos"],
        ),
        decoder_onnx,
        input_names=DECODER_INPUT_NAMES,
        output_names=DECODER_OUTPUT_NAMES,
        opset_version=opset_version,
    )

    # export OnnxSEMModel config json
    onnx_model = OnnxSEMModel(
        OnnxSEMModelConfig(
            data_preprocessor=model.cfg.data_preprocessor,
            onnx_path=output_path,
            test_noise_scheduler=model.cfg.decoder["test_noise_scheduler"],
            num_inference_timesteps=model.cfg.decoder[
                "num_inference_timesteps"
            ],
            strides=data_processor.strides,
        )
    )
    onnx_model.save_model(output_path)

    # 准备校准数据
    prepare_calib_data(dataset, data_processor, encoder)

    if validate:
        image_feature_onnx, robot_feature_onnx = (
            onnx_model.encoder_session.run(
                ENCODER_OUTPUT_NAMES,
                {key: data[key].cpu().numpy() for key in ENCODER_INPUT_NAMES},
            )
        )

        decoder_pred = decoder(
            data["noisy_action"],
            image_feature,
            robot_feature,
            timestep,
            data["joint_relative_pos"],
        )
        decoder_pred_onnx = onnx_model.decoder_session.run(
            DECODER_OUTPUT_NAMES,
            {
                "noisy_action": data["noisy_action"].cpu().numpy(),
                "image_feature": image_feature_onnx,
                "robot_feature": robot_feature_onnx,
                "timestep": timestep.cpu().numpy(),
                "joint_relative_pos": data["joint_relative_pos"].cpu().numpy(),
            },
        )[0]

        diff_image_feature = np.abs(
            image_feature_onnx - image_feature.cpu().numpy()
        )
        diff_robot_feature = np.abs(
            robot_feature_onnx - robot_feature.cpu().numpy()
        )
        diff_decoder_pred = np.abs(
            decoder_pred_onnx - decoder_pred.cpu().numpy()
        )
        logger.info(
            "Diff of image_feature: "
            f"max_diff[{diff_image_feature.max()}] "
            f"mean[{diff_image_feature.mean()}]"
        )
        logger.info(
            "Diff of robot_feature: "
            f"max_diff[{diff_robot_feature.max()}] "
            f"mean[{diff_robot_feature.mean()}]"
        )
        logger.info(
            "Diff of decoder_pred: "
            f"max_diff[{diff_decoder_pred.max()}] "
            f"mean[{diff_decoder_pred.mean()}]"
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=str)
    parser.add_argument("--model", type=str)
    parser.add_argument("--output_path", type=str, default="./onnx_files")
    parser.add_argument("--num_joint", type=int, default=14)
    parser.add_argument("--validate", action="store_true")
    args = parser.parse_args()
    log_basic_config(
        format="%(asctime)s %(levelname)s %(filename)s:%(lineno)d | %(message)s",  # noqa: E501
        level=logging.INFO,
    )
    export_onnx(
        args.config,
        args.model,
        args.num_joint,
        args.output_path,
        args.validate,
    )

[Chloe-Yan-3465]

Hi @tanjiong233,

Thanks for sharing the detailed comparison logs between ONNX and HBM models — they’re really helpful.

I have a question about the following part of your test output:

Test Results Comparison

COMPARISON SUMMARY

Output Diversity:
ONNX unique values: 895/896
HBM unique values: 3/896
Diversity loss: 892 values

Could you please clarify how this comparison was made?

• What input data were used for both ONNX and HBM inference (e.g., random noise, real dataset samples, or fixed calibration data)?
• How exactly were the “unique values” counted — was it simply np.unique(output) or some quantized metric over the tensor?
• Was the output taken from the final layer or an intermediate feature map?

Understanding this would really help me reproduce the result and diagnose the diversity collapse issue on my side.

By the way, I'm also an RMer . If it's appropriate, I’d love to get in touch and exchange some experiences. My contact information is available on my GitHub profile.

Thanks in advance!

[tanjiong233]

Hi @jin-yiwei ,Thanks for your interest in the details. Here's a comprehensive explanation of the comparison methodology:

1. Input Data Source

The input data used for both ONNX and HBM inference was calibration data generated during the ONNX export process:

Calibration Data Generation

(from export_onnx.py:295-335)

• Dataset sampling: 100 samples were uniformly sampled from the training dataset
• Preprocessing pipeline: Each sample went through:
  1. Data preprocessing (normalization, augmentation, etc.)
  2. Encoder inference to generate image_feature and robot_feature
  3. Fixed timestep assignment: timestep = 999 (int16)
  4. Data type conversion: joint_relative_pos converted to int8

Test Input

(from test_onnx_vs_hbm.py:12-28)

# Loaded calibration sample #0 from the pre-generated calibration dataset
sample_id = 0
data = {
    'noisy_action':      shape=(1, 64, 14, 8),   dtype=float32
    'image_feature':     shape=(1, 3, 400, 256), dtype=float32
    'robot_feature':     shape=(1, 14, 1, 256),  dtype=float32
    'timestep':          shape=(1,),             dtype=int16
    'joint_relative_pos': shape=(1, 14, 14),    dtype=int8
}

Key point: Both ONNX and HBM models received identical input tensors to ensure fair comparison.

2. "Unique Values" Counting Methodology

The unique value counting was performed on a specific slice of the output tensor:

Code Implementation

(from test_onnx_vs_hbm.py:43-44, 115-116)

# Extract the first action dimension (position values) from all timesteps and joints
extracted_positions = pred_action[0, :, :, 0]  # Shape: (64, 14) = 896 values
unique_count = len(np.unique(extracted_positions))

Explanation

• Output tensor shape: (1, 64, 14, 8)
  • Batch=1, Timesteps=64, Joints=14, ActionDim=8
• Extracted slice: [batch=0, all_timesteps, all_joints, first_dim]
  • This gives 64 × 14 = 896 total values
• Metric: Standard np.unique() function counts distinct floating-point values
• Not quantized: Direct count of unique values in the float32 output

Why This Metric?

In a healthy diffusion model output, predicted action values should be diverse across timesteps and joints. Collapsed output (only 3 unique values) indicates severe quantization degradation.

3. Output Layer

The comparison was performed on the final model output:

• Layer: pred_action - the decoder's final prediction
• Position in model: Final output after all decoder layers and head
• NOT an intermediate feature map: This is the actual action prediction used for robot control
• Data type: float32 (confirmed in both models)

ONNX Output

(from test_onnx_vs_hbm.py:38-40)

outputs = session.run(None, inputs)
pred_action = outputs[0]  # Final output tensor

HBM Output

(from test_onnx_vs_hbm.py:111-112)

outputs = model.run(inputs_quantized, model_name=model_name)
pred_action = outputs[model_name]['pred_action']  # Final output tensor

Summary Table

Aspect	Details
Input Data	Pre-generated calibration data (sample #0 of 100 samples), preprocessed through encoder
Unique Value Method	np.unique() on output slice [0, :, :, 0] → 896 total values
Output Layer	Final model output pred_action, not intermediate features
Data Type	float32 output, no post-processing quantization

Conclusion

The test demonstrates that under identical input conditions, the ONNX model produces diverse outputs (895/896 unique values) while the HBM quantized model suffers severe output collapse (only 3 unique values), confirming the PTQ quantization issue.

Code References

• Full test script: test_onnx_vs_hbm.py

#!/usr/bin/env python3
"""
对比测试ONNX模型和HBM模型的输出
用于验证HBM量化编译过程是否导致输出塌陷
"""

import numpy as np
import onnxruntime as ort
from hbm_runtime import HB_HBMRuntime
import os

def load_calibration_sample(sample_id=0):
    """加载校准数据样本"""
    calib_dir = "/home/sunrise/workspace/src/checkpoints/onnx_cup_42k/calibration_data_dir"

    data = {
        'noisy_action': np.load(f'{calib_dir}/noisy_action/noisy_action_{sample_id}.npy'),
        'image_feature': np.load(f'{calib_dir}/image_feature/image_feature_{sample_id}.npy'),
        'robot_feature': np.load(f'{calib_dir}/robot_feature/robot_feature_{sample_id}.npy'),
        'timestep': np.load(f'{calib_dir}/timestep/timestep_{sample_id}.npy'),
        'joint_relative_pos': np.load(f'{calib_dir}/joint_relative_pos/joint_relative_pos_{sample_id}.npy')
    }

    print(f"✅ Loaded calibration sample {sample_id}:")
    for name, arr in data.items():
        print(f"   {name}: shape={arr.shape}, dtype={arr.dtype}, range=[{arr.min():.3f}, {arr.max():.3f}]")

    return data

def test_onnx_model(inputs):
    """测试ONNX模型"""
    print("\n" + "="*70)
    print("🔷 TESTING ONNX MODEL")
    print("="*70)

    onnx_path = "/home/sunrise/workspace/src/checkpoints/onnx_cup_42k/decoder.onnx"
    session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])

    outputs = session.run(None, inputs)
    pred_action = outputs[0]

    # 分析输出
    extracted_positions = pred_action[0, :, :, 0]
    unique_count = len(np.unique(extracted_positions))

    print(f"📊 Output Analysis:")
    print(f"   Shape: {pred_action.shape}")
    print(f"   Dtype: {pred_action.dtype}")
    print(f"   Range: [{pred_action.min():.6f}, {pred_action.max():.6f}]")
    print(f"   Unique position values: {unique_count}/896")

    if unique_count < 100:
        print(f"   ⚠️  LOW DIVERSITY - All unique values:")
        print(f"   {np.unique(extracted_positions)}")
    else:
        print(f"   ✅ GOOD DIVERSITY - First 10 unique values:")
        print(f"   {np.unique(extracted_positions)[:10]}")

    return pred_action, unique_count

def test_hbm_model(inputs_float):
    """测试HBM模型"""
    print("\n" + "="*70)
    print("🔶 TESTING HBM MODEL")
    print("="*70)

    model_path = "/home/sunrise/workspace/src/hbm_model/hbm_cup_42k/decoder_opt_resize.hbm"
    model = HB_HBMRuntime(model_path)
    model_name = model.model_names[0]

    # 获取量化参数和数据类型
    print(f"📋 Input Information:")
    scales = {}
    dtypes = {}
    for inp_name in ['noisy_action', 'image_feature', 'robot_feature']:
        scale = model.input_quants[model_name][inp_name].scale[0]
        dtype = model.input_dtypes[model_name][inp_name].name
        scales[inp_name] = scale
        dtypes[inp_name] = dtype
        print(f"   {inp_name}: dtype={dtype}, scale={scale:.8e}")

    # 量化输入（自动适配int8/int16）
    print(f"\n🔢 Quantizing inputs...")
    inputs_quantized = {}
    for name in ['noisy_action', 'image_feature', 'robot_feature']:
        dtype = dtypes[name]

        # 根据数据类型选择量化范围和目标dtype
        if dtype == 'S8':
            quantized = np.clip(
                np.round(inputs_float[name] / scales[name]),
                -128, 127
            ).astype(np.int8)
        elif dtype == 'S16':
            quantized = np.clip(
                np.round(inputs_float[name] / scales[name]),
                -32768, 32767
            ).astype(np.int16)
        else:
            raise ValueError(f"Unsupported dtype: {dtype}")

        inputs_quantized[name] = quantized
        print(f"   {name} ({dtype}): range=[{quantized.min()}, {quantized.max()}]")

    # 其他输入不需要量化
    inputs_quantized['timestep'] = inputs_float['timestep']
    inputs_quantized['joint_relative_pos'] = inputs_float['joint_relative_pos']

    # 运行推理
    print(f"\n🚀 Running HBM inference...")
    outputs = model.run(inputs_quantized, model_name=model_name)
    pred_action = outputs[model_name]['pred_action']

    # 分析输出
    extracted_positions = pred_action[0, :, :, 0]
    unique_count = len(np.unique(extracted_positions))

    print(f"📊 Output Analysis:")
    print(f"   Shape: {pred_action.shape}")
    print(f"   Dtype: {pred_action.dtype}")
    print(f"   Range: [{pred_action.min():.6f}, {pred_action.max():.6f}]")
    print(f"   Unique position values: {unique_count}/896")

    if unique_count < 100:
        print(f"   ⚠️  LOW DIVERSITY - All unique values:")
        print(f"   {np.unique(extracted_positions)}")
    else:
        print(f"   ✅ GOOD DIVERSITY - First 10 unique values:")
        print(f"   {np.unique(extracted_positions)[:10]}")

    return pred_action, unique_count

def compare_outputs(onnx_output, hbm_output, onnx_unique, hbm_unique):
    """对比ONNX和HBM输出"""
    print("\n" + "="*70)
    print("📈 COMPARISON SUMMARY")
    print("="*70)

    # 计算差异
    diff = np.abs(onnx_output - hbm_output)

    print(f"📊 Output Diversity:")
    print(f"   ONNX unique values: {onnx_unique}/896 {'✅' if onnx_unique > 800 else '⚠️'}")
    print(f"   HBM  unique values: {hbm_unique}/896 {'✅' if hbm_unique > 800 else '❌'}")
    print(f"   Diversity loss: {onnx_unique - hbm_unique} values")

    print(f"\n📊 Numerical Difference:")
    print(f"   Mean absolute error: {diff.mean():.6f}")
    print(f"   Max absolute error: {diff.max():.6f}")
    print(f"   Median absolute error: {np.median(diff):.6f}")

    print(f"\n🎯 Conclusion:")
    if hbm_unique < 100:
        print(f"   ❌ HBM model has SEVERE output collapse!")
        print(f"   ❌ Only {hbm_unique} unique values vs {onnx_unique} in ONNX")
        print(f"   🔧 This confirms the quantization compilation issue")
        return False
    else:
        print(f"   ✅ HBM model produces diverse outputs")
        return True

def main():
    print("\n" + "="*70)
    print("🧪 ONNX vs HBM MODEL COMPARISON TEST")
    print("="*70)
    print("Purpose: Verify if HBM quantization causes output collapse\n")

    # 加载测试数据
    inputs = load_calibration_sample(sample_id=0)

    # 测试ONNX模型
    onnx_output, onnx_unique = test_onnx_model(inputs)

    # 测试HBM模型
    hbm_output, hbm_unique = test_hbm_model(inputs)

    # 对比结果
    is_ok = compare_outputs(onnx_output, hbm_output, onnx_unique, hbm_unique)

    if not is_ok:
        print("\n" + "="*70)
        print("💡 NEXT STEPS:")
        print("="*70)
        print("1. Try different quantization configurations (optimize_level, quant_config)")
        print("2. Test with int8 instead of int16 quantization")
        print("3. Try mixed precision (float16 for sensitive layers)")
        print("4. Check if specific layers need special handling")
        print("="*70)

    return 0 if is_ok else 1

if __name__ == "__main__":
    exit(main())