Difference between embed/convert?

[vrdn-23]

Hey @justinchuby,

First off, thanks for the amazing project! I've been looking for something like this for a while.

I had a couple of questions about the fundamental difference between the convert and embed operations for this library. I have a model trained in-house (which I can't share unfortunately), that seems to be successful when running a conversion but fails when running a embed.

I am also trying to understand what method would be helpful if I wanted to assess the memory usage of the model based on the converted safetensors file (if you know of a way to do this directly through onnx I would be very much interested to hear this). For example, would I able to use the hf-mem tool to accurately calculate VRAM usage of the onnx model after conversion?

~/dev/inferentia main ?30 > uv run onnx-safetensors convert model.onnx model.safetensors                                                                                                    10s py 3.12  14:30:02
Loading ONNX model from model.onnx...
Converting model to safetensors format...
Saving model.safetensors (onnx::MatMul_12379): 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 380/380 [00:00<00:00, 3608.57it/s]
Model saved to model.safetensors

~/dev/inferentia main ?31 > uv run onnx-safetensors embed model.onnx model.safetensors                                                                                                          py 3.12  14:39:53
Loading ONNX model from model.onnx...
Embedding model into safetensors file...
Traceback (most recent call last):
  File "/Users/vidamoda/dev/inferentia/.venv/bin/onnx-safetensors", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/Users/vidamoda/dev/inferentia/.venv/lib/python3.12/site-packages/onnx_safetensors/_cli.py", line 96, in main
    embed_command(args)
  File "/Users/vidamoda/dev/inferentia/.venv/lib/python3.12/site-packages/onnx_safetensors/_cli.py", line 53, in embed_command
    onnx_safetensors.save_safetensors_model(model, output_path)
  File "/Users/vidamoda/dev/inferentia/.venv/lib/python3.12/site-packages/onnx_safetensors/_safetensors_io.py", line 640, in save_safetensors_model
    safetensors.serialize_file(tensor_dict, safetensors_model_path, metadata=metadata)
safetensors_rust.SafetensorError: Error while serializing: header too large

Thanks again for the project!

[vrdn-23]

The converted safetensors file does not properly load for parsing metadata

from huggingface_hub import get_local_safetensors_metadata


def calculate_memory_requirements_from_hf(model_id):
    # Taken from https://bsky.app/profile/alvarobartt.me/post/3lh2esxqass2s

    dtype_bytes = {"F32": 4, "F16": 2, "BF16": 2, "F8": 1, "I64": 8, "I32": 4}

    metadata = get_local_safetensors_metadata(model_id)
    memory = (
        sum(count * dtype_bytes[key.split("_")[0]] for key, count in metadata.parameter_count.items())
        / (1024**3)
        * 1.18
    )
    return float(f"{memory:.2f}")

print(calculate_memory_requirements_from_hf("model.safetensors"))

Traceback (most recent call last):
  File "/Users/vidamoda/dev/inferentia/dummy.py", line 18, in <module>
    print(calculate_memory_requirements_from_hf("model.safetensors"))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vidamoda/dev/inferentia/dummy.py", line 9, in calculate_memory_requirements_from_hf
    metadata = get_local_safetensors_metadata(model_id)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vidamoda/dev/inferentia/.venv/lib/python3.12/site-packages/huggingface_hub/hf_api.py", line 11321, in get_local_safetensors_metadata
    file_metadata = parse_local_safetensors_file_metadata(path)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vidamoda/dev/inferentia/.venv/lib/python3.12/site-packages/huggingface_hub/hf_api.py", line 11261, in parse_local_safetensors_file_metadata
    metadata_size = _get_safetensors_metadata_size(size_bytes, filename, context_msg)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/vidamoda/dev/inferentia/.venv/lib/python3.12/site-packages/huggingface_hub/hf_api.py", line 1796, in _get_safetensors_metadata_size
    raise SafetensorsParsingError(
huggingface_hub.errors.SafetensorsParsingError: Failed to parse safetensors header for 'model.safetensors' (path 'model.safetensors'): safetensors header is too big. Maximum supported size is 25000000 bytes (got 8031177557570357512).

Process finished with exit code 1

[justinchuby]

Embed: This format is for storage/transfer only and is not compatible with ONNX Runtime, as noted in the readme. Because we embed all of the model proto data into the safetensors header. It is possible that it is over safetensor's limt. In that case embedding is not possible.

For the converted file, if the header is still too big, maybe create shards with the --max-shard-size option?

[vrdn-23]

Thanks! Sharding helped for the successful conversion.

Minor nit:

~/dev/inferentia/tmpy main ?31 > uv run onnx-safetensors convert model.onnx model.safetensors --max-shard-size 50MB                                                                             py 3.12  16:43:22
Loading ONNX model from model.onnx...
Converting model to safetensors format...
Saving model-00001-of-00008.safetensors (decoder.decode.0.conv1.parametrizations.weight.original0): 100%|█████████████████████████████████████████████████████████████████████| 116/116 [00:00<00:00, 6515.07it/s]
Saving model-00002-of-00008.safetensors (decoder.decode.1.conv2.parametrizations.weight.original0): 100%|███████████████████████████████████████████████████████████████████████| 15/15 [00:00<00:00, 1564.77it/s]
Saving model-00003-of-00008.safetensors (decoder.decode.2.conv1x1.parametrizations.weight.original0): 100%|█████████████████████████████████████████████████████████████████████| 18/18 [00:00<00:00, 2776.87it/s]
Saving model-00004-of-00008.safetensors (decoder.generator.noise_convs.6.bias): 100%|███████████████████████████████████████████████████████████████████████████████████████████| 52/52 [00:00<00:00, 5051.04it/s]
Saving model-00005-of-00008.safetensors (onnx::LSTM_12074): 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 130/130 [00:00<00:00, 8779.50it/s]
Saving model-00006-of-00008.safetensors (onnx::MatMul_12336): 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:00<00:00, 2695.65it/s]
Saving model-00007-of-00008.safetensors (onnx::MatMul_12359): 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 15/15 [00:00<00:00, 2382.41it/s]
Saving model-00008-of-00008.safetensors (onnx::MatMul_12379): 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████| 14/14 [00:00<00:00, 2473.47it/s]
Model saved to model.safetensors

It seems even after sharding there seems to be a file called model.safetensors also being saved. Is that intentional?

[justinchuby]

I don't think that should happen?

[justinchuby]

Ok. 50MB is a very small limit. Usually we use 5GB or 10GB. It is strange that the metadata size would be over limit with this little data.