yakuraku commited on
Commit
589c7b5
·
verified ·
1 Parent(s): 272da87

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

Browse files

# 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:

```jsonc
// 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|>` (151645) | 9 |
| `<|endoftext|>` (151643) | 0 |
| 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.

````markdown
## 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.
````

.gitattributes CHANGED
@@ -36,3 +36,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
36
  decoder_model.onnx_data filter=lfs diff=lfs merge=lfs -text
37
  decoder_with_past_model.onnx_data filter=lfs diff=lfs merge=lfs -text
38
  tokenizer.json filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
36
  decoder_model.onnx_data filter=lfs diff=lfs merge=lfs -text
37
  decoder_with_past_model.onnx_data filter=lfs diff=lfs merge=lfs -text
38
  tokenizer.json filter=lfs diff=lfs merge=lfs -text
39
+ onnx/model.onnx_data filter=lfs diff=lfs merge=lfs -text
40
+ onnx/model.onnx_data_1 filter=lfs diff=lfs merge=lfs -text
41
+ onnx/model.onnx_data_2 filter=lfs diff=lfs merge=lfs -text
42
+ onnx/model.onnx_data_3 filter=lfs diff=lfs merge=lfs -text
config.json CHANGED
@@ -25,5 +25,10 @@
25
  "transformers_version": "4.57.3",
26
  "use_cache": false,
27
  "use_qk_norm": true,
28
- "vocab_size": 151936
 
 
 
 
 
29
  }
 
25
  "transformers_version": "4.57.3",
26
  "use_cache": false,
27
  "use_qk_norm": true,
28
+ "vocab_size": 151936,
29
+ "transformers.js_config": {
30
+ "use_external_data_format": {
31
+ "model.onnx": 4
32
+ }
33
+ }
34
  }
generation_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "eos_token_id": [
3
+ 151643,
4
+ 151645
5
+ ],
6
+ "pad_token_id": 151643
7
+ }
onnx/model.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:44435a74a48eda79a5f1b288655c56f6c1d2e75b59e50b774445b262b5bf8e86
3
+ size 1677425
onnx/model.onnx_data ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:60377b980fce2d065236eaa816d3195b633509579947c5a841e527246d743c78
3
+ size 1602385920
onnx/model.onnx_data_1 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:540cdeba29ee065cfbf394a4e4336d2139e88e3e81f16add49a269c1a12fb70e
3
+ size 1610612736
onnx/model.onnx_data_2 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e4e4a16bd574f4df599de01e50e4a5a3718936ee32cecd326af2e3879ba20d09
3
+ size 891813888
onnx/model.onnx_data_3 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9301a5203af463f0129f995b221340a2f7daf3c39ed6cb6cbe1a0ccbfa7e5135
3
+ size 933494976
tokenizer_config.json CHANGED
@@ -204,5 +204,6 @@
204
  "pad_token": "<|endoftext|>",
205
  "split_special_tokens": false,
206
  "tokenizer_class": "Qwen2Tokenizer",
207
- "unk_token": null
 
208
  }
 
204
  "pad_token": "<|endoftext|>",
205
  "split_special_tokens": false,
206
  "tokenizer_class": "Qwen2Tokenizer",
207
+ "unk_token": null,
208
+ "chat_template": "{%- if tools %}\n {{- '<|im_start|>system\\n' }}\n {%- if messages[0]['role'] == 'system' %}\n {{- messages[0]['content'] }}\n {%- else %}\n {{- 'You are a helpful assistant.' }}\n {%- endif %}\n {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n {%- for tool in tools %}\n {{- \"\\n\" }}\n {{- tool | tojson }}\n {%- endfor %}\n {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n {%- if messages[0]['role'] == 'system' %}\n {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n {%- else %}\n {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"assistant\" %}\n {{- '<|im_start|>' + message.role }}\n {%- if message.content %}\n {{- '\\n' + message.content }}\n {%- endif %}\n {%- for tool_call in message.tool_calls %}\n {%- if tool_call.function is defined %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '\\n<tool_call>\\n{\"name\": \"' }}\n {{- tool_call.name }}\n {{- '\", \"arguments\": ' }}\n {{- tool_call.arguments | tojson }}\n {{- '}\\n</tool_call>' }}\n {%- endfor %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n {{- '<|im_start|>user' }}\n {%- endif %}\n {{- '\\n<tool_response>\\n' }}\n {{- message.content }}\n {{- '\\n</tool_response>' }}\n {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
209
  }