FredrikKarlssonSpeech/megatron-bert-base-swedish-cased-600k-onnx

megatron-bert-base-swedish-cased-600k - ONNX

ONNX export of KBLab/megatron-bert-base-swedish-cased-600k.

This package contains two task-specific exports:

• fill-mask/onnx/model.onnx for masked language modeling logits
• feature-extraction/onnx/model.onnx for last_hidden_state and cls_embedding

Available variants

Each task folder contains:

• model.onnx - FP32 baseline
• model_fp16.onnx - FP16, recommended
• model_int8.onnx - dynamic INT8
• model_uint8.onnx - dynamic UINT8
• model_q4.onnx - 4-bit MatMul quantization

File sizes

Fill-mask

• model.onnx: 738.6 MB
• model_fp16.onnx: 369.4 MB
• model_int8.onnx: 185.6 MB
• model_uint8.onnx: 185.6 MB
• model_q4.onnx: 283.8 MB

Feature-extraction

• model.onnx: 539.0 MB
• model_fp16.onnx: 269.6 MB
• model_int8.onnx: 135.4 MB
• model_uint8.onnx: 135.4 MB
• model_q4.onnx: 252.4 MB

Accuracy summary

FP32 parity vs PyTorch

• Fill-mask max logit diff: 0.000069
• Fill-mask top-5 tokens: exact match
• Feature last_hidden_state max diff: 0.000006
• Feature cls_embedding cosine similarity: 1.0

Quantized variants vs FP32 ONNX

Fill-mask

• fp16: top-5 exact match, max diff 0.0131
• int8: top-5 drift after rank 2
• uint8: top-5 drift after rank 2
• q4: top-5 drift after rank 2

Feature-extraction

• fp16: CLS cosine 0.9999997
• q4: CLS cosine 0.7404
• int8: CLS cosine 0.9731
• uint8: CLS cosine 0.9717

Recommendation

• Use model_fp16.onnx by default.
• For feature extraction, prefer fp16; int8/uint8 preserve embeddings better than q4 on this model.
• Avoid int8 and uint8 for fill-mask if token ranking fidelity matters.

Layout

megatron-bert-base-swedish-cased-600k/
├── fill-mask/
│   ├── config.json
│   ├── tokenizer.json
│   ├── tokenizer_config.json
│   ├── special_tokens_map.json
│   └── onnx/
│       ├── model.onnx
│       ├── model_fp16.onnx
│       ├── model_int8.onnx
│       ├── model_uint8.onnx
│       └── model_q4.onnx
└── feature-extraction/
    ├── config.json
    ├── tokenizer.json
    ├── tokenizer_config.json
    ├── special_tokens_map.json
    └── onnx/
        ├── model.onnx
        ├── model_fp16.onnx
        ├── model_int8.onnx
        ├── model_uint8.onnx
        └── model_q4.onnx

Export notes

• Native optimum 2.1.0 task export does not support megatron-bert.
• These graphs were exported with direct torch.onnx.export wrappers using legacy TorchScript exporter (dynamo=False).
• Feature extraction export uses cls_embedding = last_hidden_state[:, 0, :].
• No pooler output included, because loading MegatronBertModel from this MLM checkpoint would introduce randomly initialized pooler weights.