Contributing, adding convertConvInteger for ConvInteger op, cleaning up

[siherrmann]

I would like to contribute, I already forked the repository and added a working version for the ConvInteger operation. I saw, that there are some things I would also like to clean up, such as:

• some deprecated functions of graph
• some linter issues like panic(nil) which I updated to return (...)
• linter not liking the dot import for graph, I updated the code to use graph. (and yes, that's more verbose)
• missing usage of initalCell in ops.go line 1272, added in line 1286

if initialHidden != nil || initialCell != nil {
    lstmLayer = lstmLayer.InitialStates(initialHidden, initialCell)
}

• and some other linter related things

Should I just open a PR when I'm finished with the changes? Or would you like me to just add the new OP in one PR and add another PR for cleanup?

[janpfeifer]

hey Simon (@siherrmann), thanks for interest.

The ConvInteger is supported by the underlying XLA, it's just that I haven't connected yet in the stablehlo package -- it requires some designing on how to encode the zero-point (stablehlo seems to encode it in the tensor type somehow, but I haven't looked into the details). I assume you implemented by up-casting to some bigger integer, doing a normal convolution, clamping and reducing again to a small int ? I'm just guessing but I'm not sure ... But just to say that adding proper support for quantized operations is high on my list of TODOs, also for the simplego (go only) backend -- except if someone else would like to brave it 😃

Btw, that sounds very similar to the recent #35 implemented @ajroetker -- maybe he has something to add ?

Sending a PR is perfect! I tend to be very responsive for reviews. And, where possible I prefer separate PRs for different topics (in this case ConvInteger and clean ups), if that is ok.

About the clean ups: they are very welcome, thanks for doing it!

I'm a bit hesitant with removing the dot import of the gomlx/pkg/core/graph: for some functions it really helps readability (if coding graph building with vectorial math). But at the cost of "polluting" the namespace. Can we postpone this for a while ? We can disable the lint warning on that line for the time being, maybe ?

cheers!

[janpfeifer]

Definitely, move ahead with your version of the quantized convolution that works now! Add some tests, and later when it is well supported we can replace by a backend supported version (and likely much faster).

And yes, that's the simplego backend, written in Go only. I still want to add other backends: ONNX, WebGL, Vulkan, ggml (for cpu)... etc. But simplego is the one written in Go (so very portable, no need to install c++ libraries).

[siherrmann]

Sure, I would like to have it implemented fast as I'm currently also working on implementing an ImageToText pipeline in hugot and there I need MatMulInteger (which is implemented but not yet released I think) and the ConvInteger. But I would also like to look into adding QuantizeConvolve to simplehlo as this should make the whole process much faster. But I will first create a PR for the simple version and another one for the cleanup.

Regarding the simple version, the official onnx specification for ConcInteger is int8 or uint8 as input and int32 as output, in my opinion I would not have to clamp the values to int8, would I?

[janpfeifer]

Let me know if you need me to cut a new release. I'll wait your ConvInteger and can update the release then.

[janpfeifer]

Oh, if your output in int32, then you don't need to clamp indeed. Only if you are reverting back to the int8 quantized.

[ajroetker]

@timkaye11 and I have been working on adding support for a bunch of the simplego/onnx based ops (all running via hugot atm) for int8 and f16 quantized models (rerankers, chunkers, clip, and embedders) at the moment, we have a branch of onnx-gomlx and gomlx with a bunch of simd/sme optimizations for that too (although the onnx is still much faster) https://antfly.io/models we're working on infra to make it super easy to run these ollama style https://antfly.io/models hopefully open sourcing it all in the beginning of the year for the full termite pipeline, image to text and ner models on our list too. @janpfeifer

[siherrmann]

I opened a PR for the cleanup: #37
After that I would, if everything looks good with the cleanup, create another PR with the change I made.

[siherrmann]

Thanks for the PR review @janpfeifer . I'm currently looking into the gopjrt for adding support for quantized operations. I didn't find UniformQuantize and UniformDequantize in XLA, do you know whether these are available?

[janpfeifer]

Heads up: ignore the xlabuilder package, it's being deprecated.

Using XLA these days requires (1) something to generate StableHLO (the intermediary language PJRT understands), and (2) something that loads the PJRT plugin and makes the calls.

(1) Is the project https://github.com/gomlx/stablehlo, that implements a simple Go API that generates StableHLO (text). It also has to do "shape inference" (figuring out the output shape for given inputs).
(2) Is Gopjrt, in particular the pjrt package.

The StableHLO specficication is here.. There you'll find the uniform_quantize and uniform_dequantize.

Generally, adding an op in stablehlo repo is easy. The challenge in this case is that quantized types are new types, that include the zero-point and scale (?? I think, I haven't looked into details). And it will require new DTypes to represent them. The DType enum is defined in dtypes.go in Gopjrt.

If you figure this out, a contribution is most welcome -- this seems to be one of the things I've been requested a lot recently.

Just to sync up and let you know, after I finish the distributed execution release in GoMLX, my plans are to work on cleaning up stablehlo/gopjrt (but it can go in parallel to any PR):

• Merge both in one repository called gomlx/go-xla, with subpackages stableho and pjrt.
• Eliminate the deprecated xlabuilder package.
• Create a gomlx/prjt-cpu repo, with binary releases of the PJRT CPU (for linux, mac, and windows if I manage to compile it) - which today I release together with Gopjrt, but it could be separate, since they move differently.
• Add a library to automatically download the PJRT plugin on-the-fly (what today is done by gopjrt_intaller. And allow one to use this in Gomlx program (simply calling a function), so demos can work without anyone explicitly installing anything upfront. Including PJRT GPU and TPU.
• Maybe create another version of pjrt using a "pure go" approach. Should be doable ... but I expect some amount of work converting .h files to Go structs. We'll see, probably will come to next year only.

(Sorry for the laundry list, but I needed to organize my thoughts, and this was a good opportunity 😄 )

If you have any concerns or suggestions while doing your changes, let me know.

[siherrmann]

Man, that's a lot 😅 Ok, so if I understand correctly it was my correct instinct to add quantize operations to the gopjrt repo, but the xlabuilder folder should be deprecated? So adding quantization in there is not the way to go?

And yes, the new dtype needs zeropoint and yes, also scale, and both can be scalars or vectors, I would add them as two arrays. But what would be from your view be the best next step after that?