Add sharded ONNX layout, transformers.js config, and generation config

#1
by yakuraku - opened

Add sharded ONNX layout, transformers.js config, and generation config

Hi, and thanks for publishing this one in ONNX.

I tried running Maincoder-1B in a browser tab with transformers.js and it failed
after about a second, before downloading a single byte of weights. Tracing it,
the cause turned out to be the size of one file rather than anything about the
model itself, so I put a fix together and then spent most of the time checking
that none of the weights changed.

Everything here is additive, which does mean it adds about 4.7 GiB to the repo.
If you would rather these replaced the existing files than sat alongside them, or
would prefer a different approach entirely, I am glad to rework it.

Summary

This adds files. Nothing existing is removed, renamed, moved, or overwritten.
decoder_model.onnx, decoder_with_past_model.onnx, and both .onnx_data files
keep their current names, locations, and bytes, so code written against the
current README continues to work unchanged.

The additions let the model load in transformers.js with no options, in Node
and in the browser, and let generation stop at the end of a chat turn.

Added

onnx/model.onnx              1,677,425 B    re-layout of decoder_with_past_model.onnx
onnx/model.onnx_data     1,602,385,920 B
onnx/model.onnx_data_1   1,610,612,736 B
onnx/model.onnx_data_2     891,813,888 B
onnx/model.onnx_data_3     933,494,976 B
generation_config.json          77 B

Two keys added to existing files; every pre-existing key is unchanged:

// config.json
"transformers.js_config": {
  "use_external_data_format": { "model.onnx": 4 }
}

// tokenizer_config.json
"chat_template": "<contents of chat_template.jinja>"

Rationale

Sharded external data

decoder_with_past_model.onnx_data is 5,038,307,520 bytes in a single file.
transformers.js reads a model file by pre-allocating one Uint8Array of the
whole Content-Length (src/utils/hub/utils.js, readResponse). On Chrome 151 /
Windows 11 the largest single ArrayBuffer that could be allocated measured
2,145,189,888 bytes, so the read throws before any weight bytes transfer:

RangeError: Array buffer allocation failed
    at new Uint8Array (<anonymous>)
    at readResponse (transformers.web.js:6005:16)
    at loadResourceFile (transformers.web.js:6657:24)

This happens in the fetch path, above the execution provider, so WebGPU and WASM
fail identically and neither VRAM nor the storage quota is reached.

The allocation is per file. Four shards under 1.5 GiB each load where one 4.69 GiB
file cannot. Chunk names follow getExternalDataChunkNames
(src/utils/model-loader.js), and the count is declared through
use_external_data_format, which accepts a chunk count as well as a boolean
(resolveExternalDataFormat, same file).

transformers.js_config

The library defaults to subfolder: 'onnx' and a model_file_name derived from
the architecture, so the current root-level decoder_* files are not found
without explicit overrides, and external data is not fetched at all unless
declared. With onnx/model.onnx and this key,
pipeline('text-generation', ...) resolves with no options.

chat_template in tokenizer_config.json

AutoTokenizer reads chat_template from tokenizer_config.json or
chat_template.json. It does not read chat_template.jinja, so callers
currently have to fetch that file and pass it to apply_chat_template manually.
chat_template.jinja is left in place unchanged; this is a copy in a location the
tokenizer reads.

generation_config.json

The repo ships none, and config.json declares eos_token_id: 151643
(<|endoftext|>), while the chat template ends a turn with <|im_end|> (151645).
Measured over a fixed 20-prompt set, greedy, max_new_tokens=128:

Terminating condition Turns
`< im_end
`< endoftext
reached max_new_tokens 11

The declared EOS never terminated a turn; the one that does was undeclared.
Without this file those 9 turns continue past the end of the assistant turn.

This is a behaviour change: callers relying on fixed-length completions will see
shorter output. Text before the stop point is unchanged.

Bit-exactness

onnx/model.onnx holds the same weights as decoder_with_past_model.onnx. Only
the location and offset of each external tensor changed. No dtype conversion,
constant folding, operator fusion, graph optimisation, or re-export was performed.

Tensor bytes. All 356 external tensors were read from the original
.onnx_data and from the new shards and compared by SHA-256: 356/356
identical
, with byte totals reconciling to 5,038,307,520 on both sides.

356 rather than 355 because one tensor is not an initializer: the RoPE
inverse-frequency table ([1, 48, 1], 192 bytes) is stored in a Constant node
attribute. A traversal limited to graph.initializer misses it.

Graph structure. Compared with external data stripped: identical ir_version
(8), opset_import (ai.onnx:18), producer, node count (9064), node op types,
names, inputs and outputs in order, graph input and output names and types,
initializer names, shapes and dtypes, external tensor inventory, and non-external
node attributes.

Behaviour. 20 fixed prompts (12 code completion, 4 multi-turn, 4 edge cases
including a 1890-token prompt, CJK/RTL/emoji, and irregular whitespace), greedy,
max_new_tokens=128, EOS held equal on both sides: 20/20 identical token id
sequences
. Raw first-token logits (151,936 fp32 values per prompt) are
bitwise identical, max |Δ| = 0.0.

Both checks are reproducible with the scripts in this PR; see VERIFICATION.md.

Measured before / after

Node 25, @huggingface/transformers@4.2.0, fp32 CPU:

Before After
Load options required 3 0
pipeline() with no options fails works
apply_chat_template() with no template fails works
Warm load 22.3 s 20.7 s
Token ids over 20 prompts reference identical

Chrome 151, RTX 4060 Laptop, WebGPU, served over HTTP:

Before After
Session creation RangeError at ~1.2 s succeeds
Time to first token n/a 0.79 s
Throughput n/a 9.51 tok/s
Warm load n/a 16.8 s fetch + 39.7 s session init

The sharded fp32 model is 4.70 GiB. At the ~4.2 MiB/s measured from the Hub CDN
that is roughly a 19 minute first download, so this makes browser loading work
rather than making it quick.

WASM still fails on the fp32 model with
ERROR_CODE: 7 ... Failed to load model because protobuf parsing failed. The same
file parses on WebGPU, so the graph is not malformed; the cause is size against
the 4 GiB wasm32 address space. A 1.18 GiB int8 build of the same graph parses
without error on the same path.

Notes

  • tie_word_embeddings: true holds in value but not in bytes.
    model.embed_tokens.weight [151936, 1536] and onnx::MatMul_10973
    [1536, 151936] are transposes of one another with different SHA-256s. ONNX
    MatMul has no transB, so the pre-transposed copy avoids a Transpose over
    233M elements per forward pass. Hashing all 356 tensors found no duplicate byte
    ranges. Nothing here changes this.
  • tokenizer_config.json model_max_length is 32768 while config.json
    max_position_embeddings is 2048. Not changed here, since narrowing it would
    newly truncate inputs between the two values.
  • decoder_model.onnx is untouched and has no sharded counterpart. It carries no
    KV cache inputs and is not the graph used for generation.

Suggested README addition

Optional, and not included in this commit so that no existing file is
modified. Paste it wherever it fits best if you find it useful.

## Usage with transformers.js

```js
import { pipeline } from '@huggingface/transformers';

const generator = await pipeline('text-generation', 'Maincode/Maincoder-1B-ONNX');

const output = await generator(
  [{ role: 'user', content: 'Write a Python function that reverses a linked list.' }],
  { max_new_tokens: 256 },
);
console.log(output[0].generated_text.at(-1).content);
```

The chat template and stop tokens are read from the repository, so no additional
arguments are required.

### Files

`onnx/model.onnx` is a re-layout of `decoder_with_past_model.onnx` with its
external data split across four files:

| File | Bytes |
| --- | ---: |
| `onnx/model.onnx` | 1,677,425 |
| `onnx/model.onnx_data` | 1,602,385,920 |
| `onnx/model.onnx_data_1` | 1,610,612,736 |
| `onnx/model.onnx_data_2` | 891,813,888 |
| `onnx/model.onnx_data_3` | 933,494,976 |

The weights are bit-identical to `decoder_with_past_model.onnx`; only the file
each tensor is stored in has changed. `decoder_model.onnx`,
`decoder_with_past_model.onnx`, and their `.onnx_data` files are unchanged and
continue to work as documented above.

The split exists because `transformers.js` allocates a single buffer the size of
each model file it reads, and browsers cap a single allocation at roughly 2 GB.
A 4.69 GiB external data file cannot be read in a browser tab; four files under
1.5 GiB each can.

### Browser support

The fp32 model runs in Chrome under WebGPU. Measured on an RTX 4060 Laptop:
0.79 s to first token and 9.51 tok/s, after a 4.70 GiB download.

The WASM execution provider cannot load the fp32 model: 4.69 GiB of weights
exceeds the 4 GiB wasm32 address space. Use WebGPU, or a quantized build.

### Context length

`config.json` sets `max_position_embeddings` to 2048. `tokenizer_config.json`
reports `model_max_length` as 32768; inputs beyond 2048 tokens are outside the
trained context.
Ready to merge
This branch is ready to get merged automatically.

Sign up or log in to comment