[WIP] Quartet QAT support #38696
Open
[WIP] Quartet QAT support #38696
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
cc @MekkCyber |
MekkCyber
reviewed
Jun 12, 2025
There was a problem hiding this comment.
Hi @BlackSamorez ! Thanks a lot for this addition 🤗 ! Left a few comments !
| @@ -0,0 +1,49 @@ | |||
| # Copyright 2024 The HuggingFace Team. All rights reserved. | |||
There was a problem hiding this comment.
Suggested change
| # Copyright 2024 The HuggingFace Team. All rights reserved. | |
| # Copyright 2025 The HuggingFace Team. All rights reserved. |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
| "HIGGS through FLUTE (Flexible Lookup Table Engine for LUT-quantized LLMs) integration file" |
There was a problem hiding this comment.
Suggested change
| "HIGGS through FLUTE (Flexible Lookup Table Engine for LUT-quantized LLMs) integration file" | |
| "Quartet QAT integration file" |
Comment on lines
+22
to
+24
| if is_torch_available(): | ||
| pass | ||
|
|
| @@ -0,0 +1,164 @@ | |||
| # Copyright 2024 The HuggingFace Inc. team. All rights reserved. | |||
There was a problem hiding this comment.
Suggested change
| # Copyright 2024 The HuggingFace Inc. team. All rights reserved. | |
| # Copyright 2025 The HuggingFace Inc. team. All rights reserved. |
Comment on lines
+36
to
+38
| Quantizer of the HIGGS method. Enables the loading of prequantized models and in-flight quantization of full-precision models. | ||
| """ | ||
|
|
Comment on lines
+1163
to
+1165
| def is_qutlass_available(): | ||
| return _qutlass_available | ||
|
|
There was a problem hiding this comment.
I can't find a distribution for qutlass, is it not released yet ?
Comment on lines
+125
to
+128
| for name, module in tqdm(quartet_qat_modules.items(), desc="Pre-processing Quartet QAT modules", leave=False): | ||
| pass | ||
| # module.pre_forward() | ||
|
|
There was a problem hiding this comment.
What’s meant to happen here exactly ?
Comment on lines
+160
to
+164
| if isinstance(module, QuartetLinear) and tensor_name == "weight": | ||
| # Only quantize weights of QuartetLinear modules that are not already quantized | ||
| return True | ||
| else: | ||
| return False |
There was a problem hiding this comment.
is the bias quantized too ?
Comment on lines
+96
to
+97
| assert isinstance(module, QuartetLinear), f"Module {param_name} is not a QuartetLinear somehow..." | ||
|
|
There was a problem hiding this comment.
no need for assert here, or we can just raise an error instead
Comment on lines
+99
to
+100
| module.pre_forward() | ||
|
|
There was a problem hiding this comment.
- Hadamard transform matrix initialization on the correct devices.
- Since it's a QAT method, we might or might not want to keep a full-precision weight copy. If we don't need the full precision weight copy, this function also deletes the
.weightparameter after quantizing it. Here's the code.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR adds support for the Quartet QAT method.
The goal of this PR is to integrate inference and training support for the Quartet QAT method. That would allow to perform both forward and backward passes in MXFP4, allowing for very fast training on Blackwell GPUs.
Currently, we're working on the kernels here, here and here (some of the libs aren't public yet). We're planning to release the first version of the kernels this week and have optimized performance by end of June.
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.