[converter] Quantization option for `converters.convert_jax`

[josephrocca]

System information

• TensorFlow.js version (you are using): latest
• Are you willing to contribute it: Unfortunately, it is beyond me.

The title is pretty self explanatory, but one thing: I really hope if this is implemented, it will support converting and quantizing models above 2GB. Currently if I try to convert a 3GB model with converters.convert_jax, I get this error:

ValueError: Message tensorflow.GraphDef exceeds maximum protobuf size of 2GB: 3450353084

During handling of the above exception, another exception occurred:

AssertionError                            Traceback (most recent call last)
[/usr/local/lib/python3.8/dist-packages/tensorflow/python/framework/graph_util_impl.py](https://localhost:8080/#) in _assert_nodes_are_present(name_to_node, nodes)
    202   """Assert that nodes are present in the graph."""
    203   for d in nodes:
--> 204     assert d in name_to_node, "%s is not in graph" % d
    205 
    206

Of course, 3GB is a bit too big for the web, but 700mb is fine. About the same size as the combined text+image CLIP models used in this project which runs fine in the browser. So although 2GB seems like "more than enough" for any reasonable use case, please keep in mind that it actually limits models to ~500mb, in the sense that I've just described.

[mattsoulanille]

Hi @josephrocca. Thanks for the bug report! This is a bug in tfjs converter. I'd like to say you can just use the --quantize_float16 flag or similar, but I don't think this will fix the initial model loading issue.

There's no reason why it shouldn't be possible to convert larger JAX / TF models to TFJS. The TFJS model format is a JSON file with some 4MB binary files for the weights, none of which use protobuf, so this is purely a conversion issue. I ran into this as well when trying to convert this keras implementation of Stable Diffusion in December. I think the key issue is we're freezing the entire model graph (converting the variables to constants). If we can avoid this and load the variable values separately, we should be able to convert larger models. I took a look at ONNX's approach to loading large models, and I think this is somewhat similar to what they do.

[josephrocca]

Thanks, Matthew! Very much looking forward to a fix! I'm actually trying to do the same thing (Stable Diffusion) heh. One thing, just in case I misunderstand - not that it's very relevant to this issue, as you said - but:

I'd like to say you can just use the --quantize_float16 flag or similar, but I don't think this will fix the initial model loading issue.

I was under the impression that the JAX converter doesn't have a command-line approach unlike the other formats? I.e. you have to actually call it from Python? Nothing wrong with that from my perspective, but just clarifying.

[mattsoulanille]

I just checked the tfjs blog and the source code, and you're right. I forgot JAX can only be converted from python.

JAX conversion first converts to a TF model, so I think we can expose some of the options that the CLI has for configuring TF model conversion to the convert_jax function. Of course, not all of them will be applicable, but I think the quantization options would be.

[mattsoulanille]

Hi @josephrocca. I have made some progress on converting large models. This branch avoids the protobuf size limit you've run into by not converting the graph_def back into a graph (which requires serialization to send the model back to C++). This worked for a large model I converted, but I've only tested a single model. Additionally, this branch can not run Grappler optimization on the model because that requires the graph_def to be converted back into a graph. Until we have that solved, I probably won't be able to merge this branch.

You can try it out locally, though, by building the tfjs pip package on the branch:

cd tfjs-converter/python
bazel build :python3_wheel
./build-pip-package.sh --build output/

[josephrocca]

@mattsoulanille Thank you - it seems to work! Here's an example notebook that converts the 3GB unet of Stable Diffusion (note that I tried with 26GB of RAM but it crashed due to memory limitations, so I had to subscribe to Colab to use a runtime with 36GB of RAM):

https://colab.research.google.com/gist/josephrocca/a0dd5d197c7b5af7ac6b3d8e5447b776/stable_diffusion_jax-to-onnx-after-separation-of-scheduler-step.ipynb

I haven't tried to run it yet in the browser yet, since I'm quite sure it'll need to be quantized to have any chance of not crashing.

Is the jax converter able to take advantage of the existing quantization code you have in the converter? If so, would you be able to add an option for this in your branch?

[gaikwadrahul8]

Hi, @josephrocca

Apologize for the delayed response and I see this PR #7331 got merged to take care of your this issue if I'm not wrong so could you please confirm ? if yes please feel free to close this issue ? Thank you!