From LLM quantization research there is the feature of using NF4 as a low precision datatype. Paper introducing this https://arxiv.org/pdf/2305.14314
This is built upon Quantile Quantization, “which is an information-theoretically optimal data type that ensures each quantization bin has an
equal number of values assigned from the input tensor”.
Way I see it, it would be beneficial to start by implementing a QuantileQuantizedType inheriting UniformQuantizedType; where QuantileQuantizedType has the added logic for handling the quantile table.
Similar for QuantileQuantizedPerAxisType and UniformQuantizedPerAxisType.
It’s less about “NF4” being a new datatype and more about this new quantile based quantization approach, as storage type it behaves exactly as a int4 datatype.
applying the look up table to convert from NF4/INT4 to FP16/FP32
apply further scale multiply and zero point subtraction to optimize the representational range, especially in the per axis quantization case.
IR would look something like:
!quant.quantile<storageType:expressType:axis, quantiles:scales and zero points>
!qElemType = !quant.quantile<i4:f16:1, {-1.0, -0.69, … 0.72, 1.0}:{1.000000e+00:128,2.000000e+00:128,3.000000e+00:128,4.000000e+00:128}>
We’d like to contribute ourselves to such a feature and any feedback for this proposal is much welcome until we start work on it!
I’m afraid that I’m basically resigned from the quant dialect at this point and am not comfortable being a defacto owner of it. I know there are folks still interested in this form of type-based quantization, but I’m not really tracking the work to a level that would make my feedback worthwhile.
One thing I will ask: is this work planned to manifest somewhere concrete? In general, there are an endless number of such schemes and it is best for the project to take contributions for things that will receive ongoing investment or use vs everything that is possible.
While on the topic of relevance for such changes, the latest research around LLMs exploded especially on the topics of lower data types and novel quantization schemes, and data types like FP4 and NF4 over a great alternative to their integer equivalents;
One can find implementations of NF4 in different ML framerworks, like the one in OpeVino. The real exciting part of NF4 is rather the Quantile based approach and anyone can infer a optimal LUT based on their dataset and HW implementation in order to achieve better accuracy.
We at least plan to use it in our organization since we can get just as good accuracy with NF4 and basic per axis quantization routines, and expect more suits to follow : )
I think as long as it is going to be used in some MLIR based compiler and/or MLIR based frontend that can either be seen or attested to be supported, then there is a reasonable argument for upstreaming. I’m long past being excited in the abstract about new numeric schemes in the research that don’t have an effective realization in either hardware or software. Often getting over that hurdle takes a lot more than an implementation in a single mid level dialect like quant, which is why I’m probing on how real and widespread the plan is here. I’m mainly guarding against partial concepts without a long term owner or concrete use landing in tree.
As I alluded to, I’m skeptical of further propagating these schemes that encode all of the scaling parameters in the type because it doesn’t match a lot of the newer programming models (which generally need to parametrize these with runtime values). But to the extent that the quant dialect contributors would like to invest here in this way, I’m not going to push back – there are many ways to implement quantization in an MLIR based compiler, and the quant dialect is just one.
In newer work, I wouldn’t implement this concept at all in MLIR or the compiler proper but in the frontend. There are several examples of convergent evolution on this kind of thing, which show some of the different/related approaches:
The thing that all of these have in common is that they deal with encoding/layout/quantization far up the stack with runtime vs compile time parameterization. And they break the problem down at the top vs trying to preserve strong typing of a specific quantization algorithm deep in the compiler type hierarchy.
Thanks for the mention, Stella, and for the proposal, Zoran.
As a heads-up, we have currently an active PR for a revamp of the quant dialect, focusing on the !quant.uniform type with per-layer and per-axis quantization, a spec refinement with full documentation for the quant ops, lowering support for these ops, and a pass to flatten quant types. This PR involves some file restructuring to mirror other dialects, so it’d be best to get this merged before considering other dialect extensions.
I can’t speak to the relevance of quantile quantization support for the MLIR community at the moment, or specifically to the tradeoff between feature value and maintenance liability, but I’ll be happy to review your work, Zoran, if it moves forward. Your feedback on our PR could be very valuable, given your vision on a possible evolution of the dialect.
In our case, we propagate the quantized types to lower runtime stack, that way we can leverage the HW at it’s best. Still logic such as the one in mlir/lib/Dialect/Quant/Transforms/LowerQuantOps.cpp makes sense from MLIR perspective.
We’ll return with a PR for Quantile Quantization based on top of your work. Way we see it, there are already few variances on the NF4 type which use different quantile table so instead of extending the explicit datatype support we’d rather handle them as N bit low precision values mapped to M bit high precision values.
Going back on the what we could see in ML frameworks or when lowering to arithmetic operations, a QuantileQuantized type of form:
!quant.quantile<storageType:expressType:axis, quantiles:scales and zero points>
You all may at some point want to consider not carrying the scales and zero points on the type. Those really need to be values at runtime. This could be done by ensuring some structural characteristics on the IR such that you can always get back to a specific qcast* op which carries the actual values (i.e. recover the concrete quantization parameters by analysis).