Efficient Inference Kernel for SpQR

[elvircrn]

What does this PR do?

Adds support for efficient single-batch inference for SpQR.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@SunMarc @MekkCyber

[MekkCyber]

Hey @elvircrn, Thanks for adding this quantization method ! A very smooth integration 🔥! Just left a few comments

[MekkCyber]

No need to add the checks here, validate_environment in the quantizer will take care of that

[elvircrn]

Resolved by removing the checks.

[MekkCyber]

Not used, to be removed

[elvircrn]

Resolved by removing

[MekkCyber]

let's raise a warning here if you need to override the torch_dtype like in other quantizers

[elvircrn]

Since SpQR really only supports float16, I raise an error if it's None or anything other than float16. Is this the desired behavior?

[MekkCyber]

Sorry I can't find this model on the hub 🤔

[elvircrn]

Ah, didn't expect a review so soon. Thank you!

I've uploaded the model and updated this part of the test.

[MekkCyber]

No worries 😊 take your time and ping me when you finish the integration

[MekkCyber]

No need to add the requirement here, it's already verified with the decorator require_spqr

[elvircrn]

Resolved by removing.

[MekkCyber]

We can also add a small section about how to load a model quantized using SpQR using transformers

[elvircrn]

Resolved by updating the README.

[MekkCyber]

if the quantization method requires a gpu we need to check if it's availabe here

[elvircrn]

Resolved by adding an is_cuda check.

[MekkCyber]

Is it possible to add tests for multi_gpu setups ?

[elvircrn]

Resolving by bringing this back form AQLM. Will the CI be able to pick this up? I don't have the means of testing this.

[MekkCyber]

Yes it does 😄

[SunMarc]

+1 on that, you can add the test and we will check to make sure it to pass the right expected output

[elvircrn]

@MekkCyber I've removed the Draft tag from the PR.

It seems that I am unable to pass the test_raise_if_non_quantized. Could you give me some pointers on this?

[MekkCyber]

Could you rebase @elvircrn to update the branch :)

[elvircrn]

@MekkCyber Rebase done.

[MekkCyber]

Hey @elvircrn, thanks for iterating ! It looks great I just left some minor comments

[MekkCyber]

We need to add the checks here if the accelerate and spqr_quant packages are available before doing the import, what we don't need to do is raising errors, because the packages are assumed to be installed when executing the function

[elvircrn]

Not sure if I understood this, please double-check my newest change.

[MekkCyber]

We are not assuming here because we already validated the environment 😉

[MekkCyber]

Suggested change

	elif torch_dtype != torch.float16:
	raise ValueError(
	"You cannot any type other than torch.float16 for SpQR. Please either leave it None or set it to"
	"torch.float16 explicitly."
	)
	return torch_dtype
	elif torch_dtype != torch.float16:
	raise ValueError(
	"You cannot use any type other than torch.float16 for SpQR. Please either leave it None or set it to"
	"torch.float16 explicitly."
	)
	return torch_dtype

[MekkCyber]

Just FMI, are the beta and shapes parameters used solely for loading, or do they also play a role in quantization? In other words, if a model is quantized using specific beta or shapes values, can it only be loaded with those same parameters?

[elvircrn]

SpQR quantized weights are a product of tile-based (beta corresponds to the dimension of this tile where beta1, beta2 represent tile width and height respectively) bi-level quantization. You always specify beta before the quantization starts.

Each SpQR weight also comes with a sparse tensor representing the outlier weights. As this tensor is unstructured, one had to keep track of the size of this matrix during loading.

[elvircrn]

The following is a visualization of the compression format from the original publication (https://arxiv.org/pdf/2306.03078):

[MekkCyber]

Thanks for the detailed explanation

[ArthurZucker]

Let's add it to the doc!

[MekkCyber]

in the shapes attribute do we need to specify the shape for all the tensors ?

[elvircrn]

The colval size needs to be specified ahead of time for all weights. The rest can be computed from m, n and bits. Can we keep it the way it is now?

[elvircrn]

I can do it both ways easily, not sure what is preferred here.

[MekkCyber]

I'm just worried about the quantization config, I'm not sure if we can verify that shapes contains all the necessary elements in the post_init. Maybe before trying to access the key in replace_with_spqr_linear we verify if it's there or we raise an error to inform the user there is something wrong with the config.json they are using, wdyt @SunMarc

[elvircrn]

In my estimate, these configurations are directly tied to the quantized weights. This would be a fatal error in quantization serialization if the shapes don't cover the full set of quantized weights. Perhaps this is something that should be done on the side of SpQR (https://pypi.org/project/spqr-quant/)?

If we were to do it on huggingface/transformers-side:

a) during replace_with_linear, we raise a fatal error (please let me know which one would be appropriate here) if the quantized weights don't have the corresponding key in the shapes config.
b) we do a full check before conducting the replacement in order to conserve memory/device cycles.

Let me know what best suits transformers here.

[elvircrn]

I didn't see the tag to @SunMarc - I thought the wdyt referred to me.

[MekkCyber]

Yes it does 😄

[MekkCyber]

If possible it's better to use small models for the CI like Qwen 0.5, Bloom 560M, SmolLM 135M...

[elvircrn]

Will have to get back to you on this (might require a significant time investment).

[elvircrn]

Is this a blocker for the merge?

[MekkCyber]

Not really a blocker for now

[elvircrn]

@MekkCyber The latest batch of comments are addressed.

[MekkCyber]

Thanks for the quick iteration @elvircrn 🔥 ! can you re-rebase please 😅

[MekkCyber]

I'm just worried about the quantization config, I'm not sure if we can verify that shapes contains all the necessary elements in the post_init. Maybe before trying to access the key in replace_with_spqr_linear we verify if it's there or we raise an error to inform the user there is something wrong with the config.json they are using, wdyt @SunMarc

[MekkCyber]

Thanks for the detailed explanation

[MekkCyber]

Yes that's exactly what I meant, thank you for fixing it !

[MekkCyber]

Not really a blocker for now

[elvircrn]

@MekkCyber rebased.

[elvircrn]

@MekkCyber Resolved the latest batch of comments.

[elvircrn]

@SunMarc @MekkCyber
I've pushed a check for keys() in shapes in replace_with_spqr_linear. Let me know if something in this ballpark works.

[SunMarc]

pinging @ArthurZucker to have a quick look !

[SunMarc]

Could you fix the conflits @elvircrn ?

[elvircrn]

@SunMarc resolved quickly, as it was just an md conflict in the overview table.

[elvircrn]

@SunMarc @ArthurZucker Conflicts resolved.

[ArthurZucker]

LGTM!

[ArthurZucker]

we can't quantize a model on the fly with quantization_config=... using spqr as the quantize type?

[elvircrn]

We can't, this is not supported as of yet. This PR only adds Inference support.

[ArthurZucker]

Let's add it to the doc!

[ArthurZucker]

Can you just fix this:

RuntimeError: The following files are not present in the table of contents:
- quantization/spqr
Add them to ../transformers/docs/source/en/_toctree.yml.

[elvircrn]

@ArthurZucker Is the PR in order now?

[elvircrn]

@ArthurZucker Erreneous rebase issue resolved.

[MekkCyber]

Hi @elvircrn, can you just resolve the conflicts and once the CI is green we can merge ! Again sorry for the delay 🙏 and thanks for your patience

[elvircrn]

@MekkCyber Done.

[elvircrn]

@MekkCyber The failing tests, based on the logs, seem to be unrelated to my own source code.

[MekkCyber]

Yes they seem unrelated @elvircrn, i tried updating the branch I hope it works

[elvircrn]

@MekkCyber I don't think that did it - unrelated tests are still failing.

[elvircrn]

@MekkCyber Seems that everything is green now. :)

[elvircrn]

Thanks everyone!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.