joe32140 commited on
Commit
e7f3cd9
·
verified ·
1 Parent(s): f93964d

fix: Include all Dense projection layers in ONNX export (output dim 64)

Browse files

## Problem

The current ONNX exports (model.onnx and model_int8.onnx) are missing the final Dense projection layers from PyLate's modules_list pipeline. This causes the ONNX models to output the intermediate hidden dimension instead of the correct final embedding dimension.

**mxbai-edge-colbert-v0-32m:**
- Current ONNX output: **768** dimensions (intermediate)
- Expected output: **64** dimensions (after all Dense projections)

## Root Cause

The ONNX export was done using the standard HuggingFace Transformers export pipeline, which only picks up the transformer backbone and possibly the first linear head. It does not include the additional Dense modules stored in the separate `1_Dense/`, `2_Dense/`, directories that are part of PyLate's `modules_list` architecture.

## Fix

Re-exported by wrapping the full pipeline (Transformer + all Dense layers) into a single module before ONNX export:
- **model.onnx**: fp32, opset 17, output dim = 64
- **model_int8.onnx**: int8 dynamic quantization with projection layers kept in fp32 (they are small and precision-sensitive)

Both files have been verified for:
- Correct output dimensions
- ONNX model validity (onnx.checker)
- Numerical consistency with PyTorch (fp32)

Files changed (2) hide show
  1. model.onnx +2 -2
  2. model_int8.onnx +2 -2
model.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:7d23d629e62ed15828ac7cc29c950ac178f78acd479e11aed0cfb4aeb5684134
3
- size 128931530
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:01df15011c2a155eee7524613fb9125b1ffada9393c3eab02209b49ced84b368
3
+ size 131529322
model_int8.onnx CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:3d4277bc1075725da0a813e7b91bf5ebc97d86cb7d4dbcf78989eeb352113a3f
3
- size 32566895
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1f0fdd343a447c32bd6853f4162424b36bc8fe83eb5a82db18a6e555538a852f
3
+ size 36052322