diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 0774908ef997..699ec8038e6d 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -289,18 +289,24 @@ title: torchao title: Quantize on the fly - sections: + - local: quantization/blockfp8 + title: Block FP8 - local: quantization/fbgemm_fp8 title: FBGEMM - - local: quantization/finegrained_fp8 - title: Fine-grained FP8 - local: quantization/fouroversix title: Four Over Six - local: quantization/fp_quant title: FP-Quant - local: quantization/mxfp4 title: MXFP4 + - local: quantization/mxfp8 + title: MXFP8 - local: quantization/nvfp4 title: NVFP4 + - local: quantization/finegrained + title: Fine-grained + - local: quantization/finegrained_fp8 + title: Fine-grained FP8 (deprecated) title: FP4 and FP8 formats - sections: - local: quantization/auto_round diff --git a/docs/source/en/experts_interface.md b/docs/source/en/experts_interface.md index 65c45f532772..77d78f385307 100644 --- a/docs/source/en/experts_interface.md +++ b/docs/source/en/experts_interface.md @@ -95,7 +95,7 @@ model = AutoModelForCausalLM.from_pretrained( ## DeepGEMM -The `"deepgemm"` backend routes expert matmuls through the [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels distributed by [kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm). It works with unquantized `bfloat16` experts and with FP8/FP4-quantized experts loaded through [Fine-grained FP8](./quantization/finegrained_fp8). +The `"deepgemm"` backend routes expert matmuls through the [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels distributed by [kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm). It works with unquantized `bfloat16` experts, with [Block FP8](./quantization/blockfp8) experts (128x128 blocks), and with [MXFP4](./quantization/mxfp4) experts run with FP8 activations (DeepSeek-V4's W4A8, Blackwell only), including a [Fine-grained](./quantization/finegrained) plan of them. On quantized experts it quantizes activations to FP8 at run time whatever their `activation_format`, and adds no expert biases; use `grouped_mm` or `batched_mm` for MXFP8 or NVFP4 experts, biased experts, or another activation format. The `"deepgemm"` backend requires: @@ -117,7 +117,7 @@ The kernel is loaded lazily on the first forward. ### FP8 and FP4 quantized experts -DeepSeek-style checkpoints are usually pre-quantized and carry their own quantization config, so you don't need to pass a [`FineGrainedFP8Config`]. The `"deepgemm"` backend automatically picks the FP8 (or FP4 on Blackwell) grouped-GEMM kernel. DeepGEMM requires dynamic per-row activation scales (`activation_scheme="dynamic"`) and rejects static (per-tensor) activation quantization. +DeepSeek-style checkpoints are usually pre-quantized and carry their own quantization config, so you don't need to pass a [`BlockFp8Config`]. The `"deepgemm"` backend automatically picks the FP8 (or FP4 on Blackwell) grouped-GEMM kernel. DeepGEMM requires dynamic per-row activation scales (`activation_scheme="dynamic"`) and rejects static (per-tensor) activation quantization. ```py from transformers import AutoModelForCausalLM @@ -133,14 +133,14 @@ For FP4-packed expert weights (DeepSeek V4-style), the GPU must be SM100+ (Black > [!NOTE] > On Blackwell (SM100+), the `"deepgemm"` and `"deepgemm_megamoe"` experts kernels require power-of-two UE8M0 expert scales. A checkpoint quantized with plain `float32` scales (`scale_fmt="float"`) raises a `ValueError` on the first forward instead of silently corrupting the output. Load a checkpoint quantized with `scale_fmt="ue8m0"`, or switch to `grouped_mm` or `batched_mm`, which consume `float32` block scales directly. Hopper (SM90+) consumes `float32` scales on the DeepGEMM path without conversion. -The main reason to pass a [`FineGrainedFP8Config`] for a pre-quantized checkpoint is to dequantize it back to `bfloat16`, in which case the experts run in `bfloat16` rather than on the FP8/FP4 DeepGEMM path. +The main reason to pass a [`BlockFp8Config`] for a pre-quantized checkpoint is to dequantize it back to `bfloat16`, in which case the experts run in `bfloat16` rather than on the FP8/FP4 DeepGEMM path. ```py -from transformers import AutoModelForCausalLM, FineGrainedFP8Config +from transformers import AutoModelForCausalLM, BlockFp8Config model = AutoModelForCausalLM.from_pretrained( "deepseek-ai/DeepSeek-V3", - quantization_config=FineGrainedFP8Config(dequantize=True), + quantization_config=BlockFp8Config(dequantize=True), experts_implementation="deepgemm", ) ``` diff --git a/docs/source/en/main_classes/quantization.md b/docs/source/en/main_classes/quantization.md index ea2080b00e3f..b1307ff3d002 100755 --- a/docs/source/en/main_classes/quantization.md +++ b/docs/source/en/main_classes/quantization.md @@ -70,14 +70,6 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] MetalConfig -## Mxfp4Config - -[[autodoc]] Mxfp4Config - -## NVFP4Config - -[[autodoc]] NVFP4Config - ## FbgemmFp8Config [[autodoc]] FbgemmFp8Config @@ -98,8 +90,37 @@ Learn how to quantize models in the [Quantization](../quantization/overview) gui [[autodoc]] SpQRConfig +## BlockScaledConfig + +[[autodoc]] BlockScaledConfig + +## BlockFp8Config + +[[autodoc]] BlockFp8Config + +## Mxfp8Config + +[[autodoc]] Mxfp8Config + +## Mxfp4Config + +[[autodoc]] Mxfp4Config + +## NVFP4Config + +[[autodoc]] NVFP4Config + +## FineGrainedConfig + +[[autodoc]] FineGrainedConfig + ## FineGrainedFP8Config +> [!WARNING] +> `FineGrainedFP8Config` is deprecated and will be removed in a future release. Passed to `from_pretrained`, it loads +> through its own frozen integration; use [`BlockFp8Config`], or [`Mxfp8Config`] for a `[1, 32]` block with UE8M0 +> scales (see [Fine-grained FP8](../quantization/finegrained_fp8)). + [[autodoc]] FineGrainedFP8Config ## QuarkConfig diff --git a/docs/source/en/model_doc/deepseek_v32.md b/docs/source/en/model_doc/deepseek_v32.md index ce932526d48b..2bfac5f953a8 100644 --- a/docs/source/en/model_doc/deepseek_v32.md +++ b/docs/source/en/model_doc/deepseek_v32.md @@ -55,23 +55,13 @@ In DeepSeek-V3.2 **every layer runs its own indexer** — there is no cross-laye ## Usage examples -DeepSeek-V3.2-Exp is distributed as an FP8 checkpoint. The indexer projections are kept out of FP8 quantization, since the checkpoint stores them in bf16/fp32: +DeepSeek-V3.2-Exp is distributed as an FP8 checkpoint, loaded from its own quantization config. The router and the indexer's `weights_proj` stay in full precision: ```python -from transformers import FineGrainedFP8Config, AutoModelForCausalLM, AutoTokenizer -import torch +from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "deepseek-ai/DeepSeek-V3.2-Exp" -quantization_config = FineGrainedFP8Config( - modules_to_not_convert=["model.layers.*.mlp.gate.*", "*.self_attn.indexer.weights_proj.*"], - weight_block_size=(128, 128), -) -model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype="auto", - device_map="auto", - quantization_config=quantization_config, -) +model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_name) inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device) diff --git a/docs/source/en/model_doc/deepseek_v4.md b/docs/source/en/model_doc/deepseek_v4.md index a87c733721ac..efc9f6e10019 100644 --- a/docs/source/en/model_doc/deepseek_v4.md +++ b/docs/source/en/model_doc/deepseek_v4.md @@ -29,6 +29,41 @@ share the same architecture; they differ only in width / depth / expert count an > [!TIP] > Set `use_kernels=True` in [`~PreTrainedModel.from_pretrained`] to replace supported layers with optimized kernels from the Hub. Refer to [Loading kernels](../kernel_doc/loading_kernels) to learn more. +## Usage examples + +DeepSeek-V4 checkpoints store block-FP8 weights with MXFP4 routed experts (`expert_dtype="fp4"` in the model config). +They load as-is, with no quantization config: + +```python +from transformers import AutoModelForCausalLM, AutoTokenizer + +model_name = "deepseek-ai/DeepSeek-V4-Flash" +model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto", device_map="auto") +tokenizer = AutoTokenizer.from_pretrained(model_name) + +inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device) +outputs = model.generate(**inputs, max_new_tokens=20) +print(tokenizer.decode(outputs[0], skip_special_tokens=True)) +``` + +The checkpoint mixes two formats, so it loads through [`FineGrainedConfig`]. The config below is the one built +automatically; pass your own to change a format, for example `Mxfp4Config(activation_format="bf16")` to run the experts +weight-only, or `dequantize=True` to load everything in bf16: + +```python +from transformers import AutoModelForCausalLM, BlockFp8Config, FineGrainedConfig, Mxfp4Config + +quantization_config = FineGrainedConfig( + dtype_plan={ + r"\.experts($|\.)": Mxfp4Config(activation_format="mxfp4"), + ".*": BlockFp8Config(weight_block_size=(128, 128), scale_fmt="ue8m0"), + } +) +model = AutoModelForCausalLM.from_pretrained( + "deepseek-ai/DeepSeek-V4-Flash", dtype="auto", device_map="auto", quantization_config=quantization_config +) +``` + ## Architecture (paper §2) ### Hybrid attention (§2.3) diff --git a/docs/source/en/model_doc/minimax_m3_vl.md b/docs/source/en/model_doc/minimax_m3_vl.md index 2688f8c31ea2..a004b3c99921 100644 --- a/docs/source/en/model_doc/minimax_m3_vl.md +++ b/docs/source/en/model_doc/minimax_m3_vl.md @@ -153,7 +153,7 @@ from transformers import ( AutoModelForImageTextToText, AutoTokenizer, CompileConfig, - FineGrainedFP8Config, + Mxfp8Config, ) from transformers.distributed import DistributedConfig @@ -164,7 +164,7 @@ model = AutoModelForImageTextToText.from_pretrained( "MiniMaxAI/MiniMax-M3-preview", dtype=torch.bfloat16, # Dequantize the native MXFP8 weights to bf16 at load (the speed win); needs even TP/EP sharding. - quantization_config=FineGrainedFP8Config(dequantize=True), + quantization_config=Mxfp8Config(dequantize=True), distributed_config=DistributedConfig( tp_size=int(os.environ["WORLD_SIZE"]), ep_size=int(os.environ["WORLD_SIZE"]), diff --git a/docs/source/en/quantization/blockfp8.md b/docs/source/en/quantization/blockfp8.md new file mode 100644 index 000000000000..a392911efe72 --- /dev/null +++ b/docs/source/en/quantization/blockfp8.md @@ -0,0 +1,130 @@ + + +# Block FP8 + +Block FP8 stores weights in 8-bit floating point (`float8_e4m3fn`) with one scale per 2D block of the weight matrix +(`weight_block_size=(128, 128)` by default), one scale per tensor (`weight_block_size=None`), or one scale per +output channel (`weight_block_size="row"`). Rowwise FP8 uses dynamic per-token activation scales. Activations are +quantized to FP8 too, per token and block at run time (`activation_scheme="dynamic"`), or against a calibrated scale +shipped with the checkpoint (`activation_scheme="static"`). + +It is the format of [DeepSeek-V3](https://hf.co/papers/2412.19437) and DeepSeek-R1, Qwen3-FP8, and Mistral's static +FP8 checkpoints. [`BlockFp8Config`] loads these checkpoints as they ship and quantizes full-precision checkpoints on the +fly. Checkpoints declare it as `quant_method="fp8"`. + +
+ +
+ +> [!TIP] +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer (Ada, Hopper, Blackwell) or an Intel XPU, +> and the [kernels](https://github.com/huggingface/kernels) package, which loads the +> [finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels) from the Hub. Without one, a +> pre-quantized checkpoint is dequantized to full precision as it loads. + +```bash +pip install --upgrade accelerate kernels +``` + +## Load a pre-quantized checkpoint + +A pre-quantized checkpoint carries its quantization config, so no config is needed. + +```py +from transformers import AutoModelForCausalLM, AutoTokenizer + +model_id = "Qwen/Qwen3-8B-FP8" +model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto") + +tokenizer = AutoTokenizer.from_pretrained(model_id) +inputs = tokenizer("What are we having for dinner?", return_tensors="pt").to(model.device) +output = model.generate(**inputs, max_new_tokens=20) +print(tokenizer.decode(output[0], skip_special_tokens=True)) +``` + +Checkpoints exported by [compressed-tensors](./compressed_tensors) in block FP8 (`format="float-quantized"`) load as +[`BlockFp8Config`] as well. + +Rowwise FP8 from FBGEMM, ModelOpt (`FP8_PER_CHANNEL_PER_TOKEN`), and compressed-tensors (channel weights with +token activations) uses the same finegrained runtime. Producer-specific key and scale layouts are converted on load. + +## Quantize on the fly + +Pass a [`BlockFp8Config`] to [`~PreTrainedModel.from_pretrained`] to quantize a full-precision checkpoint as it loads. +Set `dtype="auto"` so the unquantized modules keep the checkpoint's dtype. + +```py +from transformers import AutoModelForCausalLM, BlockFp8Config + +quantization_config = BlockFp8Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "meta-llama/Llama-3.1-8B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +``` + +[`~PreTrainedModel.save_pretrained`] writes the quantized model back in the same layout, and +[`~PreTrainedModel.from_pretrained`] reloads it. + +```py +model.save_pretrained("llama-3.1-8b-fp8") +model = AutoModelForCausalLM.from_pretrained("llama-3.1-8b-fp8", device_map="auto") +``` + +## Activation formats + +`activation_format` picks what the activations are quantized to before each matmul. + +| `activation_format` | Scheme | Activations | +|---|---|---| +| `None` (default) or `"fp8"` | W8A8 | FP8, per token and block at run time, or against a calibrated per-tensor scale with `activation_scheme="static"` | +| `"bf16"` | W8A16 | Full precision; block-FP8 weights are decoded inside the matmul | + +## Options + +- `weight_block_size` is the `(rows, columns)` block each weight scale covers; `None` gives one scale per tensor, + and `"row"` gives one FP32 scale per output channel. +- `activation_scale_ub` caps activation amax for dynamic rowwise W8A8; values beyond it saturate. FBGEMM checkpoints + retain their configured bound (1200 by default). Other rowwise formats leave it unset. +- `activation_scheme="static"` quantizes activations against the calibrated scale a checkpoint ships (its + `input_scale`), instead of deriving one per token at run time. +- `scale_fmt="ue8m0"` holds the block scales as powers of two (`float8_e8m0fnu`, with a `uint8` exponent fallback when the dtype is unavailable), as DeepSeek-V4 ships them. + Float32 scale containers, such as DeepSeek-V4-Base's, load and save as native UE8M0. +- `modules_to_not_convert` keeps the modules it names in full precision; `modules_to_convert` names embedding tables to + hold in FP8 with one per-tensor scale (a table too large to keep in full precision). +- `dequantize=True` loads a pre-quantized checkpoint dequantized to full precision. + +## DeepGEMM fast path + +On Hopper GPUs, a dense FP8 linear with `weight_block_size=(128, 128)` and dynamic activations dispatches to the +[DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels from +[kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm). On Blackwell the Triton kernels are +faster and are always used. + +DeepGEMM JIT-compiles its kernels with the system `nvcc`, so a full CUDA toolkit must be installed (12.3 or later on +Hopper). Transformers finds the toolkit by checking `CUDA_HOME`, then `CUDA_PATH`, then `nvcc` on your `PATH`, then +`/usr/local/cuda`. If DeepGEMM cannot load, Transformers logs a warning once and uses the Triton kernels. Set +`TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to always use them. + +For MoE experts, DeepGEMM is opt-in: pass `experts_implementation="deepgemm"` (or `"deepgemm_megamoe"` on Blackwell) +to [`~PreTrainedModel.from_pretrained`]. The DeepGEMM experts kernels on Blackwell read only UE8M0 scales; with +`float32` scales, use the default `grouped_mm` and `batched_mm` experts implementations. See the +[Experts backends](../experts_interface) guide for every option. + +## Related formats + +- [MXFP8](./mxfp8) holds FP8 weights with one power-of-two scale per 32 values of a row. +- [Fine-grained](./finegrained) mixes block FP8 with the 4-bit formats in one model, as DeepSeek-V4 ships it. diff --git a/docs/source/en/quantization/compressed_tensors.md b/docs/source/en/quantization/compressed_tensors.md index 22d1b8c68c83..09e57991053c 100644 --- a/docs/source/en/quantization/compressed_tensors.md +++ b/docs/source/en/quantization/compressed_tensors.md @@ -49,6 +49,14 @@ pip install -e . Search using the compressed-tensors [tag](https://huggingface.co/models?other=compressed-tensors) to find a compatible model on the Hugging Face Hub. +> [!NOTE] +> A checkpoint in a format the [fine-grained](./finegrained) kernels serve loads through them, with +> `checkpoint_format="compressed-tensors"`: `float-quantized` FP8 with per-tensor or block weight scales as +> [`BlockFp8Config`], `nvfp4-pack-quantized` as [`NVFP4Config`], and one whose config groups mix formats as a +> [`FineGrainedConfig`]. Integer, per-channel FP8, sparse and transformed checkpoints stay with compressed-tensors, as +> do `naive-quantized` ones such as the examples on this page. Pass a [`CompressedTensorsConfig`] to +> [`~PreTrainedModel.from_pretrained`] to load any checkpoint with compressed-tensors. + Pre-quantized models can be loaded directly. To quantize a model into the compressed-tensors format, see [llm-compressor](https://github.com/vllm-project/llm-compressor). Alternatively, models can be created independently and serialized with a compressed-tensors config. ```python diff --git a/docs/source/en/quantization/concept_guide.md b/docs/source/en/quantization/concept_guide.md index 47c3b037c867..c9e429bdd89b 100644 --- a/docs/source/en/quantization/concept_guide.md +++ b/docs/source/en/quantization/concept_guide.md @@ -110,7 +110,7 @@ FP8 is used in the *A8W8* quantization scheme, which quantizes both activations While int8 has broad support, efficient FP8 computation requires specific hardware capabilities found in newer GPUs like NVIDIA H100/H200/B100, AMD Instinct MI300 series, and Intel XPU. Without native hardware acceleration, the benefits of FP8 might not be fully realized. -Transformers supports FP8 through specific backends like [FBGEMM](./fbgemm_fp8), [FineGrainedFP8](./finegrained_fp8), and [compressed-tensors](./compressed_tensors). These backends handle the underlying FP8 conversion and computation when the appropriate hardware and configurations are used. +Transformers supports FP8 through specific backends like [FBGEMM](./fbgemm_fp8), [Block FP8](./blockfp8), [MXFP8](./mxfp8), and [compressed-tensors](./compressed_tensors). These backends handle the underlying FP8 conversion and computation when the appropriate hardware and configurations are used. ## Granularity diff --git a/docs/source/en/quantization/contribute.md b/docs/source/en/quantization/contribute.md index 190f7692231b..fd5d540470f9 100644 --- a/docs/source/en/quantization/contribute.md +++ b/docs/source/en/quantization/contribute.md @@ -46,7 +46,7 @@ Some quantization methods may require "pre-quantizing" the model through data ca ## Create new HFQuantizer class -0. The best starting point would be to have a look at another quantization method such as Finegrained Fp8. You will have to update or create three files in total: the [config file](https://github.com/huggingface/transformers/blob/main/src/transformers/utils/quantization_config.py), the [integration file](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/finegrained_fp8.py) and the [quantizer file](https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_finegrained_fp8.py). +0. The best starting point would be to have a look at another quantization method such as Block FP8. You will have to update or create three files in total: the [config file](https://github.com/huggingface/transformers/blob/main/src/transformers/utils/quantization_config.py), the [integration file](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/finegrained/layers.py) and the [quantizer file](https://github.com/huggingface/transformers/blob/main/src/transformers/quantizers/quantizer_blockfp8.py). 1. Create a new quantization config class inside [src/transformers/utils/quantization_config.py](https://github.com/huggingface/transformers/blob/abbffc4525566a48a9733639797c812301218b83/src/transformers/utils/quantization_config.py). Add the new quantization config to the [\_import_structure](https://github.com/huggingface/transformers/blob/abbffc4525566a48a9733639797c812301218b83/src/transformers/__init__.py#L1088) inside Transformers' [src/transformers/\_\_init\_\_.py](https://github.com/huggingface/transformers/blob/abbffc4525566a48a9733639797c812301218b83/src/transformers/__init__.py) file. diff --git a/docs/source/en/quantization/fbgemm_fp8.md b/docs/source/en/quantization/fbgemm_fp8.md index 503eb87dd087..a7347fcefcde 100644 --- a/docs/source/en/quantization/fbgemm_fp8.md +++ b/docs/source/en/quantization/fbgemm_fp8.md @@ -16,20 +16,23 @@ rendered properly in your Markdown viewer. # FBGEMM -[FBGEMM (Facebook GEneral Matrix Multiplication)](https://github.com/pytorch/FBGEMM) is a low-precision matrix multiplication library for small batch sizes and support for accuracy-loss minimizing techniques such as row-wise quantization and outlier-aware quantization. With FBGEMM, quantize a models weights to 8-bits/channel and the activations to 8-bits/token (also known as fp8 or w8a8). +[FBGEMM](https://github.com/pytorch/FBGEMM) FP8 checkpoints store weights in FP8 with one scale per output channel +and quantize activations per token (W8A8). Transformers reads this recipe into +[`BlockFp8Config`] with `weight_block_size="row"` and executes it through the +[finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels), including Llama4 experts. +The producer's `activation_scale_ub` is preserved: it caps activation amax before FP8 quantization. > [!TIP] > You need a GPU with [compute capability 9+](https://developer.nvidia.com/cuda-gpus#collapseOne) like a H100. -Install the FBGEMM_GPU package with the command below to ensure you have the latest version. +Install the kernel loader and Accelerate: ```bash -pip install --upgrade accelerate fbgemm-gpu torch +pip install --upgrade accelerate kernels ``` -If you're having installation issues, try installing the [nightly release](https://pytorch.org/FBGEMM/fbgemm_gpu-development/InstallationInstructions.html#fbgemm-gpu-install-libraries:~:text=found%20here.-,Install%20the%20FBGEMM_GPU%20Package,-Install%20through%20PyTorch). - -Create a [`FbgemmFp8Config`] and pass it to [`~PreTrainedModel.from_pretrained`] to quantize a model to fp8. +[`FbgemmFp8Config`] remains a checkpoint-compatible entry point and normalizes to this rowwise configuration. +Pass it to [`~PreTrainedModel.from_pretrained`] to quantize a model to FP8. ```py from transformers import FbgemmFp8Config, AutoModelForCausalLM diff --git a/docs/source/en/quantization/finegrained.md b/docs/source/en/quantization/finegrained.md new file mode 100644 index 000000000000..02c57411547a --- /dev/null +++ b/docs/source/en/quantization/finegrained.md @@ -0,0 +1,118 @@ + + +# Fine-grained + +[Block FP8](./blockfp8), [MXFP8](./mxfp8), [MXFP4](./mxfp4) and [NVFP4](./nvfp4) each have their own config and +quantize a whole model on their own. [`FineGrainedConfig`] mixes them in one model: its `dtype_plan` maps module name +patterns to the format config each matching module is quantized in. A typical plan holds the MoE experts, which carry +most of the weights, in a 4-bit format and every other linear in 8 bits. + +```py +from transformers import AutoModelForCausalLM, BlockFp8Config, FineGrainedConfig, Mxfp4Config + +quantization_config = FineGrainedConfig( + dtype_plan={ + r"\.experts($|\.)": Mxfp4Config(), + ".*": BlockFp8Config(), + }, + modules_to_not_convert=["lm_head"], +) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-30B-A3B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("qwen3-30b-a3b-mixed") +``` + +The saved model records its plan, so it reloads with no config. + +> [!TIP] +> The requirements are those of the formats in the plan: a CUDA GPU with compute capability 8.9 or newer or an Intel +> XPU, and the [kernels](https://github.com/huggingface/kernels) package. + +## How a plan resolves + +- A pattern names the modules it matches with `re.search`. A targeted entry wins over the `".*"` catch-all, whatever + order the plan is written in, since a plan saved to `config.json` loses its order. Two targeted entries matching one + module raise an error. +- A module no entry matches, or that `modules_to_not_convert` names, stays in full precision. +- `modules_to_convert` names embedding tables to hold in FP8 with one per-tensor scale, whatever the plan's formats. +- `dequantize=True` expands linear and expert weights to full precision. FP8 embedding tables stay compressed; + plans containing NVFP4 embedding tables currently reject this option. + +Quantized embedding tables support column-wise tensor parallelism (`colwise` or `colwise_gather_output`). +Vocabulary sharding (`rowwise`) is currently unsupported. For NVFP4, each rank's embedding width must be divisible +by the scale group size of 16. + +## Checkpoints that mix formats + +These load as a [`FineGrainedConfig`] without one being passed: + +- DeepSeek-V4 ships block-FP8 weights with MXFP4 routed experts. Its checkpoint declares `quant_method="fp8"` and its + model config `expert_dtype="fp4"`, which together read as a plan of MXFP4 experts over block FP8 (DeepSeek-V4.1 + declares `expert_dtype` in its quantization config instead, and reads the same way). See the + [DeepSeek-V4](../model_doc/deepseek_v4) page for the plan written out. +- A ModelOpt export labelled `quant_algo="NVFP4"` whose `config_groups` give different modules different formats + reads as a plan with one entry per group target. +- ModelOpt `quant_algo="MIXED_PRECISION"` reads the `quantized_layers` mapping. Each concrete module name covers + that module and its descendants; modules absent from the mapping stay in full precision. Supported algorithms + are `FP8` (calibrated per-tensor FP8), `FP8_PER_CHANNEL_PER_TOKEN` (rowwise FP8), + `FP8_PB_WO` (block-FP8 weights with full-precision activations), `NVFP4`, `MXFP8`, `MXFP4`, and the + `W4A16_NVFP4` / `W4A16_MXFP4` weight-only variants. Fields can appear directly in the quantization config or + inside its `quantization` object. Unsupported algorithms or config groups raise an error. + +For example, ModelOpt's per-module assignment can combine weight-only block-FP8 attention and NVFP4 experts: + +```json +{ + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "layers.8.attention.wq": {"quant_algo": "FP8_PB_WO"}, + "layers.9.experts": {"quant_algo": "NVFP4", "group_size": 16} + } +} +``` + +The model must support the checkpoint's architecture and parameter names. Format parsing alone does not add +support for a new architecture. Module names are resolved through the model's checkpoint key renames. An +assignment that matches no supported linear, expert, or embedding module raises an error, including per-expert assignments +that cannot be represented by an architecture's fused expert module. + +AWQ smoothing (`pre_quant_scale`), SVD residuals, asymmetric zero points, sparse or transformed layouts, and +compressed scale encodings such as `weight_scale_encoding="csf"` are rejected. Explicit NVFP4 embedding targets +load as packed tables with E4M3 scales per 16 values and a shared FP32 multiplier. Lookup decodes only selected +rows, and saving preserves the checkpoint's row shards. These tables cannot currently be loaded with +`dequantize=True`, including the automatic CPU dequantization fallback. Supporting the embedding layout alone +does not make a checkpoint using unsupported scale encodings loadable. Weight-only block-FP8 requires an updated +kernel bundle with weight-only support for the selected dense or expert backend. + +To change one format of such a checkpoint, pass the plan you want; to dequantize it, pass +`FineGrainedConfig(dequantize=True)`. + +## Expert parallelism + +Until the EP implementation PR is merged, the default fine-grained expert-parallel path temporarily keeps tokens +replicated across ranks, masks routes to non-local experts, and sums their outputs with all-reduce. Set `tp_size` and +`ep_size` to the same value in `DistributedConfig`. DeepGEMM MegaMoE uses its own routing and communication. + +## Kernel autotuning + +Fine-grained kernels tune each new input shape and cache the chosen configuration. For correctness tests or quick +experiments, set `FINEGRAINED_AUTOTUNE_TRIALS=1` before loading the kernel bundle to reduce the Bayesian search +budget. The tuner still probes initial configurations and refines its choice, so this does not cap the total number of +benchmarks at one. Leave it unset when measuring performance so the tuner can search its full budget. See the +[kernel autotuning documentation](https://huggingface.co/kernels-community/finegrained-kernels#autotuning). diff --git a/docs/source/en/quantization/finegrained_fp8.md b/docs/source/en/quantization/finegrained_fp8.md index efc601c1cdae..11ca965e841e 100644 --- a/docs/source/en/quantization/finegrained_fp8.md +++ b/docs/source/en/quantization/finegrained_fp8.md @@ -14,73 +14,20 @@ rendered properly in your Markdown viewer. --> -# Fine-grained FP8 +# Fine-grained FP8 (deprecated) -Fine-grained FP8 quantization quantizes the weights and activations to fp8. +> [!WARNING] +> `FineGrainedFP8Config` is deprecated and will be removed in a future release. Existing code keeps working: a +> `FineGrainedFP8Config` passed to `from_pretrained` loads through its own frozen integration. A checkpoint declaring +> `quant_method="fp8"`, loaded without one, loads as [`BlockFp8Config`]. -- The weights are quantized to 8-bits for each 2D block (`weight_block_size=(128, 128)`). -- The activations are quantized to 8-bits for each group per token. The group value matches the weights in the input channel (128 by default). +New code uses the config that describes the checkpoint's format: -FP8 quantization enables support for [DeepSeek-V3](https://hf.co/papers/2412.19437) and DeepSeek-R1. +| `FineGrainedFP8Config` | Use | Guide | +|---|---|---| +| `weight_block_size=[1, 32]` with `scale_fmt="ue8m0"` | [`Mxfp8Config`] | [MXFP8](./mxfp8) | +| any other `weight_block_size` | [`BlockFp8Config`] | [Block FP8](./blockfp8) | +| on a model whose config declares `expert_dtype="fp4"` (DeepSeek-V4) | [`BlockFp8Config`], which loads as a [`FineGrainedConfig`] of MXFP4 experts over block FP8 | [Fine-grained](./finegrained) | -
- -
- -> [!TIP] -> You need a GPU with Compute Capability>=9 (H100), and install a PyTorch version compatible with the CUDA version of your GPU. - -Install Accelerate and upgrade to the latest version of PyTorch. - -```bash -pip install --upgrade accelerate torch -``` - -Create a [`FineGrainedFP8Config`] class and pass it to [`~PreTrainedModel.from_pretrained`] to quantize it. The weights are loaded in full precision (`torch.float32`) by default regardless of the actual data type the weights are stored in. Set `dtype="auto"` to load the weights in the data type defined in a models `config.json` file to automatically load the most memory-optimal data type. - -```py -from transformers import FineGrainedFP8Config, AutoModelForCausalLM, AutoTokenizer - -model_name = "meta-llama/Meta-Llama-3-8B" -quantization_config = FineGrainedFP8Config() -quantized_model = AutoModelForCausalLM.from_pretrained(model_name, dtype="auto", device_map="auto", quantization_config=quantization_config) - -tokenizer = AutoTokenizer.from_pretrained(model_name) -input_text = "What are we having for dinner?" -input_ids = tokenizer(input_text, return_tensors="pt").to(quantized_model.device.type) - -output = quantized_model.generate(**input_ids, max_new_tokens=10) -print(tokenizer.decode(output[0], skip_special_tokens=True)) -``` - -Use [`~PreTrainedModel.save_pretrained`] to save the quantized model and reload it with [`~PreTrainedModel.from_pretrained`]. - -```py -quant_path = "/path/to/save/quantized/model" -quantized_model.save_pretrained(quant_path) -model = AutoModelForCausalLM.from_pretrained(quant_path, device_map="auto") -``` - -## DeepGEMM fast path - -On Hopper (SM90+) and Blackwell (SM100+) GPUs, every FP8 linear automatically dispatches to the [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) kernels from [kernels-community/deep-gemm](https://huggingface.co/kernels-community/deep-gemm) when `weight_block_size=(128, 128)` and `activation_scheme="dynamic"`. DeepGEMM is 3-6x faster than the Triton fallback. Install or upgrade the [kernels](https://github.com/huggingface/kernels) package to enable it. - -```bash -pip install -U kernels -``` - -DeepGEMM JIT-compiles its kernels with the system `nvcc`, so a full CUDA toolkit must be installed. A runtime-only install (CUDA libraries without `nvcc`) won't work. The minimum toolkit version depends on the hardware: 12.3 or later on Hopper (SM90), 12.9 or later on Blackwell (SM100). - -Transformers finds the toolkit by checking `CUDA_HOME`, then `CUDA_PATH`, then `nvcc` on your `PATH`, then `/usr/local/cuda`. Set `CUDA_HOME` to the toolkit root if `nvcc` isn't discovered. - -If the kernel cannot load (missing `kernels`, unsupported GPU, no CUDA toolkit, or an `nvcc` older than the required version), Transformers logs a warning once and falls back to the Triton finegrained-fp8 kernel. Static activation quantization always stays on the Triton path. - -To force the Triton fallback even when DeepGEMM is available, set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`. This only affects the FP8 linear dispatch and leaves the `"deepgemm"` experts backend untouched, which you switch with [`~PreTrainedModel.set_experts_implementation`]. - -For MoE experts, the DeepGEMM path is opt-in. Pass `experts_implementation="deepgemm"` (or `"deepgemm_megamoe"` on Blackwell) at load time to route the expert matmuls through DeepGEMM. See the [Experts backends](../experts_interface) guide for the full set of options. - -## UE8M0 scale format - -DeepSeek V4-style checkpoints store FP8 weight scales in the packed `float8_e8m0fnu` format instead of `float32`. These checkpoints are pre-quantized and set `scale_fmt="ue8m0"` in their quantization config. Both the DeepGEMM and Triton kernels read UE8M0 scales, so these checkpoints run on either path. - -On Blackwell (SM100+), the DeepGEMM experts kernels only supports UE8M0 scales. A checkpoint with plain `float32` scales (`scale_fmt="float"`) raises a `ValueError`. Use a `scale_fmt="ue8m0"` checkpoint, or run the experts with `grouped_mm` or `batched_mm`, which support `float32` scales directly. Hopper (SM90+) supports `float32` scales on the DeepGEMM path without conversion. See the [Experts backends](../experts_interface) guide for the experts backend options. +The fields keep their names and meanings, so `FineGrainedFP8Config(weight_block_size=(128, 128), +activation_scheme="static")` becomes `BlockFp8Config(weight_block_size=(128, 128), activation_scheme="static")`. diff --git a/docs/source/en/quantization/mxfp4.md b/docs/source/en/quantization/mxfp4.md index f96833c0cc09..456b584ab50c 100644 --- a/docs/source/en/quantization/mxfp4.md +++ b/docs/source/en/quantization/mxfp4.md @@ -16,62 +16,70 @@ rendered properly in your Markdown viewer. # MXFP4 -Note: MXFP4 quantization currently only works for OpenAI GPT-OSS 120b and 20b. +MXFP4 is the OCP microscaling 4-bit format: weights in 4-bit floating point (E2M1, two values per byte) with one +power-of-two scale (UE8M0) per group of 32 values along a row. It cuts weight memory about 4x against bf16, so +GPT-OSS-120B fits on a single 80GB GPU and GPT-OSS-20B needs about 16GB. + +[`Mxfp4Config`] loads MXFP4 checkpoints, such as OpenAI's GPT-OSS, as they ship and quantizes full-precision +checkpoints on the fly. Checkpoints declare it as `quant_method="mxfp4"`. DeepSeek-V4 ships its routed experts in MXFP4 +beside block-FP8 weights, which loads through [Fine-grained](./finegrained). + +> [!TIP] +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer or an Intel XPU, and the +> [kernels](https://github.com/huggingface/kernels) package, which loads the +> [finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels) from the Hub. Without one, a +> pre-quantized checkpoint is dequantized to full precision as it loads. + +```bash +pip install --upgrade accelerate kernels +``` -MXFP4 is a 4-bit floating point format that dramatically reduces the memory requirements of large models. Large models (GPT-OSS-120B) can fit on a single 80GB GPU and smaller models (GPT-OSS-20B) only require 16GB of memory. It uses blockwise scaling to preserve its range and accuracy, which typically becomes degraded at lower precisions. +## Load a pre-quantized checkpoint -To use MXPF4, make sure your hardware meets the following requirements. +A checkpoint whose quantization config says `"quant_method": "mxfp4"` loads in MXFP4 with no config. -- Install Accelerate, kernels, and Triton ≥ 3.4. Only manually install Triton ≥ 3.4 if you're using PyTorch 2.7 because it is already supported in PyTorch 2.8. -- NVIDIA GPU Compute Capability ≥ 7.5 which includes Tesla GPUs and newer. Use [get_device_capability](https://docs.pytorch.org/docs/stable/generated/torch.cuda.get_device_capability.html) to check Compute Capability. +```py +from transformers import AutoModelForCausalLM, AutoTokenizer -```python -from torch import cuda -cuda.get_device_capability() +model_id = "openai/gpt-oss-20b" +model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto") -# (7, 5) +tokenizer = AutoTokenizer.from_pretrained(model_id) +inputs = tokenizer("Once upon a time", return_tensors="pt").to(model.device) +output = model.generate(**inputs, max_new_tokens=20) +print(tokenizer.decode(output[0], skip_special_tokens=True)) ``` -Check a model's quantization config as shown below to see if it supports MXFP4. If `'quant_method': 'mxfp4'`, then the model automatically uses MXFP4. +Pass `Mxfp4Config(dequantize=True)` to load the weights dequantized to bf16 instead, for example to fine-tune them. + +## Quantize on the fly ```py -from transformers import GptOssConfig - -model_id = "openai/gpt-oss-120b" -cfg = GptOssConfig.from_pretrained(model_id) -print(cfg.quantization_config) - -# Example output: -# { -# 'modules_to_not_convert': [ -# 'model.layers.*.self_attn', -# 'model.layers.*.mlp.router', -# 'model.embed_tokens', -# 'lm_head' -# ], -# 'quant_method': 'mxfp4' -# } -``` +from transformers import AutoModelForCausalLM, Mxfp4Config -## MXFP4 kernels +quantization_config = Mxfp4Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-30B-A3B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("qwen3-30b-a3b-mxfp4") +``` -Transformers automatically pulls the MXFP4-aware Triton kernels from the community repository when you load a model that needs them. The kernels are stored in your local cache and used during the forward pass. +## Activation formats -MXFP4 kernels are used by default, if available and supported, and does not require any code changes. +`activation_format` picks what the activations are quantized to before each matmul. -You can use [hf cache scan](https://huggingface.co/docs/huggingface_hub/en/guides/manage-cache#scan-your-cache) to verify the kernels are downloaded. +| `activation_format` | Scheme | Activations | +|---|---|---| +| `"bf16"` (default) | W4A16 (weight-only) | kept in full precision, as GPT-OSS runs | +| `"mxfp8"` | W4A8 | MXFP8 | +| `None` or `"mxfp4"` | W4A4 | MXFP4 | -```shell -hf cache scan -``` +With a pre-quantized checkpoint, pass the config with the `activation_format` you want; the weights load as they +ship. -```shell -REPO ID REPO TYPE SIZE ON DISK --------------------------------- --------- ------------ -kernels-community/triton_kernels model 536.2K -openai/gpt-oss-20b model 13.8G -``` +The 4-bit activation scheme uses the 4-bit block-scaled MMA of Blackwell GPUs. ## Resources -Learn more about MXFP4 quantization and how blockwise scaling works in this [blog post](https://huggingface.co/blog/faster-transformers#mxfp4-quantization). +Learn more about MXFP4 and how blockwise scaling works in this +[blog post](https://huggingface.co/blog/faster-transformers#mxfp4-quantization). diff --git a/docs/source/en/quantization/mxfp8.md b/docs/source/en/quantization/mxfp8.md new file mode 100644 index 000000000000..7f61ae5486ed --- /dev/null +++ b/docs/source/en/quantization/mxfp8.md @@ -0,0 +1,72 @@ + + +# MXFP8 + +MXFP8 is the OCP microscaling FP8 format: weights in 8-bit floating point (`float8_e4m3fn`) with one power-of-two +scale (UE8M0) per group of 32 values along a row. [`Mxfp8Config`] loads MXFP8 checkpoints, such as `MiniMaxAI/MiniMax-M3-MXFP8`, as +they ship and quantizes full-precision checkpoints on the fly. Checkpoints declare it as `quant_method="mxfp8"`. + +> [!TIP] +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer or an Intel XPU, and the +> [kernels](https://github.com/huggingface/kernels) package. Without one, a pre-quantized checkpoint is dequantized to +> full precision as it loads. On Blackwell, the group scales feed the tensor cores' block-scaled MMA directly. + +```bash +pip install --upgrade accelerate kernels +``` + +## Load a pre-quantized checkpoint + +```py +from transformers import AutoModelForImageTextToText + +model = AutoModelForImageTextToText.from_pretrained("MiniMaxAI/MiniMax-M3-MXFP8", dtype="auto", device_map="auto") +``` + +## Quantize on the fly + +```py +from transformers import AutoModelForCausalLM, Mxfp8Config + +quantization_config = Mxfp8Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen3-30B-A3B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("qwen3-30b-a3b-mxfp8") +``` + +## Activation formats + +`activation_format` picks what the activations are quantized to before each matmul. + +| `activation_format` | Scheme | Activations | +|---|---|---| +| `None` (default) or `"mxfp8"` | W8A8 | MXFP8 | +| `"mxfp4"` | W8A4 | MXFP4 | +| `"bf16"` | W8A16 (weight-only) | kept in full precision | + +With a pre-quantized checkpoint, pass the config with the `activation_format` you want; the weights load as they +ship. + +Other options are shared with the other formats: `modules_to_not_convert` keeps modules in full precision, and +`dequantize=True` loads a pre-quantized checkpoint dequantized. + +## Related formats + +- [Block FP8](./blockfp8) holds FP8 weights with one scale per 2D block or per tensor. +- [MXFP4](./mxfp4) is the 4-bit microscaling format. +- [Fine-grained](./finegrained) mixes MXFP8 with other formats in one model. diff --git a/docs/source/en/quantization/nvfp4.md b/docs/source/en/quantization/nvfp4.md index 42a2d31cfa5a..33ece19b7e0d 100644 --- a/docs/source/en/quantization/nvfp4.md +++ b/docs/source/en/quantization/nvfp4.md @@ -16,38 +16,35 @@ rendered properly in your Markdown viewer. # NVFP4 -NVFP4 quantization packs full-precision linear weights into NVIDIA's 4-bit floating-point format while a model is -loaded. [`NVFP4Config`] replaces eligible bias-free `torch.nn.Linear` modules, whose `in_features` and `out_features` are both divisible by 16, with an NVFP4 linear implementation from -the [NVFP4 Hub kernel](https://huggingface.co/kernels-community/nvfp4-gemm). The model's attention and MLP interfaces are -not replaced. +NVFP4 is NVIDIA's 4-bit floating-point format: weights in E2M1 (two values per byte) with two levels of scale, an FP8 +(E4M3) scale per group of 16 values and one FP32 global scale per tensor. [`NVFP4Config`] loads NVFP4 checkpoints as +they ship and quantizes full-precision checkpoints on the fly. Checkpoints declare it as `quant_method="nvfp4"`. > [!TIP] -> NVFP4 requires a Blackwell GPU with compute capability 10.0 or newer, a compatible CUDA-enabled PyTorch build, and -> the [kernels](https://github.com/huggingface/kernels) package. - -Install Accelerate and a compatible version of `kernels`. +> Quantized compute needs a CUDA GPU with compute capability 8.9 or newer or an Intel XPU, and the +> [kernels](https://github.com/huggingface/kernels) package, which loads the +> [finegrained kernels](https://huggingface.co/kernels-community/finegrained-kernels) from the Hub. Without one, a +> pre-quantized checkpoint is dequantized to full precision as it loads. The 4-bit activation scheme uses the NVFP4 +> MMA of Blackwell GPUs. ```bash pip install --upgrade accelerate kernels ``` -Pass [`NVFP4Config`] to [`~PreTrainedModel.from_pretrained`] with a single CUDA device. Weights are quantized as they -are loaded, so the source checkpoint should contain floating-point weights. - -```py -import torch +## Load a pre-quantized checkpoint -from transformers import AutoModelForCausalLM, AutoTokenizer, NVFP4Config +The NVFP4 exports of [NVIDIA ModelOpt](https://github.com/NVIDIA/TensorRT-Model-Optimizer) (`quant_method="modelopt"`, +`quant_algo="NVFP4"`) and of [compressed-tensors](./compressed_tensors) (`format="nvfp4-pack-quantized"`) load as +[`NVFP4Config`], with their calibrated activation scales, and save back in their own layout. A ModelOpt export +labelled `quant_algo="NVFP4"` whose `config_groups` give different modules different formats loads as a +[`FineGrainedConfig`]. ModelOpt `quant_algo="MIXED_PRECISION"` uses its explicit `quantized_layers` mapping; see +[Fine-grained](./finegrained) for supported formats and restrictions. +```py +from transformers import AutoModelForCausalLM, AutoTokenizer -model_id = "meta-llama/Llama-3.2-1B" -quantization_config = NVFP4Config() -model = AutoModelForCausalLM.from_pretrained( - model_id, - dtype=torch.bfloat16, - device_map="cuda", - quantization_config=quantization_config, -) +model_id = "nvidia/Llama-3.1-8B-Instruct-NVFP4" +model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map="auto") tokenizer = AutoTokenizer.from_pretrained(model_id) inputs = tokenizer("NVFP4 is", return_tensors="pt").to(model.device) @@ -55,19 +52,32 @@ output = model.generate(**inputs, max_new_tokens=20) print(tokenizer.decode(output[0], skip_special_tokens=True)) ``` -Use `modules_to_not_convert` to keep selected modules in their original precision. +Pass `NVFP4Config(dequantize=True)` to load linear and expert weights dequantized to bf16 instead. +Checkpoints with explicit NVFP4 embedding targets keep those tables packed and decode only selected rows during +lookup; they currently reject `dequantize=True`, including the automatic CPU fallback. + +## Quantize on the fly ```py -quantization_config = NVFP4Config(modules_to_not_convert=["vision", "lm_head"]) +from transformers import AutoModelForCausalLM, NVFP4Config + +quantization_config = NVFP4Config(modules_to_not_convert=["lm_head"]) +model = AutoModelForCausalLM.from_pretrained( + "meta-llama/Llama-3.1-8B", dtype="auto", device_map="auto", quantization_config=quantization_config +) +model.save_pretrained("llama-3.1-8b-nvfp4") ``` -NVFP4 linear modules support `torch.compile`. The first compiled invocation includes graph compilation time, so warm up -the model before measuring generation throughput. +## Activation formats + +`activation_format` picks what the activations are quantized to before each matmul. + +| `activation_format` | Scheme | Activations | +|---|---|---| +| `None` (default) or `"nvfp4"` | W4A4 | NVFP4, against the checkpoint's calibrated global scale when it ships one | +| `"bf16"` | W4A16 (weight-only) | kept in full precision | -## Current limitations +With a pre-quantized checkpoint, pass the config with the `activation_format` you want; the weights load as they +ship. -- Only one CUDA device is supported. Tensor parallelism and multi-device `device_map` configurations are rejected until - the sharding behavior of the NVFP4 scale metadata is defined. -- CPU and disk offload are not supported. -- Pre-quantized NVFP4 checkpoints are not supported. -- NVFP4 models cannot currently be serialized with [`~PreTrainedModel.save_pretrained`] or trained. +`modules_to_not_convert` keeps the modules it names in full precision. diff --git a/docs/source/en/quantization/overview.md b/docs/source/en/quantization/overview.md index a5f4ad935164..d00b5a403f7d 100644 --- a/docs/source/en/quantization/overview.md +++ b/docs/source/en/quantization/overview.md @@ -47,13 +47,16 @@ Use the [quantization picker](https://huggingface.co/spaces/stevhliu/quantizatio | [HIGGS](./higgs) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 2/4 | 🔴 | 🟢 | 🟢 | https://github.com/HanGuo97/flute | | [HQQ](./hqq) | 🟢 | 🟢 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 1/8 | 🟢 | 🔴 | 🟢 | https://github.com/mobiusml/hqq/ | | [Metal](./metal) | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 🔴 | 🔴 | 2/4/8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/mlx-quantization-metal-kernels) | -| [NVFP4](./nvfp4) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 4 | 🔴 | 🔴 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/nvfp4-gemm) | +| [Block FP8](./blockfp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [Fine-grained](./finegrained) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 4/8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [MXFP4](./mxfp4) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 4 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [MXFP8](./mxfp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 8 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | +| [NVFP4](./nvfp4) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🟢 | 4 | 🔴 | 🟢 | 🟢 | [Hub Kernels](https://huggingface.co/kernels-community/finegrained-kernels) | | [optimum-quanto](./quanto) | 🟢 | 🟢 | 🟢 | 🔴 | 🟢 | 🟢 | 🟢 | 2/4/8 | 🔴 | 🔴 | 🟢 | https://github.com/huggingface/optimum-quanto | | [SINQ](./sinq) | 🟢 | 🟢 | 🟢 | 🟡 | 🟡 | 🟡 | 🟡 | 2/3/4/6/8 | 🔴 | 🟢 | 🟢 | https://github.com/huawei-csl/SINQ | | [FBGEMM_FP8](./fbgemm_fp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🔴 | 8 | 🔴 | 🟢 | 🟢 | https://github.com/pytorch/FBGEMM | | [torchao](./torchao) | 🟢 | 🟢 | 🟢 | 🔴 | 🟡 | 🟢 | | 4/8 | | 🟢🔴 | 🟢 | https://github.com/pytorch/ao | | [VPTQ](./vptq) | 🔴 | 🔴 | 🟢 | 🟡 | 🔴 | 🔴 | 🟢 | 1/8 | 🔴 | 🟢 | 🟢 | https://github.com/microsoft/VPTQ | -| [FINEGRAINED_FP8](./finegrained_fp8) | 🟢 | 🔴 | 🟢 | 🔴 | 🔴 | 🟢 | 🔴 | 8 | 🔴 | 🟢 | 🟢 | Built-in | | [SpQR](./spqr) | 🔴 | 🔴 | 🟢 | 🔴 | 🔴 | 🔴 | 🟢 | 3 | 🔴 | 🟢 | 🟢 | https://github.com/Vahe1994/SpQR/ | | [Quark](./quark) | 🔴 | 🟢 | 🟢 | 🟢 | 🟢 | 🟢 | ? | 2/4/6/8/9/16 | 🔴 | 🔴 | 🟢 | https://quark.docs.amd.com/latest/ | diff --git a/docs/source/en/quantization/quark.md b/docs/source/en/quantization/quark.md index c493a5f365f7..a87964f9494f 100644 --- a/docs/source/en/quantization/quark.md +++ b/docs/source/en/quantization/quark.md @@ -22,7 +22,7 @@ The PyTorch support integrated through 🤗 Transformers primarily targets AMD C Users interested in Quark can refer to its [documentation](https://quark.docs.amd.com/latest/) to get started quantizing models and using them in supported open-source libraries! -Although Quark has its own checkpoint / [configuration format](https://huggingface.co/amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test/blob/main/config.json#L26), the library also supports producing models with a serialization layout compliant with other quantization/runtime implementations ([AutoAWQ](./awq), [native fp8 in 🤗 Transformers](./finegrained_fp8)). +Although Quark has its own checkpoint / [configuration format](https://huggingface.co/amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test/blob/main/config.json#L26), the library also supports producing models with a serialization layout compliant with other quantization/runtime implementations ([AutoAWQ](./awq), [native fp8 in 🤗 Transformers](./blockfp8)). To be able to load Quark quantized models in Transformers, the library first needs to be installed: @@ -51,7 +51,7 @@ The table below shows a few features supported by Quark: Public models using Quark native serialization can be found at https://huggingface.co/models?other=quark. -Although Quark also supports [models using `quant_method="fp8"`](https://huggingface.co/models?other=fp8) and [models using `quant_method="awq"`](https://huggingface.co/models?other=awq), Transformers loads these models rather through [AutoAWQ](./awq) or uses the [native fp8 support in 🤗 Transformers](./finegrained_fp8). +Although Quark also supports [models using `quant_method="fp8"`](https://huggingface.co/models?other=fp8) and [models using `quant_method="awq"`](https://huggingface.co/models?other=awq), Transformers loads these models rather through [AutoAWQ](./awq) or uses the [native fp8 support in 🤗 Transformers](./blockfp8). ## Using Quark models in Transformers diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 804b18726cca..c803a0e946e7 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -233,9 +233,12 @@ "AwqConfig", "BitNetQuantConfig", "BitsAndBytesConfig", + "BlockFp8Config", + "BlockScaledConfig", "CompressedTensorsConfig", "EetqConfig", "FbgemmFp8Config", + "FineGrainedConfig", "FineGrainedFP8Config", "FourOverSixConfig", "FPQuantConfig", @@ -246,6 +249,7 @@ "HqqConfig", "MetalConfig", "Mxfp4Config", + "Mxfp8Config", "NVFP4Config", "QuantoConfig", "QuarkConfig", @@ -734,9 +738,12 @@ from .utils.quantization_config import AwqConfig as AwqConfig from .utils.quantization_config import BitNetQuantConfig as BitNetQuantConfig from .utils.quantization_config import BitsAndBytesConfig as BitsAndBytesConfig + from .utils.quantization_config import BlockFp8Config as BlockFp8Config + from .utils.quantization_config import BlockScaledConfig as BlockScaledConfig from .utils.quantization_config import CompressedTensorsConfig as CompressedTensorsConfig from .utils.quantization_config import EetqConfig as EetqConfig from .utils.quantization_config import FbgemmFp8Config as FbgemmFp8Config + from .utils.quantization_config import FineGrainedConfig as FineGrainedConfig from .utils.quantization_config import FineGrainedFP8Config as FineGrainedFP8Config from .utils.quantization_config import FourOverSixConfig as FourOverSixConfig from .utils.quantization_config import FPQuantConfig as FPQuantConfig @@ -746,6 +753,8 @@ from .utils.quantization_config import HiggsConfig as HiggsConfig from .utils.quantization_config import HqqConfig as HqqConfig from .utils.quantization_config import MetalConfig as MetalConfig + from .utils.quantization_config import Mxfp4Config as Mxfp4Config + from .utils.quantization_config import Mxfp8Config as Mxfp8Config from .utils.quantization_config import NVFP4Config as NVFP4Config from .utils.quantization_config import QuantoConfig as QuantoConfig from .utils.quantization_config import QuarkConfig as QuarkConfig diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 28be4395b949..1a3fbb8c9f2d 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -546,17 +546,27 @@ def _build_checkpoint_conversion_mapping(): ], ), ], + # a pre-quantized checkpoint ships a weight's block scales as `.scale`; ahead of the `hc_*_scale` renames, which + # give the model's own `.scale` parameters their dot + "deepseek_v4_quantized": [ + WeightRenaming(source_patterns=r"^(.+)\.scale", target_patterns=r"\1.weight_scale_inv"), + ], "deepseek_v4": [ + # checkpoint-form renames the general `.attn.` one below would otherwise take first + WeightRenaming(source_patterns=r"\.attn\.wq_b\.", target_patterns=r".attn.q_b_proj."), # General mid-key renames acting on several keys WeightRenaming(source_patterns=r"\.attn\.", target_patterns=r".self_attn."), WeightRenaming(source_patterns=r"\.ffn\.", target_patterns=r".mlp."), WeightRenaming(source_patterns=r"\.indexer\.compressor\.", target_patterns=r"\.compressor\.indexer\."), # Leaf patterns - WeightRenaming(source_patterns=r"^embed\.weight", target_patterns="embed_tokens.weight"), + # The checkpoint has no `model.` prefix: named here, so a save takes it back off + WeightRenaming(source_patterns=r"^layers\.", target_patterns="model.layers."), + WeightRenaming(source_patterns=r"^norm\.weight", target_patterns="model.norm.weight"), + WeightRenaming(source_patterns=r"^embed\.weight", target_patterns="model.embed_tokens.weight"), WeightRenaming(source_patterns=r"^head\.weight", target_patterns="lm_head.weight"), - WeightRenaming(source_patterns=r"^hc_head_fn", target_patterns="hc_head.hc_fn"), - WeightRenaming(source_patterns=r"^hc_head_base", target_patterns="hc_head.hc_base"), - WeightRenaming(source_patterns=r"^hc_head_scale", target_patterns="hc_head.hc_scale"), + WeightRenaming(source_patterns=r"^hc_head_fn", target_patterns="model.hc_head.hc_fn"), + WeightRenaming(source_patterns=r"^hc_head_base", target_patterns="model.hc_head.hc_base"), + WeightRenaming(source_patterns=r"^hc_head_scale", target_patterns="model.hc_head.hc_scale"), WeightRenaming(source_patterns=r"\.attn_norm\.", target_patterns=r".input_layernorm."), WeightRenaming(source_patterns=r"\.ffn_norm\.", target_patterns=r".post_attention_layernorm."), WeightRenaming(source_patterns=r"\.hc_attn_fn", target_patterns=r".attn_hc.fn"), @@ -571,10 +581,13 @@ def _build_checkpoint_conversion_mapping(): target_patterns=r"\.compressor\.indexer\.scorer\.weights_proj\.", ), WeightRenaming(source_patterns=r"\.indexer\.wq_b\.", target_patterns=r"\.compressor\.indexer\.q_b_proj\."), - WeightRenaming(source_patterns=r"\.norm\.", target_patterns=r"\.kv_norm\."), + # the compressors' norms only: the attention's own is `kv_norm` already, and a save reversing a + # bare `.kv_norm.` would rename it too (`[.]`: the reverse unescapes a group's `\.`) + WeightRenaming( + source_patterns=r"\.compressor\.(|indexer[.])norm\.", target_patterns=r"\.compressor\.\1kv_norm\." + ), WeightRenaming(source_patterns=r"\.ape", target_patterns=r"\.position_bias"), WeightRenaming(source_patterns=r"\.wq_a\.", target_patterns=r".q_a_proj."), - WeightRenaming(source_patterns=r"\.self_attn\.wq_b\.", target_patterns=r"\.self_attn\.q_b_proj\."), WeightRenaming(source_patterns=r"\.wkv\.", target_patterns=r".kv_proj."), WeightRenaming(source_patterns=r"\.wgate\.", target_patterns=r".gate_proj."), WeightRenaming(source_patterns=r"\.wo_a\.", target_patterns=r".o_a_proj."), @@ -684,7 +697,12 @@ def _build_checkpoint_conversion_mapping(): source_patterns=r"^patch_merge_mlp\.linear_2\.", target_patterns="model.multi_modal_projector.merge_linear_2.", ), - WeightRenaming(source_patterns=r"\.block_sparse_moe\.", target_patterns=r"\.mlp\."), + # the MoE block only, so a save does not reverse the dense and vision `.mlp.` too (`[.]`: the + # reverse unescapes a group's `\.`) + WeightRenaming( + source_patterns=r"\.block_sparse_moe\.(experts[.]|gate[.]|shared_experts[.]|e_score_correction_bias$)", + target_patterns=r"\.mlp\.\1", + ), WeightRenaming( source_patterns=r"\.e_score_correction_bias", target_patterns=r"\.gate\.e_score_correction_bias" ), @@ -1945,6 +1963,18 @@ def _build_checkpoint_conversion_mapping(): "UltrabertForQuestionAnswering": [ WeightRenaming(r"^ultrabert\.", "model."), ], + "llama4_text": [ + WeightConverter( + source_patterns="feed_forward.experts.gate_up_proj", + target_patterns="feed_forward.experts.gate_up_proj", + operations=[Transpose(1, 2)], + ), + WeightConverter( + source_patterns="feed_forward.experts.down_proj", + target_patterns="feed_forward.experts.down_proj", + operations=[Transpose(1, 2)], + ), + ], } mapping["Tipsv2DptForNormalEstimation"] = [ @@ -2288,6 +2318,10 @@ def get_model_conversion_mapping( if model_type is not None: conversions = get_checkpoint_conversion_mapping(model_type) + if hf_quantizer is not None and hf_quantizer.pre_quantized and model_type is not None: + if (quantized := get_checkpoint_conversion_mapping(f"{model_type}_quantized")) is not None: + conversions = quantized + (conversions or []) + if conversions is None: continue diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 9ff820a656e4..56ec9e48c02f 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -81,6 +81,8 @@ def build_glob_alternation( class ConversionOps(ABC): """Base class for weight conversion operations.""" + supports_many_to_many = False + def __repr__(self): if hasattr(self, "dim"): return f"{self.__class__.__name__}(dim={self.dim})" @@ -320,8 +322,9 @@ def convert( sizes = tensor.size(self.dim) targets = self.get_target_patterns(input_dict, source_pattern, target_patterns, sizes) chunks = torch.chunk(tensor, sizes, dim=self.dim) - # We squeeze each chunk here as well to make sure to give them their original shape - all_tensors.update({target: chunk.squeeze() for target, chunk in zip(targets, chunks)}) + # We squeeze each chunk here as well to make sure to give them their original shape; only the split + # dim, so a tensor that holds a size-1 dim of its own (a single column of block scales) keeps it + all_tensors.update({target: chunk.squeeze(self.dim) for target, chunk in zip(targets, chunks)}) return all_tensors def get_target_patterns( @@ -646,6 +649,8 @@ class ErnieFuseAndSplitTextVisionExperts(ConversionOps): The final fusions are defined by the amount of original module lists. """ + supports_many_to_many = True + def __init__(self, stack_dim: int = 0, concat_dim: int = 1): self.stack_dim = stack_dim self.concat_dim = concat_dim @@ -708,6 +713,8 @@ class ErnieSplitAndDecoupleTextVisionExperts(ConversionOps): The final decoupled module lists are defined by the amount of keys. """ + supports_many_to_many = True + def __init__(self, stack_dim: int = 0, concat_dim: int = 1): self.stack_dim = stack_dim self.concat_dim = concat_dim @@ -981,6 +988,8 @@ def reverse_transform(self) -> WeightTransform: def copy(self, source_patterns, target_patterns, **kwargs) -> WeightTransform: """A transform of the same type and scope, with new patterns and e.g. new `operations`.""" + if hasattr(self, "force_cpu"): + kwargs.setdefault("force_cpu", self.force_cpu) copy = self.__class__(source_patterns=source_patterns, target_patterns=target_patterns, **kwargs) copy.scope_prefix, copy.base_model_prefix = self.scope_prefix, self.base_model_prefix return copy @@ -1177,13 +1186,6 @@ def reverse_transform(self) -> WeightTransform: return result -# List of classes that are known to be able to use m:n -_INTERNAL_MANY_TO_MANY_CONVERSIONS = ( - ErnieFuseAndSplitTextVisionExperts, - ErnieSplitAndDecoupleTextVisionExperts, -) - - class WeightConverter(WeightTransform): __slots__ = ("operations", "force_cpu") @@ -1200,7 +1202,7 @@ def __init__( if bool(len(self.source_patterns) - 1) + bool(len(self.target_patterns) - 1) >= 2: # We allow many-to-many only if we use an internal operation that can handle it - if not any(isinstance(op, _INTERNAL_MANY_TO_MANY_CONVERSIONS) for op in self.operations): + if not any(op.supports_many_to_many for op in self.operations): raise ValueError( f"source keys={self.source_patterns}, target_patterns={self.target_patterns} but you can only have one to many, one to one or many to one." ) @@ -1718,8 +1720,10 @@ def convert_and_load_state_dict_in_model( matched_dtype_pattern = dtype_policy_alt.search(renamed_key) if matched_dtype_pattern is not None: _dtype = dtype_plan[dtype_policy_by_group_name[matched_dtype_pattern.lastgroup]] - elif empty_param is not None and empty_param.dtype != _dtype: - _dtype = empty_param.dtype # usually correct when initializing + elif empty_param is not None and empty_param.dtype != _dtype and not needs_quantization: + # usually correct when initializing; only exception can be quants (int8 storage + # would zero it, float8 would double-round it) + _dtype = empty_param.dtype if hf_quantizer and hf_quantizer.pre_quantized: _dtype = hf_quantizer.get_new_dtype(mapping, tensor, _dtype) @@ -1738,7 +1742,28 @@ def convert_and_load_state_dict_in_model( param_device = get_device(device_map, renamed_key, valid_torch_device=True) sharding_op = None if is_dtensor(empty_param): - sharding_op = DtensorShardOperation(empty_param) + source_axes = list(range(empty_param.ndim)) + if ( + isinstance(mapping, WeightConverter) + and mapping.operations + and isinstance(mapping.operations[0], Transpose) + and not mapping.operations[0].check_dims + ): + source_shape = tensor.shape if isinstance(tensor, torch.Tensor) else tensor.get_shape() + transpose_count = 0 + for op in mapping.operations: + if not isinstance(op, Transpose) or op.check_dims: + break + transpose_count += 1 + preserves_axes = all( + getattr(op, "preserves_sharding_axes", lambda *_: False)(model, renamed_key) + for op in mapping.operations[transpose_count:] + ) + if len(source_shape) == empty_param.ndim and preserves_axes: + # Read the source axes corresponding to the target placements, then transpose the local shard. + for op in mapping.operations[:transpose_count]: + source_axes[op.dim0], source_axes[op.dim1] = source_axes[op.dim1], source_axes[op.dim0] + sharding_op = DtensorShardOperation(empty_param, source_axis_order=tuple(source_axes)) # Some parameters are so large (qwen4_exp ple_embedding is about ~95 GiB) that we cannot afford to perform the Operations # directly on the device, as it will completely blow up the memory during the ops memory spike. So defer to "cpu", then @@ -1857,23 +1882,26 @@ def revert_weight_conversion(model: PreTrainedModel, state_dict: dict[str, torch state_dict = sorted(state_dict.items(), key=lambda kv: dot_natural_key(kv[0])) for original_key, tensor in state_dict: # Rename the key according to all renaming pattern and optional weight converter patterns - renamed_key, source_pattern = rename_source_key(original_key, renamings, converters, reverse=True) + converted_key, source_pattern = rename_source_key(original_key, [], converters, reverse=True) if source_pattern is not None: - new_converter = deepcopy(pattern_to_converter[source_pattern]) - # each target key gets its own converter instance - mapping = conversion_mapping.setdefault(renamed_key, new_converter) + # Apply checkpoint renames after conversion, while the converter can still resolve the model's names. + if converted_key not in conversion_mapping: + conversion_mapping[converted_key] = deepcopy(pattern_to_converter[source_pattern]) + mapping = conversion_mapping[converted_key] + mapping.add_tensor(converted_key, original_key, source_pattern, tensor) else: + renamed_key, _ = rename_source_key(original_key, renamings, [], reverse=True) mapping = conversion_mapping.setdefault(renamed_key, WeightRenaming(original_key, renamed_key)) - source_pattern = original_key - - mapping.add_tensor(renamed_key, original_key, source_pattern, tensor) + mapping.add_tensor(renamed_key, original_key, original_key, tensor) new_state_dict = {} - for first_param_name, reversed_converter in conversion_mapping.items(): + for layer_name, reversed_transform in conversion_mapping.items(): # Apply the reverse converter - realized_value = reversed_converter.convert(first_param_name, model=model, config=model.config) + realized_value = reversed_transform.convert(layer_name, model=model, config=model.config) for target_name, param in realized_value.items(): param = param[0] if isinstance(param, list) else param + if isinstance(reversed_transform, WeightConverter): + target_name, _ = rename_source_key(target_name, renamings, [], reverse=True) new_state_dict[target_name] = param return new_state_dict diff --git a/src/transformers/distributed/sharding_utils.py b/src/transformers/distributed/sharding_utils.py index 666fdb82b10c..0faa21b28588 100644 --- a/src/transformers/distributed/sharding_utils.py +++ b/src/transformers/distributed/sharding_utils.py @@ -88,10 +88,11 @@ class DtensorShardOperation: we skip it. Later, `MergeModulelist` will stack the owned expert we kept to create the rank's local shard """ - def __init__(self, param: DTensor): + def __init__(self, param: DTensor, source_axis_order: tuple[int, ...] | None = None): self.device_mesh = param.device_mesh self.placements = tuple(param.placements) self.param_ndim = param.ndim + self.source_axis_order = source_axis_order or tuple(range(param.ndim)) local_shape, offsets = compute_local_shape_and_global_offset(param.shape, self.device_mesh, self.placements) # Axis-0 range owned by this rank (used to filter per-expert pieces) # [_axis0_offset, _axis0_offset + _axis0_local_size) @@ -129,7 +130,7 @@ def shard_tensor( for mesh_dim, placement in dim_placements: sub_mesh = self._get_sub_mesh(mesh_dim) rank, world_size = sub_mesh.get_local_rank(), sub_mesh.size() - dim_idx = self._normalize_param_dim(placement.dim) + dim_idx = self.source_axis_order[self._normalize_param_dim(placement.dim)] planned_ops_by_dim[dim_idx].append((placement, rank, world_size)) # prepare the slices to fetch on disk for each tensor dimension. diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 59c3670d594f..adab74502cd3 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -199,6 +199,14 @@ def should_use_local_tensors(self, module): def validate_param(self, module, param, mesh, parameter_name=None): meta = module._parameters.get(param) + group_size = getattr(module, "_hf_quantized_embedding_group_size", None) + if group_size is not None and ( + module.embedding_dim % mesh.size() or module.embedding_dim // mesh.size() % group_size + ): + raise ValueError( + f"Quantized embedding width {module.embedding_dim} must split into whole groups of {group_size} " + f"for tensor parallel size {mesh.size()}." + ) gathers_output = isinstance(self.output_layouts, Replicate) if meta is None or not gathers_output: return @@ -285,6 +293,17 @@ def should_use_local_tensors(self, module): uses_local_inference_kernel = isinstance(module, torch.nn.Linear) and not torch.is_grad_enabled() return use_local_quantized_path or uses_local_inference_kernel + def validate_param(self, module, param, mesh, parameter_name=None): + if ( + isinstance(module, torch.nn.Embedding) + and getattr(module, "_hf_quantized_needs_local_tp", False) + and mesh.size() > 1 + ): + raise ValueError( + "Rowwise tensor parallelism is not supported for quantized embeddings. " + "Use colwise or colwise_gather_output to shard the embedding width." + ) + def shard_param(self, module, param, mesh): meta = module._parameters.get(param) if meta is None: @@ -956,6 +975,9 @@ class ParallelInterface(GeneralInterface): "packed_rowwise": PackedRowwiseParallel(), "sequence_parallel": SequenceParallel(use_local_output=True), "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True), + "moe_experts_colwise": MoEParamShard(Shard(1)), + "moe_experts_packed_colwise": MoEParamShard(_StridedShard(dim=1, split_factor=2)), + "moe_experts_rowwise": MoEParamShard(Shard(2)), "ep_router": EpRouterParallel(), "ep_dispatch_experts": EpDispatchExpertsParallel(), "megamoe_router": RouterParallelMegaMoe(), diff --git a/src/transformers/integrations/__init__.py b/src/transformers/integrations/__init__.py index 1cb6e0302172..77dcefbc8e42 100755 --- a/src/transformers/integrations/__init__.py +++ b/src/transformers/integrations/__init__.py @@ -50,6 +50,11 @@ ], "eetq": ["replace_with_eetq_linear"], "fbgemm_fp8": ["FbgemmFp8Linear", "FbgemmFp8Llama4TextExperts", "replace_with_fbgemm_fp8_linear"], + "finegrained": [ + "FineGrainedEmbedding", + "FineGrainedExperts", + "FineGrainedLinear", + ], "finegrained_fp8": ["FP8Embedding", "FP8Linear", "replace_with_fp8_embedding", "replace_with_fp8_linear"], "fsdp": ["is_fsdp_enabled", "is_fsdp_managed_module"], "gemma_quant": [ @@ -207,6 +212,11 @@ ) from .eetq import replace_with_eetq_linear from .fbgemm_fp8 import FbgemmFp8Linear, FbgemmFp8Llama4TextExperts, replace_with_fbgemm_fp8_linear + from .finegrained import ( + FineGrainedEmbedding, + FineGrainedExperts, + FineGrainedLinear, + ) from .finegrained_fp8 import FP8Embedding, FP8Linear, replace_with_fp8_embedding, replace_with_fp8_linear from .fsdp import is_fsdp_enabled, is_fsdp_managed_module from .gemma_quant import ( diff --git a/src/transformers/integrations/deepgemm.py b/src/transformers/integrations/deepgemm.py index 76810552213f..1396a9206017 100644 --- a/src/transformers/integrations/deepgemm.py +++ b/src/transformers/integrations/deepgemm.py @@ -36,8 +36,7 @@ import torch -from ..utils import logging -from ..utils.deprecation import deprecate_kwarg +from ..quantizers.quantizers_utils import FUSED_STACKS from ..utils.import_utils import ( is_kernels_available, maybe_import_error, @@ -46,8 +45,6 @@ from .hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel -logger = logging.get_logger(__name__) - # ── Kernel loading ───────────────────────────────────────────────────────────── @@ -65,12 +62,9 @@ class DeepGEMM: transform_weights_for_mega_moe: Callable get_symm_buffer_for_mega_moe: Callable fp8_fp4_mega_moe: Callable - # M/K-dimension alignment for TMA-based contiguous grouped GEMM. Sourced from - # `get_mk_alignment_for_contiguous_layout()` at load time. The kernel exposes a - # `set_mk_alignment_for_contiguous_layout` setter, but we don't call it: the - # build-time default (128) was empirically the best across MoE workloads - # (bench showed kernel-recommended 240 is slower and 256 doesn't even compile). - # Same stance as vLLM, which caches and never sets it. + # M/K alignment for the TMA contiguous grouped GEMM, read at load time. The kernel's setter + # stays unused: the build-time default (128) measured best across MoE workloads (its own + # recommendation, 240, is slower; 256 does not compile). vLLM takes the same stance. m_alignment: int @@ -146,9 +140,10 @@ def is_sm100() -> bool: evaluate it once at trace time and inline the bool — keeping the untraceable `torch.cuda.get_device_capability()` pybind out of the graph — so the branches on it (SF layout, cast kwargs, psum layout, and the arch guards) stay compile-safe. Re-queries each eager call, so a faked - capability in tests is honoured. + capability in tests is honoured — which means faking BOTH, since `get_device_capability()` asserts + rather than answering on a CPU-only build and the availability check is what keeps it unreached. """ - return torch.cuda.get_device_capability()[0] >= 10 + return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 10 @torch._dynamo.assume_constant_result @@ -161,7 +156,7 @@ def is_deepgemm_loadable(raise_error: bool = False) -> bool: `@assume_constant_result` makes dynamo evaluate this once at trace time and inline the bool (its probe is untraceable — an `lru_cache`'d import check + the `get_device_capability` pybind + nvcc lookup), so - the check can sit in a compiled dispatch gate (e.g. `fp8_linear`) without breaking the graph. + the check can sit in a compiled dispatch gate (e.g. `prefers_deepgemm_linear`) without breaking the graph. FP4 / Mega MoE additionally need Blackwell (SM100); the relevant forwards enforce that via `_assert_sm100_requirements` / `is_sm100()`, since this env check can't see which dtype path the @@ -221,9 +216,8 @@ def is_deepgemm_loadable(raise_error: bool = False) -> bool: return True -# Cache the loaded bundle but not failures: re-checking each call is cheap and intended, since the env -# can change between attempts. A module global (not `@functools.cache`) avoids Dynamo warning about -# tracing a cache-wrapped function on every compile. +# Successes are cached, failures are not: the env can change between attempts. A module global +# rather than `@functools.cache`, which Dynamo warns about tracing on every compile. _DEEPGEMM: DeepGEMM | None = None @@ -307,42 +301,6 @@ def load_deepgemm_kernel() -> DeepGEMM: # ── Scale-factor helpers ─────────────────────────────────────────────────────── -def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> None: - """Before-load guard for DeepGEMM's FP8/FP4 arch constraints on the given weight/scale dtypes: - - - FP4 (``int8``-packed) weights have no Hopper (SM90) kernel — they need Blackwell (SM100+). - - Blackwell has no float32 scale-factor path: ``_coerce_sf_for_kernel`` would silently round a - ``float32`` scale to UE8M0 and corrupt the output. UE8M0 scales load as ``float8_e8m0fnu`` (the - loader normalizes even float32-container checkpoints like dsv4-flash-base), so a ``float32`` scale - on SM100 means a genuine non-UE8M0 checkpoint. - - Uses `is_sm100()` (compile-safe via `assume_constant_result`), so the whole guard folds to a constant under - ``torch.compile``: the valid case compiles away to nothing, while an unsupported combo fails loud - rather than letting the hot path silently corrupt (unlike an ``is_compiling`` skip, which would miss a - model compiled from cold with no eager warmup). Both raise ``NotImplementedError``, which - ``fp8_linear`` treats as "DeepGEMM declined" and falls back to Triton (SM90 consuming float32 SFs - directly is fine, so those cases are no-ops). - """ - if not is_sm100(): - # SM90: DeepGEMM has no FP4 (int8-packed) kernel, but consumes float32 SFs directly. - if weight.dtype == torch.int8: - raise NotImplementedError( - "DeepGEMM's FP4 (int8-packed) path requires a Blackwell (SM100+) GPU; FP4 weights have no " - "Hopper (SM90) kernel. Use an FP8 checkpoint, or run on a Blackwell GPU." - ) - return - - # SM100: DeepGEMM has no float32 scale-factor path. - if scale.dtype == torch.float32: - raise NotImplementedError( - "DeepGEMM has no float32 scale-factor path on Blackwell (SM100): these scales are plain float32 " - "(quantization_config.scale_fmt='float'), and rounding them to UE8M0 would silently corrupt the " - "output. Use a checkpoint quantized with scale_fmt='ue8m0', or a path that consumes float32 block " - "scales directly — the FP8 linear falls back to Triton automatically; for experts use " - "`model.set_experts_implementation('grouped_mm')`." - ) - - def _ceil_to_ue8m0(sf: torch.Tensor) -> torch.Tensor: """Round each fp32 SF up to the nearest power of 2 (zero mantissa). @@ -384,14 +342,14 @@ def _coerce_sf_for_kernel(sf: torch.Tensor, is_sm100: bool, expected_mn: int | N DeepGEMM kernel branch is the only UE8M0 path on SM100; for `gran_mn > 1` the kernel only handles FP32 SFs and would otherwise reject our INT SF here. """ - if sf.dtype == torch.float8_e8m0fnu: + if sf.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)): if expected_mn is not None and sf.size(-2) < expected_mn: gran_mn = expected_mn // sf.size(-2) sf = sf.repeat_interleave(gran_mn, dim=-2) if is_sm100: sf = sf.contiguous().view(torch.int32) else: - sf = sf.float() + sf = (sf.float() - 127).exp2() if sf.dtype == torch.uint8 else sf.float() elif sf.dtype == torch.float32 and is_sm100: sf = _ceil_to_ue8m0(sf) @@ -439,7 +397,7 @@ def _select_fp8_cast_kwargs( block_size = tuple(block_size) if block_size not in ((128, 128), (1, 128)): raise ValueError(f"DeepGEMM requires `block_size` ∈ {{(128, 128), (1, 128)}}, got {block_size}.") - if weight_scale_inv.dtype == torch.float8_e8m0fnu and is_sm100: + if weight_scale_inv.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) and is_sm100: return {"use_ue8m0": True, "gran_k": 128, "use_packed_ue8m0": True} return {"use_ue8m0": False, "gran_k": 128} @@ -517,23 +475,19 @@ def _dispatch_routed_input( expert_ids = top_k_index.reshape(-1) # (S,) sample_weights = top_k_weights.reshape(-1) # (S,) - # Sort by expert for grouped processing expert_ids_g, perm = torch.sort(expert_ids) sorted_hidden_states_g = hidden_states[perm // num_top_k] sample_weights_g = sample_weights[perm] - # Build the M-grouped padded layout (DeepGEMM contract: each expert's rows - # start on the kernel's M-alignment boundary, sentinels routed past valid - # expert blocks). + # DeepGEMM's contract: each expert's rows start on the kernel's M-alignment boundary, and + # sentinels route past every valid expert block. sorted_to_padded, grouped_layout, total_padded_rows = _build_deepgemm_contiguous_layout( expert_ids_g, num_experts, m_alignment, use_psum_layout ) - # EP sentinel mask is captured before the in-place clamp; used by the post-mask in - # `_combine_routed_output` to zero sentinel rows before the per-token reduction. The clamp - # keeps any per-row gather (e.g. bias) in-bounds — bias added at sentinel positions falls - # in rows the kernel skips, so harmless. Safe to mutate now: the layout was built from the - # unclamped tensor and nothing downstream needs the sentinel info from `expert_ids_g` itself. + # Only EP routes sentinels. Captured before the in-place clamp, which keeps a per-row gather + # (bias) in bounds — those rows the kernel skips anyway. `_combine_routed_output` zeroes them + # before the per-token reduction; the layout above was built from the unclamped ids. sentinel_mask = None if is_expert_parallel: sentinel_mask = (expert_ids_g >= num_experts).unsqueeze(-1) @@ -577,29 +531,242 @@ def _combine_routed_output( # ── Public dispatches ────────────────────────────────────────────────────────── -@deprecate_kwarg("output_dtype", version="v5.16") -def deepgemm_fp8_fp4_linear( +def prefers_deepgemm_linear( input: torch.Tensor, weight: torch.Tensor, weight_scale_inv: torch.Tensor, - bias: torch.Tensor | None = None, - block_size: tuple[int, int] | None = None, - output_dtype: torch.dtype | None = None, - activation_scale: torch.Tensor | None = None, -) -> torch.Tensor: - """End-to-end DeepGEMM linear: per-token activation quant + FP8/FP4 matmul. + *, + block_size: list[int] | None, + activation_scale: torch.Tensor | None, + weight_global_scale: torch.Tensor | None, + input_global_scale: torch.Tensor | None, + activation_format: str | None, + allow_deepgemm: bool, +) -> bool: + """Whether a dense fine-grained linear should go to DeepGEMM rather than Triton. + + DeepGEMM is CUDA-only, dynamic-only, SM90+, FP4 or 128x128-block FP8; the combos it would + silently corrupt are refused here rather than attempted. + ``TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1`` forces Triton for this dispatcher only.""" + return ( + # SM100 perf: Triton measures 2.9x (decode) / 1.45x (prefill) over DeepGEMM on the DSV4 + # qkv linear — the 128x128 block-FP8 shape DeepGEMM is otherwise preferred for. + not is_sm100() + # False when the model spans devices: DeepGEMM's kernels are bound to one CUDA context and + # corrupt across them, and the Triton fallback is context-free. + and allow_deepgemm + # DeepGEMM has no backward pass; a call that needs a gradient takes Triton, which has one + and not (torch.is_grad_enabled() and input.requires_grad) + # A pre-swizzled (SWIZZLE_32_4_4) scale is not readable as row-major — DeepGEMM would + # consume the permuted buffer as affine and return garbage. Correctness gate, any arch. + and weight_scale_inv.ndim <= 2 + and activation_scale is None + # DeepGEMM serves neither the NVFP4 two-level globals nor an explicit activation format + and weight_global_scale is None + and input_global_scale is None + and activation_format is None + and (weight.dtype == torch.int8 or (block_size is not None and block_size[0] == block_size[1] == 128)) + and os.environ.get("TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR", "0") != "1" + # last: an environment probe (arch, `CUDA_HOME/bin/nvcc`, nvcc version) rather than a + # tensor check, so the cheap gates above short-circuit before it + and is_deepgemm_loadable() + ) + - Static (per-tensor) activation quantization is rejected — DeepGEMM needs - per-row SFs. Callers should route static activations through the Triton fallback. +def _assert_sm100_requirements(weight: torch.Tensor, scale: torch.Tensor) -> None: + """Before-load guard for DeepGEMM's FP8/FP4 arch constraints on the given weight/scale dtypes: + + - FP4 (``int8``-packed) weights have no Hopper (SM90) kernel — they need Blackwell (SM100+). + - Blackwell has no float32 scale-factor path: ``_coerce_sf_for_kernel`` would silently round a + ``float32`` scale to UE8M0 and corrupt the output. UE8M0 scales load as ``float8_e8m0fnu`` (the + loader normalizes even float32-container checkpoints like dsv4-flash-base), so a ``float32`` scale + on SM100 means a genuine non-UE8M0 checkpoint. + + Uses `is_sm100()` (compile-safe via `assume_constant_result`), so the whole guard folds to a constant under + ``torch.compile``: the valid case compiles away to nothing, while an unsupported combo fails loud + rather than letting the hot path silently corrupt (unlike an ``is_compiling`` skip, which would miss a + model compiled from cold with no eager warmup). Both raise ``NotImplementedError``, which + ``finegrained_linear`` treats as "DeepGEMM declined" and falls back to Triton (SM90 consuming float32 SFs + directly is fine, so those cases are no-ops). """ + if not is_sm100(): + # SM90: DeepGEMM has no FP4 (int8-packed) kernel, but consumes float32 SFs directly. + if weight.dtype == torch.int8: + raise NotImplementedError( + "DeepGEMM's FP4 (int8-packed) path requires a Blackwell (SM100+) GPU; FP4 weights have no " + "Hopper (SM90) kernel. Use an FP8 checkpoint, or run on a Blackwell GPU." + ) + return + + # SM100: DeepGEMM has no float32 scale-factor path. + if scale.dtype == torch.float32: + raise NotImplementedError( + "DeepGEMM has no float32 scale-factor path on Blackwell (SM100): these scales are plain float32 " + "(quantization_config.scale_fmt='float'), and rounding them to UE8M0 would silently corrupt the " + "output. Use a checkpoint quantized with scale_fmt='ue8m0', or a path that consumes float32 block " + "scales directly — the finegrained linear falls back to Triton automatically; for experts load with " + "`experts_implementation='grouped_mm'`." + ) + + +def _assert_no_gradient(hidden_states: torch.Tensor) -> None: + """DeepGEMM has no backward pass, so a call that needs a gradient fails rather than drop it.""" + if torch.is_grad_enabled() and hidden_states.requires_grad: + raise NotImplementedError( + "DeepGEMM does not support training; train through the Triton kernels " + "(for experts, `experts_implementation='grouped_mm'` or 'batched_mm')." + ) + + +def _assert_dynamic_activation_scale(activation_scale: torch.Tensor | None) -> None: + """DeepGEMM quantizes activations per row, so a calibrated per-tensor scale has nowhere to go.""" if activation_scale is not None: raise NotImplementedError("DeepGEMM linear does not support static activation quantization.") + + +def _assert_half_precision_input(input: torch.Tensor) -> None: if input.dtype not in (torch.bfloat16, torch.float16): raise ValueError(f"DeepGEMM linear requires FP16 or BF16 activations, got {input.dtype}") - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. - _assert_sm100_requirements(weight, weight_scale_inv) +def _assert_bf16_hidden_states(hidden_states: torch.Tensor) -> None: + """Every DeepGEMM experts arm builds its intermediates and its output in bfloat16, so anything + else silently changes the dtype the caller gets back.""" + if hidden_states.dtype != torch.bfloat16: + raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") + + +def _assert_dynamic_activations(self: torch.nn.Module) -> None: + """DeepGEMM quantizes activations inside its kernels, so a checkpoint's calibrated per-tensor + scale has nowhere to be applied.""" + if getattr(self, "activation_scheme", "dynamic") == "static": + raise NotImplementedError( + "DeepGEMM experts dispatch does not support activation_scheme='static'; use " + "`experts_implementation='grouped_mm'` (or 'batched_mm')." + ) + + +def _assert_affine_scales(module: torch.nn.Module) -> None: + """DeepGEMM reads row-major block scales. A module loaded for a triton backend holds them in the + SWIZZLE_32_4_4 layout (5-D); consuming that as affine returns garbage, so a backend switched to + DeepGEMM after such a load fails here instead.""" + for name in ("gate_up_proj_scale_inv", "up_proj_scale_inv", "down_proj_scale_inv"): + scale = getattr(module, name, None) + if scale is not None and scale.ndim == 5: + raise NotImplementedError( + f"DeepGEMM experts need row-major block scales, but `{name}` is held swizzled for the " + "triton experts backends. Switching to a DeepGEMM backend after load is not supported — " + "pass `experts_implementation=...` to `from_pretrained` so the model loads for it." + ) + + +def _assert_stacked_gate_up(module: torch.nn.Module) -> None: + """The quantized DeepGEMM arms take the stacked ``[gate; up]`` rows: Mega MoE's + ``transform_weights_for_mega_moe`` does its own gate/up interleave, and the grouped GEMM reads a + block-scale grid row-major, whose blocks an interleaved stack splits between the halves. A module + loaded for a triton backend holds them interleaved; switched to this backend after load, silently + re-permuting would be a wrong answer rather than a failure.""" + if getattr(module, "has_gate", True) and getattr(module, "holds_interleaved_gate_up", False): + raise NotImplementedError( + "DeepGEMM experts need gate|up stacked, but this module was loaded interleaved for the triton " + "experts backends. Switching to a DeepGEMM backend after load is not supported — pass " + "`experts_implementation=...` to `from_pretrained` so the model loads for it." + ) + + +def _assert_block_scales_only(module: torch.nn.Module) -> None: + """The quantized DeepGEMM arms read the weights and their UE8M0/float32 block scales, and quantize the + activations to FP8 themselves. An expert bias, NVFP4's second-level globals and E4M3 block scales, or + weight-only (bf16) activations have nowhere to go, so they are refused rather than dropped.""" + scales = (getattr(module, f"{proj}_scale_inv", None) for proj in FUSED_STACKS) + if getattr(module, "has_bias", False): + raise NotImplementedError("DeepGEMM quantized experts do not apply expert biases; use 'grouped_mm'.") + if getattr(module, "global_scale_dtype", None) is not None or any( + scale is not None and scale.dtype == torch.float8_e4m3fn for scale in scales + ): + raise NotImplementedError("DeepGEMM serves no NVFP4 experts (E4M3 block scales, globals); use 'grouped_mm'.") + if getattr(module, "activation_format", None) == "bf16": + raise NotImplementedError( + "DeepGEMM quantizes the activations, so it cannot serve weight-only experts " + "(activation_format='bf16'); use 'grouped_mm'." + ) + + +def _assert_no_post_expert_norm(self: torch.nn.Module) -> None: + """Mega MoE fuses the routing-weighted reduce into its kernel, so a model's per-expert output + norm has nowhere to go — it would be dropped silently, so refuse instead. The other DeepGEMM + arms reduce in `_combine_routed_output` and apply the norm on the rows just before it.""" + if self.has_post_expert_norm: + raise NotImplementedError( + "DeepGEMM Mega MoE cannot apply this model's per-expert output norm; use " + "`experts_implementation='deepgemm'`, 'grouped_mm' or 'batched_mm'." + ) + + +def deepgemm_linear_guards(forward=None): + """State a linear's requirements on its operands, checked before any kernel work; each assert + documents its own condition. Usable bare (`@deepgemm_linear_guards`), like `deepgemm_experts_guards`.""" + + def decorate(forward): + @functools.wraps(forward) + def guarded(input, weight, weight_scale_inv, *args, activation_scale=None, **kwargs): + _assert_no_gradient(input) + _assert_half_precision_input(input) + _assert_dynamic_activation_scale(activation_scale) + _assert_sm100_requirements(weight, weight_scale_inv) + return forward(input, weight, weight_scale_inv, *args, **kwargs) + + return guarded + + return decorate(forward) if forward is not None else decorate + + +def deepgemm_experts_guards( + forward=None, + *, + affine_scales: bool = False, + sm100: bool = False, + supports_post_expert_norm: bool = True, + stacked_gate_up: bool = False, + block_scales_only: bool = False, +): + """State an experts forward's requirements on the module, checked before any kernel work; each + assert documents its own condition, and each flag enables one (``supports_post_expert_norm=False`` + the post-expert-norm refusal). Usable bare (`@deepgemm_experts_guards`) for an arm that only needs + the common checks.""" + + def decorate(forward): + @functools.wraps(forward) + def guarded(self, hidden_states, *args, **kwargs): + _assert_dynamic_activations(self) + _assert_no_gradient(hidden_states) + _assert_bf16_hidden_states(hidden_states) + if not supports_post_expert_norm: + _assert_no_post_expert_norm(self) + if affine_scales: + _assert_affine_scales(self) + if stacked_gate_up: + _assert_stacked_gate_up(self) + if block_scales_only: + _assert_block_scales_only(self) + if sm100: + _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) + return forward(self, hidden_states, *args, **kwargs) + + return guarded + + return decorate(forward) if forward is not None else decorate + + +@deepgemm_linear_guards +def deepgemm_fp8_fp4_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + bias: torch.Tensor | None = None, + block_size: tuple[int, int] | None = None, +) -> torch.Tensor: + """End-to-end DeepGEMM linear: per-token activation quant + FP8/FP4 matmul.""" deepgemm = load_deepgemm_kernel() cast_kwargs = _select_fp8_cast_kwargs(weight, weight_scale_inv, block_size, is_sm100()) @@ -622,6 +789,7 @@ def deepgemm_fp8_fp4_linear( return output +@deepgemm_experts_guards def deepgemm_bf16_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -666,7 +834,6 @@ def deepgemm_bf16_experts_forward( up_bias = (self.gate_up_proj_bias if self.has_gate else self.up_proj_bias) if self.has_bias else None down_bias = self.down_proj_bias if self.has_bias else None - # Up projection. up_out_dim = weight_up.shape[-1] if self.is_transposed else weight_up.shape[1] act = _pad_for_deepgemm(sorted_hidden, sorted_to_padded, total_padded_rows) proj_out = torch.empty(total_padded_rows, up_out_dim, device=device, dtype=hidden_states.dtype) @@ -676,12 +843,14 @@ def deepgemm_bf16_experts_forward( proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - # Down projection. out = torch.empty(total_padded_rows, hidden_dim, device=device, dtype=hidden_states.dtype) grouped_bf16_matmul(proj_out, weight_down, out, grouped_layout, use_psum_layout=is_sm100()) if self.has_bias: out.index_add_(0, sorted_to_padded, down_bias[expert_ids_g]) + if self.has_post_expert_norm: + out = self.post_expert_norm(out) + return _combine_routed_output( out, sorted_weights, @@ -695,6 +864,7 @@ def deepgemm_bf16_experts_forward( ) +@deepgemm_experts_guards(affine_scales=True, sm100=True, stacked_gate_up=True, block_scales_only=True) def deepgemm_fp8_fp4_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -703,22 +873,15 @@ def deepgemm_fp8_fp4_experts_forward( ) -> torch.Tensor: if self._deepgemm_disabled: # Set at load when the model spans >1 CUDA device in this process, where DeepGEMM's - # context-bound kernels corrupt across devices (see `quantizer_finegrained_fp8.py`). + # context-bound kernels corrupt across devices (see `disable_deepgemm_on_multi_device`). raise RuntimeError( "DeepGEMM experts selected on a model spanning multiple CUDA devices in one process; " "its kernels are bound to a single CUDA context and corrupt across devices. Use " "`experts_implementation='grouped_mm'`, or run one device per process (TP/EP)." ) - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. - _assert_sm100_requirements(self.down_proj, self.down_proj_scale_inv) - deepgemm = load_deepgemm_kernel() - if self.activation_scheme == "static": - raise NotImplementedError("DeepGEMM experts dispatch does not support static activation quantization.") - if hidden_states.dtype != torch.bfloat16: - raise ValueError(f"DeepGEMM experts path requires bfloat16 hidden states, got {hidden_states.dtype}") grouped_fp8_fp4_matmul = ( deepgemm.grouped_fp8_fp4_matmul_nn if self.is_transposed else deepgemm.grouped_fp8_fp4_matmul_nt ) @@ -754,7 +917,6 @@ def deepgemm_fp8_fp4_experts_forward( ) sf_recipe = (1, 1, cast_kwargs["gran_k"]) if cast_kwargs.get("use_packed_ue8m0") else None - # Up projection. act_fp8, act_scales = deepgemm.per_token_cast_to_fp8(sorted_hidden, **cast_kwargs) act_fp8 = _pad_for_deepgemm(act_fp8, sorted_to_padded, total_padded_rows) act_scales = _pad_for_deepgemm(act_scales, sorted_to_padded, total_padded_rows) @@ -769,7 +931,6 @@ def deepgemm_fp8_fp4_experts_forward( ) proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) - # Down projection. proj_fp8, proj_scales = deepgemm.per_token_cast_to_fp8(proj_out, **cast_kwargs) out = torch.empty(total_padded_rows, hidden_dim, device=device, dtype=torch.bfloat16) grouped_fp8_fp4_matmul( @@ -781,6 +942,9 @@ def deepgemm_fp8_fp4_experts_forward( use_psum_layout=is_sm100(), ) + if self.has_post_expert_norm: + out = self.post_expert_norm(out) + return _combine_routed_output( out, sorted_weights, @@ -795,7 +959,7 @@ def deepgemm_fp8_fp4_experts_forward( def setup_megamoe_weights(module: torch.nn.Module) -> None: - """One-shot pack + permute of an FP8Experts module's L1/L2 weights into the + """One-shot pack + permute of a `FineGrainedExperts` module's L1/L2 weights into the Mega MoE UTCCP layout. Called lazily on the first megamoe forward; idempotent via the caller's ``_megamoe_transformed`` flag. @@ -808,26 +972,25 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: the ``[E_local, 2*I, *]`` leading dims so downstream ``.size(...)`` reads stay valid. - Unwraps any ``DTensor`` wrappers FSDP2/EP may have placed around the loader- - side Parameters — the kernel takes raw pointers. + Runs inside the first forward, where the parameters are already local. """ - deepgemm = load_deepgemm_kernel() - gate_up_sf_raw = module.gate_up_proj_scale_inv.data - down_sf_raw = module.down_proj_scale_inv.data - # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. - gate_up_w = module.gate_up_proj.data.view(torch.int8).contiguous() - down_w = module.down_proj.data.view(torch.int8).contiguous() - + # The UE8M0 scale grid is group-32 along both dims, so the kernel's SF layout needs both divisible. intermediate_hidden = module.intermediate_dim num_local_experts = module.num_experts hidden_dim = module.hidden_dim - if hidden_dim % 32 != 0 or intermediate_hidden % 32 != 0: raise ValueError( f"DeepGEMM Mega MoE requires `hidden_dim` and `intermediate_hidden` divisible by 32 " f"(FP8 SF granularity); got hidden_dim={hidden_dim}, intermediate_hidden={intermediate_hidden}." ) + deepgemm = load_deepgemm_kernel() + gate_up_sf_raw = module.gate_up_proj_scale_inv.data + down_sf_raw = module.down_proj_scale_inv.data + # Force int8 view: the kernel's interleave reshape/empty_like/copy_ is bit-level. + gate_up_w = module.gate_up_proj.data.view(torch.int8).contiguous() + down_w = module.down_proj.data.view(torch.int8).contiguous() + gate_up_sf = deepgemm.transform_sf_into_required_layout( gate_up_sf_raw.float(), 2 * intermediate_hidden, @@ -852,6 +1015,9 @@ def setup_megamoe_weights(module: torch.nn.Module) -> None: module.down_proj_scale_inv = torch.nn.Parameter(down_sf, requires_grad=False) +@deepgemm_experts_guards( + affine_scales=True, sm100=True, supports_post_expert_norm=False, stacked_gate_up=True, block_scales_only=True +) def deepgemm_fp8_fp4_megamoe_experts_forward( self: torch.nn.Module, hidden_states: torch.Tensor, @@ -877,17 +1043,12 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( `transform_weights_for_mega_moe((gate_up, gate_up_sf), (down, down_sf))`. - `config.swiglu_limit` (optional): SwiGLU clamp; absent → unclamped. """ - # Fail before the (hub-download + JIT) load if this device can't serve these dtypes. Mega MoE is - # Blackwell-only, and its weights are always FP4 (int8) — so the FP4 arch check doubles as the - # SM100 gate; the explicit `!= int8` check below covers a non-FP4 (misconfigured) checkpoint. - _assert_sm100_requirements(self.gate_up_proj, self.down_proj_scale_inv) if self.gate_up_proj.dtype != torch.int8: raise NotImplementedError( f"DeepGEMM Mega MoE requires FP4-packed expert weights (dtype=`int8`), got " f"`{self.gate_up_proj.dtype}`. Use the 'deepgemm' dispatch for FP8 experts." ) - if process_group is None: raise ValueError( "DeepGEMM Mega MoE requires a `process_group` for the EP group. The TP wrapping " @@ -896,10 +1057,8 @@ def deepgemm_fp8_fp4_megamoe_experts_forward( deepgemm = load_deepgemm_kernel() - # First-forward one-shot: pack UE8M0 SFs and interleave the L1/L2 weights for UTCCP. - # Kept lazy here (instead of in a quantizer load-time hook) so the megamoe-specific - # setup lives alongside the megamoe forward — `set_experts_implementation` refuses - # to flip in/out of `deepgemm_megamoe` at runtime, so the flag won't go stale. + # One-shot on the first forward, kept here rather than in a load-time hook so the setup sits + # with the forward it serves; the flag cannot go stale because the backend cannot be flipped. if not getattr(self, "_megamoe_transformed", False): setup_megamoe_weights(self) self._megamoe_transformed = True diff --git a/src/transformers/integrations/finegrained/__init__.py b/src/transformers/integrations/finegrained/__init__.py new file mode 100644 index 000000000000..08160a809afd --- /dev/null +++ b/src/transformers/integrations/finegrained/__init__.py @@ -0,0 +1,45 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""The fine-grained quantization integration: `layers` holds the modules and their forwards, +`conversions` the loader ops that fill them. `conversions` is imported by path, so that reading +a checkpoint's layout never drags the kernels in.""" + +from .layers import ( + ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + SUPPORTED_FUSED_NORMS, + FineGrainedEmbedding, + FineGrainedExperts, + FineGrainedGroupedLinear, + FineGrainedLinear, + assert_modules_are_quantized, + disable_deepgemm_on_multi_device, + finegrained_linear, + load_finegrained_kernel, + register_shared_gate_up_input_globals, +) + + +__all__ = [ + "ALL_FINEGRAINED_EXPERTS_FUNCTIONS", + "SUPPORTED_FUSED_NORMS", + "FineGrainedEmbedding", + "FineGrainedExperts", + "FineGrainedGroupedLinear", + "FineGrainedLinear", + "assert_modules_are_quantized", + "disable_deepgemm_on_multi_device", + "register_shared_gate_up_input_globals", + "finegrained_linear", + "load_finegrained_kernel", +] diff --git a/src/transformers/integrations/finegrained/conversions.py b/src/transformers/integrations/finegrained/conversions.py new file mode 100644 index 000000000000..5b8675eee563 --- /dev/null +++ b/src/transformers/integrations/finegrained/conversions.py @@ -0,0 +1,826 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Loader/saver conversion ops for the block-scaled formats. + +Each op moves ONE checkpoint-side layout to the one the `finegrained` modules hold, and carries a +`reverse_op` so `save_pretrained` writes the checkpoint's layout back. Nothing here runs in a +forward. +""" + +from __future__ import annotations + +from collections import defaultdict + +import torch + +from ...core_model_loading import ConversionOps, Interleave +from ...quantizers.quantizers_utils import FUSED_STACKS, get_module_from_name, plan_entry_for +from .layers import ( + _FP8_DTYPE, + _FP8_MAX, + _FP8_MIN, + FineGrainedExperts, + WeightFormat, + _cdiv, + _FineGrainedModule, + _get_ue8m0_dtype, + get_weight_format, + load_finegrained_kernel, +) + + +def single_tensor(value): + """A converter's tensor, which the loader hands in a list: one source's is a one-element list. Several are + sources this op does not merge, so they are refused rather than cut down to the first.""" + if not isinstance(value, list): + return value + if len(value) != 1: + raise ValueError(f"expected a single tensor, got a list of {len(value)}") + return value[0] + + +def keyed_by_target(input_dict: dict, target_patterns) -> dict: + """A converter's ops see their single tensor under the SOURCE pattern until an op re-keys it to the + target (the core ops do; the loader then expands the target into the full name) — so a + layout op that is first in its converter re-keys the same way. Multi-tensor dicts (a + deserializer's fully named outputs) pass through.""" + if target_patterns and len(input_dict) == 1: + value = next(iter(input_dict.values())) + return {target_patterns[0]: single_tensor(value)} + return input_dict + + +def held_scale(model, full_layer_name, key): + """The finegrained module and its ``*_scale_inv`` Parameter a converter output fills, else + ``(None, None)``. A converter emitting several tensors names them fully (key); a single target's + name arrives as the pattern, its full name (with the ``_scale_inv`` suffix) in ``full_layer_name``.""" + param_name = (key if key.endswith("_scale_inv") else full_layer_name).rpartition(".")[-1] + if not param_name.endswith("_scale_inv") or model is None: + return None, None + module = model.get_submodule(full_layer_name.rpartition(".")[0]) + if not isinstance(module, _FineGrainedModule): + return None, None + return module, getattr(module, param_name, None) + + +def finegrained_module(model, full_layer_name): + """The name and finegrained module owning a tensor, or ``(None, None)``. + + On save the name can run past the module (``experts.*.gate_proj.weight``). + """ + name = full_layer_name or "" + while model is not None and name: + name = name.rpartition(".")[0] + try: + module = model.get_submodule(name) + except AttributeError: + continue + return (name, module) if isinstance(module, _FineGrainedModule) else (None, None) + return None, None + + +def finegrained_experts(model, full_layer_name): + """The ``FineGrainedExperts`` the tensor at ``full_layer_name`` fills, else ``None``.""" + _, module = finegrained_module(model, full_layer_name) + return module if isinstance(module, FineGrainedExperts) else None + + +def as_container(scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + """Reinterpret UE8M0 bytes or convert its power-of-two values between containers.""" + if scale.element_size() == dtype.itemsize: + return scale.view(dtype) + if dtype == torch.uint8: + # UE8M0 is a biased exponent, never a numeric float-to-integer cast. + return (scale.float().log2().ceil() + 127).clamp(0, 255).masked_fill(torch.isnan(scale), 255).to(torch.uint8) + return scale.to(dtype) + + +class _FineGrainedOp(ConversionOps): + """Base for the ops below. Each carries the quantizer that built it, which is how an op reaches + the config (the activation format, the block size) while it runs. + + The default reverse is the op ITSELF with `inverse` flipped: it applies the layout the modules + hold on load and puts the checkpoint's back on save. The pairs that reverse into a *different* + class — merge/split, quantize/dequantize — override `reverse_op` and ignore the flag. + """ + + def __init__(self, hf_quantizer=None, inverse: bool = False): + self.hf_quantizer = hf_quantizer + self.inverse = inverse + + @property + def reverse_op(self) -> ConversionOps: + return type(self)(self.hf_quantizer, inverse=not self.inverse) + + +class FineGrainedDeserialize(_FineGrainedOp): + """Apply a module's reversible layouts after its checkpoint tensors have been merged. + + On-the-fly quantization applies weight and scale layouts to its outputs instead, after + quantizing the full-precision rows. Biases only need their row layout. Saving reverses the layouts before splitting or renaming the tensors. + """ + + def preserves_sharding_axes(self, model, full_layer_name): + # Row-FP8 only reorders gate/up rows within the existing packed-shard intervals. + # Other formats can reshape packed weights or swizzle scale grids across axes. + _, module = finegrained_module(model, full_layer_name) + return module is not None and module.quantization_config.weight_block_size == "row" + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + if not self.inverse: + input_dict = keyed_by_target(input_dict, target_patterns) + if not self.hf_quantizer.pre_quantized and not (full_layer_name or "").endswith("_bias"): + return input_dict + module_name, module = finegrained_module(model, full_layer_name) + if module is None: + return input_dict + result = {} + for key, value in input_dict.items(): + name = key.rpartition(".")[2] + member = self.hf_quantizer.member_of(module) + ops = member.get_param_ops(model, f"{module_name}.{name}") if member is not None else [] + if self.inverse: + ops = [op.reverse_op for op in ops[::-1]] + tensors = {key: value} + for op in ops: + tensors = op.convert( + tensors, model=model, full_layer_name=f"{module_name}.{name}", target_patterns=[key], **kwargs + ) + result.update(tensors) + return result + + +class FineGrainedInterleaveGateUp(_FineGrainedOp): + """Expert gate|up rows between the checkpoint's order and the one the experts hold: stacked + ``[gate; up]`` or the kernels' ``[g0, u0, g1, u1, ...]`` — the core ``Interleave`` along dim 1, + over the weight, scale grid and bias alike. Runs on load and save wherever the two differ.""" + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + input_dict = keyed_by_target(input_dict, target_patterns) + experts = finegrained_experts(model, full_layer_name) + # `is_concatenated` is the checkpoint's stacked order + if experts is None or not experts.has_gate or experts.is_concatenated != experts.holds_interleaved_gate_up: + return input_dict + # inverse=True is stacked -> interleaved + interleave = Interleave(dim=1, inverse=experts.holds_interleaved_gate_up != self.inverse) + out = {} + for key, value in input_dict.items(): + value = single_tensor(value) + # a per-tensor scale is ONE value covering the whole gate|up stack — no row axis to + # reorder, and reshaping its length-1 axis into pairs is what would fail + if getattr(value, "ndim", 0) < 2 or value.shape[1] < 2: + out[key] = value + else: + out[key] = interleave.convert({key: value}, None, [key])[key] + return out + + +class FineGrainedSwizzleScales(_FineGrainedOp): + """Expert block scales into the ``SWIZZLE_32_4_4`` layout their module holds (a 5-D Parameter): + one triton launch per rank-local shard, values unchanged. The swizzle covers whole 128-row + blocks and 4-column groups, so the reverse reads the affine grid back off the 5-D shape.""" + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + out = {} + for key, value in keyed_by_target(input_dict, target_patterns).items(): + value = single_tensor(value) + out[key] = self._unswizzle(value) if self.inverse else self._swizzle(key, value, model, full_layer_name) + return out + + @staticmethod + def _unswizzle(value): + if value.ndim != 5: + return value + experts, row_blocks, col_groups = value.shape[:3] + kernel = load_finegrained_kernel() + return kernel.unswizzle_mx_scales(value, row_blocks * 128, col_groups * 4, num_experts=experts) + + @staticmethod + def _swizzle(key, value, model, full_layer_name): + _, held = held_scale(model, full_layer_name, key) + if held is None or held.ndim != 5 or value.ndim == 5: + return value + return load_finegrained_kernel().swizzle_mx_scales(value) + + +class FineGrainedPackedBlocks(_FineGrainedOp): + """GPT-OSS's ``{proj}_blocks`` ``(E, N, K/32, 16)`` uint8 as the packed ``(E, N, K/2)`` int8 + the kernels read: the same bytes regrouped per row, folded back into 16-byte groups on the way + out. Two low-nibble-first E2M1 values per byte.""" + + def convert(self, input_dict, target_patterns=None, **kwargs): + out = {} + for key, value in keyed_by_target(input_dict, target_patterns).items(): + value = single_tensor(value) + if "_scales" in key: # a converter that carries the scale alongside (the dequant chain) + out[key] = value + elif self.inverse: + out[key] = value.view(torch.uint8).reshape(*value.shape[:-1], value.shape[-1] // 16, 16) + else: + out[key] = value.reshape(*value.shape[:-2], -1).view(torch.int8) + return out + + +class FineGrainedScalarScales(_FineGrainedOp): + """Give per-expert scalar scales a 1x1 grid before stacking and concatenating gate/up scales.""" + + def convert(self, input_dict, **kwargs): + def reshape(tensor): + if self.inverse: + return tensor.reshape(()) + return tensor.reshape(1, 1) + + return { + key: [reshape(tensor) for tensor in value] if isinstance(value, list) else reshape(value) + for key, value in input_dict.items() + } + + +class FineGrainedRowScales(_FineGrainedOp): + """ModelOpt exports channel scales as vectors; the GEMMs hold a column grid.""" + + def convert(self, input_dict, **kwargs): + return { + key: single_tensor(value).squeeze(-1) if self.inverse else single_tensor(value).unsqueeze(-1) + for key, value in input_dict.items() + } + + +class FineGrainedViewPackedInt8(_FineGrainedOp): + """Bitcast packed-FP4 uint8 checkpoint bytes to the int8 view the finegrained modules store: + ``copy_`` into an int8 param would numerically CONVERT and corrupt values >= 128. Non-uint8 + tensors pass through, so this can ride converters that also match unquantized modules.""" + + def convert(self, input_dict, target_patterns=None, **kwargs): + held, want = (torch.int8, torch.uint8) if self.inverse else (torch.uint8, torch.int8) + out = {} + for key, value in keyed_by_target(input_dict, target_patterns).items(): + value = single_tensor(value) + out[key] = value.view(want) if torch.is_tensor(value) and value.dtype == held else value + return out + + +class FineGrainedReciprocal(_FineGrainedOp): + """compressed-tensors stores NVFP4's globals as ``448 * 6 / amax``, the reciprocal of the multiplier the kernels + (and modelopt) take. Its own inverse, so a save writes them back as shipped.""" + + def convert(self, input_dict, target_patterns=None, **kwargs): + return { + key: [1 / tensor for tensor in value] if isinstance(value, list) else 1 / value + for key, value in keyed_by_target(input_dict, target_patterns).items() + } + + +class FineGrainedWeightGlobals(_FineGrainedOp): + """The second-level NVFP4 globals and the calibrated ``input_scale``, as the kernels index them: + one fp32 global per expert per projection, from either the per-expert checkpoint keys + (``experts.*.gate_proj.weight_scale_2``) or one stacked ``(E, 2)`` tensor per layer. + + A gate|up stack has two per expert and merging them is not local, so one op owns a layer's + whole set: the stack keeps the gate's and the up half's leaves as a ratio. ``inter = + silu(gate) * up`` is linear in the up half, so that only scales the expert output by + ``1 / ratio`` — which the down projection's weight global takes back, and its calibrated input + global keeps the requant on the range the e4m3 block scales were chosen for. Rescaling the up + half's block scales instead would re-round them against codes chosen for the old global.""" + + supports_many_to_many = True + + @staticmethod + def role(key: str) -> tuple[str, str]: + """What a global-scale key holds, for a checkpoint key and a module target alike: which + projection it belongs to, and which of modelopt's two levels it is — the weight's + second-level global (``weight_scale_2``) or the calibrated activation one (``input_scale``). + + This runs while the converters are being BUILT, against raw checkpoint keys, so the + vLLM-style `w1`/`w2`/`w3` names are still in play and have not been renamed to + `gate/down/up_proj` yet. Matched on path SEGMENTS: the default is gate_up, so a loose + `"w2" in key` would silently bucket a down global as one.""" + down = any(seg == "w2" or seg.startswith("down_proj") for seg in key.split(".")) + level = "input" if "input_scale" in key or "activation_scale" in key else "weight" + return ("down" if down else "gate_up"), level + + def convert(self, input_dict, target_patterns=None, model=None, **kwargs): + sources = defaultdict(dict) + for key, value in input_dict.items(): + # the loader hands a converter's tensors in a list; unwrapped, a fused `(E, 2)` pair + # reads as `(1, 2E)` and the per-half fold below silently does not fire + sources[self.role(key)][key] = value[0] if isinstance(value, list) and len(value) == 1 else value + globals_ = {} + for target in target_patterns: + if role_sources := sources.get(self.role(target)): + # one (E, calibrated projections) fp32 tensor, gate column first: two sources are + # the stack's halves, one is the fused (E, 2) pair or a single projection's (E,) + columns = [] + for key, value in sorted(role_sources.items(), key=lambda kv: 0 if "gate_proj" in kv[0] else 1): + value = torch.stack(value, dim=0) if isinstance(value, list) else value + columns.append(value.float().reshape(value.shape[0], -1)) + globals_[target] = torch.cat(columns, dim=1) + targets = {self.role(target): target for target in globals_} + stack = targets.get(("gate_up", "weight")) + + # a stack calibrated per half keeps the gate's global; the up half's leaves as a ratio, + # which the down's weight global takes on (scaling the expert output back up) and its + # input global gives back (keeping the requantized intermediate on the calibrated range) + if stack is not None and globals_[stack].shape[1] == 2: + gate, up = globals_[stack].chunk(2, dim=1) + gate, up = gate.squeeze(1), up.squeeze(1) + if not torch.equal(gate, up): + self._assert_halves_foldable(model) + globals_[stack] = gate + ratio = (up / gate)[:, None] + for role, factor in ((("down", "weight"), ratio), (("down", "input"), 1 / ratio)): + if target := targets.get(role): + globals_[target] = globals_[target] * factor + return {target: value.reshape(-1).contiguous() for target, value in globals_.items()} + + @staticmethod + def _assert_halves_foldable(model) -> None: + """Merging the halves' globals moves the up half's past the GLU. A gate|up bias is added AFTER the global, so + its up rows would have to be divided by the same ratio, in the bias converter the layout ops own; no + checkpoint pairs the two. And the GLU must be linear in `up`: a clamped SwiGLU (`swiglu_limit`) clips the + rescaled values, and SwiGLU-OAI (`swiglu_alpha`) multiplies by `up + 1`.""" + experts = [module for module in model.modules() if isinstance(module, FineGrainedExperts)] if model else [] + if any(module.has_bias for module in experts): + raise NotImplementedError( + "this checkpoint calibrates the gate|up halves separately AND carries an expert " + "bias; the up half's bias would have to be folded with the globals." + ) + if any(module.swiglu_limit is not None or module.swiglu_alpha is not None for module in experts): + raise NotImplementedError( + "this checkpoint calibrates the gate|up halves separately, and this model's GLU " + "(`swiglu_limit` / `swiglu_alpha`) is not linear in the up half, so the halves' globals " + "cannot be merged into one per expert." + ) + + +class FineGrainedInputScales(_FineGrainedOp): + """A calibrated checkpoint's ``input_scale`` in the layout the module holds: one value per + quantized module, so a MoE brings one per expert, the gate|up pair reducing to their max since + both halves read the same rows.""" + + def convert(self, input_dict, full_layer_name=None, **kwargs): + values = [] + for value in input_dict.values(): + # unwrapped, a fused `(E,)` scale in the loader's list reads as `(1, E)` and reduces across experts + if isinstance(value, list): + value = value[0] if len(value) == 1 else torch.stack(value, dim=0) + values.append(value.float()) + stacked = torch.stack([v.reshape(v.shape[0], -1) if v.ndim > 1 else v.reshape(-1, 1) for v in values], dim=-1) + # a scale is a magnitude, so the reduce is over absolute values + per_expert = stacked.reshape(stacked.shape[0], -1).abs().amax(dim=1) + return {full_layer_name: per_expert.contiguous()} + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInputScalesSplit(self.hf_quantizer) + + +class FineGrainedInputScalesSplit(_FineGrainedOp): + """Save reverse of :class:`FineGrainedInputScales`: the module's scale back onto every + projection key the checkpoint calibrated one for. The merge took a max, so the folded halves + are gone and each projection is written the value that covers it. A collapsed global + re-expands per expert, the shape the checkpoint holds.""" + + def convert(self, input_dict, model=None, full_layer_name=None, target_patterns=None, **kwargs): + value = single_tensor(next(iter(input_dict.values()))) + if value.numel() == 1: + experts = finegrained_experts(model, full_layer_name) + if experts is not None: + value = value.reshape(1).expand(experts.num_experts).contiguous() + # one tensor per key: a save refuses keys that share storage + return {key: value.clone() for key in target_patterns or list(input_dict)} + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedInputScales(self.hf_quantizer) + + +class FineGrainedQuantize(_FineGrainedOp): + """Quantize a full-precision weight on load into the format its module holds, emitting the + scale (and the NVFP4 global) in its held dtype. Each full stack is quantized before applying + the module's layouts (gate|up interleave, scale swizzle), so a block covers one projection. + + Block-FP8 is computed in torch; the group formats run the kernels' row-wise quantizers, NVFP4 + after normalizing by the canonical global ``amax / (6 * 448)``. A tensor that is not a + finegrained module's weight passes through — `_weight_holder` decides, since rank alone does + not (an expert bias is 2-D).""" + + def convert(self, input_dict: dict[str, torch.Tensor], model=None, **kwargs) -> dict[str, torch.Tensor]: + result: dict[str, torch.Tensor] = {} + for key, value in input_dict.items(): + tensor = single_tensor(value) + holder = self._weight_holder(model, key) + if holder is None and model is not None: + # not a weight, whatever its rank: an expert bias is 2-D and would otherwise be + # quantized into a `_bias_scale_inv` no module holds. Only a direct + # invocation, with no model to ask, quantizes without a holder. + result[key] = tensor + continue + quantized = self._quantize_one(key, tensor, holder) + for name, value in quantized.items(): + tensors = {name: value} + if holder is not None and hasattr(self.hf_quantizer, "member_of"): + member = self.hf_quantizer.member_of(holder[0]) + for op in member.get_param_ops(model, name) if member is not None else []: + tensors = op.convert(tensors, model=model, full_layer_name=name, target_patterns=[name]) + result.update(tensors) + return result + + @staticmethod + def _as_expert_rows(module, name: str, value: torch.Tensor) -> torch.Tensor: + """An expert stack in the ``(E, rows, K)`` the module's slot and the kernels take. A model may store its + experts transposed (GPT-OSS: ``(E, H, 2I)``), which the module it replaced declares: quantizing before the + swap packs the wrong axis, and a square stack (GPT-OSS's down projection) cannot tell from its shape.""" + if module is None or value.ndim != 3 or name not in FUSED_STACKS: + return value + return value.transpose(1, 2).contiguous() if getattr(module, "ships_transposed", False) else value + + @staticmethod + def _weight_holder(model, key: str): + """``(module, scale_param_name)`` when ``key`` is a finegrained module's weight, else ``None``.""" + if model is None: + return None + module, name = get_module_from_name(model, key) + if not isinstance(module, _FineGrainedModule) or getattr(module, name, None) is None: + return None + if name == "weight": + return module, "weight_scale_inv" + if name in FUSED_STACKS: + return module, f"{name}_scale_inv" + return None + + def _quantize_one(self, key: str, value: torch.Tensor, holder) -> dict[str, torch.Tensor]: + if value.ndim < 2: + return {key: value} + prefix = key.rsplit(".", 1)[0] + ".weight" if key.endswith(".weight") else key + module, scale_name = holder if holder is not None else (None, None) + value = self._as_expert_rows(module, key.rsplit(".", 1)[1], value) + held = getattr(module, scale_name) if module is not None else None + format_spec = get_weight_format(module.weight_format) if module is not None else None + if format_spec is None or format_spec.scale_group is None: + if module is not None and module.quantization_config.weight_block_size == "row": + block = (1, value.shape[-1]) + elif module is not None and module.block_size: + block = tuple(module.block_size) + else: + # no module (direct invocation): take the block from the config, and failing that + # put one scale over the whole matrix + config = self.hf_quantizer.quantization_config if self.hf_quantizer is not None else None + config_block = ( + config.get("weight_block_size") + if isinstance(config, dict) + else getattr(config, "weight_block_size", None) + ) + block = ( + (1, value.shape[-1]) + if config_block == "row" + else tuple(config_block) + if config_block + else (value.shape[-2], value.shape[-1]) + ) + ue8m0 = ( + held.dtype in (torch.uint8, getattr(torch, "float8_e8m0fnu", None)) + if held is not None + else self.hf_quantizer.quantization_config.scale_fmt == "ue8m0" + ) + weight, scale = self._quantize_block_fp8(value, block, ue8m0) + return { + key: weight, + f"{prefix}_scale_inv": scale, + **self._identity_activation_scales(module, prefix, scale_name, value.device), + } + + if value.device.type not in ("cuda", "xpu"): + # rocm counts: torch reports it as "cuda" + raise ValueError( + f"on-the-fly {module.weight_format} quantization runs the kernels' quantizers on an " + f"accelerator, but the weight is on {value.device.type}" + ) + weight, scale, global_scale = self._quantize_group(value, format_spec) + scale = as_container(scale, held.dtype) + out = {key: weight, f"{prefix}_scale_inv": scale} + if global_scale is not None: + # the experts name their global `_weight_global_scale`, a dense linear just + # `weight_global_scale` — both sit beside a `_scale_inv`, so ask for the slot + stem = scale_name.removesuffix("_scale_inv") + suffix = "_global_scale" if hasattr(module, f"{stem}_global_scale") else "_weight_global_scale" + out[f"{prefix}{suffix}"] = global_scale + out.update(self._identity_activation_scales(module, prefix, scale_name, value.device)) + return out + + @staticmethod + def _identity_activation_scales(module, prefix: str, scale_name: str | None, device) -> dict: + """This module's activation-side scales, as the identity, because quantizing on the fly + is not a calibration pass. They have to be WRITTEN: the loader materializes a key no + checkpoint supplies with `torch.empty_like` and `_init_weights` has no branch for a + scale, so a slot left out here reaches the kernels as uninitialized memory.""" + if module is None: + return {} + stem = scale_name.removesuffix("_scale_inv") + # the experts prefix each slot with the projection; a dense linear, whose stem IS + # `weight`, names them bare + proj = "" if stem == "weight" else f"{stem}_" + out = {} + slot = f"{proj}activation_scale" + held = getattr(module, slot, None) + if held is not None: + # on the weight's device, not the slot's: the module is still on meta here, and + # `ones_like` would inherit that and leave the parameter unmaterialized + out[prefix.removesuffix(stem) + slot] = torch.ones(held.shape, dtype=held.dtype, device=device) + return out + + @staticmethod + def _quantize_block_fp8(value: torch.Tensor, block: tuple[int, int], ue8m0: bool): + """``(E4M3 weight, inverse scale grid)`` for a ``(..., rows, cols)`` tensor at ``block``. + + A trailing PARTIAL block is padded, not refused — DeepSeek-V3 ships `kv_a_proj_with_mqa` + as `(576, 7168)` against a 128x128 block with a `(5, 56)` grid, so the format's own + producers round up and the module allocates to match. The zeros cannot move a block's + amax, so the short block is scaled by its real values. UE8M0 rounds the inverse scale up + to a power of two before quantizing, so dequant multiplies by exactly what it divided by. + """ + block_m, block_n = block + rows, cols = value.shape[-2], value.shape[-1] + padded = torch.nn.functional.pad(value.float(), (0, -cols % block_n, 0, -rows % block_m)) + grid_m, grid_n = padded.shape[-2] // block_m, padded.shape[-1] // block_n + tiles = padded.reshape(*padded.shape[:-2], grid_m, block_m, grid_n, block_n) + max_abs = tiles.abs().amax(dim=(-3, -1)) + inv_scale = torch.where(max_abs > 0, max_abs / _FP8_MAX, torch.ones_like(max_abs)) + if ue8m0: + inv_scale = torch.pow(2.0, torch.ceil(torch.log2(inv_scale.clamp(min=torch.finfo(torch.float32).tiny)))) + scaled = tiles / inv_scale.unsqueeze(-1).unsqueeze(-3) + quantized = torch.clamp(scaled, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE).reshape(padded.shape) + # a no-op when the shape tiled + return quantized[..., :rows, :cols].contiguous(), as_container( + inv_scale, _get_ue8m0_dtype() + ) if ue8m0 else inv_scale + + @staticmethod + def _quantize_group(value: torch.Tensor, format_spec: WeightFormat): + """``(weight, scale, global)`` for a ``(..., rows, K)`` tensor in a group-scaled format through the + kernels' row-wise quantizers, one launch per tensor; ``global`` is ``None`` unless the format has + one. NVFP4's global is per matrix (per expert for a stack, a scalar for a dense weight): the + smallest that keeps every E4M3 block scale in range, divided out before the block quant.""" + kernel = load_finegrained_kernel() + rows = value.reshape(-1, value.shape[-1]) + global_scale = None + if format_spec.global_scale_dtype is not None: + per_matrix = value.float().reshape(-1, value.shape[-2] * value.shape[-1]).abs().amax(-1) / (6.0 * 448.0) + global_scale = per_matrix.clamp(min=torch.finfo(torch.float32).tiny) + rows = (value.float() / global_scale.reshape(-1, *[1] * (value.ndim - 1))).reshape_as(rows) + packed, scale = kernel.nvfp4_act_quant(rows.contiguous()) + global_scale = global_scale.reshape(()) if value.ndim == 2 else global_scale + else: + quantize = kernel.mxfp4_act_quant if format_spec.values_per_byte == 2 else kernel.mxfp8_act_quant + packed, scale = quantize(rows.to(torch.bfloat16).contiguous()) + return ( + packed.view(format_spec.weight_dtype).reshape(*value.shape[:-1], -1), + scale.reshape(*value.shape[:-1], -1), + global_scale, + ) + + @property + def reverse_op(self) -> ConversionOps: + return FineGrainedDequantize(self.hf_quantizer) + + +class FineGrainedDequantize(_FineGrainedOp): + """A quantized weight folded back to full precision against its per-block scale grid. + + Runs FIRST in its converter under ``dequantize=True``, which is why + :meth:`update_weight_conversions` attaches it to each of the model's own converters: the + merge / concat ops after it collapse the per-expert structure the (weight, scale) pairing + needs. It pairs each weight pattern with its sibling scale pattern by index and emits the + result under the weight key, dropping the scales so the rest of the chain sees weights only. + """ + + def __init__(self, hf_quantizer=None, inverse=False, *, transpose_scales=False): + super().__init__(hf_quantizer, inverse) + self.transpose_scales = transpose_scales + + def _scale_pattern_for(self, weight_pattern: str) -> str: + anchor = "(?=$)" if weight_pattern.endswith("(?=$)") else "$" if weight_pattern.endswith("$") else "" + base = weight_pattern.removesuffix(anchor) if anchor else weight_pattern + if base.endswith("_blocks"): + # GPT-OSS packs its experts as `{proj}_blocks` + `{proj}_scales` + scale = base[: -len("_blocks")] + "_scales" + elif base.endswith(".weight"): + scale = base[: -len(".weight")] + ".weight_scale_inv" + elif base == "weight": + scale = "weight_scale_inv" + else: + scale = base + "_scale_inv" + return scale + anchor + + def _config_for(self, full_layer_name: str | None): + """Resolve this module's format independently of other entries in a mixed plan.""" + if self.hf_quantizer is None: + return None + config = self.hf_quantizer.quantization_config + if full_layer_name is not None: + module_name = full_layer_name.rpartition(".")[0] + config = plan_entry_for(self.hf_quantizer.plan(), module_name) + return config + + # the weights a converter hands one at a time: a dense `weight`, or an experts module's fused stack + _SINGLE_WEIGHTS = ("weight", *FUSED_STACKS) + + @staticmethod + def _name(pattern: str) -> str: + """The name a source pattern matches: `(? `gate_up_proj`.""" + return pattern.removeprefix(r"(? str | None: + """The source pattern in `input_dict` matching the same name as `pattern`, if the converter collected one.""" + return next((key for key in input_dict if self._name(key) == self._name(pattern)), None) + + @staticmethod + def _global_pattern_for(weight_pattern: str) -> str: + """NVFP4's second-level global beside `weight_pattern`: `x.weight` -> `x.weight_global_scale`, and a fused + stack's `gate_up_proj` -> `gate_up_proj_weight_global_scale`.""" + anchor = "(?=$)" if weight_pattern.endswith("(?=$)") else "$" if weight_pattern.endswith("$") else "" + base = weight_pattern.removesuffix(anchor) if anchor else weight_pattern + suffix = "_global_scale" if base.endswith("weight") else "_weight_global_scale" + return base + suffix + anchor + + # packed-FP4 experts arrive as int8 / float4_e2m1fn_x2, two e2m1 nibbles per byte + _FP4_E2M1_LUT = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0) + + def _unpack_fp4(self, packed: torch.Tensor) -> torch.Tensor: + """Two ``e2m1`` FP4 values per byte → float32 tensor twice as wide on the last dim.""" + lut = torch.tensor(self._FP4_E2M1_LUT, dtype=torch.float32, device=packed.device) + u8 = packed.contiguous().view(torch.uint8) + low = (u8 & 0xF).long() + high = ((u8 >> 4) & 0xF).long() + unpacked = torch.stack([lut[low], lut[high]], dim=-1) + return unpacked.reshape(*packed.shape[:-1], 2 * packed.shape[-1]) + + def _dequantize_one( + self, + quantized: torch.Tensor, + scales: torch.Tensor, + output_dtype: torch.dtype | None = None, + global_scale: torch.Tensor | None = None, + *, + reciprocal_globals: bool = False, + weight_block_size: tuple[int, int] | str | None = None, + ) -> torch.Tensor: + # unpacked first, so the rest of the routine sees a normal (rows, cols) float matrix + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None) + # modelopt ships its packed bytes as uint8 + if quantized.dtype in (torch.int8, torch.uint8) or (fp4_dtype is not None and quantized.dtype == fp4_dtype): + quantized_fp32 = self._unpack_fp4(quantized) + else: + quantized_fp32 = quantized.to(torch.float32) + rows, cols = quantized_fp32.shape[-2:] + if weight_block_size is None: + weight_block_size = getattr(self._config_for(None), "weight_block_size", None) + if weight_block_size == "row": + if scales.ndim < quantized_fp32.ndim: + scales = scales.unsqueeze(-1) + if self.transpose_scales: + scales = scales.transpose(-1, -2) + # A per-tensor scale has no grid: one block over the whole matrix. + scale_rows, scale_cols = scales.shape[-2:] if scales.ndim >= 2 else (1, 1) + # Ceil-rounded grids do not uniquely identify a block size. Use this module's config, + # including in mixed plans; infer from the grid only when no configured block matches. + if ( + weight_block_size + and weight_block_size != "row" + and tuple(_cdiv(size, block) for size, block in zip((rows, cols), weight_block_size)) + == ( + scale_rows, + scale_cols, + ) + ): + block_m, block_n = weight_block_size + elif rows % scale_rows or cols % scale_cols: + raise ValueError( + f"Weight shape ({rows}, {cols}) not divisible by scale grid ({scale_rows}, {scale_cols})." + ) + else: + block_m = rows // scale_rows + block_n = cols // scale_cols + # the math runs in fp32 either way (``float8_e8m0fnu`` has no ``mul`` kernel); the + # destination parameter's dtype wins when known, so an eager module keeps the model's + if output_dtype is None: + output_dtype = ( + scales.dtype if scales.dtype.is_floating_point and scales.element_size() >= 2 else torch.bfloat16 + ) + # an MXFP8 checkpoint's uint8 scale is a biased E8M0 exponent, not a multiplier + if scales.dtype == torch.uint8: + s_fp32 = (scales.to(torch.float32) - 127.0).exp2() + else: + s_fp32 = scales.to(torch.float32) + original_shape = quantized_fp32.shape + # pad out to whole blocks so a ceil-rounded grid lines up, then crop back + pad_m, pad_n = -rows % block_m, -cols % block_n + padded = torch.nn.functional.pad(quantized_fp32, (0, pad_n, 0, pad_m)) + blocked = padded.reshape(-1, scale_rows, block_m, scale_cols, block_n) + per_block = s_fp32.reshape(-1, scale_rows, scale_cols).unsqueeze(-1).unsqueeze(2) + out = (blocked * per_block).reshape(*original_shape[:-2], rows + pad_m, cols + pad_n)[..., :rows, :cols] + if global_scale is not None: + global_scale = global_scale.to(torch.float32) + if out.ndim > 2 and global_scale.ndim: + # a stack's globals, one per expert (E,) or one per stacked gate|up half (E, 2), over its rows + global_scale = global_scale.reshape(global_scale.shape[0], -1) + global_scale = global_scale.repeat_interleave(rows // global_scale.shape[1], dim=1)[..., None] + out = out / global_scale if reciprocal_globals else out * global_scale + return out.to(output_dtype) + + def _get_target_dtype(self, model: torch.nn.Module | None, full_layer_name: str | None) -> torch.dtype | None: + if model is None or full_layer_name is None: + return None + module, tensor_name = get_module_from_name(model, full_layer_name) + param = getattr(module, tensor_name, None) + return getattr(param, "dtype", None) + + def convert( + self, + input_dict: dict[str, list[torch.Tensor] | torch.Tensor], + full_layer_name: str | None = None, + model: torch.nn.Module | None = None, + **kwargs, + ) -> dict[str, list[torch.Tensor] | torch.Tensor]: + output_dtype = self._get_target_dtype(model, full_layer_name) + config = self._config_for(full_layer_name) + reciprocal_globals = getattr(config, "checkpoint_format", None) == "compressed-tensors" + weight_block_size = getattr(config, "weight_block_size", None) + # The dense converter (``["weight$", companions...] -> weight``) and a fused stack's (``["gate_up_proj$", its + # companions...] -> gate_up_proj``) each hand one weight, with or without a scale (RMSNorm weights match + # ``weight$`` too). + weight_pattern = next((p for p in input_dict if self._name(p) in self._SINGLE_WEIGHTS), None) + if weight_pattern is not None: + # the loader derives prefix/suffix from the output KEY; without a `full_layer_name` + # (direct invocation) key it by the weight's own name + target_key = full_layer_name if full_layer_name is not None else self._name(weight_pattern) + quantized = single_tensor(input_dict[weight_pattern]) + scale, global_scale = ( + single_tensor(input_dict.get(self._companion(input_dict, companion(weight_pattern)))) + for companion in (self._scale_pattern_for, self._global_pattern_for) + ) + if scale is None: + return {target_key: quantized.to(output_dtype) if output_dtype is not None else quantized} + return { + target_key: self._dequantize_one( + quantized, + scale, + output_dtype, + global_scale, + reciprocal_globals=reciprocal_globals, + weight_block_size=weight_block_size, + ) + } + + # Generic chain path: dequantize every weight pattern that has a sibling scale. + consumed = { + pattern(key) for key in input_dict for pattern in (self._scale_pattern_for, self._global_pattern_for) + } + result: dict[str, list[torch.Tensor] | torch.Tensor] = {} + for key, value in input_dict.items(): + if "activation_scale" in key or key in consumed: + continue # consumed by the dequant; drop from the chain + scale_key = self._scale_pattern_for(key) + if scale_key not in input_dict: + if output_dtype is not None: + value = ( + [tensor.to(output_dtype) for tensor in value] + if isinstance(value, list) + else value.to(output_dtype) + ) + result[key] = value + continue + weights = value if isinstance(value, list) else [value] + scales = input_dict[scale_key] + scales = scales if isinstance(scales, list) else [scales] + if len(weights) != len(scales): + raise ValueError( + f"FineGrainedDequantize: weight/scale count mismatch for {key} " + f"({len(weights)} weights vs {len(scales)} scales)." + ) + global_scales = input_dict.get(self._global_pattern_for(key), [None] * len(weights)) + global_scales = global_scales if isinstance(global_scales, list) else [global_scales] + result[key] = [ + self._dequantize_one( + w, s, output_dtype, g, reciprocal_globals=reciprocal_globals, weight_block_size=weight_block_size + ) + for w, s, g in zip(weights, scales, global_scales) + ] + return result + + @property + def reverse_op(self) -> ConversionOps: + # a dequantized model saves in full precision: none of its weights fills a finegrained module, + # so the quantize passes each through + return FineGrainedQuantize(self.hf_quantizer) diff --git a/src/transformers/integrations/finegrained/layers.py b/src/transformers/integrations/finegrained/layers.py new file mode 100644 index 000000000000..74352440e4ed --- /dev/null +++ b/src/transformers/integrations/finegrained/layers.py @@ -0,0 +1,1051 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Quantized modules and forwards for the block-scaled formats — block-FP8, MXFP8, MXFP4, NVFP4 +and weight-only — served by `kernels-community/finegrained-kernels`. The checkpoint's layout is +moved onto these modules by the ops in `conversions`. +""" + +from __future__ import annotations + +import functools +import os +import sys +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from ...activations import ACT2FN +from ...quantizers.quantizers_utils import FUSED_STACKS, STACK_COMPANIONS +from ...utils import logging +from ...utils.import_utils import is_kernels_available +from ...utils.quantization_config import BlockFp8Config, BlockScaledConfig +from ..deepgemm import ( + deepgemm_fp8_fp4_experts_forward, + deepgemm_fp8_fp4_linear, + deepgemm_fp8_fp4_megamoe_experts_forward, + is_sm100, + prefers_deepgemm_linear, +) +from ..hub_kernels import _MISSING_KERNELS_MESSAGE, lazy_load_kernel +from ..moe import ExpertsInterface, use_experts_implementation + + +logger = logging.get_logger(__name__) + + +_FP8_DTYPE = torch.float8_e4m3fn +_FP8_MIN = torch.finfo(_FP8_DTYPE).min +_FP8_MAX = torch.finfo(_FP8_DTYPE).max + + +@functools.cache +def _get_ue8m0_dtype() -> torch.dtype: + """Use native UE8M0 when available, otherwise store its biased exponent bytes.""" + return getattr(torch, "float8_e8m0fnu", torch.uint8) + + +def _first_attr(obj, *names, raise_error: bool = True): + """The first of `names` the object defines; `None` when it defines none and the caller passes + `raise_error=False` because it has its own fallback.""" + for name in names: + if hasattr(obj, name): + return getattr(obj, name) + if raise_error: + raise AttributeError(f"{type(obj).__name__} has none of: {names}") + return None + + +@dataclass(frozen=True) +class FineGrained: + """Entry points exposed by the `kernels-community/finegrained-kernels` Triton kernel. + + Every format (block-FP8, MXFP8, MXFP4, NVFP4, weight-only) flows through the matmuls and + the two MoE forwards, with the weight format resolved off the weight/scale dtypes inside the + kernel and the module's ``activation_format`` passed through unchanged — the kernels speak the + same vocabulary (``None`` = the weights' format, ``"bf16"`` = weight-only). The experts + forwards hand the kernels' MoE forwards the module's tensors and let them run the whole chain: + scheduling, fused GLU + intermediate requant, biases, the routing-weighted reduce. What they + cannot fuse runs in between — an activation outside ``get_supported_act_fns()`` as the module's + own callable, a per-expert output norm by name where ``get_supported_norms()`` has it. The + quantizers serve on-the-fly quantization into the group formats. All symbols are required; a + build missing any raises at load with the full list. + """ + + matmul_2d: Callable + matmul_batched: Callable + matmul_grouped: Callable + moe_fused_batched: Callable + moe_fused_grouped: Callable + swizzle_mx_scales: Callable + unswizzle_mx_scales: Callable + mxfp8_act_quant: Callable + mxfp4_act_quant: Callable + nvfp4_act_quant: Callable + get_supported_act_fns: Callable + get_supported_norms: Callable + rms_norm_rows: Callable + embedding_lookup: Callable + + +# Cache the loaded kernel but not failures: re-checking each call is cheap and intended, since the env +# can change between attempts. A module global (not `@functools.cache`) avoids Dynamo warning about +# tracing a cache-wrapped function on every compile. +_FINEGRAINED: FineGrained | None = None + + +@torch._dynamo.allow_in_graph +def _load_finegrained_kernel() -> None: + """ + Load the finegrained-kernels Triton kernel once into the `_FINEGRAINED` module global. + + `@allow_in_graph` makes `torch.compile` treat the untraceable hub download + dynamic import as a + single opaque node instead of tracing into it; it returns `None` (proxyable) and populates the + global, which `load_finegrained_kernel` then returns. + + Under NO circumstances may this function return a value: an `@allow_in_graph` fx node's + return must be proxyable, and returning the bundle (e.g. from the warm-cache + short-circuit) breaks torch.compile with `Unsupported: torch.* op returned non-Tensor`. + + Raises `ImportError` if the `kernels` package is missing, or the kernel or required + symbols cannot be found. + """ + global _FINEGRAINED + if _FINEGRAINED is not None: + return + + if not is_kernels_available(): + raise ImportError(f"finegrained-kernels kernel unavailable: {_MISSING_KERNELS_MESSAGE}") + + kernel = lazy_load_kernel("finegrained-kernels") + if kernel is None: + raise ImportError( + "Failed to load the finegrained-kernels kernel — check that `kernels-community/finegrained-kernels` " + "has a build matching the current torch/CUDA." + ) + + # the bundle's fields ARE the required symbols, under the kernel's own names + symbols = {name: getattr(kernel, name, None) for name in FineGrained.__dataclass_fields__} + missing = [name for name, attr in symbols.items() if attr is None] + if missing: + raise ImportError( + f"finegrained-kernels build is missing required symbols: {', '.join(missing)}. {_MISSING_KERNELS_MESSAGE}" + ) + _FINEGRAINED = FineGrained(**symbols) + + +def load_finegrained_kernel() -> FineGrained: + _load_finegrained_kernel() + return _FINEGRAINED + + +def _cdiv(a: int, b: int) -> int: + """Ceiling division.""" + return (a + b - 1) // b + + +@dataclass(frozen=True) +class WeightFormat: + """Storage layout of one quantized weight format — the single source both module classes + derive their parameter shapes from. ``scale_dtype`` ``None`` defers to the config's + ``scale_fmt`` (block-FP8 ships fp32 or UE8M0 containers; the group formats pin theirs). + ``scale_group`` ``None`` means the block comes from the quant config's ``weight_block_size``. + ``global_scale_dtype`` ``None`` means the format has no second level; NVFP4 names the dtype its + per-matrix global is held in.""" + + weight_dtype: torch.dtype + values_per_byte: int = 1 + scale_dtype: torch.dtype | None = None + scale_group: tuple[int, int] | None = None + global_scale_dtype: torch.dtype | None = None + + +def get_weight_format(weight_format: str) -> WeightFormat: + """Resolve only the requested format, so FP8 and NVFP4 do not require UE8M0 support.""" + if weight_format == "fp8": + return WeightFormat(weight_dtype=_FP8_DTYPE) + if weight_format in ("mxfp8", "mxfp4"): + packed = weight_format == "mxfp4" + return WeightFormat( + weight_dtype=torch.int8 if packed else _FP8_DTYPE, + values_per_byte=2 if packed else 1, + scale_dtype=_get_ue8m0_dtype(), + scale_group=(1, 32), + ) + if weight_format == "nvfp4": + return WeightFormat( + weight_dtype=torch.int8, + values_per_byte=2, + scale_dtype=_FP8_DTYPE, + scale_group=(1, 16), + global_scale_dtype=torch.float32, + ) + raise ValueError(f"unknown weight_format {weight_format!r}; expected fp8, mxfp8, mxfp4 or nvfp4") + + +def resolve_weight_format( + weight_format: str, + scale_fmt: str = "float", + block_size: tuple[int, int] | str | None = None, +) -> tuple[WeightFormat, torch.dtype, tuple[int, int | None] | None]: + """``(format, scale_dtype, (sf_gran_n, sf_gran_k))`` for one format name, with the config's + ``scale_fmt``/``weight_block_size`` filling the slots the format leaves open.""" + format_spec = get_weight_format(weight_format) + scale_dtype = format_spec.scale_dtype + if scale_dtype is None: + scale_dtype = _get_ue8m0_dtype() if scale_fmt == "ue8m0" else torch.float32 + scale_group = (1, None) if isinstance(block_size, str) else block_size + return format_spec, scale_dtype, format_spec.scale_group if format_spec.scale_group is not None else scale_group + + +def _set_optional_parameter(module: nn.Module, name: str, tensor: torch.Tensor | None) -> None: + """Register `name` as a Parameter over `tensor`, or as an absent one when `tensor` is None — + a slot a checkpoint may or may not fill (a bias, an NVFP4 global, a static activation scale).""" + module.register_parameter(name, None if tensor is None else nn.Parameter(tensor)) + + +def _alloc_expert_proj( + num_experts: int, + proj_out: int, + proj_in: int, + format_spec: WeightFormat, + scale_dtype: torch.dtype, + scale_group: tuple[int, int | None] | None, + min_scale_out: int = 1, + swizzled: bool = False, +) -> tuple[nn.Parameter, nn.Parameter]: + """``(weight, weight_scale_inv)`` Parameters for one expert projection: the weight in the + format's storage (FP4 packs two values per byte along K), the scale grid at the format's + granularity (``None`` = one scale per expert). ``min_scale_out`` floors the grid's output dim + so a fused gate_up keeps room for both halves. ``swizzled`` allocates the grid in the + ``SWIZZLE_32_4_4`` layout instead — expert axis leading, so EP shards it on dim 0 — which needs + whole blocks; a projection that does not tile stays affine.""" + weight = torch.empty(num_experts, proj_out, proj_in // format_spec.values_per_byte, dtype=format_spec.weight_dtype) + if scale_group is None: + scale_shape = (num_experts, max(1, min_scale_out), 1) + else: + rows = max(_cdiv(proj_out, scale_group[0]), min_scale_out) + cols = _cdiv(proj_in, scale_group[1] or proj_in) + scale_shape = (num_experts, rows, cols) + if swizzled and rows % 128 == 0 and cols % 4 == 0: + scale_shape = (num_experts, rows // 128, cols // 4, 2, 256) + scale = torch.empty(scale_shape, dtype=scale_dtype) + return ( + nn.Parameter(weight, requires_grad=weight.is_floating_point()), + nn.Parameter(scale, requires_grad=scale.is_floating_point() and not swizzled), + ) + + +# the experts backends that read gate|up stacked and block scales affine +_DEEPGEMM_EXPERTS_IMPLEMENTATIONS = ("deepgemm", "deepgemm_megamoe") + + +def _holds_swizzled_scales(config, format_spec: WeightFormat, activation_format: str | None) -> bool: + """Whether an experts module holds its block scales in the ``SWIZZLE_32_4_4`` layout the + Blackwell tcgen05 scaled-MMA reads directly (row-major forces a per-tile gather). Needs SM100, + a triton dispatch (the DeepGEMM backends read affine scales), a group-scaled format and a chain + that quantizes activations. Block-FP8 reaches the scaled-MMA under UE8M0 by broadcasting one + 128-block scalar in-register, so it has no per-group grid to lay out.""" + return ( + getattr(config, "_experts_implementation", None) not in _DEEPGEMM_EXPERTS_IMPLEMENTATIONS + and format_spec.scale_group is not None + and activation_format != "bf16" + and is_sm100() + ) + + +def finegrained_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale_inv: torch.Tensor, + block_size: list[int] | None = None, + bias: torch.Tensor | None = None, + activation_scale: torch.Tensor | None = None, + allow_deepgemm: bool = True, + weight_global_scale: torch.Tensor | None = None, + input_global_scale: torch.Tensor | None = None, + activation_format: str | None = None, + activation_scale_ub: float | None = None, +) -> torch.Tensor: + """End-to-end FP8/FP4 linear used by `FineGrainedLinear` and the eager `FineGrainedExperts` loop. + + Dispatch order — both backends handle FP8 and FP4 weights with fp32 or UE8M0 scales: + 1. DeepGEMM (`deepgemm_fp8_fp4_linear`) where it can actually run and is faster — the + pre-SM100 shapes it supports (FP4, UE8M0 SFs, 128×128 block FP8). Never on SM100, + where Triton is tuned per shape and is the only path reading pre-swizzled scales. + 2. Triton finegrained fallback everywhere else: SM100, an ``activation_scale`` + (DeepGEMM is dynamic-only), a call that needs a gradient (DeepGEMM has no backward), + or any shape DeepGEMM declined. + + Args: + input: (..., K) bf16/fp16 activations. + weight: (N, K) `float8_e4m3fn` or (N, K // 2) `int8` (FP4-packed). + weight_scale_inv: per-block weight scales — `float32` (as DeepSeek-V3 ships them) or + `float8_e8m0fnu` (DeepSeek-V4; reinterpreted as int32 at the DeepGEMM kernel boundary). + block_size: [block_n, block_k] for FP8 block-wise quant, or None/[N, K] for per-tensor. + Ignored for FP4 weights (the kernel infers SF granularity from the dtype). + bias: optional bias added to the matmul output. + activation_scale: pass a per-tensor scalar to use static activation quant; leave `None` + for dynamic (per-token) quant. + allow_deepgemm: set ``False`` to force the Triton fallback for this call. Used when the + model spans multiple CUDA devices in one process — DeepGEMM's cached kernels are bound + to a single CUDA context and produce garbage across devices (see + ``disable_deepgemm_on_multi_device``). + weight_global_scale: the NVFP4 two-level weight global, one value for this matrix + (None for other formats). + input_global_scale: the NVFP4 calibrated activation global, the checkpoint's + ``input_scale`` (None = quantize against the block scales alone). + activation_format: the activation format where the weights leave it open + (``"bf16"`` = weight-only). + """ + if prefers_deepgemm_linear( + input, + weight, + weight_scale_inv, + block_size=block_size, + activation_scale=activation_scale, + weight_global_scale=weight_global_scale, + input_global_scale=input_global_scale, + activation_format=activation_format, + allow_deepgemm=allow_deepgemm, + ): + try: + return deepgemm_fp8_fp4_linear( + input, + weight, + weight_scale_inv, + block_size=block_size, + activation_scale=activation_scale, + bias=bias, + ) + except (ImportError, NotImplementedError, ValueError) as e: + # Triton is the more permissive of the two, so a decline here is not fatal: it serves + # arch/input combos DeepGEMM has no kernel for, and raises its own error if it cannot. + logger.warning_once( + f"DeepGEMM declined this call, falling back to Triton. Reason: {e} " + "Set `TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR=1` to skip DeepGEMM for finegrained linear entirely." + ) + + kernel = load_finegrained_kernel() + original_shape = input.shape + output = kernel.matmul_2d( + input.reshape(-1, original_shape[-1]), + weight, + activation_scale, + weight_scale_inv, + activation_format=activation_format, + output_dtype=input.dtype, + a_global_scale=input_global_scale, + b_global_scale=weight_global_scale, + activation_scale_ub=activation_scale_ub, + ) + output = output.reshape(*original_shape[:-1], output.shape[-1]) + if bias is not None: + output.add_(bias) + return output + + +# the kernels' fused form of a model's post-expert norm, by the norm's exact class: a subclass may compute +# something else, so a model whose own norm one of them implements registers that class here +SUPPORTED_FUSED_NORMS: dict[type[nn.Module], str] = {nn.RMSNorm: "rms_norm"} + + +class _FineGrainedModule: + """What every finegrained module (dense linear or experts) carries besides its parameters.""" + + # the kernels take raw pointers, so the TP layer must hand local shards and not DTensors, whatever + # `should_use_local_tensors` would decide from grad mode alone + _hf_quantized_needs_local_tp = True + + # Internal, temporary flag — not public API, don't set it directly. `disable_deepgemm_on_multi_device` + # flips it True at load when the model spans >1 CUDA device in one process (DeepGEMM's context-bound + # kernels corrupt across devices); removable once the kernel ships a context-free loader. + _deepgemm_disabled = False + + quantization_config: BlockScaledConfig + + @property + def weight_format(self) -> str: + return self.quantization_config.method.value + + @property + def activation_format(self) -> str | None: + return self.quantization_config.activation_format + + @property + def activation_scheme(self) -> str: + return self.quantization_config.activation_scheme + + @property + def block_size(self) -> tuple[int, int] | None: + block = self.quantization_config.weight_block_size + return None if isinstance(block, str) else block + + +class FineGrainedLinear(_FineGrainedModule, nn.Linear): + def __init__( + self, + in_features: int, + out_features: int, + *, + quantization_config: BlockScaledConfig | None = None, + has_bias: bool = False, + ): + super().__init__(in_features, out_features) + + self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) + self.has_bias = has_bias + activation_format = self.activation_format + + # the format table decides storage: value dtype (packed E2M1 = 2 values per int8 byte + # along K), scale dtype/granularity, and whether a two-level global exists + format_spec, sf_dtype, scale_group = resolve_weight_format( + self.weight_format, self.quantization_config.scale_fmt, self.quantization_config.weight_block_size + ) + # the format's own dtype both decides whether there is a second level and says what it + # is held in + global_scale_dtype = format_spec.global_scale_dtype + in_storage = in_features // format_spec.values_per_byte + self.weight = torch.nn.Parameter( + torch.empty(out_features, in_storage, dtype=format_spec.weight_dtype), + requires_grad=format_spec.weight_dtype.is_floating_point, + ) + # NVFP4 two-level: the per-tensor fp32 global the kernel recovers on the accumulator + _set_optional_parameter( + self, + "weight_global_scale", + torch.tensor(1.0, dtype=global_scale_dtype) if global_scale_dtype is not None else None, + ) + + # no group and no block: one per-tensor scale + scale = ( + torch.tensor(1.0, dtype=torch.float32) + if scale_group is None + else torch.empty( + _cdiv(out_features, scale_group[0]), _cdiv(in_features, scale_group[1] or in_features), dtype=sf_dtype + ) + ) + self.weight_scale_inv = nn.Parameter(scale, requires_grad=scale.dtype.is_floating_point) + + # the checkpoint's calibrated activation scale: a static FP8 scale, or NVFP4's activation global + calibrated = self.activation_scheme == "static" or ( + global_scale_dtype is not None and activation_format != "bf16" + ) + _set_optional_parameter( + self, "activation_scale", torch.tensor(1.0, dtype=torch.float32) if calibrated else None + ) + _set_optional_parameter(self, "bias", torch.empty(self.out_features) if self.has_bias else None) + + @classmethod + def from_module( + cls, module: nn.Linear, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedLinear: + return cls( + module.in_features, + module.out_features, + quantization_config=quantization_config, + has_bias=module.bias is not None, + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + two_level = self.weight_global_scale is not None + return finegrained_linear( + input, + self.weight, + self.weight_scale_inv, + block_size=self.block_size, + activation_scale=None if two_level else self.activation_scale, + bias=self.bias, + allow_deepgemm=not self._deepgemm_disabled, + weight_global_scale=self.weight_global_scale, + input_global_scale=self.activation_scale if two_level else None, + activation_scale_ub=getattr(self.quantization_config, "activation_scale_ub", None), + activation_format=self.activation_format, + ) + + +class FineGrainedEmbedding(_FineGrainedModule, nn.Embedding): + """FP8 or NVFP4 table decoded only for the rows selected by a lookup.""" + + def __init__( + self, + num_embeddings: int, + embedding_dim: int, + padding_idx: int | None = None, + *, + quantization_config: BlockScaledConfig | None = None, + dtype: torch.dtype | None = None, + ): + nn.Module.__init__(self) + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + if padding_idx is not None: + if not -num_embeddings <= padding_idx < num_embeddings: + raise ValueError("padding_idx must be within num_embeddings.") + padding_idx %= num_embeddings + self.padding_idx = padding_idx + self.max_norm = None + self.norm_type = 2.0 + self.scale_grad_by_freq = False + self.sparse = False + self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) + if self.weight_format not in ("fp8", "nvfp4") or ( + self.weight_format == "fp8" + and self.quantization_config.weight_block_size is not None + and not self.quantization_config.modules_to_convert + ): + raise ValueError("Fine-grained embeddings support per-tensor FP8 or NVFP4 weights.") + self.compute_dtype = dtype or torch.get_default_dtype() + if self.weight_format == "nvfp4": + self._hf_quantized_embedding_group_size = 16 + if embedding_dim % 16: + raise ValueError("NVFP4 embedding width must be divisible by its scale group size (16).") + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim // 2, dtype=torch.int8), requires_grad=False + ) + self.weight_scale_inv = nn.Parameter( + torch.empty(num_embeddings, embedding_dim // 16, dtype=_FP8_DTYPE), requires_grad=False + ) + self.weight_global_scale = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) + else: + self.weight = nn.Parameter( + torch.empty(num_embeddings, embedding_dim, dtype=_FP8_DTYPE), requires_grad=False + ) + self.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32), requires_grad=False) + + @classmethod + def from_module( + cls, module: nn.Embedding, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedEmbedding: + return cls( + module.num_embeddings, + module.embedding_dim, + module.padding_idx, + quantization_config=quantization_config, + dtype=module.weight.dtype, + ) + + def forward(self, input: torch.Tensor) -> torch.Tensor: + if self.weight.device.type in ("cuda", "xpu") and not ( + self.weight_format == "fp8" and torch.compiler.is_compiling() + ): + scale = self.weight_scale if self.weight_format == "fp8" else self.weight_scale_inv + global_scale = None if self.weight_format == "fp8" else self.weight_global_scale + if input.device == self.weight.device == scale.device and ( + global_scale is None or global_scale.device == self.weight.device + ): + dtype = scale.dtype if global_scale is None else self.compute_dtype + return load_finegrained_kernel().embedding_lookup(self.weight, input, scale, global_scale, dtype) + if self.weight_format == "fp8": + rows = super().forward(input) + return rows.to(self.weight_scale.dtype) * self.weight_scale + packed = nn.functional.embedding(input, self.weight).view(torch.uint8) + scales = nn.functional.embedding(input.to(self.weight_scale_inv.device), self.weight_scale_inv).to( + device=packed.device, dtype=torch.float32 + ) + low, high = packed & 15, packed >> 4 + codes = torch.stack((low, high), dim=-1).flatten(-2).long() + values = torch.tensor( + (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0), + device=packed.device, + ) + rows = values[codes] * scales.repeat_interleave(16, dim=-1) * self.weight_global_scale.to(packed.device) + return rows.to(self.compute_dtype) + + +class FineGrainedGroupedLinear(FineGrainedLinear): + """Quantized drop-in for block-diagonal grouped linears. NOT the MoE experts path + (`FineGrainedExperts`): every group runs on every token, with no routing. + + The underlying nn.Linear stores a single `(n_groups * out_per_group, in_per_group)` weight; + logically that is `n_groups` independent `(out_per_group, in_per_group)` sub-matrices, each + consuming a disjoint slice of the input's last-but-one dim. Forward takes `(..., n_groups, + in_per_group)` and returns `(..., n_groups, out_per_group)` — the same contract as the bf16 + grouped linear it replaces.""" + + def __init__( + self, + in_features_per_group: int, + out_features: int, + n_groups: int, + *, + quantization_config: BlockScaledConfig | None = None, + has_bias: bool = False, + ): + super().__init__( + in_features_per_group, out_features, quantization_config=quantization_config, has_bias=has_bias + ) + self.n_groups = n_groups + + @classmethod + def from_module( + cls, module: nn.Linear, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedGroupedLinear: + return cls( + in_features_per_group=module.in_features, + out_features=module.out_features, + n_groups=module.n_groups, + quantization_config=quantization_config, + has_bias=module.bias is not None, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + input_shape = x.shape[:-2] + hidden_dim = x.shape[-1] + + # the stored K, which packed FP4 halves + w = self.weight.view(self.n_groups, -1, self.weight.shape[-1]) + scale_inv = self.weight_scale_inv + x = x.movedim(-2, 0).reshape(-1, hidden_dim) + if scale_inv.numel() == 1: + # A per-tensor checkpoint calibrates the whole stored matrix, so every + # group uses the same scale. The grouped kernel expects an expert axis. + scale_inv = scale_inv.reshape(1, 1, 1).expand(self.n_groups, 1, 1).contiguous() + else: + scale_inv = scale_inv.view(self.n_groups, scale_inv.size(0) // self.n_groups, scale_inv.size(1)) + + tokens_per_group = x.size(0) // self.n_groups + # Kernels pad the expert schedule to a power of two. Empty trailing groups + # point at the final row, including when n_groups itself is not a power of two. + padded_groups = 1 << (self.n_groups - 1).bit_length() + expert_start = ( + torch.arange(padded_groups + 1, device=x.device, dtype=torch.int32).clamp_max(self.n_groups) + * tokens_per_group + ) + + two_level = self.weight_global_scale is not None + kernel = load_finegrained_kernel() + y = kernel.matmul_grouped( + x, + w, + None if two_level else self.activation_scale, + scale_inv, + expert_start=expert_start, + activation_format=self.activation_format, + a_global_scale=self.activation_scale if two_level else None, + b_global_scale=self.weight_global_scale, + activation_scale_ub=getattr(self.quantization_config, "activation_scale_ub", None), + ) + y = y.reshape(self.n_groups, *input_shape, -1).movedim(0, -2) + if self.has_bias: + y.add_(self.bias.view(self.n_groups, -1)) + return y + + +def _assert_interleaved_gate_up(module: nn.Module) -> None: + """The kernels' fused chains read gate|up interleaved. A module loaded for a DeepGEMM backend holds it + stacked; switched to a triton backend after load, the GLU would pair the wrong rows.""" + if module.has_gate and not module.holds_interleaved_gate_up: + raise NotImplementedError( + "The triton experts backends need gate|up interleaved, but this module was loaded stacked for a " + "DeepGEMM backend. Switching backends after load is not supported — pass `experts_implementation=...` " + "to `from_pretrained` so the model loads for it." + ) + + +def finegrained_batched_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: + """Batched (decode) experts forward: one program per routed row.""" + _assert_interleaved_gate_up(self) + kernel = load_finegrained_kernel() + return kernel.moe_fused_batched(hidden_states, top_k_index, top_k_weights, **self._moe_operands(kernel)) + + +def finegrained_grouped_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) -> torch.Tensor: + """Grouped (prefill) experts forward: one on-device routing pass, then the chain over the + expert-sorted schedule, scattering back to routed rows (EP-sentinel rows skipped).""" + _assert_interleaved_gate_up(self) + kernel = load_finegrained_kernel() + return kernel.moe_fused_grouped(hidden_states, top_k_index, top_k_weights, **self._moe_operands(kernel)) + + +class FineGrainedExperts(_FineGrainedModule, nn.Module): + # the plan kinds an `_experts_implementation` needs in place of the impl-agnostic default + # (megamoe wants no gradient-sync hooks and an EP `process_group`); extend when adding an impl + _impl_tp_layer_overrides: dict[str, dict[str, str]] = { + "deepgemm_megamoe": { + "moe_tp_experts": "megamoe_experts", + "ep_dispatch_experts": "megamoe_experts", + "ep_router": "megamoe_router", + }, + } + # a projection's tensors, by the suffix they carry on the module: the weight, its scale grid, the NVFP4 + # weight global, the calibrated activation scale, a bias + _projection_slots = ("", *STACK_COMPANIONS, "_bias") + + def __init__( + self, + config, + *, + quantization_config: BlockScaledConfig | None = None, + has_gate: bool = True, + has_bias: bool = False, + is_concatenated: bool = True, + ): + super().__init__() + + self.config = config + self.swiglu_alpha = getattr(config, "swiglu_alpha", None) + self.swiglu_limit = getattr(config, "swiglu_limit", None) + self.act_fn_name = _first_attr(config, "hidden_activation", "hidden_act") + self.hidden_dim = _first_attr(config, "moe_hidden_size", "hidden_size") + self.num_experts = _first_attr(config, "num_local_experts", "num_experts") + self.intermediate_dim = _first_attr(config, "moe_intermediate_size", "intermediate_size") + self.act_fn = ACT2FN[self.act_fn_name] + + self.quantization_config = quantization_config or BlockFp8Config(weight_block_size=None) + activation_format = self.activation_format + format_spec, scale_dtype, scale_group = resolve_weight_format( + self.weight_format, self.quantization_config.scale_fmt, self.quantization_config.weight_block_size + ) + self.global_scale_dtype = format_spec.global_scale_dtype + + # Checkpoint and module layouts may differ: conversion ops normalize stacked or interleaved + # gate/up rows and swizzle scales at load. The triton and eager paths use interleaved rows; + # DeepGEMM paths use stacked rows. + self.has_gate = has_gate + self.has_bias = has_bias + self.is_concatenated = is_concatenated + impl = getattr(config, "_experts_implementation", None) + self.holds_interleaved_gate_up = self.has_gate and impl not in _DEEPGEMM_EXPERTS_IMPLEMENTATIONS + swizzled = _holds_swizzled_scales(config, format_spec, activation_format) + self.gate_up_name = "gate_up_proj" if self.has_gate else "up_proj" + up_rows = (2 if self.has_gate else 1) * self.intermediate_dim + storage = (format_spec, scale_dtype, scale_group, swizzled) + self._register_projection(self.gate_up_name, up_rows, self.hidden_dim, 2 if self.has_gate else 1, storage) + self._register_projection("down_proj", self.hidden_dim, self.intermediate_dim, 1, storage) + + self.post_expert_norm = None + self.has_post_expert_norm = False + self.ships_transposed = False + + # optimization knob, set by `register_shared_gate_up_input_globals`: one gate_up activation global + self.register_buffer("shared_gate_up_input_global", None, persistent=False) + + @classmethod + def from_module( + cls, module: nn.Module, quantization_config: BlockScaledConfig, model: nn.Module + ) -> FineGrainedExperts: + """A `FineGrainedExperts` holding `module`'s weights, in the model's own gate|up convention. + + The flags travel twice over: `use_experts_implementation` stamps them on the instance after + `__init__`, so the class it builds needs them as well as the constructor. The model's per-expert + output norm is the model's own module, so it is taken from the experts this replaces. + """ + flags = { + "has_gate": getattr(module, "has_gate", True), + "has_bias": getattr(module, "has_bias", False), + "is_concatenated": getattr(module, "is_concatenated", True), + } + experts_class = _experts_class_with(cls, **flags) + config = getattr(module, "config", None) or model.config.get_text_config() + experts = experts_class(config=config, quantization_config=quantization_config, **flags) + experts.post_expert_norm = getattr(module, "post_expert_norm", None) + experts.has_post_expert_norm = getattr(module, "has_post_expert_norm", False) + # a full-precision stack quantized on load arrives in the layout of the module this replaces + experts.ships_transposed = getattr(module, "is_transposed", False) + return experts + + def _register_projection(self, proj, rows, in_dim, min_scale_out, storage): + """One expert projection: the packed weight and its scale grid, then the slots a + checkpoint may or may not fill — a bias, the NVFP4 second-level globals, a calibrated + activation scale. `storage` is the layout `resolve_weight_format` settled on.""" + format_spec, scale_dtype, scale_group, swizzled = storage + weight, scale = _alloc_expert_proj( + self.num_experts, rows, in_dim, format_spec, scale_dtype, scale_group, min_scale_out, swizzled + ) + self.register_parameter(proj, weight) + self.register_parameter(f"{proj}_scale_inv", scale) + + # the model dtype (the default dtype under `from_pretrained`), like the bf16 experts it + # replaces: the kernels add it on the fp32 accumulator, and a save keeps the checkpoint dtype + bias = torch.empty(self.num_experts, rows) if self.has_bias else None + # NVFP4 two-level: the fp32 globals the kernels recover on the accumulator, one per expert + # for the weight (`FineGrainedWeightGlobals` merges a separately calibrated gate|up stack + # down to that). The calibrated activation scale is the checkpoint's `input_scale`: a static FP8 scale, + # or NVFP4's activation global. + two_level = self.global_scale_dtype is not None + calibrated = self.activation_scheme == "static" or (two_level and self.activation_format != "bf16") + weight_global = torch.ones(self.num_experts, dtype=self.global_scale_dtype) if two_level else None + activation = torch.ones(self.num_experts, dtype=torch.float32) if calibrated else None + _set_optional_parameter(self, f"{proj}_bias", bias) + _set_optional_parameter(self, f"{proj}_weight_global_scale", weight_global) + _set_optional_parameter(self, f"{proj}_activation_scale", activation) + + def _apply_gate(self, gate_up: torch.Tensor) -> torch.Tensor: + # interleaved rows give alternating output columns, the split the fused epilogue does; stacked rows two halves + if self.holds_interleaved_gate_up: + gate, up = gate_up[..., 0::2], gate_up[..., 1::2] + else: + gate, up = gate_up.chunk(2, dim=-1) + if self.swiglu_limit is not None: + gate = gate.clamp(max=self.swiglu_limit) + up = up.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + if self.swiglu_alpha is not None: + # SwiGLU-OAI gate (same math as the model's non-quantized experts) + glu = gate * torch.sigmoid(gate * self.swiglu_alpha) + return (up + 1.0) * glu + return self.act_fn(gate) * up + + def forward( + self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor + ) -> torch.Tensor: + # index_add_ will accumulate using the dtype of the tensor we write into + # so we use float32 for the accumulation to avoid numerical issues in bf16/fp16 + final_hidden_states = torch.zeros_like(hidden_states, dtype=torch.float32) + + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts + 1) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero(as_tuple=False).view(-1) + + # one sync here instead of one per iteration: indexing with a 0-dim cuda tensor syncs + for expert_idx in expert_hit.tolist(): + if expert_idx == self.num_experts: + continue + + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + proj_out = self.expert_linear(hidden_states[token_idx], self.gate_up_name, expert_idx) + proj_out = self._apply_gate(proj_out) if self.has_gate else self.act_fn(proj_out) + proj_out = self.expert_linear(proj_out, "down_proj", expert_idx) + if self.has_post_expert_norm: + proj_out = self.post_expert_norm(proj_out) + routing_weights = top_k_weights[token_idx, top_k_pos, None] + weighted_out = proj_out * routing_weights.to(proj_out.dtype) + final_hidden_states.index_add_(0, token_idx, weighted_out.to(final_hidden_states.dtype)) + return final_hidden_states.to(hidden_states.dtype) + + def _projection_operands(self, proj: str) -> dict[str, torch.Tensor | None]: + """``proj``'s tensors by the kernels' slot names: the calibrated activation scale goes in as NVFP4's + activation global (the gate_up's shared one where it is set) or as a static scale.""" + operands = {slot: getattr(self, f"{proj}{slot}") for slot in self._projection_slots} + activation = operands.pop("_activation_scale") + if proj == self.gate_up_name and self.shared_gate_up_input_global is not None: + activation = self.shared_gate_up_input_global + two_level = self.global_scale_dtype is not None + operands["_input_global_scale"] = activation if two_level else None + operands["_activation_scale"] = None if two_level else activation + return operands + + def expert_linear(self, input: torch.Tensor, proj: str, expert_idx: int) -> torch.Tensor: + """One expert's ``proj`` as a dense linear, over that expert's slice of the projection's operands.""" + operands = self._projection_operands(proj) + weight, scale, weight_global, input_global, activation_scale, bias = ( + operands[slot] + for slot in ("", "_scale_inv", "_weight_global_scale", "_input_global_scale", "_activation_scale", "_bias") + ) + weight = weight[expert_idx] + # one expert's slice of a swizzled stack is the `(1, ...)` artifact `matmul_2d` reads + scale = scale[expert_idx : expert_idx + 1] if scale.ndim == 5 else scale[expert_idx] + # None where the checkpoint fills no such slot + weight_global, activation_scale, bias = ( + operand[expert_idx] if operand is not None else None for operand in (weight_global, activation_scale, bias) + ) + # one per expert, or one shared by all + if input_global is not None and input_global.numel() > 1: + input_global = input_global[expert_idx : expert_idx + 1] + + linear = functools.partial( + finegrained_linear, + block_size=self.block_size, + activation_scale=activation_scale, + allow_deepgemm=not self._deepgemm_disabled, + weight_global_scale=weight_global, + input_global_scale=input_global, + activation_format=self.activation_format, + activation_scale_ub=getattr(self.quantization_config, "activation_scale_ub", None), + ) + # Separately calibrated gate/up grids interleave per block, rather than per weight row: + # run each half with its own blocks, then interleave the output columns. A single scale row + # is shared by both halves and can be read directly. + if ( + proj == self.gate_up_name + and self.holds_interleaved_gate_up + and scale.ndim == 2 + and 1 < scale.shape[0] < weight.shape[0] + ): + halves = [ + linear( + input, + weight[half::2].contiguous(), + scale[half::2].contiguous(), + bias=bias[half::2] if bias is not None else None, + ) + for half in (0, 1) + ] + return torch.stack(halves, dim=-1).flatten(-2) + return linear(input, weight, scale, bias=bias) + + def _moe_operands(self, kernel) -> dict: + """Everything the kernels' MoE forwards take from the module: the two projections with their + scales (held swizzled for dot_scaled chains), the per-expert globals and biases, + and the activation — a ``get_supported_act_fns()`` name the gate_up epilogue fuses, else the + module's own GLU as a callable the kernels run on the host, so any activation works without a + kernel change. + + A method rather than a function so an adapter library can extend it per module: PEFT overrides + it to add its LoRA factors, which the kernels then apply on the routed rows between the GEMMs + (`expert_linear` is the eager loop's counterpart).""" + act_name = self.act_fn_name + fused = act_name in kernel.get_supported_act_fns() and (self.swiglu_alpha is None or act_name == "silu") + if fused: + act_fn = act_name + else: + act_fn = self._apply_gate if self.has_gate else self.act_fn + + # A named norm is fused into the reduce; anything else runs as a callable on the routed rows. + # Fusing is off under intra-expert TP: those rows are a partial sum, so the all-reduce the + # module's wrapped forward does has to happen first (`_hf_tp_input_reduce` marks it). + post_expert_norm, norm_weight, norm_eps = None, None, 1e-6 + if self.has_post_expert_norm: + norm_name = SUPPORTED_FUSED_NORMS.get(type(self.post_expert_norm)) + fusable = norm_name in kernel.get_supported_norms() + if not fusable or getattr(self.post_expert_norm, "_hf_tp_input_reduce", False): + post_expert_norm = self.post_expert_norm + else: + norm = self.post_expert_norm + post_expert_norm, norm_weight = norm_name, norm.weight + # the model owns this module, so take epsilon under either name; `nn.RMSNorm` leaves + # it None and resolves to the dtype's own + eps = _first_attr(norm, "eps", "variance_epsilon", raise_error=False) + norm_eps = float(eps) if eps is not None else torch.finfo(norm.weight.dtype).eps + + return { + **{ + f"{kernel_proj}{slot}": operand + for kernel_proj, held_proj in (("gate_up_proj", self.gate_up_name), ("down_proj", "down_proj")) + for slot, operand in self._projection_operands(held_proj).items() + }, + # a supported activation NAME is fused into the gate_up epilogue; any other callable + # leaves that GEMM plain and runs on the host between the two + "act_fn": act_fn, + "gate": self.has_gate, + "swiglu_alpha": self.swiglu_alpha, + "swiglu_limit": self.swiglu_limit, + "activation_format": self.activation_format, + "activation_scale_ub": getattr(self.quantization_config, "activation_scale_ub", None), + # a supported norm NAME is folded into the routing-weighted reduce, one pass ahead of it + # computing the rows' rsqrt; any other callable is applied to the routed rows before it + "post_expert_norm": post_expert_norm, + "post_expert_norm_weight": norm_weight, + "post_expert_norm_eps": norm_eps, + } + + +@functools.cache +def _experts_class_with(cls, *, has_gate: bool, has_bias: bool, is_concatenated: bool) -> type: + """`cls` with the experts dispatch and these flags: `use_experts_implementation` wraps the class it is handed in + place, so each flag set wraps a subclass of its own, once, and `cls` itself stays as defined. The subclass is + registered on `cls`'s module under its own name, so a pickled model finds it.""" + qualname = f"{cls.__qualname__}_gate{int(has_gate)}_bias{int(has_bias)}_concatenated{int(is_concatenated)}" + subclass = type(cls.__name__, (cls,), {"__module__": cls.__module__, "__qualname__": qualname}) + setattr(sys.modules[cls.__module__], qualname, subclass) + return use_experts_implementation( + experts_class=subclass, + experts_interface=ALL_FINEGRAINED_EXPERTS_FUNCTIONS, + has_gate=has_gate, + has_bias=has_bias, + is_concatenated=is_concatenated, + ) + + +class FineGrainedExpertsInterface(ExpertsInterface): + """Interface for registering custom finegrained experts forward functions.""" + + _global_mapping = { + "deepgemm": deepgemm_fp8_fp4_experts_forward, + "batched_mm": finegrained_batched_mm_experts_forward, + "grouped_mm": finegrained_grouped_mm_experts_forward, + "deepgemm_megamoe": deepgemm_fp8_fp4_megamoe_experts_forward, + } + + +ALL_FINEGRAINED_EXPERTS_FUNCTIONS = FineGrainedExpertsInterface() + + +def assert_modules_are_quantized(model: nn.Module) -> None: + """Every finegrained module must hold a quantized weight once the checkpoint is in. + + A full-precision weight in a quantized slot means the swap and the checkpoint disagree: the + module was allocated quantized storage and the loader replaced it with what the checkpoint + actually ships. Serving that in full precision would hand back a model the caller believes is + quantized and is not, so it is a load failure. Checked once here rather than per forward — + the dtype cannot change between them. + """ + for name, module in model.named_modules(): + if not isinstance(module, _FineGrainedModule): + continue + for attr in ("weight", *FUSED_STACKS): + weight = getattr(module, attr, None) + if weight is not None and weight.element_size() > 1: + raise ValueError( + f"{name}.{attr} was converted for quantized compute but holds a {weight.dtype} " + "weight — this checkpoint does not quantize it. Name the module in the " + "quantization config's `modules_to_not_convert` to leave it as it is." + ) + + +def register_shared_gate_up_input_globals(model: nn.Module) -> None: + """Share the gate_up's NVFP4 activation global across experts, as the max of their calibrated ones, so + its rows quantize once per token before routing and the GEMM gathers them. Per-expert globals, the + calibration exactly, quantize per routed row (top_k times the rows), which the GEMM then reads + materialized: 3-11% slower on a GLM NVFP4 MoE layer at prefill, the most at small batches. Decode is + unaffected. A max never clips, but experts below it lose block-scale precision. + + ``TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL`` turns it on and picks the experts the max + runs over under expert parallelism: ``global`` all of them, so the numerics do not depend on the EP + size, or ``rank`` this rank's, a tighter global. Unset keeps every expert's own.""" + scope = os.environ.get("TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL") + if scope is None: + return + if scope not in ("global", "rank"): + raise ValueError( + f"TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL takes `global` or `rank`, got {scope!r}" + ) + for module in model.modules(): + if not isinstance(module, FineGrainedExperts): + continue + input_global = getattr(module, f"{module.gate_up_name}_activation_scale") + if input_global is not None: + if hasattr(input_global, "full_tensor"): # sharded over experts + input_global = input_global.full_tensor() if scope == "global" else input_global.to_local() + module.shared_gate_up_input_global = input_global.detach().amax().reshape(1) + + +def disable_deepgemm_on_multi_device(model: nn.Module) -> None: + """Flag every quantized module to skip DeepGEMM when the model spans >1 CUDA device in one + process. DeepGEMM loads each kernel via `cuKernelGetFunction`, which binds the `CUfunction` + handle to the CUDA context live at load time; driving that cached handle from another device + launches it against the wrong context and produces garbage. (Build-time fix: compile DeepGEMM + with `DG_JIT_USE_RUNTIME_API=1` for a context-free `cudaKernel_t` loader; until our wheel picks + that up we avoid single-process multi-device.) Setting `_deepgemm_disabled` routes both the + linear and experts paths through our triton kernels. A model that fits on one device keeps + DeepGEMM even with other GPUs visible; TP/EP put one device per process, so this is a no-op + there.""" + quantized_modules = [m for m in model.modules() if isinstance(m, _FineGrainedModule)] + cuda_devices = set() + for m in quantized_modules: + param = next(m.parameters(), None) + if param is not None and param.device.type == "cuda": + cuda_devices.add(param.device.index) + if len(cuda_devices) <= 1: + return + for m in quantized_modules: + m._deepgemm_disabled = True + logger.warning_once( + "This finegrained quantized model spans multiple CUDA devices in one process; routing its linear and experts " + "layers through our triton kernels instead of DeepGEMM (DeepGEMM's cached kernels are bound to a " + "single CUDA context and corrupt across devices). Run tensor/expert parallel (one device per " + "process) to use the faster DeepGEMM path." + ) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 9b9a11e52fb3..5222dcb38d0f 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -15,6 +15,7 @@ import functools import os +import warnings from collections.abc import Callable from dataclasses import dataclass @@ -38,6 +39,14 @@ from .moe import ExpertsInterface, use_experts_implementation +warnings.warn( + "`transformers.integrations.finegrained_fp8` is frozen and will be removed in a future release; " + "`BlockFp8Config` and `Mxfp8Config` load through `transformers.integrations.finegrained`.", + DeprecationWarning, + stacklevel=2, +) + + logger = logging.get_logger(__name__) diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index 7aba365f3fcd..10c653f10d99 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -799,9 +799,15 @@ def register_kernel_mapping_transformers(*args, **kwargs): _HUB_KERNEL_MAPPING: dict[str, dict[str, str]] = { - "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, - "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2}, + # DeepGEMM JIT-compiles its CUDA kernels for the running device; the build metadata declares only + # the arch it was packaged on (9.0a), so the `kernels` arch check would wrongly reject sm_100. + "deep-gemm": {"repo_id": "kernels-community/deep-gemm", "version": 2, "check_arch": False}, "sonic-moe": {"repo_id": "kernels-community/sonic-moe", "revision": "ep-support"}, + "finegrained-fp8": {"repo_id": "kernels-community/finegrained-fp8", "version": 4}, + # the multi-recipe superset of finegrained-fp8 (block-FP8, MXFP8, MXFP4, NVFP4, weight-only); + # exports the same three matmuls plus the fused-MoE forwards and per-recipe quant helpers + "finegrained-kernels": {"repo_id": "kernels-community/finegrained-kernels", "version": 0}, + # dense NVFP4 GEMM behind `NVFP4Linear`; MoE experts route to finegrained-kernels instead "nvfp4": {"repo_id": "kernels-community/nvfp4-gemm", "version": 1}, } @@ -912,11 +918,18 @@ def lazy_load_kernel(kernel_name: str, mapping: dict[str, ModuleType | None] = _ repo_id = _HUB_KERNEL_MAPPING[kernel_name]["repo_id"] revision = _HUB_KERNEL_MAPPING[kernel_name].get("revision", None) version = _HUB_KERNEL_MAPPING[kernel_name].get("version", None) + check_arch = _HUB_KERNEL_MAPPING[kernel_name].get("check_arch", True) # Default version as it's mandatory if version is None and revision is None: version = 1 - kernel = get_kernel(repo_id, revision=revision, version=version, allow_all_kernels=ALLOW_ALL_KERNELS) + kernel = get_kernel( + repo_id, + revision=revision, + version=version, + allow_all_kernels=ALLOW_ALL_KERNELS, + check_arch=check_arch, + ) mapping[kernel_name] = kernel except FileNotFoundError as e: mapping[kernel_name] = None @@ -975,6 +988,7 @@ def get_kernel( revision: str | None = None, version: int | str | None = None, allow_all_kernels: bool = False, + check_arch: bool = True, ) -> ModuleType: from .. import __version__ @@ -983,7 +997,12 @@ def get_kernel( user_agent = {"framework": "transformers", "version": __version__, "repo_id": kernel_name} return get_kernel_hub( - kernel_name, revision=revision, version=version, user_agent=user_agent, trust_remote_code=allow_all_kernels + kernel_name, + revision=revision, + version=version, + user_agent=user_agent, + trust_remote_code=allow_all_kernels, + check_arch=check_arch, ) diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 749e40117f78..ff2f36efcbe9 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -12,6 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. + +import warnings + from ..utils import get_current_accelerator, is_torch_available, logging @@ -24,8 +27,17 @@ from ..quantizers.quantizers_utils import get_module_from_name, on_device, should_convert_module +warnings.warn( + "`transformers.integrations.mxfp4` is frozen and will be removed in a future release; " + "`Mxfp4Config` now loads through `transformers.integrations.finegrained`.", + DeprecationWarning, + stacklevel=2, +) + + logger = logging.get_logger(__name__) + FP4_VALUES = [ +0.0, +0.5, diff --git a/src/transformers/integrations/nvfp4.py b/src/transformers/integrations/nvfp4.py index 476e459df405..daa10c4bb43a 100644 --- a/src/transformers/integrations/nvfp4.py +++ b/src/transformers/integrations/nvfp4.py @@ -14,6 +14,7 @@ from __future__ import annotations import functools +import warnings from collections.abc import Callable from dataclasses import dataclass @@ -28,6 +29,14 @@ from .hub_kernels import lazy_load_kernel +warnings.warn( + "`transformers.integrations.nvfp4` is frozen and will be removed in a future release; " + "`NVFP4Config` now loads through `transformers.integrations.finegrained`.", + DeprecationWarning, + stacklevel=2, +) + + if is_torch_available(): import torch import torch.nn as nn diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 26458230c6c4..f09ad00df01a 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -74,7 +74,7 @@ ) from .integrations.deepspeed import _load_state_dict_into_zero3_model from .integrations.eager_paged import eager_paged_attention_forward -from .integrations.finegrained_fp8 import ALL_FP8_EXPERTS_FUNCTIONS +from .integrations.finegrained import ALL_FINEGRAINED_EXPERTS_FUNCTIONS from .integrations.flash_attention import flash_attention_forward from .integrations.flex_attention import flex_attention_forward from .integrations.hub_kernels import allow_all_hub_kernels, is_kernel, kernelize @@ -1881,7 +1881,9 @@ def get_correct_attn_implementation(self, requested_attention: str | None, is_in def get_correct_experts_implementation(self, requested_experts: str | None) -> str: applicable_experts = "grouped_mm" if requested_experts is None else requested_experts - base_experts_fns = ["eager"] + list(set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FP8_EXPERTS_FUNCTIONS.keys())) + base_experts_fns = ["eager"] + list( + set(ALL_EXPERTS_FUNCTIONS.keys()) | set(ALL_FINEGRAINED_EXPERTS_FUNCTIONS.keys()) + ) valid_experts_str_list = [f'`experts_implementation="{fn}"`' for fn in base_experts_fns] valid_experts_str_list[-1] = "and " + valid_experts_str_list[-1] valid_experts_str = ", ".join(valid_experts_str_list) diff --git a/src/transformers/models/llama4/configuration_llama4.py b/src/transformers/models/llama4/configuration_llama4.py index 020bceaa7249..5c33c4dc6e60 100644 --- a/src/transformers/models/llama4/configuration_llama4.py +++ b/src/transformers/models/llama4/configuration_llama4.py @@ -117,8 +117,8 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.feed_forward.shared_expert.gate_proj": "colwise", "layers.*.feed_forward.shared_expert.up_proj": "colwise", "layers.*.feed_forward.shared_expert.down_proj": "rowwise", - "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise", # row because not linear - "layers.*.feed_forward.experts.down_proj": "colwise", # col because not linear + "layers.*.feed_forward.experts.gate_up_proj": "packed_colwise", + "layers.*.feed_forward.experts.down_proj": "rowwise", "layers.*.feed_forward.experts": "moe_tp_experts", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", @@ -129,8 +129,9 @@ class Llama4TextConfig(PreTrainedConfig): "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", - "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", # row because not linear - "layers.*.feed_forward.experts.down_proj": "grouped_gemm", # col because not linear + "layers.*.feed_forward.experts.gate_up_proj": "grouped_gemm", + "layers.*.feed_forward.experts.down_proj": "grouped_gemm", + "layers.*.feed_forward.experts": "moe_tp_experts", "layers.*.feed_forward.gate_proj": "colwise", "layers.*.feed_forward.up_proj": "colwise", "layers.*.feed_forward.down_proj": "rowwise", diff --git a/src/transformers/models/llama4/modeling_llama4.py b/src/transformers/models/llama4/modeling_llama4.py index 81ea67db7260..7d48966f2336 100644 --- a/src/transformers/models/llama4/modeling_llama4.py +++ b/src/transformers/models/llama4/modeling_llama4.py @@ -26,6 +26,8 @@ from ...activations import ACT2FN from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin +from ...integrations import use_experts_implementation +from ...integrations.moe import batched_mm_experts_forward from ...masking_utils import create_causal_mask, create_chunked_causal_mask from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_layers import GradientCheckpointingLayer @@ -51,36 +53,24 @@ logger = logging.get_logger(__name__) +@use_experts_implementation class Llama4TextExperts(nn.Module): def __init__(self, config: Llama4TextConfig): super().__init__() self.num_experts = config.num_local_experts - self.intermediate_size = config.intermediate_size - self.hidden_size = config.hidden_size - self.expert_dim = self.intermediate_size - self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim)) - self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size))) + self.intermediate_dim = config.intermediate_size + self.hidden_dim = config.hidden_size + self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim)) self.act_fn = ACT2FN[config.hidden_act] - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """ - This should really not be run on a single machine, as we are reaching compute bound: - - the inputs are expected to be "sorted" per expert already. - - the weights are viewed with another dim, to match num_expert, 1, shape * num_tokens, shape - - Args: - hidden_states (torch.Tensor): (batch_size * token_num, hidden_size) - selected_experts (torch.Tensor): (batch_size * token_num, top_k) - routing_weights (torch.Tensor): (batch_size * token_num, top_k) - Returns: - torch.Tensor - """ - hidden_states = hidden_states.view(self.gate_up_proj.shape[0], -1, self.hidden_size) - gate_up = torch.bmm(hidden_states, self.gate_up_proj) - gate, up = gate_up.chunk(2, dim=-1) # not supported for DTensors - next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj) - next_states = next_states.view(-1, self.hidden_size) - return next_states + def forward( + self, + hidden_states: torch.Tensor, + top_k_index: torch.Tensor, + top_k_weights: torch.Tensor, + ) -> torch.Tensor: + return batched_mm_experts_forward(self, hidden_states, top_k_index, top_k_weights) # Phi3MLP @@ -145,13 +135,11 @@ def __init__(self, config): def forward(self, hidden_states): router_logits = super().forward(hidden_states) - router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=1) - router_scores = torch.full_like(router_logits, float("-inf")).scatter_(1, router_indices, router_top_value) - router_scores = torch.nn.functional.sigmoid(router_scores.float()).to(router_scores.dtype) - return router_scores, router_logits + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) + router_scores = torch.sigmoid(router_top_value.float()).to(router_logits.dtype) + return router_logits, router_scores, router_indices -# @use_kernel_forward_from_hub("Llama4TextMoe") class Llama4TextMoe(nn.Module): def __init__(self, config): super().__init__() @@ -164,13 +152,14 @@ def __init__(self, config): def forward(self, hidden_states): hidden_states = hidden_states.reshape(-1, self.hidden_dim) - router_scores, router_logits = self.router(hidden_states) - routed_in = hidden_states.repeat(router_scores.shape[1], 1) - routed_in = routed_in * router_scores.transpose(0, 1).reshape(-1, 1) - routed_out = self.experts(routed_in) - out = self.shared_expert(hidden_states) - out.add_(routed_out.reshape(router_scores.shape[1], -1, routed_out.shape[-1]).sum(dim=0)) - return out, router_logits + router_logits, router_scores, router_indices = self.router(hidden_states) + # Llama4 weights each selected expert's input, before the nonlinear GLU. + routed_in = (hidden_states[:, None, :] * router_scores[:, :, None]).reshape(-1, self.hidden_dim) + routed_out = self.experts( + routed_in, router_indices.reshape(-1, 1), torch.ones_like(router_scores).reshape(-1, 1) + ) + routed_out = routed_out.reshape(-1, self.top_k, self.hidden_dim).sum(dim=1) + return self.shared_expert(hidden_states) + routed_out, router_logits # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Llama4Text @@ -480,7 +469,7 @@ class Llama4TextModel(Llama4PreTrainedModel): _can_record_outputs = { "attentions": Llama4TextAttention, "hidden_states": Llama4TextDecoderLayer, - "router_logits": OutputRecorder(Llama4Router, index=1), + "router_logits": OutputRecorder(Llama4Router, index=0), } def __init__(self, config: Llama4TextConfig): diff --git a/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py b/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py index 5f97f6bf0cfd..ec15ed4f7217 100644 --- a/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py +++ b/src/transformers/models/ministral3/convert_ministral3_weights_to_hf.py @@ -29,9 +29,8 @@ PixtralProcessor, PixtralVisionConfig, ) -from transformers.integrations.finegrained_fp8 import replace_with_fp8_linear from transformers.integrations.mistral import convert_tekken_tokenizer -from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig # fmt: off @@ -254,11 +253,11 @@ def convert_and_write_model(input_dir: str, output_dir: str, max_position_embedd else: raise ValueError(f"Unknown config type {type(config)}.") - # let's swap nn.Linear to FP8 Linear before loading + # swap nn.Linear for the quantized module before loading if hasattr(model.config, "quantization_config"): - model = replace_with_fp8_linear( - model, model.config.quantization_config.modules_to_not_convert, model.config.quantization_config - ) + qconfig = model.config.quantization_config + quantizer = AutoHfQuantizer.from_config(qconfig, pre_quantized=True) + quantizer.replace_modules(model, modules_to_not_convert=qconfig.modules_to_not_convert) model.load_state_dict(full_state_dict, strict=True, assign=True) model.save_pretrained(output_dir) diff --git a/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py b/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py index 87c212f06ec2..fcd59234b90e 100644 --- a/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py +++ b/src/transformers/models/mistral4/convert_mistral4_weight_to_hf.py @@ -37,10 +37,9 @@ MergeModulelist, WeightRenaming, ) -from transformers.integrations.finegrained_fp8 import replace_with_fp8_linear from transformers.integrations.mistral import convert_tekken_tokenizer from transformers.models.mistral4.modeling_mistral4 import Mistral4ForCausalLM -from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig _FP8_DTYPE = torch.float8_e4m3fn @@ -536,7 +535,8 @@ def convert_and_write_model( if output_fp8 and hasattr(model.config, "quantization_config"): qconfig = model.config.quantization_config - model = replace_with_fp8_linear(model, qconfig.modules_to_not_convert, qconfig) + quantizer = AutoHfQuantizer.from_config(qconfig, pre_quantized=True) + quantizer.replace_modules(model, modules_to_not_convert=qconfig.modules_to_not_convert) model.load_state_dict(full_state_dict, strict=True, assign=True) model.save_pretrained(str(output_dir)) diff --git a/src/transformers/models/muse_spark/configuration_muse_spark.py b/src/transformers/models/muse_spark/configuration_muse_spark.py index f5e607c2348a..46e8bd9f4442 100644 --- a/src/transformers/models/muse_spark/configuration_muse_spark.py +++ b/src/transformers/models/muse_spark/configuration_muse_spark.py @@ -105,6 +105,12 @@ class MuseSparkTextConfig(PreTrainedConfig): "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", + # Post-expert normalization requires complete expert outputs before ranks combine them. + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", + "layers.*.mlp.experts": "moe_tp_experts", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), @@ -140,6 +146,7 @@ class MuseSparkTextConfig(PreTrainedConfig): post_norm_eps: float = 1e-8 layer_rope_theta: list[float | int] | None = None base_model_ep_plan = { + "layers.*.mlp.gate": "replicate", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", diff --git a/src/transformers/models/muse_spark/modular_muse_spark.py b/src/transformers/models/muse_spark/modular_muse_spark.py index e8f5a997f869..81ecda602c14 100644 --- a/src/transformers/models/muse_spark/modular_muse_spark.py +++ b/src/transformers/models/muse_spark/modular_muse_spark.py @@ -110,8 +110,15 @@ class MuseSparkTextConfig(MuseGlimmerTextConfig): "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", + # Post-expert normalization requires complete expert outputs before ranks combine them. + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", + "layers.*.mlp.experts": "moe_tp_experts", } base_model_ep_plan = { + "layers.*.mlp.gate": "replicate", "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", "layers.*.mlp.experts.down_proj": "grouped_gemm", "layers.*.mlp.experts.post_expert_norm": "replicated_with_grad_allreduce", diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index f5dc3fe5d5d8..be8605c0c512 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -12,6 +12,8 @@ # 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. +import copy +import json import warnings from ..models.auto.configuration_auto import AutoConfig @@ -22,9 +24,12 @@ AwqConfig, BitNetQuantConfig, BitsAndBytesConfig, + BlockFp8Config, + BlockScaledConfig, CompressedTensorsConfig, EetqConfig, FbgemmFp8Config, + FineGrainedConfig, FineGrainedFP8Config, FourOverSixConfig, FPQuantConfig, @@ -35,6 +40,7 @@ HqqConfig, MetalConfig, Mxfp4Config, + Mxfp8Config, NVFP4Config, QuantizationConfigMixin, QuantizationMethod, @@ -44,18 +50,21 @@ SpQRConfig, TorchAoConfig, VptqConfig, + read_compressed_tensors, + read_modelopt, ) from .base import HfQuantizer from .quantizer_aqlm import AqlmHfQuantizer from .quantizer_auto_round import AutoRoundQuantizer from .quantizer_awq import AwqQuantizer from .quantizer_bitnet import BitNetHfQuantizer +from .quantizer_blockfp8 import BlockFp8HfQuantizer from .quantizer_bnb_4bit import Bnb4BitHfQuantizer from .quantizer_bnb_8bit import Bnb8BitHfQuantizer from .quantizer_compressed_tensors import CompressedTensorsHfQuantizer from .quantizer_eetq import EetqHfQuantizer from .quantizer_fbgemm_fp8 import FbgemmFp8HfQuantizer -from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer +from .quantizer_finegrained import FineGrainedHfQuantizer from .quantizer_fouroversix import FourOverSixHfQuantizer from .quantizer_fp_quant import FPQuantHfQuantizer from .quantizer_gemma import GemmaQuantizer @@ -65,6 +74,7 @@ from .quantizer_hqq import HqqHfQuantizer from .quantizer_metal import MetalHfQuantizer from .quantizer_mxfp4 import Mxfp4HfQuantizer +from .quantizer_mxfp8 import Mxfp8HfQuantizer from .quantizer_nvfp4 import NVFP4HfQuantizer from .quantizer_quanto import QuantoHfQuantizer from .quantizer_quark import QuarkHfQuantizer @@ -72,6 +82,7 @@ from .quantizer_spqr import SpQRHfQuantizer from .quantizer_torchao import TorchAoHfQuantizer from .quantizer_vptq import VptqHfQuantizer +from .quantizers_utils import EXPERTS_PATTERN AUTO_QUANTIZER_MAPPING = { @@ -93,16 +104,15 @@ "bitnet": BitNetHfQuantizer, "vptq": VptqHfQuantizer, "spqr": SpQRHfQuantizer, - "fp8": FineGrainedFP8HfQuantizer, "gguf": GgufHfQuantizer, - "nvfp4": NVFP4HfQuantizer, - # MXFP8 = FP8 (E4M3 weights) with per-block ``[1, 32]`` E8M0 (uint8) scales — - # reuses the FineGrainedFP8 dequant path, with the E8M0 byte→exponent - # unpacking handled inside ``Fp8Dequantize._dequantize_one``. - "mxfp8": FineGrainedFP8HfQuantizer, + "metal": MetalHfQuantizer, "auto-round": AutoRoundQuantizer, + # fine-grained: one quantizer per weight format, and one hosting several on the parts of a model + "fp8": BlockFp8HfQuantizer, + "mxfp8": Mxfp8HfQuantizer, "mxfp4": Mxfp4HfQuantizer, - "metal": MetalHfQuantizer, + "nvfp4": NVFP4HfQuantizer, + "finegrained": FineGrainedHfQuantizer, "sinq": SinqHfQuantizer, "gemma": GemmaQuantizer, } @@ -127,24 +137,33 @@ "bitnet": BitNetQuantConfig, "vptq": VptqConfig, "spqr": SpQRConfig, - "fp8": FineGrainedFP8Config, - "nvfp4": NVFP4Config, - "mxfp8": FineGrainedFP8Config, + "metal": MetalConfig, "auto-round": AutoRoundConfig, + "fp8": BlockFp8Config, + "mxfp8": Mxfp8Config, "mxfp4": Mxfp4Config, - "metal": MetalConfig, + "nvfp4": NVFP4Config, + "finegrained": FineGrainedConfig, + # a modelopt export reads through `read_modelopt` (see `AutoQuantizationConfig.from_dict`) as its format's config + # or a `FineGrainedConfig`; listed so the method counts as supported + "modelopt": FineGrainedConfig, "sinq": SinqConfig, "gemma": GemmaQuantizationConfig, } +# a passed config of these types reads a pre-quantized checkpoint's config as its own; the deprecated +# `FineGrainedFP8Config` so keeps loading through its frozen integration +CHECKPOINT_READING_CONFIG_TYPES = (AutoRoundConfig, CompressedTensorsConfig, FineGrainedFP8Config) + LOADING_ATTRIBUTES_CONFIG_TYPES = ( GPTQConfig, AwqConfig, AutoRoundConfig, FbgemmFp8Config, CompressedTensorsConfig, - Mxfp4Config, MetalConfig, + FineGrainedConfig, + BlockScaledConfig, FineGrainedFP8Config, ) @@ -169,6 +188,21 @@ def from_dict(cls, quantization_config_dict: dict): "The model's quantization config from the arguments has no `quant_method` attribute. Make sure that the model has been correctly quantized" ) + # a modelopt export reads as its one format's config, or a `FineGrainedConfig` when it mixes them + if quant_method == QuantizationMethod.MODELOPT: + return read_modelopt(quantization_config_dict) + if quant_method == QuantizationMethod.FBGEMM_FP8: + return BlockFp8Config( + weight_block_size="row", + activation_scale_ub=quantization_config_dict.get("activation_scale_ub", 1200.0), + modules_to_not_convert=quantization_config_dict.get("modules_to_not_convert"), + checkpoint_format="fbgemm", + ) + # so does a compressed-tensors checkpoint in formats these kernels serve + if quant_method == QuantizationMethod.COMPRESSED_TENSORS: + if (config := read_compressed_tensors(quantization_config_dict)) is not None: + return config + if quant_method not in AUTO_QUANTIZATION_CONFIG_MAPPING: raise ValueError( f"Unknown quantization type, got {quant_method} - supported types are:" @@ -203,6 +237,13 @@ def from_config(cls, quantization_config: QuantizationConfigMixin | dict, **kwar # Convert it to a QuantizationConfig if the q_config is a dict if isinstance(quantization_config, dict): quantization_config = AutoQuantizationConfig.from_dict(quantization_config) + elif isinstance(quantization_config, FbgemmFp8Config): + quantization_config = AutoQuantizationConfig.from_dict(quantization_config.to_dict()) + # the deprecated config keeps its own, frozen integration + if isinstance(quantization_config, FineGrainedFP8Config): + from .quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + + return FineGrainedFP8HfQuantizer(quantization_config, **kwargs) quant_method = quantization_config.quant_method @@ -251,8 +292,8 @@ def merge_quantization_configs( if isinstance(quantization_config, dict): # Convert the config based on the type of quantization_config_from_args (e.g., AutoRoundConfig), which takes priority before automatic configuration dispatch. - if isinstance(quantization_config_from_args, AutoRoundConfig): - quantization_config = AutoRoundConfig.from_dict(quantization_config) + if isinstance(quantization_config_from_args, CHECKPOINT_READING_CONFIG_TYPES): + quantization_config = type(quantization_config_from_args).from_dict(quantization_config) else: quantization_config = AutoQuantizationConfig.from_dict(quantization_config) @@ -265,20 +306,33 @@ def merge_quantization_configs( "Please make sure to pass the same quantization config class to `from_pretrained` with different loading attributes." ) + ignored = [] if isinstance(quantization_config, LOADING_ATTRIBUTES_CONFIG_TYPES) and isinstance( quantization_config_from_args, LOADING_ATTRIBUTES_CONFIG_TYPES ): loading_attr_dict = quantization_config_from_args.get_loading_attributes() + if isinstance(quantization_config_from_args, (BlockScaledConfig, FineGrainedConfig)): + ignored = _ignored_options(quantization_config, quantization_config_from_args, loading_attr_dict) for attr, val in loading_attr_dict.items(): + if attr == "dtype_plan" and isinstance(quantization_config, FineGrainedConfig): + val = _with_activation_formats(quantization_config.dtype_plan, val) + elif attr == "activation_format": + _check_calibrated(quantization_config, val) setattr(quantization_config, attr, val) if loading_attr_dict: warning_msg += f"However, loading attributes (e.g. {list(loading_attr_dict.keys())}) will be overwritten with the one you passed to `from_pretrained`. The rest will be ignored." + if ignored: + warning_msg += f" Ignored: {ignored}." - if warning_msg != "" and not isinstance(quantization_config, (Mxfp4Config, MetalConfig, FineGrainedFP8Config)): + # the block-scaled formats take a passed config for its loading attributes, so that alone is not worth a warning + lends_only = ( + isinstance(quantization_config, (MetalConfig, FineGrainedConfig, BlockScaledConfig, FineGrainedFP8Config)) + and not ignored + ) + if warning_msg != "" and not lends_only: warnings.warn(warning_msg) else: - # in the case of mxfp4, we don't want to print the warning message, bit confusing for users logger.info(warning_msg) return quantization_config @@ -335,6 +389,107 @@ def register_quantizer_fn(cls): return register_quantizer_fn +def _as_json(config) -> dict: + """`config` as `config.json` holds it, so a tuple and the list it serializes to compare equal.""" + return json.loads(config.to_json_string(use_diff=False)) + + +def _set_options(config) -> dict: + """The options `config` sets away from its class's defaults.""" + defaults = _as_json(type(config)()) + return {name: value for name, value in _as_json(config).items() if defaults.get(name) != value} + + +def _ignored_options(quantization_config, quantization_config_from_args, loading_attributes: dict) -> list[str]: + """The options a config passed for a pre-quantized checkpoint sets away from its defaults, differently from the + checkpoint's own, and that are not loading attributes, so the checkpoint's stand; a passed plan's, per entry.""" + own, passed = _as_json(quantization_config), _set_options(quantization_config_from_args) + ignored = [ + name + for name, value in passed.items() + if name not in loading_attributes and name != "quant_method" and own.get(name) != value + ] + for pattern, entry in loading_attributes.get("dtype_plan", {}).items(): + own_entry = quantization_config.dtype_plan.get(pattern) + if type(own_entry) is not type(entry): + ignored.append(f"dtype_plan[{pattern!r}]") + continue + own_options = _as_json(own_entry) + ignored += [ + f"dtype_plan[{pattern!r}].{name}" + for name, value in _set_options(entry).items() + if name != "activation_format" and own_options.get(name) != value + ] + return ignored + + +def _with_activation_formats(dtype_plan: dict, passed_plan: dict) -> dict: + """`dtype_plan` with the activation format of each `passed_plan` entry naming the same format under the same + pattern: the kernels the same weights run on.""" + merged = dict(dtype_plan) + for pattern, entry in passed_plan.items(): + if type(merged.get(pattern)) is type(entry) and entry._names_activation_format: + _check_calibrated(merged[pattern], entry.activation_format) + merged[pattern] = copy.copy(merged[pattern]) + merged[pattern].activation_format = entry.activation_format + return merged + + +def _check_calibrated(checkpoint_config, activation_format): + """NVFP4 activations quantize against the checkpoint's calibrated activation global, which a weight-only + NVFP4 checkpoint does not ship: such a checkpoint can only run weight-only.""" + if ( + isinstance(checkpoint_config, NVFP4Config) + and checkpoint_config.activation_format == "bf16" + and activation_format != "bf16" + ): + raise ValueError( + "This NVFP4 checkpoint is weight-only: it ships no calibrated activation globals (`input_scale`), which " + f"activation_format={activation_format!r} quantizes against. Load it with activation_format='bf16'." + ) + + +def _with_mxfp4_experts(block_fp8_config): + """Block-FP8 over W4A4 MXFP4 experts: what a checkpoint declaring `expert_dtype="fp4"` ships. The loading options + are the plan's.""" + dense = BlockFp8Config( + weight_block_size=block_fp8_config.weight_block_size, + scale_fmt=block_fp8_config.scale_fmt, + activation_scheme=block_fp8_config.activation_scheme, + activation_format=block_fp8_config.activation_format, + checkpoint_format=block_fp8_config.checkpoint_format, + ) + return FineGrainedConfig( + dtype_plan={EXPERTS_PATTERN: Mxfp4Config(activation_format=None), ".*": dense}, + modules_to_not_convert=block_fp8_config.modules_to_not_convert, + modules_to_convert=block_fp8_config.modules_to_convert, + dequantize=block_fp8_config.dequantize, + ) + + +def _declared_expert_dtype(model_config): + """`expert_dtype` where a checkpoint declares it: on its model config (DeepSeek-V4), or in the quantization config + it ships (DeepSeek-V4.1), read from the raw dict since the parsed config has no such field.""" + text_config = model_config.get_text_config() + shipped = getattr(model_config, "quantization_config", None) or getattr(text_config, "quantization_config", None) + return getattr(text_config, "expert_dtype", None) or ( + shipped.get("expert_dtype") if isinstance(shipped, dict) else None + ) + + +def finegrained_quantization_config(quantization_config, model_config): + """The config a checkpoint really describes, with its model config in hand: block-FP8 whose checkpoint declares + FP4 experts (`expert_dtype`, DeepSeek-V4 and V4.1) as block-FP8 over W4A4 MXFP4 experts. Any other config + unchanged.""" + if isinstance(quantization_config, FbgemmFp8Config): + quantization_config = AutoQuantizationConfig.from_dict(quantization_config.to_dict()) + if not isinstance(quantization_config, BlockFp8Config): + return quantization_config + if _declared_expert_dtype(model_config) != "fp4": + return quantization_config + return _with_mxfp4_experts(quantization_config) + + def get_hf_quantizer(config, quantization_config, device_map, weights_only, user_agent, gguf_file=None): if gguf_file is not None: if quantization_config is None: @@ -355,9 +510,16 @@ def get_hf_quantizer(config, quantization_config, device_map, weights_only, user pre_quantized = False if pre_quantized or quantization_config is not None: + # read as what it describes for this model, the model config's `expert_dtype` included: on the fly too, as + # the checkpoint it writes declares it + quantization_config = finegrained_quantization_config(quantization_config, config) if pre_quantized: + # a passed config of `CHECKPOINT_READING_CONFIG_TYPES` reads the checkpoint as its own (the merge does) + routed = isinstance(quantization_config, CHECKPOINT_READING_CONFIG_TYPES) + if isinstance(quantization_params_from_config, dict) and not routed: + quantization_params_from_config = AutoQuantizationConfig.from_dict(quantization_params_from_config) config.quantization_config = AutoHfQuantizer.merge_quantization_configs( - quantization_params_from_config, quantization_config + finegrained_quantization_config(quantization_params_from_config, config), quantization_config ) else: config.quantization_config = quantization_config diff --git a/src/transformers/quantizers/block_scaled_utils.py b/src/transformers/quantizers/block_scaled_utils.py new file mode 100644 index 000000000000..6935ae317273 --- /dev/null +++ b/src/transformers/quantizers/block_scaled_utils.py @@ -0,0 +1,556 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""What the block-scaled weight formats' quantizers share: block-FP8, MXFP8, MXFP4 and NVFP4 are each a quantizer with +its own config, served by the modules in `integrations/finegrained`, and each quantizes a whole model alone; +`FineGrainedHfQuantizer` hosts several of them, each on the modules its `dtype_plan` gives it. +""" + +import re +from copy import deepcopy +from typing import TYPE_CHECKING + +from ..utils import ( + is_accelerate_available, + is_torch_available, + is_torch_xpu_available, + logging, +) +from .base import HfQuantizer +from .quantizers_utils import ( + FUSED_STACKS, + STACK_COMPANIONS, + get_module_from_name, + module_kind, + plan_entry_for, + replace_with_quantized_modules, + should_convert_module, +) + + +if is_torch_available(): + import torch + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..utils.quantization_config import BlockScaledConfig, FineGrainedConfig + +logger = logging.get_logger(__name__) + + +class BlockScaledQuantizer(HfQuantizer): + """The quantizer of a block-scaled weight format (block-FP8, MXFP8, MXFP4, NVFP4), written over `plan()`: the + modules it covers and the format config each loads in. Alone, a format covers every module with its own config; + `FineGrainedHfQuantizer` covers each part of a model with the format quantizer its plan entry names.""" + + requires_calibration = False + quantization_config: "BlockScaledConfig | FineGrainedConfig" + # checkpoint keys this format's producers ship beside the weights that no module holds + unexpected_keys: tuple[str, ...] = () + # narrowed once the modules are built: whether this format holds experts, and the up projection they stack + # (`gate_up_proj` gated, `up_proj` not) + holds_experts = True + holds_embeddings = False + up_projections: tuple[str, ...] = ("gate_up_proj",) + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + # the embedding tables the plan quantizes, by name suffix: each holds its own per-tensor `weight_scale` + self.embeddings = list(getattr(quantization_config, "modules_to_convert", None) or []) + + def plan(self) -> dict[str, "BlockScaledConfig"]: + """`{module regex: format config}` for the modules this quantizer covers.""" + return {".*": self.quantization_config} + + def members(self) -> list["BlockScaledQuantizer"]: + """The quantizer of each format config in `plan()`.""" + return [self] + + def member_of(self, module) -> "BlockScaledQuantizer | None": + """The quantizer whose plan entry `module` was built from; `None` for a module left as it was.""" + config = getattr(module, "quantization_config", None) + if config is None: + return None + return next( + ( + member + for member in self.members() + if member.quantization_config is config or member.quantization_config.to_dict() == config.to_dict() + ), + None, + ) + + # what a format brings + + def format_conversions(self, dequantize: bool) -> list: + """The converters this format's producers' key layouts need on top of the model's plan. Under `dequantize`, + what folding this format's keys needs beyond the shared `weight` + `weight_scale_inv` chain.""" + return [] + + def format_weight_conversions(self, converters, model): + """Adjust the model's companion converters after the base quantizer has cloned them.""" + + def format_parallel_plan(self, model, plan): + """Set this format's companion placements after the shared parallel plan is built.""" + + def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: + """The ops a tensor loading into `param_name` runs, for the layout this format's modules hold it in: + an expert's stacked gate|up rows interleaved (weight, scale grid and bias share the row axis), a scale + in the swizzled layout. Each op has a reverse, so a save + restores the checkpoint layout.""" + from ..integrations.finegrained import FineGrainedExperts + from ..integrations.finegrained.conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedSwizzleScales, + ) + + module, name = get_module_from_name(model, param_name) + if isinstance(module, FineGrainedExperts) and re.fullmatch( + r"(gate_up_proj|up_proj|down_proj)(_scale_inv|_bias)?", name + ): + ops = [FineGrainedInterleaveGateUp(self)] if name.startswith("gate_up_proj") else [] + if not name.endswith("_bias"): + ops += [FineGrainedSwizzleScales(self)] + return ops + return [] + + def finish_modules(self, model): + """What this format does to its loaded modules beyond what every format does.""" + + def _process_model_after_weight_loading(self, model, **kwargs): + from ..integrations.finegrained import assert_modules_are_quantized, disable_deepgemm_on_multi_device + + for member in self.members(): + member.finish_modules(model) + disable_deepgemm_on_multi_device(model) + assert_modules_are_quantized(model) + if self.pre_quantized and self.quantization_config.dequantize: + # The loading plan also consumed scales. Saving full-precision weights must reverse only + # the original model transforms, preserving expert split/merge and user key mappings. + model._weight_conversions = [ + original for original, loaded in self._dequantized_save_conversions if loaded.was_used() + ] + if not self.pre_quantized and getattr(model, "_weight_conversions", None) is not None: + # Quantization emitted scales that were absent on disk. Give them the same + # reverse conversions as their weights when saving. + model._weight_conversions = self.update_weight_conversions(model._weight_conversions, model, for_save=True) + return model + + @staticmethod + def _per_expert_sources(proj: str, slot: str) -> list[str]: + """The per-expert checkpoint keys of `slot` that stack into the experts' `proj`, under any parent + (`mlp.experts`, `mixer.experts`, `feed_forward.experts`): the gate|up stack reads both halves.""" + halves = ("gate_proj", "up_proj") if proj == "gate_up_proj" else (proj,) + return [f"experts.*.{half}.{slot}" for half in halves] + + def _block_scale_rename(self): + """Rename block scales while leaving the host's scalar embedding scales untouched.""" + from ..core_model_loading import WeightRenaming + + kept = "".join(rf"(? list: + """A calibrated checkpoint's `input_scale` onto the `activation_scale` its module holds: renamed for every + module first, then one value per expert of each of `projections` the experts hold, the gate|up pair reduced + to the one both halves read; `ops` run on each value before that reduction. Both expert layouts are covered, a + checkpoint matching one while the other never fires; neither when this format holds no experts.""" + from ..core_model_loading import MergeModulelist, WeightConverter, WeightRenaming + from ..integrations.finegrained.conversions import FineGrainedInputScales + + rename = WeightRenaming(source_patterns=r"^(.+)\.input_scale", target_patterns=r"\1.activation_scale") + if not self.holds_experts: + return [rename] + held = [proj for proj in projections if proj in (*self.up_projections, "down_proj")] + return [ + rename, + *( + WeightConverter( + source_patterns=self._per_expert_sources(proj, "activation_scale"), + target_patterns=f"experts.{proj}_activation_scale", + operations=[MergeModulelist(dim=0), *ops, FineGrainedInputScales(self)], + ) + for proj in held + ), + *( + WeightConverter( + source_patterns=rf"experts\.{proj}_input_scale", + target_patterns=f"experts.{proj}_activation_scale", + operations=[*ops, FineGrainedInputScales(self)], + ) + for proj in held + ), + ] + + # what every format shares + + def validate_environment(self, *args, **kwargs): + if not is_accelerate_available(): + raise ImportError("Loading a fine-grained quantized model requires accelerate (`pip install accelerate`)") + + if self.quantization_config.dequantize: + return + + if not torch.cuda.is_available() and not is_torch_xpu_available(): + if self.pre_quantized: + logger.warning_once( + "Fine-grained quantized models require a GPU or XPU; dequantizing to bf16 since neither is " + "available." + ) + self.quantization_config.dequantize = True + return + else: + raise RuntimeError("No GPU or XPU found. A GPU or XPU is needed for fine-grained quantization.") + + if torch.cuda.is_available(): + compute_capability = torch.cuda.get_device_capability() + major, minor = compute_capability + # W4A16 decodes packed FP4 weights before a regular dot product; it does not + # need the FP8 instructions used by the other activation/weight formats. + weight_only_mxfp4 = all( + config.quant_method == "mxfp4" and config.activation_format == "bf16" + for config in self.plan().values() + ) + minimum_capability = (7, 5) if weight_only_mxfp4 else (8, 9) + if compute_capability < minimum_capability: + logger.warning_once( + "This fine-grained quantization plan requires a GPU with compute capability >= " + f"{minimum_capability[0]}.{minimum_capability[1]}, actual = `{major}.{minor}`. " + "We will default to dequantizing the model to bf16." + ) + self.quantization_config.dequantize = True + return + + device_map = kwargs.get("device_map") + if device_map is None: + logger.warning_once( + "You have loaded a fine-grained quantized model on CPU and have a CUDA or XPU device available, " + "make sure to set your model on a GPU or XPU device in order to run your model. To remove this " + "warning, pass device_map = 'cuda' or 'xpu'." + ) + elif isinstance(device_map, dict): + if ( + not self.pre_quantized + and len(device_map) > 1 + and ("cpu" in device_map.values() or "disk" in device_map.values()) + ): + raise ValueError( + "You are attempting to load a fine-grained quantized model with a device_map that contains " + "a cpu/disk device." + "This is not supported when the model is quantized on the fly. " + "Please use a quantized checkpoint or remove the cpu/disk device from the device_map." + ) + + def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: + # `FineGrainedGroupedLinear` subclasses `FineGrainedLinear`, so the tuple covers it implicitly. + from ..integrations.finegrained import FineGrainedExperts, FineGrainedLinear + + module, tensor_name = get_module_from_name(model, param_name) + if self.pre_quantized: + return False + if isinstance(module, FineGrainedExperts): + return tensor_name in FUSED_STACKS + return isinstance(module, FineGrainedLinear) and tensor_name == "weight" + + @property + def quantized_element_size(self) -> float: + # FP4 modules already hold two values per byte in a half-width weight. + return 1 + + def get_new_dtype(self, converter, param, dtype): + if any(name.endswith(("activation_scale", "_global_scale")) for name in converter.target_patterns): + return None + # The initialized scale dtype is authoritative for UE8M0, including float32 containers + # (DeepSeek-V4-Base). Integer exponent containers retain their checkpoint dtype. + is_float = ( + param.get_dtype().startswith(("F", "BF")) if hasattr(param, "get_dtype") else param.is_floating_point() + ) + if dtype == getattr(torch, "float8_e8m0fnu", None) and is_float: + return dtype + return super().get_new_dtype(converter, param, dtype) + + def _normalize_modules_to_not_convert(self, model: "PreTrainedModel"): + """Extend the skip-list, which names modules in the checkpoint's namespace, with the model's + modules it covers there. A released model renames its modules on load, so each one is named back + through the reversed renames (written on keys, so on the module's key prefix) and taken with its + sub-tree when the list covers it: a VLM's bare `vision_tower` matches no rename on its own.""" + skip = self.quantization_config.modules_to_not_convert + if not skip: + return + + from ..conversion_mapping import get_model_conversion_mapping + from ..core_model_loading import WeightRenaming + + renamings = get_model_conversion_mapping(model) + reverse = [r.reverse_transform() for r in renamings[::-1] if isinstance(r, WeightRenaming)] + covered = [] + for module_name, _ in model.named_modules(): + if any(module_name.startswith(f"{parent}.") for parent in covered): + continue + checkpoint_name = f"{module_name}." + for rename in reverse: + checkpoint_name, _ = rename.rename_source_key(checkpoint_name) + checkpoint_name = checkpoint_name.removesuffix(".") + if checkpoint_name != module_name and not should_convert_module(checkpoint_name, skip): + covered.append(module_name) + self.quantization_config.modules_to_not_convert = [*skip, *covered] + + def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): + self._normalize_modules_to_not_convert(model) + self.modules_to_not_convert = self.get_modules_to_not_convert( + model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules + ) + # an embedding table worth quantizing is too big to expand, so it stays FP8 under `dequantize=True` + # too; quantizing one on the fly has no path + plan = self.plan() + embeddings = { + rf"^{re.escape(name)}$": config + for name, module in model.named_modules() + if module_kind(name, module) == "embedding" + and (config := plan_entry_for(plan, name, catch_all=False)) is not None + } + if not self.pre_quantized: + self.modules_to_not_convert = [*self.modules_to_not_convert, *embeddings] + elif self.quantization_config.dequantize: + plan = embeddings + replaced = self.replace_modules(model, plan, self.modules_to_not_convert) + from ..integrations.finegrained import FineGrainedEmbedding, FineGrainedExperts + + if excluded := getattr(model, "_no_placement_params", None): + excluded = set(excluded) + for name, module in model.named_modules(): + if isinstance(module, FineGrainedEmbedding): + for prefix in list(excluded): + if prefix.endswith(".weight") and (name + ".weight").endswith(prefix): + excluded.update( + prefix.removesuffix(".weight") + "." + param + for param, _ in module.named_parameters(recurse=False) + ) + setattr(model, "_no_placement_params", sorted(excluded)) + + experts = [module for module in model.modules() if isinstance(module, FineGrainedExperts)] + ignored = set(model._keys_to_ignore_on_load_unexpected or []) + for member in self.members(): + # a format's expert converters exist only where its experts do, for the stacks they hold + held = [module for module in experts if module.quantization_config is member.quantization_config] + member.holds_embeddings = any( + isinstance(module, FineGrainedEmbedding) and self.member_of(module) is member + for module in model.modules() + ) + member.holds_experts = bool(held) + member.up_projections = tuple(sorted({module.gate_up_name for module in held})) + ignored.update(member.unexpected_keys) + # with no module holding an activation scale (weight-only, dynamic activations), a calibrated + # `input_scale` is one this run has no use for, not an unexpected key: a later run may read it + if not any(name.endswith("activation_scale") for name, _ in model.named_parameters()): + ignored.add(r"input(_global)?_scale$") + model._keys_to_ignore_on_load_unexpected = ignored + if not replaced and not self.quantization_config.dequantize: + logger.warning( + "You are loading a fine-grained quantized model but no linear or experts modules were " + "found to convert. Please double check your model architecture." + ) + + def replace_modules(self, model, plan=None, modules_to_not_convert=None) -> bool: + """Swap the modules `plan` (this quantizer's own by default) covers for their quantized classes.""" + from ..integrations.finegrained import ( + FineGrainedEmbedding, + FineGrainedExperts, + FineGrainedGroupedLinear, + FineGrainedLinear, + ) + + module_classes = { + "linear": FineGrainedLinear, + "grouped_linear": FineGrainedGroupedLinear, + "experts": FineGrainedExperts, + "embedding": FineGrainedEmbedding, + } + plan = self.plan() if plan is None else plan + return replace_with_quantized_modules(model, plan, modules_to_not_convert, module_classes) + + def update_parallel_plan(self, model, plan): + from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard + from ..integrations.finegrained import FineGrainedExperts + + original = dict(plan) + super().update_parallel_plan(model, plan) + overrides = FineGrainedExperts._impl_tp_layer_overrides.get( + getattr(model.config, "_experts_implementation", None), {} + ) + plan.update({key: overrides.get(style, style) for key, style in plan.items()}) + # Temporary until the EP implementation PR is merged: use replicated routing, + # local expert masking, and all-reduce to keep token shapes fixed across ranks. + if "ep_dispatch_experts" in plan.values(): + modules = {replace_layer_number_by_wildcard(name) for name, _ in model.named_modules()} + for name, style in list(plan.items()): + if style != "ep_dispatch_experts": + continue + parent = name.rpartition(".")[0] + router = next((f"{parent}.{key}" for key in ("gate", "router") if f"{parent}.{key}" in modules), None) + if router is None: + raise ValueError(f"Cannot find a gate or router for masked expert parallelism at `{name}`.") + plan[name] = "moe_tp_experts" + plan[router] = "ep_router" + grids = {} + for name, module in model.named_modules(): + if not isinstance(module, FineGrainedExperts): + continue + glob = replace_layer_number_by_wildcard(name) + if _get_parameter_plan(name, plan, is_weight=False) not in ("moe_tp_experts", "megamoe_experts"): + continue + blocked = _has_scale_grid(module.quantization_config) + if glob in grids and grids[glob] != blocked: + raise ValueError(f"Tensor parallelism requires every `{glob}` to use the same scale grid layout.") + grids[glob] = blocked + for projection in (module.gate_up_name, "down_proj"): + key = f"{glob}.{projection}" + style = _get_parameter_plan(f"{name}.{projection}", plan) + rows = { + "packed_colwise": "moe_experts_packed_colwise", + "colwise": "moe_experts_colwise", + "rowwise": "moe_experts_rowwise", + "packed_rowwise": "moe_experts_rowwise", + }.get(style) + if rows is None: + continue + for suffix in ("_bias", *STACK_COMPANIONS): + companion = key + suffix + if companion in original or not hasattr(module, projection + suffix): + continue + if (suffix == "_scale_inv" and blocked) or (suffix == "_bias" and projection != "down_proj"): + plan[companion] = rows + else: + # Per-expert globals and down biases stay whole under intra-expert TP. + plan.pop(companion, None) + plan.update( + { + replace_layer_number_by_wildcard(name): "replicate" + for name, param in model.named_parameters() + if param.ndim == 0 and _get_parameter_plan(name, plan) is not None + } + ) + + for member in self.members(): + member.format_parallel_plan(model, plan) + + def is_serializable(self): + return True + + @property + def is_trainable(self) -> bool: + return False + + @property + def is_compileable(self) -> bool: + return True + + def get_quantize_ops(self): + from ..integrations.finegrained.conversions import FineGrainedQuantize + + return FineGrainedQuantize(self) + + @property + def serialized_names(self): + dequantize = self.quantization_config.dequantize + dense = ["weight", "weight_scale_inv"] + if dequantize: + dense += ["weight_global_scale", "activation_scale"] + return { + "weight": dense, + **{ + proj: [ + proj, + *(proj + suffix for suffix in STACK_COMPANIONS), + *([] if dequantize else [proj + "_bias"]), + ] + for proj in FUSED_STACKS + }, + } + + def get_load_operations(self): + from ..integrations.finegrained.conversions import FineGrainedDequantize + + return {"dequantize": FineGrainedDequantize} + + def get_weight_conversions(self): + conversions, claimed = [], {} + if not self.pre_quantized: + return conversions + # Fused stacks precede the dense fallback, whose companion names are suffixes of theirs. + general = super().get_weight_conversions() + for converter in [*general[1:], *general[:1]]: + conversions.append( + converter.copy( + [rf"(? bool: + """Whether a weight in `config`'s format carries a scale grid, which shards along with it under tensor parallelism, + rather than one scale for the whole tensor.""" + from ..integrations.finegrained.layers import resolve_weight_format + + return resolve_weight_format(config.method.value, config.scale_fmt, config.weight_block_size)[2] is not None + + +def _signature(conversion) -> tuple: + """What a converter does: which keys it reads, where they go and through which ops.""" + operations = tuple(type(op).__name__ for op in getattr(conversion, "operations", ())) + return type(conversion), tuple(conversion.source_patterns), tuple(conversion.target_patterns), operations diff --git a/src/transformers/quantizers/quantizer_blockfp8.py b/src/transformers/quantizers/quantizer_blockfp8.py new file mode 100644 index 000000000000..ab4dd4edc7cc --- /dev/null +++ b/src/transformers/quantizers/quantizer_blockfp8.py @@ -0,0 +1,145 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Block-FP8: `float8_e4m3fn` weights with fp32 or UE8M0 scales over blocks (or one per tensor), W8A8.""" + +from typing import TYPE_CHECKING + +from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import embedding_pattern + + +if TYPE_CHECKING: + from ..utils.quantization_config import BlockFp8Config + + +class BlockFp8HfQuantizer(BlockScaledQuantizer): + """Block-FP8 and per-tensor FP8, with static or dynamic activations and scalar-scaled embedding tables.""" + + quantization_config: "BlockFp8Config" + + def plan(self): + embeddings = self.quantization_config.modules_to_convert or [] + return {**dict.fromkeys(map(embedding_pattern, embeddings), self.quantization_config), **super().plan()} + + def get_param_ops(self, model, param_name): + from ..integrations.finegrained.conversions import FineGrainedRowScales + + ops = super().get_param_ops(model, param_name) + if ( + self.quantization_config.weight_block_size == "row" + and self.quantization_config.checkpoint_format == "modelopt" + and param_name.endswith("_scale_inv") + ): + ops = [FineGrainedRowScales(self), *ops] + return ops + + def format_parallel_plan(self, model, plan): + from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard + from ..integrations.finegrained import FineGrainedEmbedding, FineGrainedExperts, FineGrainedLinear + + for name, module in model.named_modules(): + if ( + self.quantization_config.weight_block_size == "row" + and isinstance(module, FineGrainedExperts) + and self.member_of(module) is not None + and _get_parameter_plan(f"{name}.down_proj", plan) in ("rowwise", "packed_rowwise") + ): + plan.pop(f"{replace_layer_number_by_wildcard(name)}.down_proj_scale_inv", None) + if ( + self.quantization_config.weight_block_size == "row" + and isinstance(module, FineGrainedLinear) + and self.member_of(module) is not None + and _get_parameter_plan(f"{name}.weight", plan) + in ("rowwise", "rowwise_rep", "rowwise_split_input", "packed_rowwise") + ): + plan[f"{replace_layer_number_by_wildcard(name)}.weight_scale_inv"] = "replicate" + if ( + isinstance(module, FineGrainedEmbedding) + and self.member_of(module) is not None + and _get_parameter_plan(name, plan, is_weight=False) is not None + ): + plan.setdefault(f"{replace_layer_number_by_wildcard(name)}.weight_scale", "replicate") + + def format_weight_conversions(self, converters, model): + from ..core_model_loading import MergeModulelist, Transpose + from ..integrations.finegrained import FineGrainedExperts + from ..integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedScalarScales + + producer = self.quantization_config.checkpoint_format + if self.quantization_config.weight_block_size == "row" and producer in ("fbgemm", "modelopt"): + for converter in converters: + dequantize = next((op for op in converter.operations if isinstance(op, FineGrainedDequantize)), None) + if dequantize is None or not any( + isinstance(op, Transpose) and not op.check_dims and {op.dim0, op.dim1} == {1, 2} + for op in converter.operations + ): + continue + reverse = converter.reverse_transform() + if any( + reverse.rename_source_key(name)[1] is not None + and dequantize._config_for(name) is self.quantization_config + and (producer == "modelopt" or name.endswith(".down_proj")) + for name, _ in model.named_parameters() + ): + # Dequantization precedes the BMM transpose; these producers' scales already name output rows. + converter.operations = [ + FineGrainedDequantize(op.hf_quantizer, transpose_scales=True) if op is dequantize else op + for op in converter.operations + ] + if producer == "fbgemm" or (producer == "modelopt" and self.quantization_config.weight_block_size == "row"): + # ModelOpt leaves BMM channel vectors untransposed; FBGEMM's down grid is already rowwise. + scales = [ + f"{name}.{param}" + for name, module in model.named_modules() + if self.member_of(module) is not None + for param, _ in module.named_parameters(recurse=False) + if param.endswith("_scale_inv") and (producer == "modelopt" or param == "down_proj_scale_inv") + ] + for converter in converters: + reverse = converter.reverse_transform() + if any(reverse.rename_source_key(name)[1] is not None for name in scales): + converter.operations = [op for op in converter.operations if not isinstance(op, Transpose)] + + if ( + self.quantization_config.checkpoint_format != "modelopt" + or self.quantization_config.weight_block_size is not None + ): + return + scalar_scales = [ + f"{name}.{proj}_scale_inv" + for name, module in model.named_modules() + if isinstance(module, FineGrainedExperts) and self.member_of(module) is not None + for proj in (module.gate_up_name, "down_proj") + ] + for converter in converters: + if any(isinstance(op, MergeModulelist) for op in converter.operations) and not any( + isinstance(op, FineGrainedScalarScales) for op in converter.operations + ): + reverse = converter.reverse_transform() + if any(reverse.rename_source_key(name)[1] is not None for name in scalar_scales): + converter.operations = [FineGrainedScalarScales(self), *converter.operations] + + def format_conversions(self, dequantize): + conversions = [] + if self.quantization_config.checkpoint_format in ("compressed-tensors", "modelopt", "fbgemm"): + # its block scale ships as `weight_scale`, and rides its weight's converter once renamed + conversions.append(self._block_scale_rename()) + if self.quantization_config.checkpoint_format in ("fbgemm", "modelopt"): + from ..core_model_loading import WeightRenaming + + suffix = "weight_scale" if self.quantization_config.checkpoint_format == "modelopt" else "scale" + conversions.append(WeightRenaming(rf"(gate_up_proj|up_proj|down_proj)_{suffix}(?=$)", r"\1_scale_inv")) + if not dequantize and self.quantization_config.activation_scheme == "static": + conversions += self._activation_scale_conversions() + return conversions diff --git a/src/transformers/quantizers/quantizer_finegrained.py b/src/transformers/quantizers/quantizer_finegrained.py new file mode 100644 index 000000000000..83a8a1a60dd6 --- /dev/null +++ b/src/transformers/quantizers/quantizer_finegrained.py @@ -0,0 +1,103 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""The fine-grained quantizer: a model whose modules are in different block-scaled formats, each part loaded by its +own format's quantizer.""" + +import re +from typing import TYPE_CHECKING + +from ..utils.quantization_config import BlockFp8Config +from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import embedding_pattern, get_module_from_name, module_kind, plan_entry_for + + +if TYPE_CHECKING: + from ..modeling_utils import PreTrainedModel + from ..utils.quantization_config import FineGrainedConfig + + +class FineGrainedHfQuantizer(BlockScaledQuantizer): + """Hosts one quantizer per entry of the config's `dtype_plan`, each the one its format config names in + `AUTO_QUANTIZER_MAPPING`: what the formats share runs once over the whole plan, and each module's tensors take + the ops of the quantizer its entry belongs to. DeepSeek-V4 is MXFP4 experts over block-FP8 linears; a modelopt + export can mix NVFP4 and MXFP8.""" + + quantization_config: "FineGrainedConfig" + + def __init__(self, quantization_config: "FineGrainedConfig", **kwargs): + super().__init__(quantization_config, **kwargs) + from .auto import AUTO_QUANTIZER_MAPPING + + # the embedding tables `modules_to_convert` names hold FP8 rows with one per-tensor scale + self._embedding_config = BlockFp8Config(weight_block_size=None) + self._resolved_plan = None + configs = {id(config): config for config in self.plan().values()}.values() + self._members = [AUTO_QUANTIZER_MAPPING[config.method.value](config, **kwargs) for config in configs] + # the plan's embedding tables are the host's to name + for member in self._members: + member.embeddings = self.embeddings + + def plan(self): + return { + **dict.fromkeys(map(embedding_pattern, self.embeddings), self._embedding_config), + **(self.quantization_config.dtype_plan if self._resolved_plan is None else self._resolved_plan), + } + + def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): + from ..conversion_mapping import get_model_conversion_mapping + from ..core_model_loading import WeightRenaming + + # Match checkpoint patterns against the source namespace without rewriting the saved config. + # A real leaf also exercises rename rules anchored to a parameter name. + reverse = [ + rule.reverse_transform() + for rule in get_model_conversion_mapping(model)[::-1] + if isinstance(rule, WeightRenaming) + ] + plan = self.quantization_config.dtype_plan + resolved = dict(plan) + matched = set() + for name, module in model.named_modules(): + if module_kind(name, module) not in ("linear", "grouped_linear", "experts", "embedding"): + continue + checkpoint_key = name + ".weight" + for rule in reverse: + checkpoint_key, _ = rule.rename_source_key(checkpoint_key) + checkpoint_name = checkpoint_key.removesuffix(".weight") + native = plan_entry_for(plan, name, catch_all=False) + source = plan_entry_for(plan, checkpoint_name, catch_all=False) + if native is not None and source is not None and native is not source: + raise ValueError( + f"Conflicting quantization entries for {name!r} and checkpoint module {checkpoint_name!r}" + ) + if source is not None and native is None: + resolved[rf"^{re.escape(name)}$"] = source + matched.update( + pattern for pattern in plan if re.search(pattern, name) or re.search(pattern, checkpoint_name) + ) + for pattern, config in plan.items(): + if config.checkpoint_format == "modelopt" and pattern not in matched: + raise ValueError( + f"ModelOpt target {pattern!r} matches no supported linear, expert, or embedding module. " + "Per-expert assignments cannot target an architecture that stores fused expert tensors." + ) + self._resolved_plan = resolved + return super()._process_model_before_weight_loading(model, **kwargs) + + def members(self): + return self._members + + def get_param_ops(self, model: "PreTrainedModel", param_name: str) -> list: + member = self.member_of(get_module_from_name(model, param_name)[0]) + return member.get_param_ops(model, param_name) if member is not None else [] diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index 0714aaf198cd..68ea88ab9387 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -1,4 +1,4 @@ -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -11,218 +11,78 @@ # 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. +"""MXFP4: E2M1 weights two per byte with a UE8M0 scale per 32 — GPT-OSS's experts, which ship each projection as +`{proj}_blocks` + `{proj}_scales`, and DeepSeek-V4's, which ship a plain weight/scale pair.""" + from typing import TYPE_CHECKING -from .base import HfQuantizer +from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import FUSED_STACKS, get_module_from_name if TYPE_CHECKING: - from ..modeling_utils import PreTrainedModel from ..utils.quantization_config import Mxfp4Config -from ..utils import ( - get_current_accelerator, - is_accelerate_available, - is_kernels_available, - is_torch_available, - is_triton_available, - logging, -) -from ..utils.import_utils import KERNELS_MAX_VERSION, KERNELS_MIN_VERSION -from .quantizers_utils import get_module_from_name - - -if is_torch_available(): - import torch - - -logger = logging.get_logger(__name__) +class Mxfp4HfQuantizer(BlockScaledQuantizer): + """MXFP4 weights, W4A16 (the config's default, what GPT-OSS runs), W4A8 or W4A4 (DeepSeek-V4's experts). -class Mxfp4HfQuantizer(HfQuantizer): - """ - FP4 quantization using fbgemm kernels + One thing is MXFP4's alone: GPT-OSS names each projection `{proj}_blocks` + `{proj}_scales` rather than a + weight/scale pair. """ - requires_calibration = False quantization_config: "Mxfp4Config" - def validate_environment(self, *args, **kwargs): - if not is_torch_available(): - raise ImportError( - "Using mxfp4 quantization requires torch" - "Please install the latest version of torch ( pip install --upgrade torch )" - ) - - if self.quantization_config.dequantize: - return - - if not is_accelerate_available(): - raise ImportError("Using mxfp4 requires Accelerate: `pip install accelerate`") - - device = get_current_accelerator() or torch.device("cpu") - if device.type not in ["cuda", "xpu", "cpu"]: - if self.pre_quantized: - logger.warning_once( - f"Using MXFP4 quantized models requires model on cuda/xpu/cpu, but found {device}, we will default to dequantizing the model to bf16. To use mxfp4, please disable the current accelerator." - ) - self.quantization_config.dequantize = True - return - else: - raise RuntimeError( - f"Quantizing a model using MXFP4 requires model on cuda/xpu/cpu, but found {device}. To use mxfp4, please disable the current accelerator." - ) - - if torch.xpu.is_available(): - is_device_supported_mxfp4 = True - triton_available = is_triton_available("3.5.0") - kernels_installed = is_kernels_available() - elif torch.cuda.is_available(): - compute_capability = torch.cuda.get_device_capability() - is_device_supported_mxfp4 = compute_capability >= (7, 5) - triton_available = is_triton_available("3.4.0") - kernels_installed = is_kernels_available() - elif device.type == "cpu": - is_device_supported_mxfp4 = True - triton_available = is_triton_available("3.5.0") - kernels_installed = is_kernels_available() - else: - is_device_supported_mxfp4 = False - triton_available = False - kernels_installed = False - - if self.pre_quantized: - if not is_device_supported_mxfp4: - logger.warning_once( - "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 " - "(e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series). " - "We will default to dequantizing the model to bf16." - ) - self.quantization_config.dequantize = True - return - - if not triton_available: - logger.warning_once( - "MXFP4 quantization requires Triton: CUDA requires Triton >= 3.4.0, " - "XPU/CPU requires Triton >= 3.5.0. Please install triton: `pip install triton`. " - "We will default to dequantizing the model to bf16." - ) - self.quantization_config.dequantize = True - return - - if not kernels_installed: - logger.warning_once( - "MXFP4 quantization requires the `kernels` package: " - f"Please install a compatible version ({KERNELS_MIN_VERSION} <= version < {KERNELS_MAX_VERSION}), " - f"e.g. `pip install kernels=={KERNELS_MIN_VERSION}`" - "We will default to dequantizing the model to bf16." - ) - self.quantization_config.dequantize = True - return - elif not is_device_supported_mxfp4: - raise ValueError( - "MXFP4 quantization is only supported on GPUs with compute capability >= 7.5 " - "(e.g T4, A100, L4, H100, or B200) or XPUs (e.g Intel® Data Center GPU Max Series) or CPU" + def format_conversions(self, dequantize): + conversions = self._dequantize_conversions() if dequantize else self._mxfp4_conversions() + if self.quantization_config.checkpoint_format == "modelopt": + conversions = [self._block_scale_rename(), *conversions] + return conversions + + def get_param_ops(self, model, param_name): + from ..integrations.finegrained.conversions import FineGrainedViewPackedInt8 + from ..integrations.finegrained.layers import _FineGrainedModule + + ops = super().get_param_ops(model, param_name) + module, name = get_module_from_name(model, param_name) + if ( + self.quantization_config.checkpoint_format == "modelopt" + and isinstance(module, _FineGrainedModule) + and name in ("weight", *FUSED_STACKS) + ): + ops = [FineGrainedViewPackedInt8(self), *ops] + return ops + + def _dequantize_conversions(self): + """The blocks regroup into the packed rows and the exponent-byte scales come back out of + them as bf16, in the `(E, hidden, 2I)` orientation the unquantized experts hold.""" + from ..core_model_loading import Transpose, WeightConverter + from ..integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks + + return [ + WeightConverter( + source_patterns=[rf"{proj}_blocks", rf"{proj}_scales"], + target_patterns=f"{proj}", + operations=[FineGrainedPackedBlocks(self), FineGrainedDequantize(self), Transpose(1, 2)], ) - elif not triton_available: - raise ValueError( - "MXFP4 quantization requires Triton: CUDA requires Triton >= 3.4.0, " - "XPU/CPU requires Triton >= 3.5.0. Please install triton: `pip install triton`" - ) - elif not kernels_installed: - raise ValueError("MXFP4 quantization requires the `kernels` package: `pip install kernels>=0.12.0`") - - device_map = kwargs.get("device_map") - if device_map is not None and isinstance(device_map, dict): - if not self.pre_quantized and "disk" in device_map.values(): - raise ValueError( - "You are attempting to load an FP4 model with a device_map that contains a disk device." - "This is not supported when the model is quantized on the fly. " - "Please use a quantized checkpoint or remove the disk device from the device_map." - ) - - def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: - from ..integrations import Mxfp4GptOssExperts - - module, tensor_name = get_module_from_name(model, param_name) - if isinstance(module, Mxfp4GptOssExperts): - if tensor_name in ["down_proj_bias", "gate_up_proj_bias"]: - return False - return True - return False - - def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): - # clean cache due to triton ops - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif torch.xpu.is_available(): - torch.xpu.empty_cache() - - def _process_model_before_weight_loading( - self, - model: "PreTrainedModel", - use_kernels: bool = False, - **kwargs, - ): - from ..integrations import replace_with_mxfp4_linear - - # if we are using kernels, we can't use the quantized model, since the forward pass is different and needs special handling - # only CPU and XPU kernels can work with pre-quantized models - device = get_current_accelerator() or torch.device("cpu") - if use_kernels and device.type not in ["cpu", "xpu"]: - logger.warning_once( - "You are using full precision kernels, we will dequantize the model to bf16. " - "To use the quantized model with quantization kernels, please set use_kernels=False" - ) - self.quantization_config.dequantize = True - - if not use_kernels and device.type in ["cpu"]: - logger.warning_once( - "MXFP4 inference on CPU requires use_kernels=True, but use_kernels is disabled. " - "We will dequantize the model to bf16. To run MXFP4 natively on CPU, please set use_kernels=True." - ) - self.quantization_config.dequantize = True - - self.modules_to_not_convert = self.get_modules_to_not_convert( - model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules - ) - - model = replace_with_mxfp4_linear( - model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config - ) - - def get_state_dict_and_metadata(self, model): - from ..integrations.mxfp4 import Mxfp4GptOssExperts, Mxfp4ReverseDeserialize - - state_dict = model.state_dict() - reverse_op = Mxfp4ReverseDeserialize(self) - for name, module in model.named_modules(): - if isinstance(module, Mxfp4GptOssExperts): - for proj in ("gate_up_proj", "down_proj"): - state_dict.update(reverse_op.convert({}, model=model, full_layer_name=f"{name}.{proj}_blocks")) - return state_dict, {} - - def is_serializable(self): - return True - - @property - def is_trainable(self) -> bool: - logger.warning_once( - "MXFP4 quantization don't support training, please consider dequantizing the model first by passing quantization_config=Mxfp4Config(dequantize=True) to .from_pretrained()" - ) - return False - - def get_quantize_ops(self): - from ..integrations.mxfp4 import Mxfp4Quantize - - return Mxfp4Quantize(self) - - serialized_names = { - "gate_up_proj": ["gate_up_proj_blocks", "gate_up_proj_scales"], - "down_proj": ["down_proj_blocks", "down_proj_scales"], - } - - def get_load_operations(self): - from ..integrations.mxfp4 import Mxfp4Dequantize, Mxfp4Deserialize - - return {"deserialize": Mxfp4Deserialize, "dequantize": Mxfp4Dequantize} + for proj in ("gate_up_proj", "down_proj") + ] + + def _mxfp4_conversions(self): + """{proj}_blocks + {proj}_scales checkpoints (GPT-OSS): the blocks regroup into the packed weight and + the exponent bytes are the scale grid. Bare patterns, with no `experts.` scope; the layout each lands + in is the module's, from `get_param_ops`.""" + from ..core_model_loading import WeightConverter, WeightRenaming + from ..integrations.finegrained.conversions import FineGrainedPackedBlocks + + conversions = [] + for proj in ("gate_up_proj", "down_proj"): + conversions += [ + WeightConverter( + source_patterns=rf"{proj}_blocks", + target_patterns=f"{proj}", + operations=[FineGrainedPackedBlocks(self)], + ), + WeightRenaming(source_patterns=rf"{proj}_scales", target_patterns=f"{proj}_scale_inv"), + ] + return conversions diff --git a/src/transformers/quantizers/quantizer_mxfp8.py b/src/transformers/quantizers/quantizer_mxfp8.py new file mode 100644 index 000000000000..06808f4ac6fb --- /dev/null +++ b/src/transformers/quantizers/quantizer_mxfp8.py @@ -0,0 +1,36 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""MXFP8: `float8_e4m3fn` weights with a UE8M0 scale per 32. It ships the plain `weight` / `weight_scale_inv` +pair; ModelOpt exports name that scale `weight_scale` instead.""" + +from typing import TYPE_CHECKING + +from .block_scaled_utils import BlockScaledQuantizer + + +if TYPE_CHECKING: + from ..utils.quantization_config import Mxfp8Config + + +class Mxfp8HfQuantizer(BlockScaledQuantizer): + """MXFP8 weights, W8A8, W8A4 or weight-only.""" + + quantization_config: "Mxfp8Config" + + def format_conversions(self, dequantize): + return ( + [self._block_scale_rename()] + if self.quantization_config.checkpoint_format in ("modelopt", "compressed-tensors") + else [] + ) diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index 0779c9d60ebf..15725172a9d7 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -11,109 +11,182 @@ # 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. -from __future__ import annotations +"""NVFP4: E2M1 weights two per byte with an E4M3 scale per 16 and an fp32 global, in the layouts NVIDIA modelopt +exports (per expert per projection, or stacked per layer); quantizes a full-precision checkpoint on the fly too.""" from typing import TYPE_CHECKING -from ..utils import is_kernels_available, is_torch_available -from ..utils.import_utils import KERNELS_MAX_VERSION, KERNELS_MIN_VERSION -from .base import HfQuantizer -from .quantizers_utils import get_module_from_name +from .block_scaled_utils import BlockScaledQuantizer +from .quantizers_utils import FUSED_STACKS, get_module_from_name -if is_torch_available(): - import torch - if TYPE_CHECKING: - from ..modeling_utils import PreTrainedModel from ..utils.quantization_config import NVFP4Config -def _as_cuda_device(device) -> torch.device | None: - if isinstance(device, int): - return torch.device("cuda", device) - try: - device = torch.device(device) - except (TypeError, RuntimeError): - return None - return device if device.type == "cuda" else None +class NVFP4HfQuantizer(BlockScaledQuantizer): + """NVFP4, in both key layouts modelopt ships: per expert per projection, or stacked per layer.""" + quantization_config: "NVFP4Config" + # modelopt calibrates an FP8 KV cache beside the weights; the cache here stays in the model dtype + unexpected_keys = (r"\.(k_proj\.k|v_proj\.v)_scale$",) -class NVFP4HfQuantizer(HfQuantizer): - """Quantize eligible linear weights to NVFP4 while loading a full-precision checkpoint.""" + def finish_modules(self, model): + from ..integrations.finegrained import register_shared_gate_up_input_globals - requires_calibration = False - quantization_config: NVFP4Config + register_shared_gate_up_input_globals(model) - @property - def quantized_element_size(self) -> float: - return 0.5 - - def validate_environment(self, device_map, **kwargs): - if self.pre_quantized: - raise ValueError("Loading pre-quantized NVFP4 checkpoints is not supported yet.") - if not is_kernels_available(): - raise ImportError( - "NVFP4 quantization requires the `kernels` package. " - f"Install a compatible version ({KERNELS_MIN_VERSION} <= version < {KERNELS_MAX_VERSION}), " - f"for example with `pip install kernels=={KERNELS_MIN_VERSION}`." - ) - if not torch.cuda.is_available(): - raise RuntimeError("NVFP4 quantization requires a CUDA GPU.") - - if device_map is None: - raise ValueError("NVFP4 quantization requires a CUDA `device_map`.") - devices = set(device_map.values()) if isinstance(device_map, dict) else {device_map} - cuda_devices = {_as_cuda_device(device) for device in devices} - if None in cuda_devices: - raise ValueError("NVFP4 quantization does not support CPU or disk offload.") - if len(cuda_devices) != 1: + def format_conversions(self, dequantize): + if dequantize and self.holds_embeddings: raise ValueError( - "NVFP4 quantization currently supports one CUDA device. Tensor parallelism and multi-device " - "`device_map` configurations are not supported yet." + "dequantize=True cannot expand an NVFP4 embedding table. Load it compressed on a supported device; " + "lookup decodes only the selected rows." ) - self.quantization_device = cuda_devices.pop() - - major, minor = torch.cuda.get_device_capability(self.quantization_device) - if major < 10: - raise RuntimeError(f"NVFP4 requires Blackwell (sm100+) block-scaled tensor cores; found sm{major}{minor}.") - - def update_parallel_plan(self, model, plan): - distributed_config = getattr(model.config, "distributed_config", None) - tp_size = getattr(distributed_config, "tp_size", None) - if tp_size is not None and tp_size > 1: - raise ValueError( - "NVFP4 quantization does not support tensor parallelism yet because its scale metadata does not have " - "a sharding plan." + # weight-only keeps activations bf16, so there is no activation global for `input_scale` to land on; nor + # does a dequantized model, which leaves those keys to the unexpected-key filter + calibrated = not dequantize and self.quantization_config.activation_format != "bf16" + conversions = self._nvfp4_renames(calibrated) + if calibrated: + # the down projection's is read with its globals + conversions += self._activation_scale_conversions( + projections=self.up_projections, ops=self._reciprocal_ops ) - - def param_needs_quantization(self, model: PreTrainedModel, param_name: str, **kwargs) -> bool: - from ..integrations.nvfp4 import NVFP4Linear - - module, tensor_name = get_module_from_name(model, param_name) - return isinstance(module, NVFP4Linear) and tensor_name == "weight" - - def _process_model_before_weight_loading(self, model: PreTrainedModel, **kwargs): - from ..integrations.nvfp4 import replace_with_nvfp4_linear - - self.modules_to_not_convert = self.get_modules_to_not_convert( - model, self.quantization_config.modules_to_not_convert, model._keep_in_fp32_modules - ) - replace_with_nvfp4_linear(model, modules_to_not_convert=self.modules_to_not_convert) - - def get_quantize_ops(self): - from ..integrations.nvfp4 import NVFP4Quantize - - return NVFP4Quantize(self.quantization_device) + if dequantize: + # a stacked layout's scales reach its stack's dequantize converter under the slot names it collects; the + # experts are left as they are, so every stack a model may hold is covered + conversions += self._nvfp4_fused_renames( + FUSED_STACKS, (("_weight_scale", "_scale_inv"), ("_weight_scale_2", "_weight_global_scale")) + ) + elif self.holds_experts: + conversions += self._nvfp4_per_expert_conversions(calibrated) + self._nvfp4_fused_conversions(calibrated) + return conversions + + def format_parallel_plan(self, model, plan): + from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard + from ..integrations.finegrained import FineGrainedEmbedding + + for name, module in model.named_modules(): + if not isinstance(module, FineGrainedEmbedding) or self.member_of(module) is None: + continue + style = _get_parameter_plan(name, plan, is_weight=False) + if style is None: + continue + glob = replace_layer_number_by_wildcard(name) + # A scale grid follows its weight's axis; output gather/reduce hooks belong to the module. + scale_style = {"colwise_gather_output": "colwise", "rowwise_split_input": "rowwise"}.get(style, style) + plan.setdefault(f"{glob}.weight_scale_inv", scale_style) + plan.setdefault(f"{glob}.weight_global_scale", "replicate") + + def get_param_ops(self, model, param_name): + """The packed weights first take the int8 view the modules hold: the loader keeps a pre-quantized + non-float tensor in the checkpoint's dtype, modelopt's uint8, and the reverse saves it back as that.""" + from ..integrations.finegrained import FineGrainedLinear + from ..integrations.finegrained.conversions import FineGrainedViewPackedInt8 + from ..integrations.finegrained.layers import _FineGrainedModule + + ops = super().get_param_ops(model, param_name) + module, name = get_module_from_name(model, param_name) + if isinstance(module, _FineGrainedModule) and name in ("weight", *FUSED_STACKS): + ops = [FineGrainedViewPackedInt8(self), *ops] + # an experts module's globals are inverted in their converters, before the per-expert reductions + elif isinstance(module, FineGrainedLinear) and name in ("weight_global_scale", "activation_scale"): + ops = [*self._reciprocal_ops, *ops] + return ops @property - def is_serializable(self) -> bool: - return False + def _reciprocal_ops(self) -> list: + """compressed-tensors ships the globals as their reciprocal; modelopt as the multiplier the kernels take.""" + from ..integrations.finegrained.conversions import FineGrainedReciprocal - @property - def is_trainable(self) -> bool: - return False + return ( + [FineGrainedReciprocal(self)] if self.quantization_config.checkpoint_format == "compressed-tensors" else [] + ) - @property - def is_compileable(self) -> bool: - return True + def _nvfp4_renames(self, calibrated: bool): + """modelopt's scales onto the slots every module holds them in, experts included: renamed before converter + collection, a block scale then rides its weight's converter as `weight_scale_inv` (the merge that stacks + the experts' weights stacks their grids too), and the globals reach the converters below by slot name.""" + from ..core_model_loading import WeightRenaming + + renames = [ + self._block_scale_rename(), + WeightRenaming(source_patterns=r"^(.+)\.weight_scale_2", target_patterns=r"\1.weight_global_scale"), + ] + if self.quantization_config.checkpoint_format == "compressed-tensors": + # its packed values ship as `weight_packed` and its activation global as `input_global_scale`; its weight + # global already carries the slot's name + renames.append(WeightRenaming(source_patterns=r"^(.+)\.weight_packed", target_patterns=r"\1.weight")) + if calibrated: + renames.append( + WeightRenaming( + source_patterns=r"^(.+)\.input_global_scale", target_patterns=r"\1.activation_scale" + ) + ) + return renames + + def _nvfp4_per_expert_conversions(self, calibrated: bool): + """One key per expert per projection, as GLM-5.2-NVFP4 ships them. + + `MergeModulelist` stamps each source with its expert index, which is what lets expert + parallelism keep only this rank's experts (`tensor_idx` in `core_model_loading`); without + it every rank collects all E values and the forward asserts on the per-expert count. + """ + from ..core_model_loading import MergeModulelist, WeightConverter + from ..integrations.finegrained.conversions import FineGrainedWeightGlobals + + return [ + # every second-level global of a layer in ONE converter: folding the gate|up stack's + # two halves into one moves the up half's onto the down projection, so they decide + # together (`FineGrainedWeightGlobals`) + WeightConverter( + source_patterns=[ + *self._per_expert_sources(up, "weight_global_scale"), + *self._per_expert_sources("down_proj", "weight_global_scale"), + *(self._per_expert_sources("down_proj", "activation_scale") if calibrated else []), + ], + target_patterns=[ + f"experts.{up}_weight_global_scale", + "experts.down_proj_weight_global_scale", + *(["experts.down_proj_activation_scale"] if calibrated else []), + ], + operations=[MergeModulelist(dim=0), *self._reciprocal_ops, FineGrainedWeightGlobals(self)], + ) + for up in self.up_projections + ] + + @staticmethod + def _nvfp4_fused_renames(projections, slots) -> list: + """The stacked layout's `{proj}{shipped}` scales onto the `{proj}{slot}` slots, for each `(shipped, slot)`.""" + from ..core_model_loading import WeightRenaming + + return [ + WeightRenaming(source_patterns=rf"experts\.{proj}{shipped}", target_patterns=f"experts.{proj}{slot}") + for proj in projections + for shipped, slot in slots + ] + + def _nvfp4_fused_conversions(self, calibrated: bool): + """Already stacked per layer, as the vLLM fused layout ships them — so no `MergeModulelist`: + the scales only take their slot's name, and the module's ops give every tensor its layout.""" + from ..core_model_loading import WeightConverter + from ..integrations.finegrained.conversions import FineGrainedWeightGlobals + + converters = self._nvfp4_fused_renames((*self.up_projections, "down_proj"), (("_weight_scale", "_scale_inv"),)) + converters += [ + # the same one-converter rule as the per-expert layout, over the fused keys + WeightConverter( + source_patterns=[ + rf"experts\.{up}_weight_scale_2", + r"experts\.down_proj_weight_scale_2", + *([r"experts\.down_proj_input_scale"] if calibrated else []), + ], + target_patterns=[ + f"experts.{up}_weight_global_scale", + "experts.down_proj_weight_global_scale", + *(["experts.down_proj_activation_scale"] if calibrated else []), + ], + operations=[*self._reciprocal_ops, FineGrainedWeightGlobals(self)], + ) + for up in self.up_projections + ] + return converters diff --git a/src/transformers/quantizers/quantizers_utils.py b/src/transformers/quantizers/quantizers_utils.py index 42036ff4416f..4932a7444d0f 100644 --- a/src/transformers/quantizers/quantizers_utils.py +++ b/src/transformers/quantizers/quantizers_utils.py @@ -16,6 +16,11 @@ from typing import Any +# an experts module's fused projection stacks, and what each holds beside it +FUSED_STACKS = ("gate_up_proj", "up_proj", "down_proj") +STACK_COMPANIONS = ("_scale_inv", "_weight_global_scale", "_activation_scale") + + def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: if "." in tensor_name: module_name, tensor_name = tensor_name.rsplit(".", 1) @@ -23,6 +28,72 @@ def get_module_from_name(module, tensor_name: str) -> tuple[Any, str]: return module, tensor_name +def module_kind(module_name: str, module) -> str | None: + """The kind of layer a quantized-module registry is keyed by: `"experts"`, `"grouped_linear"` (a + block-diagonal linear, matched on the attribute a swap needs rather than its class name), `"linear"` or + `"embedding"`; `None` for anything a registry does not replace.""" + from torch import nn + + if module_name.endswith(".experts"): + return "experts" + if isinstance(module, nn.Linear) and hasattr(module, "n_groups"): + return "grouped_linear" + if type(module) is nn.Linear: + return "linear" + if type(module) is nn.Embedding: + return "embedding" + return None + + +EXPERTS_PATTERN = r"\.experts($|\.)" + + +def embedding_pattern(name: str) -> str: + """The `dtype_plan` key of a `modules_to_convert` entry, which names an embedding table by suffix.""" + return rf"{re.escape(name)}$" + + +def plan_entry_for(dtype_plan: dict, module_name: str, catch_all: bool = True): + """The format config `dtype_plan` gives `module_name`, `None` if none covers it. + + A targeted entry wins over the `".*"` catch-all whatever order they were declared in, because + `to_json_string` sorts keys and a config that round-tripped through `config.json` has lost that + order. Two targeted entries claiming the same module is ambiguous, and says so. Without + `catch_all`, only an entry that names the module counts. + """ + claimed = [pattern for pattern in dtype_plan if pattern != ".*" and re.search(pattern, module_name)] + if len(claimed) > 1: + raise ValueError(f"{module_name!r} is claimed by more than one dtype_plan entry: {claimed}") + if claimed: + return dtype_plan[claimed[0]] + return dtype_plan.get(".*") if catch_all else None + + +def replace_with_quantized_modules( + model, dtype_plan: dict, modules_to_not_convert: list[str] | None, module_classes: dict[str, Any] +) -> bool: + """Swap every module `dtype_plan` gives a config for the class `module_classes` builds that kind of layer as, built + on `meta` by its `from_module(module, config, model)`. A module the plan gives no config, of a kind with no class, + or named by `modules_to_not_convert` is kept. Returns whether any module was replaced.""" + import torch + + replaced = False + for module_name, module in model.named_modules(): + kind = module_kind(module_name, module) + if kind is None or not should_convert_module(module_name, modules_to_not_convert): + continue + # the catch-all covers linears and experts; an embedding is quantized only where an entry names it + config = plan_entry_for(dtype_plan, module_name, catch_all=kind != "embedding") + quantized_class = module_classes.get(kind) if config is not None else None + if quantized_class is None: + continue + with torch.device("meta"): + quantized = quantized_class.from_module(module, config, model) + model.set_submodule(module_name, quantized) + replaced = True + return replaced + + def should_convert_module(full_name, patterns: list[str] | None = None): if patterns is None: return True diff --git a/src/transformers/utils/quantization_config.py b/src/transformers/utils/quantization_config.py index a5c3fba8a1f1..85c6da01b3c7 100644 --- a/src/transformers/utils/quantization_config.py +++ b/src/transformers/utils/quantization_config.py @@ -15,12 +15,15 @@ # See the License for the specific language governing permissions and # limitations under the License. import copy +import fnmatch import importlib.metadata import json +import math import os +import warnings from dataclasses import dataclass from enum import Enum -from typing import Any, Optional, Union +from typing import Any, Optional, Union, cast from packaging import version @@ -62,6 +65,8 @@ class QuantizationMethod(str, Enum): AUTOROUND = "auto-round" MXFP4 = "mxfp4" NVFP4 = "nvfp4" + FINEGRAINED = "finegrained" + MODELOPT = "modelopt" MXFP8 = "mxfp8" METAL = "metal" FOUR_OVER_SIX = "fouroversix" @@ -1292,7 +1297,7 @@ def is_quantization_compressed(self): class FbgemmFp8Config(QuantizationConfigMixin): """ This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using fbgemm fp8 quantization. + loaded from FBGEMM rowwise FP8 metadata. Execution uses the finegrained kernels. Args: activation_scale_ub (`float`, *optional*, defaults to 1200.0): @@ -1686,16 +1691,633 @@ def post_init(self): raise TypeError("shapes must be a dict") -@dataclass +def subtree_patterns(globs: list[str]) -> list[str]: + """Glob-style module names ("model.layers.0*", "model.layers.1.*", bare "self_attn") as + regexes that match the named module and everything under it. + + Handed on as regexes the globs are wrong twice: the dots match any character, and the star is + greedy, so "model.layers.1.*" also takes layers 10-19. Anchor each at a path boundary instead. + """ + patterns = [] + for glob in globs: + subtree = glob.removesuffix(".*").rstrip(".") + pattern = fnmatch.translate(subtree).removesuffix(r"\Z").removesuffix(r"\z") + patterns.append(r"(?:^|.*\.)" + pattern + r"(\..*)?$") + return patterns + + +# the block and scale MXFP8 and MXFP4 fix +_MX_FORMAT_OPTIONS = {"weight_block_size": (1, 32), "scale_fmt": "ue8m0"} + + +class _FormatDefault: + """The `activation_format` a format config takes when none is passed.""" + + def __repr__(self): + return "the format's default" + + +_FORMAT_DEFAULT = cast(Any, _FormatDefault()) + + +def _loading_attributes(config, **attributes) -> dict[str, Any]: + """What a config passed to `from_pretrained` lends a pre-quantized checkpoint's own config: `dequantize`, the + skip-list when it sets one (left unset, the checkpoint's stays), and `attributes`.""" + loading = {"dequantize": config.dequantize} + if config.modules_to_not_convert is not None: + loading["modules_to_not_convert"] = config.modules_to_not_convert + return {**loading, **attributes} + + +class BlockScaledConfig(QuantizationConfigMixin): + """What every block-scaled weight format's config holds: the options of the format itself, which the quantized + modules read, and the loading options. One such config quantizes a whole model; a [`FineGrainedConfig`] maps + several onto parts of one. It is built as one of its formats: [`BlockFp8Config`], [`Mxfp8Config`], + [`Mxfp4Config`] or [`NVFP4Config`]. + + Args: + activation_format (`str | None`, *optional*, defaults to `the format's default`): + Activation format used by the quantized layers. Omit it to use the weight format's default. + dequantize (`bool`, *optional*, defaults to `False`): + Whether to load a quantized checkpoint in full precision. + modules_to_not_convert (`list | None`, *optional*): + Module names or patterns to keep in full precision. + checkpoint_format (`str | None`, *optional*): + Producer layout of the checkpoint tensors, such as `"compressed-tensors"`. + """ + + # what `quant_method` this config is, the activation formats its weights' kernels serve, and the one it takes + method: QuantizationMethod + activation_formats: tuple[str | None, ...] = (None,) + default_activation_format: str | None = None + # fixed by the MX and NVFP4 formats; block-FP8 sets them + weight_block_size: tuple[int, int] | str | None = None + scale_fmt: str = "float" + activation_scheme: str = "dynamic" + # the values those fields take in this format, which a checkpoint's `config.json` may spell out + format_options: dict[str, Any] = {} + + def __init__( + self, + *, + activation_format: str | None = _FORMAT_DEFAULT, + dequantize: bool = False, + modules_to_not_convert: list | None = None, + checkpoint_format: str | None = None, + **kwargs, + ): + if not hasattr(self, "method"): + raise TypeError( + f"{type(self).__name__} is what the block-scaled format configs share; build one of them: " + f"{', '.join(config.__name__ for config in BLOCK_SCALED_CONFIGS.values())}" + ) + self.quant_method = self.method + kwargs.pop("quant_method", None) + if modules_to_not_convert is None and "ignored_layers" in kwargs: + modules_to_not_convert = kwargs.pop("ignored_layers") + # a `config.json` written before the format fixed these carries their defaults + fixed = { + name: kwargs.pop(name) + for name in ("weight_block_size", "scale_fmt", "activation_scheme") + if name in kwargs + } + changed = { + name: value + for name, value in fixed.items() + # a block size read back from `config.json` is a list + if (tuple(value) if isinstance(value, list) else value) + not in (getattr(BlockScaledConfig, name), self.format_options.get(name)) + } + if changed: + raise ValueError( + f"{self.method.value} weights fix {sorted(changed)} by their format; got {changed}. Only " + "`BlockFp8Config` sets them." + ) + if kwargs: + logger.info(f"Unused kwargs: {list(kwargs)}. These kwargs are not used in {type(self).__name__}.") + self.checkpoint_format = checkpoint_format + # a passed config lends its activation format to a pre-quantized checkpoint only where it names one + self._names_activation_format = activation_format is not _FORMAT_DEFAULT + self.activation_format = ( + self.default_activation_format if activation_format is _FORMAT_DEFAULT else activation_format + ) + self.dequantize = dequantize + self.modules_to_not_convert = modules_to_not_convert + self.post_init() + + def post_init(self): + if self.activation_format not in self.activation_formats: + raise ValueError( + f"{self.method.value} weights take activation_format in {self.activation_formats}; " + f"got {self.activation_format!r}" + ) + # one spelling for the weights' own activation format + if self.activation_format == self.method.value: + self.activation_format = None + if self.checkpoint_format not in (None, "compressed-tensors", "modelopt", "fbgemm"): + raise ValueError( + f"checkpoint_format must be None, 'compressed-tensors', 'modelopt', or 'fbgemm'; got {self.checkpoint_format!r}" + ) + + def get_loading_attributes(self): + named = {"activation_format": self.activation_format} if self._names_activation_format else {} + return _loading_attributes(self, **named) + + def to_dict(self) -> dict[str, Any]: + return {name: value for name, value in super().to_dict().items() if not name.startswith("_")} + + def __iter__(self): + yield from self.to_dict().items() + + +class BlockFp8Config(BlockScaledConfig): + """Block-FP8 weights: `float8_e4m3fn` values with fp32 or UE8M0 scales over `weight_block_size` blocks (one + per-tensor scale without a block), as DeepSeek-V3 and Qwen3 ship them; W8A8 with dynamic or calibrated activations. + + Args: + weight_block_size (`tuple[int, int] | str`, *optional*, defaults to `(128, 128)`): + The block each weight scale covers; `None` for one scale per tensor, or `"row"` for one scale + per output channel. Dynamic FP8 activations then use one scale per token. + scale_fmt (`str`, *optional*, defaults to `"float"`): + Storage dtype of the scales: `"float"` (fp32) or `"ue8m0"` (`torch.float8_e8m0fnu`, DeepSeek-V4). + activation_scheme (`str`, *optional*, defaults to `"dynamic"`): + `"dynamic"` (quantized inline, per 128-wide group) or `"static"` (a calibrated per-tensor activation + scale from the checkpoint). + modules_to_convert (`list`, *optional*): + Embedding tables to quantize, by name suffix: FP8 rows with one per-tensor scale. + activation_scale_ub (`float`, *optional*): + Maximum activation amax for rowwise FP8 quantization. Values beyond it saturate, as in FBGEMM. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.FP8 + activation_formats = (None, "fp8", "bf16") + + def __init__( + self, + *, + weight_block_size: tuple[int, int] | str | None = (128, 128), + scale_fmt: str = "float", + activation_scheme: str = "dynamic", + modules_to_convert: list | None = None, + activation_scale_ub: float | None = None, + **kwargs, + ): + self.weight_block_size = ( + weight_block_size + if isinstance(weight_block_size, str) or weight_block_size is None + else tuple(weight_block_size) + ) + self.scale_fmt = scale_fmt + self.activation_scheme = activation_scheme + self.modules_to_convert = modules_to_convert + self.activation_scale_ub = activation_scale_ub + super().__init__(**kwargs) + + def get_loading_attributes(self): + attributes = super().get_loading_attributes() + if self.activation_scale_ub is not None or ( + self._names_activation_format and self.activation_format == "bf16" + ): + attributes["activation_scale_ub"] = self.activation_scale_ub + return attributes + + def post_init(self): + super().post_init() + self.activation_scheme = self.activation_scheme.lower() + if self.activation_scheme not in ("dynamic", "static"): + raise ValueError(f"activation_scheme must be 'dynamic' or 'static'; got {self.activation_scheme!r}") + if self.activation_format == "bf16" and self.activation_scheme == "static": + raise ValueError( + "Weight-only block-FP8 has no activation quantization; activation_scheme must be 'dynamic'" + ) + if self.scale_fmt not in ("float", "ue8m0"): + raise ValueError(f"scale_fmt must be 'float' or 'ue8m0'; got {self.scale_fmt!r}") + if self.weight_block_size == "row": + if self.scale_fmt != "float": + raise ValueError("Rowwise FP8 requires FP32 weight scales") + elif isinstance(self.weight_block_size, str) or ( + self.weight_block_size is not None + and (len(self.weight_block_size) != 2 or min(self.weight_block_size) <= 0) + ): + raise ValueError( + f"weight_block_size must be 'row', None, or two positive integers; got {self.weight_block_size}" + ) + if self.activation_scale_ub is not None and ( + self.weight_block_size != "row" + or not math.isfinite(self.activation_scale_ub) + or self.activation_scale_ub <= 0 + or self.activation_scheme != "dynamic" + or self.activation_format is not None + ): + raise ValueError("activation_scale_ub must be positive and requires dynamic rowwise W8A8 FP8") + + +class Mxfp8Config(BlockScaledConfig): + """MXFP8 weights: `float8_e4m3fn` values with a UE8M0 scale per 32, as MiniMax-M3 ships them. + + Args: + activation_format (`str | None`, *optional*, defaults to `the format's default`): + Activation format for the quantized layers; `"bf16"` keeps activations in full precision. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.MXFP8 + activation_formats = (None, "mxfp8", "mxfp4", "bf16") + format_options = _MX_FORMAT_OPTIONS + + def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs): + super().__init__(activation_format=activation_format, **kwargs) + + +class Mxfp4Config(BlockScaledConfig): + """MXFP4 weights: E2M1 values two per byte with a UE8M0 scale per 32, as GPT-OSS and DeepSeek-V4's experts ship + them. + + Args: + activation_format (`str`, *optional*, defaults to `"bf16"`): + `"bf16"` for weight-only (W4A16, what GPT-OSS runs), `"mxfp8"` for W4A8, or `None` (or `"mxfp4"`) for W4A4. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.MXFP4 + activation_formats = (None, "mxfp8", "mxfp4", "bf16") + default_activation_format = "bf16" + format_options = _MX_FORMAT_OPTIONS + + def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs): + super().__init__(activation_format=activation_format, **kwargs) + + +class NVFP4Config(BlockScaledConfig): + """NVFP4 weights: E2M1 values two per byte with an E4M3 scale per 16 and an fp32 global, as NVIDIA modelopt + exports them; quantizes a full-precision checkpoint on the fly as well. + + Args: + activation_format (`str | None`, *optional*, defaults to `the format's default`): + Activation format for the quantized layers; `"bf16"` keeps activations in full precision. + kwargs (`dict[str, Any]`, *optional*): + The options every block-scaled format takes, see [`BlockScaledConfig`]. + """ + + method = QuantizationMethod.NVFP4 + activation_formats = (None, "nvfp4", "bf16") + format_options = {"weight_block_size": (1, 16)} + + def __init__(self, *, activation_format: str | None = _FORMAT_DEFAULT, **kwargs): + super().__init__(activation_format=activation_format, **kwargs) + + +# the block-scaled formats by their `quant_method`, for a plan entry read back from `config.json` +BLOCK_SCALED_CONFIGS = { + config.method.value: config for config in (BlockFp8Config, Mxfp8Config, Mxfp4Config, NVFP4Config) +} + + +def plan_from_config_groups(spec: dict) -> dict[str, BlockScaledConfig] | None: + """One `config_groups` entry as `dtype_plan` entries, one per target, or `None` if it names a + format we do not serve. ModelOpt rejects such groups; compressed-tensors keeps its own backend. + + compressed-tensors describes a format by its parameters; we name it. `group_size` is what + separates the two 4-bit formats: nvfp4's block scale is E4M3 over 16 values, mxfp4's is E8M0 + over 32. `targets` is theirs too: a bare entry is a CLASS name (GLM-5.2 ships `["Linear"]`, + i.e. every linear), and an `re:` prefix introduces a pattern. Activations are read the same way, and only where + the weights' kernels quantize to them. + """ + try: + _check_block_scaled_features(spec) + except ValueError: + return None + formats = { + (4, "float", 16): "nvfp4", + (4, "float", 32): "mxfp4", + (8, "float", 32): "mxfp8", + (8, "float", None): "fp8", + } + weights = spec.get("weights") or {} + quant_method = formats.get((weights.get("num_bits"), weights.get("type"), weights.get("group_size"))) + if quant_method is None: + return None + # MX block scales are E8M0 (some producers store its exponent bytes as uint8): group-32 values over another + # scale dtype are not an MX format + scale_dtype = str(weights.get("scale_dtype") or "e8m0") + if quant_method in ("mxfp8", "mxfp4") and not any(name in scale_dtype for name in ("e8m0", "uint8")): + return None + targets = [] + for target in spec.get("targets") or []: + if target.startswith("re:"): + targets.append(target[3:]) + elif target == "Linear": # every linear — the catch-all, in their spelling + targets.append(".*") + else: # a class name we cannot resolve to module paths here + return None + activations = spec.get("input_activations") or {} + rowwise = weights.get("strategy") == "channel" + if rowwise: + if quant_method != "fp8": + return None + if activations and not ( + (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) == (8, "float", None) + and ( + (activations.get("strategy") == "token" and activations.get("dynamic") is True) + or (activations.get("strategy") == "tensor" and activations.get("dynamic") is False) + ) + ): + return None + act_key = (activations.get("num_bits"), activations.get("type"), activations.get("group_size")) + # block-FP8's own activations are FP8 per token, or per 128-wide group as its kernels quantize them; + # no activation quantization is weight-only + activation_format = {**formats, (8, "float", 128): "fp8"}.get(act_key) if activations else "bf16" + config_class = BLOCK_SCALED_CONFIGS[quant_method] + if activation_format is None or activation_format not in config_class.activation_formats: + return None + options = {"activation_format": activation_format} + if quant_method == "fp8": + # only a per-tensor scale is calibrated: `dynamic: false` on group-scaled activations is NVFP4's + # second-level global (`input_scale`), which the loader carries on its own + static = not activations.get("dynamic", True) and act_key == (8, "float", None) + options.update( + activation_scheme="static" if static else "dynamic", + weight_block_size="row" if rowwise else weights.get("block_structure"), + ) + config = config_class(**options) + return dict.fromkeys(targets or [".*"], config) + + +def _check_block_scaled_features(spec: dict): + """Refuse runtime operations that a block-scaled GEMM alone cannot reproduce.""" + for field in ( + "pre_quant_scale", + "lora_rank", + "has_zero_point", + "transform_config", + "sparsity_config", + "weight_scale_encoding", + ): + if spec.get(field): + raise ValueError(f"Unsupported block-scaled quantization feature {field}: {spec[field]!r}") + for name in ("weights", "input_activations"): + if options := spec.get(name): + if options.get("symmetric") is False or options.get("zero_point") is not None: + raise ValueError(f"Unsupported asymmetric quantization of {name}") + + +def _modelopt_format(spec: dict) -> BlockScaledConfig: + """Translate an algorithm without inferring a different format from unsupported metadata.""" + if not isinstance(spec, dict): + raise ValueError(f"Expected a modelopt format object; got {spec!r}") + _check_block_scaled_features(spec) + algo = spec.get("quant_algo") + formats = { + "NVFP4": (NVFP4Config, 16, None), + "W4A16_NVFP4": (NVFP4Config, 16, "bf16"), + "MXFP8": (Mxfp8Config, 32, None), + "MXFP4": (Mxfp4Config, 32, None), + "W4A16_MXFP4": (Mxfp4Config, 32, "bf16"), + } + if algo == "FP8": + return BlockFp8Config(weight_block_size=None, activation_scheme="static", checkpoint_format="modelopt") + if algo == "FP8_PER_CHANNEL_PER_TOKEN": + return BlockFp8Config(weight_block_size="row", checkpoint_format="modelopt") + if algo == "FP8_PB_WO": + block = spec.get("group_size", 128) + if not isinstance(block, int) or isinstance(block, bool) or block <= 0: + raise ValueError(f"FP8_PB_WO requires a positive group_size; got {block!r}") + return BlockFp8Config(weight_block_size=(block, block), activation_format="bf16", checkpoint_format="modelopt") + if algo not in formats: + raise ValueError(f"Unsupported modelopt quant_algo: {algo!r}") + config_class, group_size, activation_format = formats[algo] + if spec.get("group_size", group_size) != group_size: + raise ValueError(f"modelopt {algo} requires group_size={group_size}; got {spec['group_size']!r}") + return config_class(activation_format=activation_format, checkpoint_format="modelopt") + + +def _merge_config_group_plan(dtype_plan: dict, entries: dict): + """Combine groups without making their order decide a target's quantization format.""" + for pattern, config in entries.items(): + if pattern in dtype_plan and dtype_plan[pattern].to_dict() != config.to_dict(): + raise ValueError(f"Conflicting quantization configs for target {pattern!r}") + dtype_plan[pattern] = config + + +def read_modelopt(config_dict: dict) -> QuantizationConfigMixin: + """Read a ModelOpt format or its explicit per-module `MIXED_PRECISION` assignment. + + ModelOpt's `quant_cfg.json` nests these fields under `quantization`; HF configs can store them directly. + Unsupported groups and algorithms raise rather than loading their tensors as blanket NVFP4. + """ + config_dict = {**config_dict, **config_dict.get("quantization", {})} + _check_block_scaled_features(config_dict) + quant_algo = config_dict.get("quant_algo") + if config_dict.get("kv_cache_scheme") or config_dict.get("kv_cache_quant_algo"): + logger.warning_once( + "This modelopt checkpoint calibrates a quantized KV cache; the KV cache stays in the model's dtype, " + "so its scales are not loaded." + ) + # Producer exclusions are globs; an explicitly supplied native skip-list already uses + # the library's regex semantics and must not be escaped or discarded. + skip_fields = ("modules_to_not_convert", "ignore", "exclude_modules") + native_skip = config_dict.get("modules_to_not_convert") or [] + producer_skip = [*(config_dict.get("ignore") or []), *(config_dict.get("exclude_modules") or [])] + modules_to_not_convert = ( + list(dict.fromkeys([*native_skip, *subtree_patterns(producer_skip)])) + if any(config_dict.get(field) is not None for field in skip_fields) + else None + ) + if quant_algo == "MIXED_PRECISION": + layers = config_dict.get("quantized_layers") + if not isinstance(layers, dict) or not layers: + raise ValueError("modelopt MIXED_PRECISION requires a non-empty quantized_layers mapping") + dtype_plan = {} + for name, spec in layers.items(): + if not isinstance(name, str) or not name or "*" in name: + raise ValueError(f"modelopt quantized_layers requires concrete module names; got {name!r}") + try: + dtype_plan[subtree_patterns([name])[0]] = _modelopt_format(spec) + except ValueError as error: + raise ValueError(f"modelopt module {name!r}: {error}") from error + config = FineGrainedConfig(dtype_plan=dtype_plan, modules_to_not_convert=modules_to_not_convert) + return config + if groups := config_dict.get("config_groups"): + dtype_plan = {} + # Validate the declared algorithm too; groups cannot disguise an unsupported export. + _modelopt_format(config_dict) + for name, spec in groups.items(): + _check_block_scaled_features(spec) + spec = dict(spec) + targets = [] + for target in spec.get("targets") or ["Linear"]: + if target == "Linear" or target.startswith("re:"): + targets.append(target) + elif "." in target or target[:1].islower(): + targets.append("re:" + subtree_patterns([target])[0]) + else: + raise ValueError(f"Unsupported modelopt target {target!r} in config group {name!r}") + spec["targets"] = targets + entries = plan_from_config_groups(spec) + if entries is None: + raise ValueError(f"Unsupported modelopt config_groups entry {name!r}: {spec}") + for config in entries.values(): + config.checkpoint_format = "modelopt" + _merge_config_group_plan(dtype_plan, entries) + return _config_for_plan(dtype_plan, modules_to_not_convert) + config = _modelopt_format(config_dict) + config.modules_to_not_convert = modules_to_not_convert + return config + + +def read_compressed_tensors(config_dict: dict) -> QuantizationConfigMixin | None: + """The block-scaled config a compressed-tensors checkpoint describes, or `None` when it is not one: integer, + sparse or transformed checkpoints stay with [`CompressedTensorsConfig`]. Each config's + `checkpoint_format` names the layout its tensors ship in, which a save writes back.""" + try: + _check_block_scaled_features(config_dict) + except ValueError: + return None + if config_dict.get("quantization_status", "compressed") != "compressed": + return None + groups = config_dict.get("config_groups") or {} + served = ("float-quantized", "nvfp4-pack-quantized") + if not groups or any(spec.get("format", config_dict.get("format")) not in served for spec in groups.values()): + return None + mapped = [plan_from_config_groups(spec) for spec in groups.values()] + if any(entries is None for entries in mapped): + return None + dtype_plan = {} + try: + for entries in mapped: + if entries is not None: + _merge_config_group_plan(dtype_plan, entries) + except ValueError: + return None + for config in dtype_plan.values(): + config.checkpoint_format = "compressed-tensors" + if config_dict.get("kv_cache_scheme"): + logger.warning_once( + "This compressed-tensors checkpoint calibrates a quantized KV cache; the KV cache is not quantized here " + "and stays in the model's dtype, so those scales are not loaded." + ) + # entries are module names, or `re:`-prefixed patterns; a bare name is that module and its sub-tree, not a prefix + ignore = config_dict.get("ignore") or [] + skip = [entry[3:] for entry in ignore if entry.startswith("re:")] + skip += subtree_patterns([entry for entry in ignore if not entry.startswith("re:")]) + return _config_for_plan(dtype_plan, skip or None) + + +def _config_for_plan(dtype_plan: dict, modules_to_not_convert: list | None) -> QuantizationConfigMixin: + """The one format config of a plan that covers everything with it, else the `FineGrainedConfig` of the plan.""" + if list(dtype_plan) == [".*"]: + (config,) = dtype_plan.values() + config.modules_to_not_convert = modules_to_not_convert + return config + return FineGrainedConfig(dtype_plan=dtype_plan, modules_to_not_convert=modules_to_not_convert) + + +class FineGrainedConfig(QuantizationConfigMixin): + """ + Configuration for a model whose modules are quantized in different fine-grained formats: which format covers + which modules, each as its own format config ([`BlockFp8Config`], [`Mxfp8Config`], [`Mxfp4Config`], + [`NVFP4Config`]), and the options that hold for the whole model. A model in one format takes that format's config + alone. + + | format | weights | block scales | a checkpoint that ships it | + |-----------|--------------------------------|-----------------------------|----------------------------| + | block-FP8 | `float8_e4m3fn` | fp32 or UE8M0, 128x128 | `deepseek-ai/DeepSeek-V3` | + | MXFP8 | `float8_e4m3fn` | UE8M0, group-32 | `MiniMaxAI/MiniMax-M3` | + | MXFP4 | E2M1, two values per `int8` | UE8M0, group-32 | `openai/gpt-oss-20b` | + | NVFP4 | E2M1, two values per `int8` | E4M3 group-16 + fp32 global | `nvidia/GLM-5.2-NVFP4` | + + Args: + dtype_plan (`dict[str, BlockScaledConfig]`, *optional*): + A module regex and the format config the modules it matches load in. A targeted entry wins over the + `".*"` catch-all; modules no entry covers stay in full precision. An entry takes no loading options + (`dequantize`, `modules_to_not_convert`, `modules_to_convert`): those are this config's. On a + pre-quantized checkpoint, a passed plan's entries replace the activation formats of the checkpoint's + entries under the same pattern. + modules_to_not_convert (`list`, *optional*): + Modules to keep in full precision. + modules_to_convert (`list`, *optional*): + Embedding tables to quantize, by name suffix: FP8 rows with one per-tensor scale. + dequantize (`bool`, *optional*, defaults to `False`): + Whether to dequantize the model to full precision while loading. + """ + + # what an entry of the plan does not hold: they are the whole model's + _entry_loading_options = ("dequantize", "modules_to_not_convert", "modules_to_convert") + + def __init__( + self, + *, + dtype_plan: dict[str, BlockScaledConfig] | None = None, + modules_to_not_convert: list | None = None, + modules_to_convert: list | None = None, + dequantize: bool = False, + **kwargs, + ): + self.quant_method = QuantizationMethod.FINEGRAINED + kwargs.pop("quant_method", None) + if kwargs: + logger.info(f"Unused kwargs: {list(kwargs)}. These kwargs are not used in {type(self).__name__}.") + self.dtype_plan = {pattern: self._plan_entry(pattern, entry) for pattern, entry in (dtype_plan or {}).items()} + self.modules_to_not_convert = modules_to_not_convert + self.modules_to_convert = modules_to_convert + self.dequantize = dequantize + + def _plan_entry(self, pattern: str, entry) -> BlockScaledConfig: + """`entry` as its format config: from `config.json`, a plain dict naming its format.""" + if isinstance(entry, dict): + method = entry.get("quant_method") + if method not in BLOCK_SCALED_CONFIGS: + raise ValueError( + f"dtype_plan entry {pattern!r} has quant_method {method!r}; an entry is one of " + f"{sorted(BLOCK_SCALED_CONFIGS)}" + ) + entry = BLOCK_SCALED_CONFIGS[method].from_dict(entry) + if not isinstance(entry, BlockScaledConfig): + raise TypeError( + f"dtype_plan entry {pattern!r} is a {type(entry).__name__}; an entry is one of " + f"{', '.join(config.__name__ for config in BLOCK_SCALED_CONFIGS.values())}" + ) + loading = [name for name in self._entry_loading_options if getattr(entry, name, None)] + if loading: + raise ValueError( + f"dtype_plan entry {pattern!r} sets {loading}, which hold for the whole model: set them on the " + "FineGrainedConfig" + ) + return entry + + def to_dict(self) -> dict[str, Any]: + return { + "quant_method": self.quant_method, + "dtype_plan": { + pattern: { + name: value for name, value in config.to_dict().items() if name not in self._entry_loading_options + } + for pattern, config in self.dtype_plan.items() + }, + "modules_to_not_convert": self.modules_to_not_convert, + "modules_to_convert": self.modules_to_convert, + "dequantize": self.dequantize, + } + + def get_loading_attributes(self): + return _loading_attributes(self, **({"dtype_plan": self.dtype_plan} if self.dtype_plan else {})) + + class FineGrainedFP8Config(QuantizationConfigMixin): """ - FineGrainedFP8Config is a configuration class for fine-grained FP8 quantization used mainly for deepseek models. + Deprecated: passed to `from_pretrained`, it loads through its own frozen integration. Use [`BlockFp8Config`], or + [`Mxfp8Config`] for a `[1, 32]` block with UE8M0 scales. Args: activation_scheme (`str`, *optional*, defaults to `"dynamic"`): - The scheme used for activation, the defaults and only support scheme for now is "dynamic". + `"dynamic"` (quantized inline) or `"static"` (a calibrated per-tensor activation scale from the + checkpoint). weight_block_size (`typing.tuple[int, int]`, *optional*, defaults to `(128, 128)`): - The size of the weight blocks for quantization, default is (128, 128). + The block each weight scale covers; `None` for one scale per tensor. dequantize (`bool`, *optional*, defaults to `False`): Whether to dequantize the model during loading. modules_to_not_convert (`list`, *optional*): @@ -1717,12 +2339,16 @@ def __init__( scale_fmt: str = "float", **kwargs, ): + warnings.warn( + "`FineGrainedFP8Config` is deprecated and will be removed in a future release: use `BlockFp8Config`, or " + "`Mxfp8Config` for a `[1, 32]` block with UE8M0 scales.", + FutureWarning, + ) self.quant_method = kwargs.pop("quant_method", QuantizationMethod.FP8) # MiniMax ships the skip-list under ``ignored_layers``; accept it as an alias. if modules_to_not_convert is None and "ignored_layers" in kwargs: modules_to_not_convert = kwargs.pop("ignored_layers") self.modules_to_not_convert = modules_to_not_convert - # TODO: check overlap with not to convert self.modules_to_convert = modules_to_convert self.activation_scheme = activation_scheme self.weight_block_size = weight_block_size @@ -1789,41 +2415,6 @@ def __init__( self.quant_method = QuantizationMethod.QUARK -@dataclass -class Mxfp4Config(QuantizationConfigMixin): - """ - This is a wrapper class about all possible attributes and features that you can play with a model that has been - loaded using mxfp4 quantization. - - Args: - modules_to_not_convert (`list`, *optional*, default to `None`): - The list of modules to not quantize, useful for quantizing models that explicitly require to have - some modules left in their original precision. - dequantize (`bool`, *optional*, default to `False`): - Whether we dequantize the model to bf16 precision or not - """ - - def __init__( - self, - modules_to_not_convert: list | None = None, - dequantize: bool = False, - **kwargs, - ): - self.quant_method = QuantizationMethod.MXFP4 - self.modules_to_not_convert = modules_to_not_convert - self.dequantize = dequantize - - def get_loading_attributes(self): - return {"dequantize": self.dequantize} - - def to_dict(self) -> dict[str, Any]: - """ - Serializes this instance to a Python dictionary. Returns: - `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. - """ - return {"quant_method": self.quant_method, "modules_to_not_convert": self.modules_to_not_convert} - - class MetalConfig(QuantizationConfigMixin): """ Configuration class for Metal affine quantization targeting Apple Silicon (MPS) devices. @@ -2066,25 +2657,6 @@ def __init__( self.modules_to_not_convert = modules_to_not_convert -class NVFP4Config(QuantizationConfigMixin): - """Configuration for on-the-fly NVFP4 weight quantization. - - Args: - modules_to_not_convert (`list[str]`, *optional*): - Module-name patterns that should remain in their original precision. - kwargs (`dict[str, Any]`, *optional*): - Additional values are ignored and reported through the Transformers logger. - """ - - def __init__(self, modules_to_not_convert: list[str] | None = None, **kwargs): - self.quant_method = QuantizationMethod.NVFP4 - self.modules_to_not_convert = modules_to_not_convert - if kwargs: - logger.info( - f"Unused kwargs: {list(kwargs.keys())}. These kwargs are not used in {self.__class__.__name__}." - ) - - @dataclass class GgufConfig(QuantizationConfigMixin): """Load a quantized GGUF checkpoint with its weights left in GGUF blocks. diff --git a/tests/integrations/finegrained/__init__.py b/tests/integrations/finegrained/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/integrations/finegrained/test_conversions.py b/tests/integrations/finegrained/test_conversions.py new file mode 100644 index 000000000000..a1a160658e49 --- /dev/null +++ b/tests/integrations/finegrained/test_conversions.py @@ -0,0 +1,750 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""The `finegrained` conversion ops a load and a save run whatever the format: the layout ops (gate|up +interleave, scale container, swizzle) and the ops each expert tensor takes from its module, and a core chain +composed with on-the-fly quantization. Each format's checkpoint layouts are tested in its own folder. Pure +torch: no kernel is loaded and no GPU is needed.""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +import transformers.integrations.finegrained.layers as fg +from transformers import PreTrainedConfig, PreTrainedModel +from transformers.core_model_loading import ( + Chunk, + Concatenate, + PermuteForRope, + Transpose, + WeightConverter, + convert_and_load_state_dict_in_model, + revert_weight_conversion, +) +from transformers.integrations.moe import use_experts_implementation +from transformers.modeling_utils import LoadStateDictConfig +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import ( + BlockFp8Config, + Mxfp4Config, + Mxfp8Config, + NVFP4Config, +) + +from .utils import ( + E, + H, + I, + _Cfg, + _format_config, + _load_and_save, + _model_mapping, + _MoeModel, + _per_expert, + _quantizer_for, + _replace, + fake_finegrained_kernel, + patch_finegrained_kernel, +) + + +@use_experts_implementation(is_transposed=True) +class _TransposedExperts(torch.nn.Module): + """GPT-OSS's layout: the stacked experts stored `(E, in, out)`.""" + + def __init__(self, config): + super().__init__() + self.gate_up_proj = torch.nn.Parameter(torch.empty(E, H, 2 * I)) + self.down_proj = torch.nn.Parameter(torch.empty(E, I, H)) + + +@require_torch +class FineGrainedLayoutOpsTest(unittest.TestCase): + """The module HOLDS its block scales swizzled (a 5-D Parameter) when its chain runs the tcgen05 + scaled-MMA under a triton dispatch on SM100 — keyed off the declared weight format, not the + scale dtype (V4-style block-FP8 ships UE8M0 scales) — and affine otherwise; the loader's + `FineGrainedSwizzleScales` op fills that layout and its reverse restores the checkpoint grid.""" + + E, HIDDEN, INTER = 4, 256, 128 # gate_up rows 2*INTER = 256 and down rows HIDDEN = 256: whole 128-row blocks + + def _experts(self, weight_format, activation_format=None, impl="grouped_mm", sm100=True): + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = self.HIDDEN, self.INTER, self.E + cfg._experts_implementation = impl + with ( + mock.patch.object(fg, "is_sm100", return_value=sm100), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts( + cfg, + quantization_config=_format_config( + weight_format, + activation_format=activation_format, + # the group formats fix their own block + **({"weight_block_size": (128, 128)} if weight_format == "fp8" else {}), + ), + ) + model = torch.nn.Module() + model.config = cfg + model.experts = experts + return model, experts + + def test_module_holds_swizzled_scales_for_the_scaled_mma_chain(self): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2, 2, 256)) # (E, 256/128, (256/32)/4, 2, 256) + self.assertEqual(experts.down_proj_scale_inv.shape, (4, 2, 1, 2, 256)) + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertTrue(experts.holds_interleaved_gate_up) + # every MX group-32 format holds UE8M0 scales whatever `scale_fmt` says (the kernels reject a + # float32 grid there; GPT-OSS's config has no scale_fmt at all) + _, experts = self._experts("mxfp4", activation_format="bf16") + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.float8_e8m0fnu) + + def test_affine_off_sm100_weight_only_deepgemm_and_block_fp8(self): + for kw in ( + {"weight_format": "mxfp8", "sm100": False}, + {"weight_format": "mxfp4", "activation_format": "bf16"}, # W4A16 reads scales per group affinely + {"weight_format": "mxfp8", "impl": "deepgemm"}, # the DeepGEMM backends read affine scales + {"weight_format": "nvfp4", "impl": "deepgemm_megamoe"}, + ): + with self.subTest(**kw): + _, experts = self._experts(**kw) + self.assertEqual( + experts.gate_up_proj_scale_inv.shape, + (4, 256, 256 // (16 if kw["weight_format"] == "nvfp4" else 32)), + ) + # block-FP8 has no per-group scale grid to swizzle: one scalar per (N/128, K/128) block + _, experts = self._experts("fp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 2, 2)) + with mock.patch.object(fg, "_holds_swizzled_scales", return_value=False): + _, experts = self._experts("mxfp8") + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (4, 256, 8)) + + def test_megamoe_holds_gate_up_stacked(self): + _, experts = self._experts("fp8", impl="deepgemm_megamoe") + self.assertFalse(experts.holds_interleaved_gate_up) + + def test_deepgemm_holds_gate_up_stacked(self): + """DeepGEMM's grouped GEMM reads a block grid row-major, which interleaved rows do not follow.""" + _, experts = self._experts("fp8", impl="deepgemm") + self.assertFalse(experts.holds_interleaved_gate_up) + + def _op_kernel(self): + kernel, _ = fake_finegrained_kernel() + kernel.swizzle_mx_scales = mock.Mock( + side_effect=lambda s: s.reshape(s.shape[0], s.shape[1] // 128, -1, 2, 256) + ) + kernel.unswizzle_mx_scales = mock.Mock( + side_effect=lambda s, rows, cols, num_experts=None: s.reshape(num_experts, rows, cols) + ) + return kernel + + def test_swizzle_op_fills_only_the_scales_the_module_holds_swizzled(self): + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + model, experts = self._experts("mxfp8") + kernel = self._op_kernel() + op = FineGrainedSwizzleScales(hf_quantizer=None) + grid = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + weight = torch.zeros(4, 256, 256, dtype=torch.float8_e4m3fn) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + # the loader hands a single target under its pattern, the full name (with suffix) alongside + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, model=model, full_layer_name="experts.gate_up_proj_scale_inv" + ) + self.assertEqual(out["mlp.experts.gate_up_proj"].shape, (4, 2, 2, 2, 256)) + # the same converter's weight passes through + out = op.convert({"mlp.experts.gate_up_proj": weight}, model=model, full_layer_name="experts.gate_up_proj") + self.assertIs(out["mlp.experts.gate_up_proj"], weight) + # a multi-tensor converter (GPT-OSS deserialize) names its outputs fully + out = op.convert( + { + "experts.down_proj": weight, + "experts.down_proj_scale_inv": torch.zeros(4, 256, 4, dtype=torch.float8_e8m0fnu), + }, + model=model, + full_layer_name="experts.down_proj", + ) + self.assertEqual(out["experts.down_proj_scale_inv"].shape, (4, 2, 1, 2, 256)) + self.assertIs(out["experts.down_proj"], weight) + # a module that holds affine scales keeps them + affine_model, _ = self._experts("mxfp4", activation_format="bf16") + out = op.convert( + {"mlp.experts.gate_up_proj": grid}, + model=affine_model, + full_layer_name="experts.gate_up_proj_scale_inv", + ) + self.assertIs(out["mlp.experts.gate_up_proj"], grid) + self.assertEqual(kernel.swizzle_mx_scales.call_count, 2) + + def test_fused_dense_scales_save_without_changing_the_runtime_dtype(self): + from transformers.integrations.finegrained.conversions import FineGrainedDeserialize + + config = Mxfp8Config() + model = torch.nn.Module() + model.config = PreTrainedConfig() + model.mlp = torch.nn.Module() + model.mlp.gate_up_proj = fg.FineGrainedLinear(64, 64, quantization_config=config) + scales = torch.arange(120, 128, dtype=torch.uint8).repeat(16, 1).reshape(64, 2) + model._weight_conversions = [ + WeightConverter( + [r"\.gate_proj.weight_scale_inv", r"\.up_proj.weight_scale_inv"], + r"\.gate_up_proj.weight_scale_inv", + [Concatenate(dim=0), FineGrainedDeserialize(_quantizer_for(config))], + ) + ] + saved = revert_weight_conversion( + model, {"mlp.gate_up_proj.weight_scale_inv": scales.view(torch.float8_e8m0fnu)} + ) + for name, expected in zip(("gate_proj", "up_proj"), scales.chunk(2)): + actual = saved[f"mlp.{name}.weight_scale_inv"] + self.assertEqual(actual.dtype, torch.float8_e8m0fnu) + self.assertTrue(torch.equal(actual.view(torch.uint8), expected)) + + def test_reverse_op_restores_the_affine_grid_from_the_artifact_alone(self): + from transformers.integrations.finegrained.conversions import FineGrainedSwizzleScales + + kernel = self._op_kernel() + reverse = FineGrainedSwizzleScales(hf_quantizer=None).reverse_op + artifact = torch.zeros(4, 2, 2, 2, 256, dtype=torch.float8_e8m0fnu) + affine = torch.zeros(4, 256, 8, dtype=torch.float8_e8m0fnu) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = reverse.convert({"experts.gate_up_proj_scale_inv": artifact}, model=None, full_layer_name="x") + self.assertEqual(out["experts.gate_up_proj_scale_inv"].shape, (4, 256, 8)) + kernel.unswizzle_mx_scales.assert_called_once_with(artifact, 256, 8, num_experts=4) + out = reverse.convert({"experts.gate_up_proj_scale_inv": affine}, model=None, full_layer_name="x") + self.assertIs(out["experts.gate_up_proj_scale_inv"], affine) + + def test_expert_tensors_take_their_layout_from_their_module(self): + """The layout an expert tensor lands in comes from the module it fills, whatever converter + brought it: gate|up interleaved (weight, scale grid and bias), then the scale + swizzle; a bias only interleaves.""" + from transformers.integrations.finegrained.conversions import ( + FineGrainedInterleaveGateUp, + FineGrainedSwizzleScales, + ) + + model, experts = self._experts("mxfp4") + experts.quantization_config = Mxfp4Config(activation_format=None) # what the replacement walk records + quantizer = _quantizer_for(Mxfp4Config()) + layout = [FineGrainedInterleaveGateUp, FineGrainedSwizzleScales] + for param, expected in ( + ("experts.gate_up_proj", layout), + ("experts.gate_up_proj_scale_inv", layout), + ("experts.gate_up_proj_bias", [FineGrainedInterleaveGateUp]), + ("experts.down_proj_scale_inv", [FineGrainedSwizzleScales]), + ("experts.down_proj_bias", []), + ): + ops = quantizer.get_param_ops(model, param) + self.assertEqual([type(op) for op in ops], expected, param) + + def test_a_key_under_its_fused_name_takes_its_modules_ops_by_its_full_name(self): + """A key already under the fused name (GPT-OSS's gate_up_proj_bias, a renamed MiniMax + down_proj_scale_inv) has no converter of its own: the loader runs its module's ops in a + one-to-one converter, whose ops receive the tensor under the key and hand it back by name.""" + from transformers.core_model_loading import WeightConverter + + model, experts = self._experts("mxfp8") + experts.quantization_config = Mxfp8Config() # what the replacement walk records + quantizer = _quantizer_for(Mxfp8Config()) + for key, shape in (("experts.gate_up_proj_bias", (4, 256)), ("experts.down_proj_scale_inv", (4, 256, 4))): + converter = WeightConverter(key, key, operations=quantizer.get_param_ops(model, key)) + tensor = torch.zeros(shape, dtype=torch.float8_e8m0fnu if "scale" in key else torch.float32) + converter.add_tensor(key, key, key, lambda t=tensor: t) + kernel = self._op_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = converter.convert(key, model=model) + self.assertEqual(list(out), [key]) + + def test_replacement_keeps_the_models_gate_up_convention(self): + """`FineGrainedExperts.from_module` goes through `use_experts_implementation`, which stamps + the layout flags on the instance after __init__ — GPT-OSS's `is_concatenated=False` must + survive it (it decides whether the loader interleaves).""" + from transformers import GptOssConfig, GptOssForCausalLM, Qwen3MoeConfig, Qwen3MoeForCausalLM + + gpt_oss = GptOssConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=32, + num_local_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + sliding_window=16, + ) + qwen = Qwen3MoeConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_experts=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + ) + for cls, cfg, quant_method, expected in ( + (GptOssForCausalLM, gpt_oss, "mxfp4", False), + (Qwen3MoeForCausalLM, qwen, "mxfp8", True), + ): + with self.subTest(model=cls.__name__), torch.device("meta"): + model = cls(cfg) + _replace(model, _format_config(quant_method, flat=True)) + experts = model.model.layers[0].mlp.experts + self.assertIsInstance(experts, fg.FineGrainedExperts) + self.assertIs(experts.is_concatenated, expected) + + def test_on_the_fly_quantization_still_interleaves_expert_biases(self): + from transformers.integrations.finegrained.conversions import FineGrainedDeserialize + + model, experts = self._experts("fp8", sm100=False) + quantizer = _quantizer_for(experts.quantization_config) + quantizer.pre_quantized = False + bias = torch.arange(4 * 256, dtype=torch.float32).reshape(4, 256) + key = "experts.gate_up_proj_bias" + actual = FineGrainedDeserialize(quantizer).convert({key: bias}, model=model, full_layer_name=key)[key] + gate, up = bias.chunk(2, dim=1) + torch.testing.assert_close(actual, torch.stack((gate, up), dim=-1).flatten(1)) + + def test_interleave_op_follows_the_layer_it_fills(self): + """The op reads the experts module a key fills: quantized experts get the interleaved order and the + reverse restores the stack, while a model whose rows already come interleaved (`is_concatenated=False`, + GPT-OSS), Mega MoE (which packs gate|up itself) and full-precision experts of a mixed checkpoint keep + the stack.""" + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp + + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) # rows [g0,g1,g2,u0,u1,u2] + interleaved = stacked[:, [0, 3, 1, 4, 2, 5]] + full_precision = torch.nn.Module() + full_precision.gate_up_proj = torch.nn.Parameter(torch.empty(2, 6, 4)) + full_precision.down_proj = torch.nn.Parameter(torch.empty(2, 4, 3)) + model = torch.nn.Module() + model.layers = torch.nn.ModuleList([torch.nn.Module(), torch.nn.Module()]) + model.layers[0].experts = full_precision + model.layers[1].experts = quantized = self._experts("mxfp8")[1] + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + + for layer, expected in ((0, stacked), (1, interleaved)): + key = f"layers.{layer}.experts.gate_up_proj" + out = op.convert({key: stacked}, model=model, full_layer_name=key)[key] + torch.testing.assert_close(out, expected) + saved_key = f"layers.{layer}.experts.*.gate_proj.weight" + back = op.reverse_op.convert({saved_key: out}, model=model, full_layer_name=saved_key)[saved_key] + torch.testing.assert_close(back, stacked) + + key = "layers.1.experts.gate_up_proj" + quantized.is_concatenated = False + self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) + quantized.is_concatenated, quantized.holds_interleaved_gate_up = True, False + self.assertIs(op.convert({key: stacked}, model=model, full_layer_name=key)[key], stacked) + + def test_interleave_op_stacks_an_interleaved_checkpoint_for_stacked_experts(self): + """An interleaved checkpoint (GPT-OSS) loaded for a backend that holds gate|up stacked (DeepGEMM) is + de-interleaved on load and interleaved back on save.""" + from transformers.integrations.finegrained.conversions import FineGrainedInterleaveGateUp + + stacked = torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4) + interleaved = stacked[:, [0, 3, 1, 4, 2, 5]] + model, experts = self._experts("mxfp8", impl="deepgemm") + experts.is_concatenated = False + op = FineGrainedInterleaveGateUp(hf_quantizer=None) + key = "experts.gate_up_proj" + loaded = op.convert({key: interleaved}, model=model, full_layer_name=key)[key] + torch.testing.assert_close(loaded, stacked) + torch.testing.assert_close( + op.reverse_op.convert({key: loaded}, model=model, full_layer_name=key)[key], interleaved + ) + + def test_ops_refuse_a_list_of_several_tensors(self): + """The loader hands a single-source tensor in a one-element list; an op that unwraps it must not + silently keep the first of several.""" + from transformers.integrations.finegrained.conversions import ( + FineGrainedInputScalesSplit, + FineGrainedPackedBlocks, + FineGrainedQuantize, + FineGrainedViewPackedInt8, + ) + + tensor = torch.zeros(2, 4, 32, 16, dtype=torch.uint8) + for op in ( + FineGrainedPackedBlocks(None), + FineGrainedViewPackedInt8(None), + FineGrainedInputScalesSplit(None), + FineGrainedQuantize(None), + ): + with self.subTest(op=type(op).__name__), self.assertRaises(ValueError): + op.convert({"experts.down_proj": [tensor, tensor]}) + # a single target, re-keyed by the op before anything reads it + with self.assertRaises(ValueError): + FineGrainedViewPackedInt8(None).convert({"x": [tensor, tensor]}, target_patterns=["experts.down_proj"]) + + def test_per_half_globals_are_not_folded_past_a_glu_nonlinear_in_up(self): + """Folding the up half's NVFP4 global onto the down projection assumes the GLU is linear in `up`; a + clamped SwiGLU or SwiGLU-OAI's `(up + 1)` is not, so differing halves are refused there. Halves that + agree have nothing to fold.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = ["experts.gate_up_proj_weight_global_scale", "experts.down_proj_weight_global_scale"] + for attribute in ("swiglu_limit", "swiglu_alpha"): + config = _Cfg() + setattr(config, attribute, 7.0) + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts(config, quantization_config=NVFP4Config()) + with self.subTest(attribute=attribute): + with self.assertRaises(NotImplementedError): + FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 3.0], [2.0, 4.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), + }, + target_patterns=targets, + model=model, + ) + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 1.0], [2.0, 2.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), + }, + target_patterns=targets, + model=model, + ) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([5.0, 6.0])) + + +@require_torch +class FineGrainedLoaderChainTest(unittest.TestCase): + def test_rowwise_fp8_checkpoint_scales_load_and_save_in_their_producer_layout(self): + """Producer vectors/grids and transposed BMM weights share one runtime row-scale layout.""" + + def tensors(rows, cols, vector): + return { + "weight": torch.randn(rows, cols).to(torch.float8_e4m3fn), + "weight_scale": torch.rand((rows,) if vector else (rows, 1)) + 0.1, + } + + for producer, transposed in ( + ("modelopt", False), + ("compressed-tensors", False), + ("fbgemm", False), + ("modelopt", True), + ("fbgemm", True), + ): + with self.subTest(producer=producer, transposed=transposed): + source = _per_expert(lambda rows, cols: tensors(rows, cols, producer == "modelopt")) + checkpoint, mapping = source, _model_mapping() + if transposed: + prefix = "model.layers.0.mlp.experts" + checkpoint = {key: value for key, value in source.items() if not key.startswith(prefix)} + mapping = [] + for fused, projections in ( + ("gate_up_proj", ("gate_proj", "up_proj")), + ("down_proj", ("down_proj",)), + ): + weight = torch.stack( + [ + torch.cat([source[f"{prefix}.{expert}.{proj}.weight"] for proj in projections]) + for expert in range(E) + ] + ) + scale = torch.stack( + [ + torch.cat( + [ + source[f"{prefix}.{expert}.{proj}.weight_scale"].reshape(-1, 1) + for proj in projections + ] + ) + for expert in range(E) + ] + ) + key = f"{prefix}.{fused}" + checkpoint[key] = weight.transpose(1, 2).contiguous() + suffix = "weight_scale" if producer == "modelopt" else "scale" + checkpoint[f"{key}_{suffix}"] = ( + scale.squeeze(-1) + if producer == "modelopt" + else scale.transpose(1, 2) + if fused == "gate_up_proj" + else scale + ) + mapping.append(WeightConverter(key, key, [Transpose(1, 2)])) + config = BlockFp8Config(weight_block_size="row", checkpoint_format=producer) + model, info, saved = _load_and_save(checkpoint, config, mapping) + self.assertEqual( + (info.missing_keys, info.unexpected_keys, info.mismatched_keys, info.conversion_errors), + (set(), set(), set(), {}), + ) + experts = model.model.layers[0].mlp.experts + self.assertEqual(experts.gate_up_proj_scale_inv.shape, (E, 2 * I, 1)) + self.assertEqual(experts.down_proj_scale_inv.shape, (E, H, 1)) + self.assertEqual(set(saved), set(checkpoint)) + for key, value in checkpoint.items(): + torch.testing.assert_close(saved[key].float(), value.float(), rtol=0, atol=0) + dequantized, info, _ = _load_and_save( + checkpoint, + BlockFp8Config(weight_block_size="row", checkpoint_format=producer, dequantize=True), + mapping, + ) + self.assertEqual((info.missing_keys, info.unexpected_keys, info.conversion_errors), (set(), set(), {})) + for fused, projections in (("gate_up_proj", ("gate_proj", "up_proj")), ("down_proj", ("down_proj",))): + expected = torch.stack( + [ + torch.cat( + [ + source[f"model.layers.0.mlp.experts.{expert}.{proj}.weight"].float() + * source[f"model.layers.0.mlp.experts.{expert}.{proj}.weight_scale"].reshape(-1, 1) + for proj in projections + ] + ) + for expert in range(E) + ] + ) + torch.testing.assert_close(getattr(dequantized.model.layers[0].mlp.experts, fused), expected) + + """A core conversion chain and an on-the-fly quantization in the same load: the permute has to + run on the full-precision tensor, and only the finegrained module of the three gets a scale.""" + + def _quantize_on_load(self, model, checkpoint, weight_mapping): + """`checkpoint`'s full-precision weights quantized on load into 32x32-block FP8, off SM100.""" + quantizer = _quantizer_for(BlockFp8Config(weight_block_size=(32, 32))) + quantizer.pre_quantized = False + with mock.patch.object(fg, "is_sm100", return_value=False): + quantizer._process_model_before_weight_loading(model) + load_config = LoadStateDictConfig( + weight_mapping=quantizer.update_weight_conversions(weight_mapping, model), + hf_quantizer=quantizer, + ) + info, _ = convert_and_load_state_dict_in_model(model, checkpoint, load_config) + self.assertEqual(info.conversion_errors, {}) + model.hf_quantizer = quantizer + quantizer._process_model_after_weight_loading(model) + return model.model.layers[0].mlp.experts + + def test_dequantizing_a_fused_nvfp4_stack_applies_each_experts_globals(self): + """The vLLM layout ships each stack's globals per expert, `(E,)`, and a gate|up stack's per half, + `(E, 2)`: the stack's dequantize scales each expert's rows, and each half's, by its own.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + + torch.manual_seed(0) + checkpoint, expected = {}, {} + for proj, rows, cols, halves in (("gate_up_proj", 2 * I, H, 2), ("down_proj", H, I, 1)): + codes = torch.randint(0, 256, (E, rows, cols // 2), dtype=torch.uint8) + blocks = (torch.rand(E, rows, cols // 16) + 0.5).to(torch.float8_e4m3fn) + globals_ = torch.pow(2.0, torch.randint(-8, 0, (E, halves)).float()) + prefix = f"model.layers.0.mlp.experts.{proj}" + checkpoint |= { + prefix: codes, + f"{prefix}_weight_scale": blocks, + f"{prefix}_weight_scale_2": globals_.squeeze(1) if halves == 1 else globals_, + } + half_rows = rows // halves + expected[proj] = torch.stack( + [ + torch.cat( + [ + FineGrainedDequantize(None)._dequantize_one( + codes[e, h * half_rows : (h + 1) * half_rows], + blocks[e, h * half_rows : (h + 1) * half_rows], + torch.float32, + globals_[e, h], + ) + for h in range(halves) + ] + ) + for e in range(E) + ] + ) + model, info, _ = _load_and_save(checkpoint, NVFP4Config(dequantize=True)) + self.assertEqual(info.conversion_errors, {}) + experts = model.model.layers[0].mlp.experts + for proj, reference in expected.items(): + torch.testing.assert_close(getattr(experts, proj).float(), reference, msg=proj) + + def test_on_the_fly_quantization_blocks_each_projection_whole(self): + """The gate|up interleave runs on the quantized weight and its grid, so every 32-row block holds rows of + ONE projection, scaled by its own magnitude: a gate four orders below its up half keeps its values. + Both stack layouts: a per-expert checkpoint concatenated by the model's converter, and one stored + transposed `(E, H, 2I)`, which the quantization reads as rows.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + + torch.manual_seed(0) + gate, up, down = torch.randn(E, I, H) * 1e-5, torch.randn(E, I, H), torch.randn(E, H, I) + dense = {"model.layers.0.self_attn.q_proj.weight": torch.randn(H, H)} + per_expert = { + f"model.layers.0.mlp.experts.{expert}.{proj}.weight": tensor[expert] + for expert in range(E) + for proj, tensor in (("gate_proj", gate), ("up_proj", up), ("down_proj", down)) + } + transposed = { + "model.layers.0.mlp.experts.gate_up_proj": torch.cat([gate, up], dim=1).transpose(1, 2).contiguous(), + "model.layers.0.mlp.experts.down_proj": down.transpose(1, 2).contiguous(), + } + config = PreTrainedConfig() + config.hidden_size, config.intermediate_size, config.num_local_experts, config.hidden_act = H, I, E, "silu" + for layout, checkpoint, mapping in ( + ("concatenated", per_expert, _model_mapping()), + ("transposed", transposed, []), + ): + with self.subTest(layout=layout): + model = _MoeModel(config) + if layout == "transposed": + model.model.layers[0].mlp.experts = _TransposedExperts(config) + experts = self._quantize_on_load(model, {**dense, **checkpoint}, mapping) + saved = revert_weight_conversion(model, model.state_dict()) + reloaded, info, _ = _load_and_save(saved, BlockFp8Config(weight_block_size=(32, 32))) + self.assertEqual((info.missing_keys, info.unexpected_keys, info.conversion_errors), (set(), set(), {})) + restored = reloaded.model.layers[0].mlp.experts + for name in ("gate_up_proj", "gate_up_proj_scale_inv", "down_proj", "down_proj_scale_inv"): + torch.testing.assert_close( + getattr(restored, name).float(), getattr(experts, name).float(), rtol=0, atol=0 + ) + + self.assertTrue(experts.holds_interleaved_gate_up) + for half, expected in ((0, gate), (1, up)): + held = FineGrainedDequantize(None)._dequantize_one( + experts.gate_up_proj[:, half::2], experts.gate_up_proj_scale_inv[:, half::2], torch.float32 + ) + error = ((held - expected).norm() / expected.norm()).item() + self.assertLess(error, 0.1, f"half {half}") + + def test_a_chunked_rope_permute_quantizes_only_the_finegrained_target(self): + from transformers.integrations.finegrained import FineGrainedLinear + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + n_heads = 2 + head_dim = 4 + in_dim = 4 + out_dim = n_heads * head_dim + block_size = (4, 4) + + class RopeProjector(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(out_dim, in_dim)) + + class RopeSelfAttn(torch.nn.Module): + def __init__(self): + super().__init__() + self.q_proj = FineGrainedLinear( + in_dim, out_dim, quantization_config=BlockFp8Config(weight_block_size=block_size) + ) + self.k_proj = RopeProjector() + self.v_proj = RopeProjector() + + class RopeLayer(torch.nn.Module): + def __init__(self): + super().__init__() + self.self_attn = RopeSelfAttn() + + class RopeModel(PreTrainedModel): + base_model_prefix = "model" + + def __init__(self, config): + super().__init__(config) + self.layers = torch.nn.ModuleList([RopeLayer()]) + self.post_init() + + def _init_weights(self, module): + pass # every weight here comes from the state dict, and `normal_` has no fp8 kernel + + config = PreTrainedConfig() + config.num_attention_heads = n_heads + model = RopeModel(config) + + raw_q = torch.tensor( + [ + [1.0, -1.0, 1.0, -1.0], + [0.5, -0.5, 0.5, -0.5], + [-1.0, 1.0, -1.0, 1.0], + [-0.5, 0.5, -0.5, 0.5], + [1.0, 1.0, -1.0, -1.0], + [0.5, 0.5, -0.5, -0.5], + [-1.0, -1.0, 1.0, 1.0], + [-0.5, -0.5, 0.5, 0.5], + ], + dtype=torch.float32, + ) + raw_k = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + raw_v = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + 100.0 + raw_qkv = torch.cat([raw_q, raw_k, raw_v], dim=0) + state_dict = {"layers.0.self_attn.qkv_proj.weight": raw_qkv.clone()} + + quantizer_cls = type( + "FineGrainedHfQuantizer", + (), + { + "__init__": lambda self, bs=block_size: setattr( + self, "quantization_config", SimpleNamespace(weight_block_size=bs, scale_fmt="float") + ), + "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), + "get_quantize_ops": lambda self: FineGrainedQuantize(self), + "get_param_ops": lambda self, _model, _param_name: [], + "pre_quantized": False, + }, + ) + quantizer = quantizer_cls() + + weight_mapping = [ + WeightConverter( + "self_attn.qkv_proj.weight", + [ + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + ], + operations=[Chunk(dim=0), PermuteForRope(permute_layer_names=["q_proj", "k_proj"])], + ) + ] + load_config = LoadStateDictConfig(weight_mapping=weight_mapping, hf_quantizer=quantizer) + loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict, load_config) + + self.assertEqual(loading_info.missing_keys, set()) + self.assertEqual(loading_info.unexpected_keys, set()) + self.assertEqual(loading_info.mismatched_keys, set()) + self.assertEqual(loading_info.conversion_errors, {}) + + permute_op = PermuteForRope(permute_layer_names=["q_proj", "k_proj"]) + permute_op.config = model.config + expected_q = permute_op._apply(raw_q) + expected_k = permute_op._apply(raw_k) + expected_v = raw_v # V carries no RoPE, so the permute names only q and k + + model_state = model.state_dict() + self.assertFalse(torch.allclose(raw_k, expected_k)) + torch.testing.assert_close(model_state["layers.0.self_attn.k_proj.weight"], expected_k) + torch.testing.assert_close(model_state["layers.0.self_attn.v_proj.weight"], expected_v) + + q_weight_key = "layers.0.self_attn.q_proj.weight" + scale_key = "layers.0.self_attn.q_proj.weight_scale_inv" + self.assertIn(scale_key, model_state) + self.assertEqual(model_state[q_weight_key].dtype, torch.float8_e4m3fn) + self.assertEqual(model_state[q_weight_key].shape, torch.Size((out_dim, in_dim))) + self.assertEqual(model_state[scale_key].dtype, torch.float32) + self.assertEqual( + model_state[scale_key].shape, + torch.Size((out_dim // block_size[0], in_dim // block_size[1])), + ) + + dequantized_q = FineGrainedDequantize(None)._dequantize_one( + model_state[q_weight_key], model_state[scale_key], output_dtype=torch.float32 + ) + torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) diff --git a/tests/integrations/finegrained/test_layers.py b/tests/integrations/finegrained/test_layers.py new file mode 100644 index 000000000000..dc9ab8e7f4b7 --- /dev/null +++ b/tests/integrations/finegrained/test_layers.py @@ -0,0 +1,728 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""The `finegrained` integration's modules, no kernel involved: what replaces what, what the modules hold and +hand the kernels, and the parallel plans a quantized model needs. The kernel call contract lives in +`tests/kernels/test_finegrained.py`, the conversion ops in `test_conversions.py`.""" + +import pickle +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +import transformers.integrations.finegrained.layers as fg +from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import BlockFp8Config, NVFP4Config + +from .utils import ( + _Cfg, + _format_config, + _quantizer_for, + _replace, + fake_finegrained_kernel, + patch_finegrained_kernel, +) + + +def _block_scaled_matmul_2d(A, B, As, Bs, *, output_dtype=None, **kwargs): + """`matmul_2d`'s contract on a block-scaled weight: the grid read row-major, its block inferred from the shapes.""" + block_n, block_k = B.shape[0] // Bs.shape[0], B.shape[1] // Bs.shape[1] + weight = B.float() * Bs.float().repeat_interleave(block_n, 0).repeat_interleave(block_k, 1) + return (A.float() @ weight.T).to(output_dtype) + + +@require_torch +class FineGrainedModulesTest(unittest.TestCase): + """The finegrained modules themselves: what they hold and what replaces what.""" + + def test_nvfp4_lookup_uses_packed_nibbles_block_scales_and_shared_global(self): + m = fg.FineGrainedEmbedding(8, 32, quantization_config=NVFP4Config(), dtype=torch.bfloat16) + m.weight.data.copy_(torch.full_like(m.weight, -78)) + m.weight_scale_inv.data.copy_(torch.tensor([1.0, 2.0]).expand(8, 2).to(torch.float8_e4m3fn)) + m.weight_global_scale.data.fill_(0.5) + ids = torch.tensor([[3, 1, 3], [7, 0, 1]]).t() + out = m(ids) + expected = torch.tensor([1.0, -1.5]).repeat(16) * torch.tensor([0.5] * 16 + [1.0] * 16) + torch.testing.assert_close(out, expected.expand(*ids.shape, 32).to(torch.bfloat16)) + self.assertEqual(m(torch.empty(0, dtype=torch.long)).shape, (0, 32)) + with self.assertRaises((IndexError, RuntimeError)): + m(torch.tensor([-1])) + with self.assertRaises((IndexError, RuntimeError)): + m(torch.tensor([8])) + + def test_nvfp4_lookup_decodes_every_e2m1_value_without_a_device_lut(self): + m = fg.FineGrainedEmbedding(2, 16, quantization_config=NVFP4Config(), dtype=torch.float32) + m.weight.data.copy_(torch.arange(8, dtype=torch.uint8).mul(17).view(torch.int8).expand(2, 8)) + m.weight_scale_inv.data.fill_(1) + m.weight_global_scale.data.fill_(1) + expected = torch.tensor([0.0, 0.0, 0.5, 0.5, 1.0, 1.0, 1.5, 1.5, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 6.0, 6.0]) + torch.testing.assert_close(m(torch.tensor([1, 0])), expected.expand(2, 16), rtol=0, atol=0) + m.weight.data.bitwise_or_(-120) # Set both sign bits without changing the magnitude. + torch.testing.assert_close(m(torch.tensor([0])), -expected.unsqueeze(0), rtol=0, atol=0) + + def test_storage_can_be_built_on_meta_without_a_full_precision_table(self): + with torch.device("meta"): + m = fg.FineGrainedEmbedding(320001536, 160, quantization_config=NVFP4Config()) + self.assertEqual(m.weight.shape, (320001536, 80)) + self.assertEqual(m.weight_scale_inv.shape, (320001536, 10)) + self.assertTrue(all(p.is_meta for p in m.parameters())) + + def test_colwise_tp_requires_complete_scale_groups(self): + from torch.distributed.tensor import Replicate + + from transformers.distributed.tensor_parallel import ColwiseParallel + + mesh = mock.Mock() + mesh.size.return_value = 2 + style = ColwiseParallel(output_layouts=Replicate()) + valid = fg.FineGrainedEmbedding(8, 64, quantization_config=NVFP4Config()) + style.validate_param(valid, "weight", mesh) + style.validate_param(valid, "weight_scale_inv", mesh) + invalid = fg.FineGrainedEmbedding(8, 48, quantization_config=NVFP4Config()) + with self.assertRaisesRegex(ValueError, "whole groups of 16"): + style.validate_param(invalid, "weight", mesh) + + def test_rowwise_tp_rejects_quantized_vocabulary_shards(self): + from transformers.distributed.tensor_parallel import RowwiseParallel + + mesh = mock.Mock() + mesh.size.return_value = 2 + style = RowwiseParallel() + style.validate_param(torch.nn.Embedding(8, 32), "weight", mesh) + for config in (BlockFp8Config(weight_block_size=None), NVFP4Config()): + with self.subTest(format=config.method): + table = fg.FineGrainedEmbedding(8, 32, quantization_config=config) + with self.assertRaisesRegex(ValueError, "Rowwise.*quantized embeddings"): + style.validate_param(table, "weight", mesh) + mesh.size.return_value = 1 + style.validate_param(table, "weight", mesh) + + def test_non_mx_formats_do_not_require_ue8m0(self): + from transformers.integrations.finegrained import conversions + + with ( + mock.patch.object(fg, "_get_ue8m0_dtype", side_effect=RuntimeError("UE8M0 unavailable")), + mock.patch.object(conversions, "_get_ue8m0_dtype", side_effect=RuntimeError("UE8M0 unavailable")), + ): + for config in (BlockFp8Config(), NVFP4Config()): + with self.subTest(format=config.method): + linear = FineGrainedLinear(64, 64, quantization_config=config) + self.assertEqual(linear.weight.element_size(), 1) + model = torch.nn.Module() + config = BlockFp8Config() + model.linear = FineGrainedLinear(64, 64, quantization_config=config) + quantizer = _quantizer_for(config) + quantizer.pre_quantized = False + result = quantizer.get_quantize_ops().convert({"linear.weight": torch.randn(64, 64)}, model=model) + self.assertEqual(result["linear.weight"].dtype, torch.float8_e4m3fn) + self.assertEqual(result["linear.weight_scale_inv"].dtype, torch.float32) + + def test_ue8m0_falls_back_to_exponent_bytes(self): + from transformers.integrations.finegrained.conversions import FineGrainedQuantize, as_container + from transformers.utils.quantization_config import Mxfp8Config + + fg._get_ue8m0_dtype.cache_clear() + try: + with mock.patch.object(torch, "float8_e8m0fnu", create=True): + del torch.float8_e8m0fnu + self.assertEqual(fg._get_ue8m0_dtype(), torch.uint8) + linear = FineGrainedLinear(64, 64, quantization_config=Mxfp8Config()) + self.assertEqual(linear.weight_scale_inv.dtype, torch.uint8) + scales = torch.tensor([0.5, 1.0, 2.0, 4.0]) + encoded = as_container(scales, torch.uint8) + torch.testing.assert_close(encoded, torch.tensor([126, 127, 128, 129], dtype=torch.uint8)) + _, scale = FineGrainedQuantize._quantize_block_fp8(torch.full((2, 2), 448.0), (2, 2), True) + torch.testing.assert_close(scale, torch.tensor([[127]], dtype=torch.uint8)) + finally: + fg._get_ue8m0_dtype.cache_clear() + + def test_only_weights_need_on_the_fly_quantization(self): + config = BlockFp8Config() + quantizer = _quantizer_for(config) + quantizer.pre_quantized = False + model = torch.nn.Module() + model.linear = FineGrainedLinear(64, 64, has_bias=True, quantization_config=config) + model.experts = FineGrainedExperts(_Cfg(), quantization_config=config, has_bias=True) + for module_name, module in model.named_children(): + weights = {"weight"} if module_name == "linear" else {"gate_up_proj", "down_proj"} + for name, _ in module.named_parameters(recurse=False): + with self.subTest(parameter=f"{module_name}.{name}"): + self.assertEqual( + quantizer.param_needs_quantization(model, f"{module_name}.{name}"), name in weights + ) + + def test_replacement_respects_a_disabled_post_expert_norm(self): + from transformers.integrations.moe import use_experts_implementation + + @use_experts_implementation(has_post_expert_norm=False) + class Experts(torch.nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.post_expert_norm = torch.nn.RMSNorm(config.hidden_size) + + def forward(self, hidden_states, top_k_index, top_k_weights): + return hidden_states + + config = _Cfg() + source = Experts(config) + self.assertFalse(source.has_post_expert_norm) + model = SimpleNamespace(config=SimpleNamespace(get_text_config=lambda: config)) + quantized = FineGrainedExperts.from_module(source, BlockFp8Config(), model) + self.assertIs(quantized.post_expert_norm, source.post_expert_norm) + self.assertFalse(quantized.has_post_expert_norm) + + def test_a_grouped_linear_is_recognised_by_its_groups_not_its_name(self): + """`replace_with_quantized_modules` picks `FineGrainedGroupedLinear` for a block-diagonal + linear — a plain one would collapse the groups into a single giant matmul and return the + wrong output dim. The branch keys on `n_groups`, the attribute the swap consumes: a class + NAME substring would miss a grouped linear named anything else, and would match (then + `AttributeError` on) one that carries its group count under another name. + """ + from transformers.integrations.finegrained import FineGrainedGroupedLinear + + class OddlyNamedBlockDiagonal(torch.nn.Linear): + def __init__(self): + super().__init__(8, 16, bias=False) + self.n_groups = 2 + + model = torch.nn.Module() + model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) + model.grouped = OddlyNamedBlockDiagonal() + model.plain = torch.nn.Linear(8, 16, bias=False) + + with torch.device("meta"): + _replace(model, BlockFp8Config()) + self.assertIsInstance(model.grouped, FineGrainedGroupedLinear) + self.assertNotIsInstance(model.plain, FineGrainedGroupedLinear) + + # A quantized embedding TABLE (Qwen4-Exp's n-gram table): FP8 rows with one per-tensor scale, + # rescaled on the rows a lookup gathers; swapped in for the names in `modules_to_convert`. + def test_replacing_experts_leaves_the_experts_class_as_defined(self): + """Each layer's flags come with its own class, shared by the layers that agree on them; a load modifies + neither `FineGrainedExperts` nor what a later direct construction gets.""" + from transformers import Qwen3MoeConfig, Qwen3MoeForCausalLM + + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_hidden_layers=3, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + num_experts=2, + num_experts_per_tok=1, + ) + with torch.device("meta"): + model = Qwen3MoeForCausalLM(config) + _replace(model, BlockFp8Config()) + experts = {type(layer.mlp.experts) for layer in model.model.layers} + self.assertEqual(len(experts), 1) + self.assertTrue(issubclass(experts.pop(), FineGrainedExperts)) + self.assertFalse(hasattr(FineGrainedExperts.forward, "__wrapped__")) + direct = FineGrainedExperts(_Cfg(), quantization_config=_format_config("fp8"), has_bias=True) + self.assertTrue(direct.has_bias) + # a replaced experts class is found by name, so a pickled model loads back + experts = model.model.layers[0].mlp.experts + self.assertIs(pickle.loads(pickle.dumps(type(experts))), type(experts)) + + def test_lookup_rescales_the_gathered_rows(self): + torch.manual_seed(0) + table = torch.randn(64, 16) + scale = table.abs().max() / 448.0 + embedding = fg.FineGrainedEmbedding(64, 16) + embedding.weight = torch.nn.Parameter((table / scale).to(torch.float8_e4m3fn), requires_grad=False) + embedding.weight_scale = torch.nn.Parameter(scale.to(torch.bfloat16).reshape(1), requires_grad=False) + ids = torch.tensor([[3, 7], [60, 0]]) + out = embedding(ids) + self.assertEqual(out.dtype, torch.bfloat16) + expected = (table / scale).to(torch.float8_e4m3fn)[ids].to(torch.bfloat16) * scale.to(torch.bfloat16) + torch.testing.assert_close(out, expected) + + def test_a_skip_list_names_modules_by_pattern(self): + """`should_convert_module`: everything converts by default, a skip-list entry names its modules.""" + from transformers.quantizers.quantizers_utils import should_convert_module + + self.assertTrue(should_convert_module("model", None)) + self.assertTrue(should_convert_module("model", [])) + patterns = ["model.layers.*.self_attn", "lm_head"] + self.assertFalse(should_convert_module("lm_head", patterns)) + self.assertTrue(should_convert_module("experts", patterns)) + + def test_an_empty_plan_replaces_nothing(self): + """A dequantized load walks an empty plan; it must not fall back to the quantizer's own.""" + model = torch.nn.Sequential(torch.nn.Linear(256, 256)) + model.config = SimpleNamespace(get_text_config=lambda: SimpleNamespace()) + self.assertFalse(_quantizer_for(BlockFp8Config()).replace_modules(model, plan={})) + self.assertIsInstance(model[0], torch.nn.Linear) + + def test_replacement_follows_the_patterns_and_the_skip_list(self): + """A table is quantized where `modules_to_convert` names it, never through the plan's catch-all.""" + + model = torch.nn.Module() + model.ngram_embedding = torch.nn.Embedding(8, 4) + model.embed_tokens = torch.nn.Embedding(8, 4) + model.other = torch.nn.Embedding(8, 4) + _replace(model, BlockFp8Config(modules_to_convert=["ngram_embedding", "other"]), ["other"]) + self.assertIsInstance(model.ngram_embedding, fg.FineGrainedEmbedding) + self.assertTrue(model.ngram_embedding._hf_quantized_needs_local_tp) + self.assertIs(type(model.embed_tokens), torch.nn.Embedding) + self.assertIs(type(model.other), torch.nn.Embedding) + + # `from_pretrained` sets the default dtype to the checkpoint's during construction, so a scale + # allocated as `torch.ones(n)` comes out bf16 and the kernels read it as fp32 (NaN logits, no + # error). Building under both defaults catches it. + def _dtypes(self, build, default): + previous = torch.get_default_dtype() + torch.set_default_dtype(default) + try: + with torch.device("meta"): + module = build() + return {name: p.dtype for name, p in module.named_parameters() if p is not None} + finally: + torch.set_default_dtype(previous) + + def _assert_default_dtype_invariant(self, build): + under_fp32 = self._dtypes(build, torch.float32) + under_bf16 = self._dtypes(build, torch.bfloat16) + self.assertTrue(under_fp32) + # the weight and bias DO follow the model dtype; every scale beside them must not + scales = {name for name in under_fp32 if "scale" in name} + self.assertTrue(scales) + for name in sorted(scales): + self.assertEqual(under_bf16[name], under_fp32[name], f"{name} followed the default dtype") + + def test_linear_scales_ignore_the_default_dtype(self): + for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): + with self.subTest(weight_format=weight_format, activation_scheme=scheme): + self._assert_default_dtype_invariant( + lambda f=weight_format, s=scheme: FineGrainedLinear( + quantization_config=_format_config( + f, **({"weight_block_size": (4, 4), "activation_scheme": s} if f == "fp8" else {}) + ), + in_features=64, + out_features=32, + ) + ) + + def test_embedding_scale_ignores_the_default_dtype(self): + self._assert_default_dtype_invariant(lambda: fg.FineGrainedEmbedding(64, 16)) + + def test_experts_scales_ignore_the_default_dtype(self): + for weight_format, scheme in (("fp8", "static"), ("fp8", "dynamic"), ("nvfp4", "dynamic")): + with self.subTest(weight_format=weight_format, activation_scheme=scheme): + self._assert_default_dtype_invariant( + lambda f=weight_format, s=scheme: FineGrainedExperts( + _Cfg(), + quantization_config=_format_config( + f, **({"weight_block_size": (4, 4), "activation_scheme": s} if f == "fp8" else {}) + ), + ) + ) + + def test_static_activation_scheme_reaches_the_fused_chain(self): + """Mistral-3 ships `activation_scheme="static"`: a calibrated scale per activation, and a + MoE calibrates each expert separately, so the module holds one per expert. Both fused + chains take them per projection and the kernels apply each tile's own expert scale in + register — a host pre-quant could not, since top-k routes one row to several experts + whose scales differ. A dynamic module holds no such slot and passes `None`.""" + static, dynamic = ( + FineGrainedExperts( + _Cfg(), quantization_config=BlockFp8Config(weight_block_size=(128, 128), activation_scheme=scheme) + ) + for scheme in ("static", "dynamic") + ) + held = static.gate_up_proj_activation_scale + self.assertIsNotNone(held, "the calibrated scale the chain needs") + self.assertEqual(held.numel(), _Cfg.num_local_experts, "one per expert, which the kernels index") + # a dynamic module holds the slot empty, like every other optional parameter + self.assertIsNone(dynamic.gate_up_proj_activation_scale) + + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = () + for module, calibrated in ((static, True), (dynamic, False)): + operands = module._moe_operands(kernel) + for projection in ("gate_up_proj", "down_proj"): + with self.subTest(scheme=module.activation_scheme, projection=projection): + scale = operands[f"{projection}_activation_scale"] + self.assertEqual(scale is not None, calibrated) + if calibrated: + self.assertEqual(scale.numel(), _Cfg.num_local_experts) + + def test_eager_gated_block_scales_pair_each_half_with_its_own_blocks(self): + """The loader interleaves a gate|up block grid per BLOCK, so one scale covers rows of one half while + the rows alternate between halves: the eager loop runs each half on its own rows and blocks rather + than hand `matmul_2d` the stack as a row-major grid. Halves of very different magnitudes, as real + gate and up projections have, so a scale read from the wrong half shows.""" + experts = FineGrainedExperts(_Cfg(), quantization_config=BlockFp8Config(weight_block_size=(32, 32))) + experts._deepgemm_disabled = True + num_experts, hidden, inter = _Cfg.num_local_experts, _Cfg.hidden_size, _Cfg.intermediate_size + torch.manual_seed(0) + codes = torch.randn(num_experts, 2, inter, hidden).to(torch.float8_e4m3fn) # [gate, up] + exponents = torch.tensor([-8.0, 0.0]).reshape(1, 2, 1, 1) + torch.randint(-2, 3, (num_experts, 2, 1, 2)) + grid = torch.pow(2.0, exponents) # (experts, half, row blocks, column blocks) + # interleaved: rows [g0, u0, g1, u1, ...], grid rows [gate block, up block] + experts.gate_up_proj = torch.nn.Parameter(codes.transpose(1, 2).reshape(num_experts, 2 * inter, hidden)) + experts.gate_up_proj_scale_inv = torch.nn.Parameter(grid.transpose(1, 2).reshape(num_experts, 2, 2)) + + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = _block_scaled_matmul_2d + x = torch.randn(3, hidden) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = experts.expert_linear(x, "gate_up_proj", 1) + halves = [ + x @ (codes[1, half].float() * grid[1, half].repeat_interleave(32, 0).repeat_interleave(32, 1)).T + for half in (0, 1) + ] + torch.testing.assert_close(out, torch.stack(halves, dim=-1).flatten(-2)) + + def test_eager_gated_per_tensor_scale_is_shared_by_gate_and_up(self): + experts = FineGrainedExperts(_Cfg(), quantization_config=BlockFp8Config(weight_block_size=None)) + experts._deepgemm_disabled = True + torch.manual_seed(0) + codes = torch.randn_like(experts.gate_up_proj, dtype=torch.float32).to(torch.float8_e4m3fn) + experts.gate_up_proj = torch.nn.Parameter(codes) + # On-the-fly quantization of a fused projection produces one scale per expert, + # whereas separately calibrated gate/up checkpoints can hold two. + scales = torch.tensor([0.125, 0.5, 2.0, 8.0]).view(-1, 1, 1) + experts.gate_up_proj_scale_inv = torch.nn.Parameter(scales) + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = _block_scaled_matmul_2d + x = torch.randn(3, _Cfg.hidden_size) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + for expert in range(_Cfg.num_local_experts): + with self.subTest(expert=expert): + actual = experts.expert_linear(x, "gate_up_proj", expert) + expected = x @ (codes[expert].float() * scales[expert]).T + torch.testing.assert_close(actual, expected) + + def test_stacked_gate_up_is_gated_by_halves_and_kept_from_the_triton_chains(self): + """The DeepGEMM backends load gate|up stacked: the host GLU splits the halves, and the triton chains, + which read the rows interleaved, refuse such a module rather than pair the wrong rows.""" + config = _Cfg() + config._experts_implementation = "deepgemm" + experts = FineGrainedExperts(config, quantization_config=BlockFp8Config(weight_block_size=(32, 32))) + self.assertFalse(experts.holds_interleaved_gate_up) + gate, up = torch.randn(3, _Cfg.intermediate_size), torch.randn(3, _Cfg.intermediate_size) + torch.testing.assert_close( + experts._apply_gate(torch.cat([gate, up], dim=-1)), torch.nn.functional.silu(gate) * up + ) + + hidden = torch.randn(2, _Cfg.hidden_size, dtype=torch.bfloat16) + routing = torch.zeros(2, 1, dtype=torch.long), torch.ones(2, 1) + with mock.patch.object(fg, "load_finegrained_kernel") as load: + for forward in (fg.finegrained_grouped_mm_experts_forward, fg.finegrained_batched_mm_experts_forward): + with ( + self.subTest(forward=forward.__name__), + self.assertRaisesRegex(NotImplementedError, "interleaved"), + ): + forward(experts, hidden, *routing) + load.assert_not_called() + + def test_swiglu_alpha_without_a_limit_runs_unclamped(self): + config = _Cfg() + config.swiglu_alpha = 1.702 + experts = FineGrainedExperts(config, quantization_config=BlockFp8Config(weight_block_size=(32, 32))) + gate_up = torch.randn(3, 2 * _Cfg.intermediate_size) + gate, up = gate_up[..., 0::2], gate_up[..., 1::2] + torch.testing.assert_close(experts._apply_gate(gate_up), (up + 1) * gate * torch.sigmoid(gate * 1.702)) + + def test_grouped_linear_broadcasts_a_shared_per_tensor_scale(self): + from transformers.integrations.finegrained import FineGrainedGroupedLinear + + linear = FineGrainedGroupedLinear( + 8, 12, 3, quantization_config=BlockFp8Config(weight_block_size=None), has_bias=True + ) + torch.manual_seed(0) + linear.weight = torch.nn.Parameter(torch.randn(12, 8).to(torch.float8_e4m3fn)) + linear.weight_scale_inv.data.fill_(0.25) + linear.bias.data.copy_(torch.randn(12)) + kernel, _ = fake_finegrained_kernel() + + def grouped_matmul(a, b, activation_scale, scales, *, expert_start, **kwargs): + self.assertEqual(expert_start.tolist(), [0, 10, 20, 30, 30]) + return torch.cat( + [ + _block_scaled_matmul_2d( + a[expert_start[group] : expert_start[group + 1]], b[group], None, scales[group] + ) + for group in range(3) + ] + ) + + kernel.matmul_grouped = grouped_matmul + x = torch.randn(2, 5, 3, 8) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + actual = linear(x) + weight = linear.weight.float().view(3, 4, 8) * 0.25 + expected = torch.einsum("...gk,gnk->...gn", x, weight) + linear.bias.view(3, 4) + torch.testing.assert_close(actual, expected) + + def test_grouped_linear_hands_over_its_stored_weight_and_activation_scales(self): + """The groups split the weight as stored (packed FP4 halves K), and the calibrated activation scale goes + where `FineGrainedLinear` hands it: NVFP4's as the activation global, a static FP8 one as `As`.""" + from transformers.integrations.finegrained import FineGrainedGroupedLinear + + for config, packed in ( + (NVFP4Config(), True), + (BlockFp8Config(weight_block_size=(16, 16), activation_scheme="static"), False), + (BlockFp8Config(weight_block_size="row", activation_scale_ub=1200.0), False), + ): + with self.subTest(weight_format=config.method.value): + linear = FineGrainedGroupedLinear(64, 32, 2, quantization_config=config) + kernel, recorder = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + linear(torch.randn(3, 2, 64, dtype=torch.bfloat16)) + (call,) = recorder.calls["matmul_grouped"] + self.assertEqual(call.args[1].shape, (2, 16, 32 if packed else 64)) + two_level = linear.weight_global_scale is not None + self.assertIs(call.args[2], None if two_level else linear.activation_scale) + self.assertIs(call.kwargs["a_global_scale"], linear.activation_scale if two_level else None) + self.assertEqual(call.kwargs["activation_scale_ub"], getattr(config, "activation_scale_ub", None)) + if config.weight_block_size == "row": + self.assertEqual(call.args[3].shape, (2, 16, 1)) + + +@require_torch +class FineGrainedFusedNormTest(unittest.TestCase): + """A norm the kernels can fuse is not fused when a forward collective sits on its input: under intra-expert + TP the rows are a partial sum over the sharded intermediate, which the wrapped module all-reduces in its own + forward.""" + + def _operands(self, *, input_reduce): + experts = FineGrainedExperts(_Cfg(), quantization_config=_format_config("fp8", weight_block_size=(4, 4))) + experts.post_expert_norm, experts.has_post_expert_norm = torch.nn.RMSNorm(_Cfg.hidden_size), True + if input_reduce: + experts.post_expert_norm._hf_tp_input_reduce = True # what `ReplicatedWithInputAllReduce` marks + kernel = mock.Mock() + kernel.get_supported_act_fns.return_value = ("silu",) + kernel.get_supported_norms.return_value = ("rms_norm",) + with mock.patch.dict(fg.SUPPORTED_FUSED_NORMS, {torch.nn.RMSNorm: "rms_norm"}): + return experts._moe_operands(kernel), experts + + def test_a_fusable_norm_fuses_when_nothing_reduces_its_input(self): + operands, experts = self._operands(input_reduce=False) + self.assertEqual(operands["post_expert_norm"], "rms_norm") + self.assertIs(operands["post_expert_norm_weight"], experts.post_expert_norm.weight) + + def test_a_fusable_norm_falls_back_when_its_input_is_all_reduced(self): + operands, experts = self._operands(input_reduce=True) + self.assertIs(operands["post_expert_norm"], experts.post_expert_norm) + self.assertIsNone(operands["post_expert_norm_weight"]) + + +class FineGrainedParallelPlanTest(unittest.TestCase): + """Update the built model's TP/EP plans without mutating config plans.""" + + BASE_EP = { + "layers.*.mlp.gate": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.self_attn.q_proj": "colwise", + } + + def _planned(self, impl=None, raw=None, config=None): + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + quantizer = BlockFp8HfQuantizer(config or BlockFp8Config()) + model = torch.nn.Module() + model.config = SimpleNamespace(_experts_implementation=impl) + model.layers = torch.nn.ModuleList([torch.nn.Module() for _ in range(4)]) + with torch.device("meta"), mock.patch.object(fg, "is_sm100", return_value=False): + for layer in model.layers: + layer.mlp = torch.nn.Module() + layer.mlp.experts = FineGrainedExperts( + _Cfg(), quantization_config=config or NVFP4Config(activation_format="nvfp4"), has_bias=True + ) + plan = dict(raw or self.BASE_EP) + quantizer.update_parallel_plan(model, plan) + return plan + + def test_embedding_companions_follow_the_format_and_preserve_explicit_placements(self): + for config in (NVFP4Config(), BlockFp8Config(modules_to_convert=["table"])): + quantizer = _quantizer_for(config) + model = torch.nn.Module() + model.config = SimpleNamespace(_experts_implementation=None) + model.table = fg.FineGrainedEmbedding.from_module(torch.nn.Embedding(8, 64), config, model) + for style in ("colwise_gather_output", "rowwise", "replicate"): + with self.subTest(format=config.method, style=style): + plan = {"table": style} + quantizer.update_parallel_plan(model, plan) + if isinstance(config, NVFP4Config): + self.assertEqual( + plan["table.weight_scale_inv"], "colwise" if style.startswith("colwise") else style + ) + self.assertEqual(plan["table.weight_global_scale"], "replicate") + explicit = {"table": style, "table.weight_scale_inv": "replicate"} + quantizer.update_parallel_plan(model, explicit) + self.assertEqual(explicit["table.weight_scale_inv"], "replicate") + else: + self.assertEqual(plan["table.weight_scale"], "replicate") + + def test_every_companion_shards_with_its_expert_weight(self): + """The plan matcher keys on the exact parameter name and falls back only to the owning + module, whose entry shards nothing — so a scale with no entry of its own stays whole while + its weight is this rank's expert slice, and every rank but the first reads another rank's + scales.""" + from transformers.distributed.tensor_parallel import _get_parameter_plan + + plan = self._planned() + + def style(name): + return _get_parameter_plan(f"layers.3.mlp.experts.{name}", plan) + + for proj in ("gate_up_proj", "down_proj"): + for name in (proj, f"{proj}_scale_inv", f"{proj}_weight_global_scale", f"{proj}_activation_scale"): + self.assertEqual(style(name), "grouped_gemm", name) + for name in (f"{proj}_bias",): + self.assertEqual(style(name), "grouped_gemm", name) + self.assertEqual(plan["layers.*.self_attn.q_proj"], "colwise") + + def test_intra_expert_entries_follow_the_projection_axis(self): + """Under intra-expert TP the experts stay whole and the projection's own axis splits, so the scale grid + follows its weight while the per-expert globals stay replicated.""" + from transformers.distributed.tensor_parallel import _get_parameter_plan + + base = { + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + plan = self._planned(raw=base) + + def style(name): + return _get_parameter_plan(f"layers.3.mlp.experts.{name}", plan) + + # the weight keeps the style the MODEL declared — `packed_colwise` is the strided split a + # concatenated `[gate; up]` checkpoint needs, since the interleave into the kernels' row + # order runs on each rank's shard, after sharding. The companions follow it. + self.assertEqual(style("gate_up_proj"), "packed_colwise") + for name in ("gate_up_proj_scale_inv", "gate_up_proj_bias"): + self.assertEqual(style(name), "moe_experts_packed_colwise", name) + # the down scale splits on dim 2 — the reduce axis affine and 5-D swizzled alike + self.assertEqual(style("down_proj"), "rowwise") + self.assertEqual(style("down_proj_scale_inv"), "moe_experts_rowwise") + # TP does not split experts, so anything indexed per expert stays whole + for name in ("gate_up_proj_weight_global_scale", "down_proj_weight_global_scale", "down_proj_bias"): + self.assertEqual(style(name), "moe_tp_experts", name) + + # a natively INTERLEAVED layout declares `colwise`, and its companions split contiguously + inter = dict(base, **{"layers.*.mlp.experts.gate_up_proj": "colwise"}) + plan_i = self._planned(raw=inter) + self.assertEqual(_get_parameter_plan("layers.3.mlp.experts.gate_up_proj", plan_i), "colwise") + self.assertEqual( + _get_parameter_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan_i), + "moe_experts_colwise", + ) + rowwise = self._planned(raw=base, config=BlockFp8Config(weight_block_size="row")) + self.assertEqual( + _get_parameter_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", rowwise), "moe_experts_packed_colwise" + ) + self.assertEqual(_get_parameter_plan("layers.3.mlp.experts.down_proj_scale_inv", rowwise), "moe_tp_experts") + + def test_dense_projections_that_share_a_name_are_left_alone(self): + """`mlp.up_proj` and `mlp.shared_experts.down_proj` end in the same words as the stacked expert + projections but are plain 2-D linears that colwise/rowwise already shard right; an expert-axis style + would shard a dim they do not have.""" + from transformers.distributed.tensor_parallel import _get_parameter_plan + + plan = self._planned( + raw={ + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + "layers.*.mlp.shared_experts.up_proj": "colwise", + "layers.*.mlp.shared_experts.down_proj": "rowwise", + } + ) + for dense, expected in ( + ("layers.*.mlp.up_proj", "colwise"), + ("layers.*.mlp.down_proj", "rowwise"), + ("layers.*.mlp.shared_experts.up_proj", "colwise"), + ("layers.*.mlp.shared_experts.down_proj", "rowwise"), + ): + self.assertEqual(plan[dense], expected, dense) + # and no scale entry was invented for a dense linear, whose scale `rowwise`/`colwise` + # already reaches through the module fallback + strays = [k for k in plan if k.endswith("_scale_inv") and ".experts." not in k] + self.assertEqual(strays, [], f"invented dense scale entries: {strays}") + # the stacked experts still get theirs + self.assertEqual( + _get_parameter_plan("layers.3.mlp.experts.gate_up_proj_scale_inv", plan), + "moe_experts_packed_colwise", + ) + + def test_a_multimodal_models_expert_plans_are_reached(self): + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + from .utils import _conversion_model + + quantizer = BlockFp8HfQuantizer(BlockFp8Config()) + text_model = _conversion_model(quantizer) + model = torch.nn.Module() + model.config = SimpleNamespace(_experts_implementation=None) + model.language_model = text_model + plan = {"language_model.model.layers.*.mlp.experts.gate_up_proj": "grouped_gemm"} + quantizer.update_parallel_plan(model, plan) + self.assertEqual(plan["language_model.model.layers.*.mlp.experts.gate_up_proj_scale_inv"], "grouped_gemm") + self.assertFalse(text_model.config.base_model_ep_plan) + + def test_the_impl_rewrites_the_layer_kinds(self): + megamoe = self._planned("deepgemm_megamoe") + self.assertEqual(megamoe["layers.*.mlp.gate"], "megamoe_router") + self.assertEqual(megamoe["layers.*.mlp.experts"], "megamoe_experts") + # Mega MoE dispatches across ranks itself, in place of the EP dispatch a plan may default to + dispatched = self._planned("deepgemm_megamoe", raw={"layers.*.mlp.experts": "ep_dispatch_experts"}) + self.assertEqual(dispatched["layers.*.mlp.experts"], "megamoe_experts") + + def test_a_real_models_plan_keeps_its_entries_and_gains_the_expert_companions(self): + from transformers import AutoModelForCausalLM, Qwen3Config, Qwen3MoeConfig + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + for config_cls in (Qwen3Config, Qwen3MoeConfig): + with self.subTest(config=config_cls.__name__), torch.device("meta"): + config = config_cls( + vocab_size=64, + hidden_size=64, + intermediate_size=64, + moe_intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=32, + num_experts=2, + num_experts_per_tok=1, + ) + shipped = dict(config.base_model_tp_plan) + model = AutoModelForCausalLM.from_config(config) + before = dict(model.tp_plan) + quantizer = BlockFp8HfQuantizer(BlockFp8Config()) + quantizer.preprocess_model(model) + for key, style in before.items(): + self.assertEqual(model.tp_plan[key], style, key) + self.assertEqual(config.base_model_tp_plan, shipped) + for companion in ("gate_up_proj_scale_inv", "down_proj_scale_inv"): + self.assertIn(f"model.layers.*.mlp.experts.{companion}", model.tp_plan) diff --git a/tests/integrations/finegrained/test_loading_regressions.py b/tests/integrations/finegrained/test_loading_regressions.py new file mode 100644 index 000000000000..67a897f94ddb --- /dev/null +++ b/tests/integrations/finegrained/test_loading_regressions.py @@ -0,0 +1,420 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""CPU loading, inference, and saving against local quantized checkpoints with real model mappings.""" + +import os +import tempfile +import unittest +from unittest import mock + +import torch +from safetensors.torch import load_file, save_file + +import transformers.conversion_mapping as conversion_mapping +from transformers import ( + AutoModelForCausalLM, + LlamaConfig, + NVFP4Config, + PreTrainedConfig, + PreTrainedModel, + Qwen2MoeConfig, +) +from transformers.core_model_loading import Concatenate, WeightConverter, WeightRenaming +from transformers.integrations.finegrained import FineGrainedEmbedding +from transformers.integrations.finegrained.conversions import FineGrainedDequantize +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.quantizers.block_scaled_utils import BlockScaledQuantizer +from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import BlockFp8Config, FineGrainedConfig, Mxfp4Config, Mxfp8Config + +from .utils import _load_and_save, _per_expert + + +class EmbeddingModel(PreTrainedModel): + config_class = PreTrainedConfig + + def __init__(self, config): + super().__init__(config) + self.embed_tokens = torch.nn.Embedding(config.vocab_size, config.hidden_size) + self.ngram_embedding = torch.nn.Embedding(config.vocab_size, config.hidden_size) + self.proj = torch.nn.Linear(config.hidden_size, config.hidden_size, bias=False) + self.lm_head = torch.nn.Linear(config.hidden_size, 1, bias=False) + self.post_init() + + def forward(self, ids): + return self.proj(self.ngram_embedding(ids)) + + +@require_torch +class FineGrainedLoadingRegressionTest(unittest.TestCase): + def test_sharded_checkpoint_roundtrips_and_refuses_full_table_dequantization(self): + conversion_mapping.get_checkpoint_conversion_mapping("qwen4_exp_text") + conversion = WeightConverter( + "ngram_embedding.shard_*.weight", + "ngram_embedding.weight", + [Concatenate(dim=0, num_shards_attribute="split_ngram_parts")], + force_cpu=True, + ) + with ( + mock.patch.dict(conversion_mapping._checkpoint_conversion_mapping_cache, {"EmbeddingModel": [conversion]}), + mock.patch.object( + conversion_mapping, + "USER_REGISTERED_MAPPINGS", + conversion_mapping.USER_REGISTERED_MAPPINGS | {"EmbeddingModel"}, + ), + tempfile.TemporaryDirectory() as directory, + tempfile.TemporaryDirectory() as saved, + ): + config = PreTrainedConfig(vocab_size=8, hidden_size=32, split_ngram_parts=2, dtype="float32") + config.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "ngram_embedding": {"quant_algo": "NVFP4", "group_size": 16}, + "proj": {"quant_algo": "NVFP4", "group_size": 16}, + }, + } + config.save_pretrained(directory) + state = { + "embed_tokens.weight": torch.zeros(8, 32), + "lm_head.weight": torch.zeros(1, 32), + "ngram_embedding.weight_scale_2": torch.tensor([0.5]), + "proj.weight": torch.full((32, 16), 0x22, dtype=torch.uint8), + "proj.weight_scale": torch.ones(32, 2).to(torch.float8_e4m3fn), + "proj.weight_scale_2": torch.tensor(0.25), + "proj.input_scale": torch.tensor(0.125), + } + for shard in range(2): + state[f"ngram_embedding.shard_{shard}.weight"] = torch.full((4, 16), 0xB2, dtype=torch.uint8) + state[f"ngram_embedding.shard_{shard}.weight_scale"] = torch.full((4, 2), shard + 1.0).to( + torch.float8_e4m3fn + ) + save_file(state, os.path.join(directory, "model.safetensors"), metadata={"format": "pt"}) + with mock.patch.object(BlockScaledQuantizer, "validate_environment"): + packed, info = EmbeddingModel.from_pretrained(directory, output_loading_info=True) + self.assertFalse(info["missing_keys"]) + self.assertFalse(info["unexpected_keys"]) + self.assertIsInstance(packed.ngram_embedding, FineGrainedEmbedding) + packed.save_pretrained(saved) + written = load_file(os.path.join(saved, "model.safetensors")) + self.assertEqual(set(state), set(written)) + for key, tensor in state.items(): + torch.testing.assert_close( + written[key].reshape(-1).view(torch.uint8), tensor.reshape(-1).view(torch.uint8) + ) + with mock.patch.object(BlockScaledQuantizer, "validate_environment"): + reloaded = EmbeddingModel.from_pretrained(saved) + ids = torch.tensor([[0, 5, 0]]) + torch.testing.assert_close(packed.ngram_embedding(ids), reloaded.ngram_embedding(ids)) + with self.assertRaisesRegex(ValueError, "dequantize=True cannot expand an NVFP4 embedding"): + EmbeddingModel.from_pretrained(directory, quantization_config=FineGrainedConfig(dequantize=True)) + + def _checkpoint(self, directory, moe=False, mixed=False, modelopt_mixed=False): + config_class = Qwen2MoeConfig if moe else LlamaConfig + config = config_class( + vocab_size=32, + hidden_size=64, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + **({"num_experts": 2, "num_experts_per_tok": 1, "moe_intermediate_size": 32} if moe else {}), + ) + model = AutoModelForCausalLM.from_config(config).to(torch.bfloat16).eval() + # Saving the model first supplies the architecture's checkpoint key layout, including split experts. + model.save_pretrained(directory) + filename = os.path.join(directory, "model.safetensors") + state = load_file(filename) + quantized = {} + for name, tensor in state.items(): + if tensor.ndim != 2 or name.endswith(("embed_tokens.weight", "lm_head.weight", "mlp.gate.weight")): + quantized[name] = tensor + continue + if modelopt_mixed and ".q_proj." not in name and ".mlp.experts." not in name: + quantized[name] = tensor + continue + stem = name.removesuffix("weight") + if modelopt_mixed and ".q_proj." in name: + quantized[name] = torch.full_like(tensor, 2, dtype=torch.float8_e4m3fn) + quantized[stem + "weight_scale"] = torch.full((1, 1), 1 / 128) + continue + ct = mixed and ".q_proj." in name + quantized[stem + ("weight_packed" if ct else "weight")] = torch.full( + (tensor.shape[0], tensor.shape[1] // 2), 0x22, dtype=torch.uint8 + ) + quantized[stem + "weight_scale"] = torch.ones( + tensor.shape[0], tensor.shape[1] // 16, dtype=torch.float8_e4m3fn + ) + quantized[stem + ("weight_global_scale" if ct else "weight_scale_2")] = torch.tensor( + 64.0 if ct else 1 / 64 + ) + config.quantization_config = ( + FineGrainedConfig( + dtype_plan={r"\.q_proj$": NVFP4Config(checkpoint_format="compressed-tensors"), ".*": NVFP4Config()}, + modules_to_not_convert=["lm_head", "mlp.gate"], + ) + if mixed + else NVFP4Config(modules_to_not_convert=["lm_head", "mlp.gate"]) + ) + if modelopt_mixed: + config.quantization_config = { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8_PB_WO"}, + "model.layers.0.mlp.experts": {"quant_algo": "NVFP4", "group_size": 16}, + }, + } + config.save_pretrained(directory) + save_file(quantized, filename, metadata={"format": "pt"}) + + def test_dequantized_dense_model_respects_requested_dtype(self): + with tempfile.TemporaryDirectory() as directory: + self._checkpoint(directory) + model = AutoModelForCausalLM.from_pretrained( + directory, quantization_config=NVFP4Config(dequantize=True), dtype=torch.float16, device_map="cpu" + ).eval() + self.assertEqual({parameter.dtype for parameter in model.parameters()}, {torch.float16}) + with torch.no_grad(): + self.assertTrue(torch.isfinite(model(torch.tensor([[1, 2, 3]])).logits).all()) + + def test_dequantized_moe_can_save_and_reload(self): + with tempfile.TemporaryDirectory() as directory, tempfile.TemporaryDirectory() as saved: + self._checkpoint(directory, moe=True) + model = AutoModelForCausalLM.from_pretrained( + directory, quantization_config=NVFP4Config(dequantize=True), dtype=torch.float32, device_map="cpu" + ).eval() + model.save_pretrained(saved) + reloaded, info = AutoModelForCausalLM.from_pretrained(saved, output_loading_info=True) + self.assertEqual(info["missing_keys"], set()) + self.assertEqual(info["unexpected_keys"], set()) + self.assertFalse(hasattr(reloaded.config, "quantization_config")) + for key, value in model.state_dict().items(): + torch.testing.assert_close(reloaded.state_dict()[key], value) + with torch.no_grad(): + inputs = torch.tensor([[1, 2, 3]]) + torch.testing.assert_close(model(inputs).logits, reloaded(inputs).logits) + + def test_mixed_checkpoint_global_conventions_are_scoped_per_module(self): + with tempfile.TemporaryDirectory() as modelopt, tempfile.TemporaryDirectory() as mixed: + self._checkpoint(modelopt, moe=True) + self._checkpoint(mixed, moe=True, mixed=True) + plain = AutoModelForCausalLM.from_pretrained( + modelopt, quantization_config=NVFP4Config(dequantize=True), dtype=torch.float32, device_map="cpu" + ) + loaded = AutoModelForCausalLM.from_pretrained( + mixed, + quantization_config=FineGrainedConfig( + dtype_plan={ + r"\.q_proj$": NVFP4Config(checkpoint_format="compressed-tensors"), + ".*": NVFP4Config(), + }, + modules_to_not_convert=["lm_head", "mlp.gate"], + dequantize=True, + ), + dtype=torch.float32, + device_map="cpu", + ) + for name, tensor in plain.state_dict().items(): + if name.endswith(("q_proj.weight", "experts.gate_up_proj", "experts.down_proj")): + torch.testing.assert_close(loaded.state_dict()[name], tensor) + torch.testing.assert_close(tensor, torch.full_like(tensor, 1 / 64)) + + def test_mixed_block_sizes_use_the_current_modules_grid(self): + # Both block sizes produce a 2x2 scale grid for 192x192 weights, but different boundaries. + quantizer = FineGrainedHfQuantizer( + FineGrainedConfig( + dtype_plan={ + "a$": BlockFp8Config(weight_block_size=(96, 96)), + "b$": BlockFp8Config(weight_block_size=(128, 128)), + }, + dequantize=True, + ) + ) + model = torch.nn.Module() + model.a = torch.nn.Linear(192, 192, bias=False) + model.b = torch.nn.Linear(192, 192, bias=False) + scales = torch.tensor([[1.0, 2.0], [4.0, 8.0]]) + for name, block in (("a", 96), ("b", 128)): + with self.subTest(module=name): + result = FineGrainedDequantize(quantizer).convert( + {"weight": torch.ones(192, 192).to(torch.float8_e4m3fn), "weight_scale_inv": scales}, + full_layer_name=f"{name}.weight", + model=model, + )[f"{name}.weight"] + expected = scales.repeat_interleave(block, dim=0).repeat_interleave(block, dim=1)[:192, :192] + torch.testing.assert_close(result, expected) + + def test_modelopt_mixed_precision_checkpoint_loads_and_saves(self): + for rename_namespace in (False, True): + with ( + self.subTest(rename_namespace=rename_namespace), + tempfile.TemporaryDirectory() as directory, + tempfile.TemporaryDirectory() as saved, + ): + self._checkpoint(directory, moe=True, modelopt_mixed=True) + mapping = conversion_mapping.get_checkpoint_conversion_mapping("qwen2_moe") + if rename_namespace: + config = Qwen2MoeConfig.from_pretrained(directory) + layers = config.quantization_config["quantized_layers"] + config.quantization_config["quantized_layers"] = { + name.replace("model.layers.", "decoder.blocks."): spec for name, spec in layers.items() + } + config.quantization_config = AutoQuantizationConfig.from_dict(config.quantization_config).to_dict() + config.save_pretrained(directory) + filename = os.path.join(directory, "model.safetensors") + save_file( + { + name.replace("model.layers.", "decoder.blocks."): value + for name, value in load_file(filename).items() + }, + filename, + metadata={"format": "pt"}, + ) + mapping = [WeightRenaming(r"^decoder\.blocks\.", "model.layers."), *mapping] + with mock.patch.dict(conversion_mapping._checkpoint_conversion_mapping_cache, {"qwen2_moe": mapping}): + model, info = AutoModelForCausalLM.from_pretrained( + directory, + quantization_config=FineGrainedConfig(dequantize=True), + dtype=torch.float16, + device_map="cpu", + output_loading_info=True, + ) + self.assertEqual(info["missing_keys"], set()) + self.assertEqual(info["unexpected_keys"], set()) + self.assertEqual({parameter.dtype for parameter in model.parameters()}, {torch.float16}) + state = model.state_dict() + for name, tensor in state.items(): + if name.endswith(("q_proj.weight", "experts.gate_up_proj", "experts.down_proj")): + torch.testing.assert_close(tensor, torch.full_like(tensor, 1 / 64)) + original = load_file(os.path.join(directory, "model.safetensors")) + unlisted = "model.layers.0.self_attn.k_proj.weight" + checkpoint_key = ( + unlisted.replace("model.layers.", "decoder.blocks.") if rename_namespace else unlisted + ) + torch.testing.assert_close(state[unlisted], original[checkpoint_key].to(torch.float16)) + with torch.no_grad(): + inputs = torch.tensor([[1, 2, 3]]) + expected = model(inputs).logits + self.assertTrue(torch.isfinite(expected).all()) + model.save_pretrained(saved) + reloaded = AutoModelForCausalLM.from_pretrained(saved).eval() + with torch.no_grad(): + torch.testing.assert_close(reloaded(inputs).logits, expected) + + def test_modelopt_mx_checkpoint_loads_and_saves_dense_and_expert_weights(self): + for format_name, config_class in (("mxfp8", Mxfp8Config), ("mxfp4", Mxfp4Config)): + with self.subTest(format=format_name): + packed = format_name == "mxfp4" + checkpoint = _per_expert( + lambda rows, cols: { + "weight": ( + torch.full((rows, cols // 2), 0xB2, dtype=torch.uint8) + if packed + else torch.full((rows, cols), 2, dtype=torch.float8_e4m3fn) + ), + "weight_scale": torch.full((rows, cols // 32), 121, dtype=torch.uint8), + } + ) + model, info, saved = _load_and_save(checkpoint, config_class(checkpoint_format="modelopt")) + self.assertFalse(info.missing_keys) + self.assertFalse(info.unexpected_keys) + self.assertFalse(info.conversion_errors) + q_proj = model.model.layers[0].self_attn.q_proj + experts = model.model.layers[0].mlp.experts + if packed: + self.assertEqual(q_proj.weight.dtype, torch.int8) + self.assertEqual(experts.gate_up_proj.dtype, torch.int8) + torch.testing.assert_close( + q_proj.weight.view(torch.uint8), checkpoint["model.layers.0.self_attn.q_proj.weight"] + ) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + actual = saved[name] + if tensor.element_size() == 1: + torch.testing.assert_close(actual.view(torch.uint8), tensor.view(torch.uint8)) + else: + torch.testing.assert_close(actual, tensor) + + def test_modelopt_static_fp8_checkpoint_loads_and_saves_dense_and_experts(self): + checkpoint = _per_expert( + lambda rows, cols: { + "weight": torch.full((rows, cols), 2, dtype=torch.float8_e4m3fn), + "weight_scale": torch.tensor(1 / 128), + "input_scale": torch.tensor(1 / 256), + } + ) + for name in checkpoint: + if name.endswith(".up_proj.weight_scale"): + checkpoint[name] = torch.tensor(1 / 64) + config = AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "FP8"}) + model, info, saved = _load_and_save(checkpoint, config) + self.assertFalse(info.missing_keys) + self.assertFalse(info.unexpected_keys) + self.assertFalse(info.conversion_errors) + q_proj = model.model.layers[0].self_attn.q_proj + self.assertEqual(q_proj.weight.dtype, torch.float8_e4m3fn) + torch.testing.assert_close(q_proj.weight_scale_inv, torch.tensor(1 / 128)) + torch.testing.assert_close(q_proj.activation_scale, torch.tensor(1 / 256)) + experts = model.model.layers[0].mlp.experts + torch.testing.assert_close( + experts.gate_up_proj_scale_inv, + torch.tensor([[[1 / 128], [1 / 64]], [[1 / 128], [1 / 64]]]), + ) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + if name.endswith(".weight"): + torch.testing.assert_close(saved[name].view(torch.uint8), tensor.view(torch.uint8)) + else: + torch.testing.assert_close(saved[name], tensor) + + def test_compressed_tensors_mxfp8_scales_reach_dense_and_expert_modules(self): + config = AutoQuantizationConfig.from_dict( + { + "quant_method": "compressed-tensors", + "format": "float-quantized", + "config_groups": { + "group_0": { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 32}, + "input_activations": { + "num_bits": 8, + "type": "float", + "strategy": "group", + "group_size": 32, + "dynamic": True, + }, + } + }, + } + ) + self.assertIsInstance(config, Mxfp8Config) + checkpoint = _per_expert( + lambda rows, cols: { + "weight": torch.full((rows, cols), 2, dtype=torch.float8_e4m3fn), + "weight_scale": torch.full((rows, cols // 32), 121, dtype=torch.uint8), + } + ) + model, info, saved = _load_and_save(checkpoint, config) + self.assertFalse(info.missing_keys) + self.assertFalse(info.unexpected_keys) + self.assertFalse(info.conversion_errors) + torch.testing.assert_close( + model.model.layers[0].self_attn.q_proj.weight_scale_inv, + checkpoint["model.layers.0.self_attn.q_proj.weight_scale"], + ) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + torch.testing.assert_close(saved[name].view(torch.uint8), tensor.view(torch.uint8)) diff --git a/tests/integrations/finegrained/utils.py b/tests/integrations/finegrained/utils.py new file mode 100644 index 000000000000..a38456903a94 --- /dev/null +++ b/tests/integrations/finegrained/utils.py @@ -0,0 +1,271 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Helpers shared by the fine-grained quantization tests: the format configs, the quantizer a load picks, the +checkpoint-config builders, the mocked kernel bundle, and the tiny MoE checkpoint the conversion chains load +and save back.""" + +from copy import deepcopy +from dataclasses import dataclass +from types import SimpleNamespace +from unittest import mock + +import torch +from torch import nn + +import transformers.integrations.finegrained.layers as fg +from transformers import PreTrainedConfig, PreTrainedModel +from transformers.conversion_mapping import get_checkpoint_conversion_mapping +from transformers.core_model_loading import ( + convert_and_load_state_dict_in_model, + revert_weight_conversion, +) +from transformers.integrations.moe import use_experts_implementation +from transformers.modeling_utils import LoadStateDictConfig +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.utils.quantization_config import BLOCK_SCALED_CONFIGS + + +def _format_config(weight_format, flat=False, **options): + """`weight_format`'s config. Built for a module directly, per-tensor FP8 and W4A4 MXFP4 unless asked otherwise; + `flat`, the defaults a config loading a whole checkpoint takes.""" + defaults = ( + {} + if flat + else {"fp8": {"weight_block_size": None}, "mxfp4": {"activation_format": None}}.get(weight_format, {}) + ) + return BLOCK_SCALED_CONFIGS[weight_format](**{**defaults, **options}) + + +class _Cfg: + hidden_size = 64 + num_local_experts = 4 + intermediate_size = 32 + hidden_act = "silu" + _experts_implementation = None + + +def _quantizer_for(config): + """The quantizer a `from_pretrained` picks for this config.""" + from transformers.quantizers.auto import AUTO_QUANTIZER_MAPPING + + method = getattr(config.quant_method, "value", config.quant_method) + return AUTO_QUANTIZER_MAPPING[method](config) + + +def _replace(model, config, modules_to_not_convert=None): + """The module swap a load with `config` makes, through the quantizer it picks.""" + return _quantizer_for(config).replace_modules(model, modules_to_not_convert=modules_to_not_convert) + + +def _targets(conv): + t = conv.target_patterns + return t if isinstance(t, (list, tuple)) else [t] + + +def _modelopt(**fields): + return AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "NVFP4", **fields}) + + +def _model_config(**text_fields): + text = SimpleNamespace(**text_fields) + return SimpleNamespace(get_text_config=lambda: text) + + +def _compressed_tensors(weights, activations=None, fmt="nvfp4-pack-quantized", **fields): + group = { + "weights": weights, + "targets": ["Linear"], + **({"input_activations": activations} if activations else {}), + } + return {"quant_method": "compressed-tensors", "format": fmt, "config_groups": {"group_0": group}, **fields} + + +@dataclass +class _Call: + args: tuple + kwargs: dict + + +class _Recorder: + """Records every call and returns a correctly-shaped zero tensor for matmuls.""" + + def __init__(self): + self.calls = {} + + def _op(self, name): + def op(*args, **kwargs): + self.calls.setdefault(name, []).append(_Call(args, kwargs)) + if name == "swizzle_mx_scales": + return args[0] + if name.startswith("moe_fused"): # (hidden, top_k_index, top_k_weights, gate_up, down, ...) + hidden, down = args[0], kwargs["down_proj"] + return hidden.new_zeros(hidden.shape[0], down.shape[1], dtype=torch.bfloat16) + a, b = args[0], args[1] + n = b.shape[-2] + gather = kwargs.get("gather_idx") + rows = gather.shape[0] if gather is not None else a.shape[0] + return a.new_zeros(rows, n, dtype=torch.bfloat16) + + return op + + +def fake_finegrained_kernel(): + """A `finegrained-kernels` stand-in whose ops record their calls and return correctly shaped zeros.""" + rec = _Recorder() + kernel = mock.Mock() + for name in ( + "matmul_2d", + "matmul_batched", + "matmul_grouped", + "moe_fused_batched", + "moe_fused_grouped", + "swizzle_mx_scales", + "unswizzle_mx_scales", + "mxfp8_act_quant", + "mxfp4_act_quant", + "nvfp4_act_quant", + ): + setattr(kernel, name, rec._op(name)) + kernel.get_supported_act_fns = lambda: ("silu", "gelu", "relu") + kernel.get_supported_norms = lambda: ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm") + return kernel, rec + + +def patch_finegrained_kernel(kernel): + """Patches that make the finegrained integration load `kernel` as its bundle.""" + return ( + mock.patch.object(fg, "_FINEGRAINED", None), + mock.patch.object(fg, "is_kernels_available", return_value=True), + mock.patch.object(fg, "lazy_load_kernel", return_value=kernel), + ) + + +E, H, I = 2, 64, 32 # experts, hidden, intermediate: one 32-group along each input + + +@use_experts_implementation +class _Experts(nn.Module): + """A model's unquantized experts, which the quantizer replaces with `FineGrainedExperts`.""" + + def __init__(self, config): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.empty(E, 2 * I, H)) + self.down_proj = nn.Parameter(torch.empty(E, H, I)) + + +class _Layer(nn.Module): + def __init__(self, config): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(H, H, bias=False) + self.mlp = nn.Module() + self.mlp.experts = _Experts(config) + + +class _MoeModel(PreTrainedModel): + """One decoder layer: a dense `q_proj` and a MoE, the two module kinds a checkpoint quantizes.""" + + base_model_prefix = "model" + config: PreTrainedConfig + + def __init__(self, config): + super().__init__(config) + self.model = nn.Module() + self.model.layers = nn.ModuleList([_Layer(config)]) + self.post_init() + + def _init_weights(self, module): + pass # every weight here comes from the state dict + + +def _model_mapping(): + """The per-expert merge a Qwen3-MoE / DeepSeek model declares for its experts.""" + return get_checkpoint_conversion_mapping("qwen2_moe") + + +def _conversion_model(quantizer): + """Build the tiny model's parameter tree before asking the quantizer for its conversions.""" + config = PreTrainedConfig() + config.hidden_size, config.intermediate_size, config.num_local_experts, config.hidden_act = H, I, E, "silu" + with torch.device("meta"), mock.patch.object(fg, "is_sm100", return_value=False): + model = _MoeModel(config) + quantizer._process_model_before_weight_loading(model) + return model + + +def _load_and_save(checkpoint, quantization_config, weight_mapping=None): + """`checkpoint` through the real quantizer's chain, as `from_pretrained` runs it (off SM100, so the + scales stay affine and no kernel is needed), then saved back. Returns `(model, loading_info, saved)`.""" + config = PreTrainedConfig() + config.hidden_size, config.intermediate_size, config.num_local_experts, config.hidden_act = H, I, E, "silu" + quantizer = _quantizer_for(quantization_config) + quantizer.pre_quantized = True + model = _MoeModel(config) + with mock.patch.object(fg, "is_sm100", return_value=False): + quantizer._process_model_before_weight_loading(model) + load_config = LoadStateDictConfig( + weight_mapping=quantizer.update_weight_conversions( + _model_mapping() if weight_mapping is None else deepcopy(weight_mapping), model + ), + hf_quantizer=quantizer, + ) + info, _ = convert_and_load_state_dict_in_model( + model, {k: v.clone() for k, v in checkpoint.items()}, load_config + ) + quantizer._process_model_after_weight_loading(model) + model.hf_quantizer = quantizer + saved = revert_weight_conversion(model, model.state_dict()) + return model, info, saved + + +def _raw(tensor): + return tensor.reshape(-1).view(torch.uint8) if tensor.element_size() == 1 else tensor + + +def _per_expert(make): + """A per-expert checkpoint layout: `make(n, k)` gives one projection's `{suffix: tensor}`.""" + checkpoint = {} + for expert in range(E): + for proj, shape in (("gate_proj", (I, H)), ("up_proj", (I, H)), ("down_proj", (H, I))): + for suffix, tensor in make(*shape).items(): + checkpoint[f"model.layers.0.mlp.experts.{expert}.{proj}.{suffix}"] = tensor + for suffix, tensor in make(H, H).items(): + checkpoint[f"model.layers.0.self_attn.q_proj.{suffix}"] = tensor + return checkpoint + + +class _ChainAssertions: + def assertLoadsCleanly(self, info): + self.assertEqual( + (info.missing_keys, info.unexpected_keys, info.mismatched_keys, info.conversion_errors), + (set(), set(), set(), {}), + ) + + def assertSavesBack(self, saved, checkpoint): + """The save writes the checkpoint's own keys, dtypes and bytes.""" + self.assertEqual(sorted(saved), sorted(checkpoint)) + for key, tensor in checkpoint.items(): + with self.subTest(key=key): + self.assertEqual((saved[key].dtype, saved[key].shape), (tensor.dtype, tensor.shape)) + self.assertTrue(torch.equal(_raw(saved[key]), _raw(tensor))) + + def assertHoldsInterleaved(self, experts, checkpoint, suffix, held): + """The experts hold gate|up stacked and row-interleaved: `[g0, u0, g1, u1, ...]`.""" + self.assertTrue(experts.holds_interleaved_gate_up) + for expert in range(E): + gate = checkpoint[f"model.layers.0.mlp.experts.{expert}.gate_proj.{suffix}"] + up = checkpoint[f"model.layers.0.mlp.experts.{expert}.up_proj.{suffix}"] + expected = torch.stack([_raw(gate).view(gate.shape), _raw(up).view(up.shape)], dim=1) + got = _raw(held[expert]).view(held[expert].shape) + self.assertTrue(torch.equal(got, expected.reshape(-1, gate.shape[-1])), f"expert {expert} {suffix}") diff --git a/tests/integrations/test_moe.py b/tests/integrations/test_moe.py index bdf2681350ca..efef759ae22f 100644 --- a/tests/integrations/test_moe.py +++ b/tests/integrations/test_moe.py @@ -41,6 +41,7 @@ def apply_gate(gate_up): num_experts=num_experts, has_gate=True, has_bias=False, + has_post_expert_norm=False, is_transposed=False, act_fn=act_fn, _apply_gate=apply_gate, @@ -49,7 +50,6 @@ def apply_gate(gate_up): down_proj=torch.randn(num_experts, hidden, inter, device=torch_device), down_proj_bias=None, _is_expert_parallel=is_expert_parallel, - has_post_expert_norm=False, ) diff --git a/tests/kernels/__init__.py b/tests/kernels/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/kernels/conftest.py b/tests/kernels/conftest.py new file mode 100644 index 000000000000..571e4a374495 --- /dev/null +++ b/tests/kernels/conftest.py @@ -0,0 +1,35 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. + +import os +from unittest import mock + +import pytest + + +@pytest.fixture(scope="session", autouse=True) +def cheap_finegrained_tuning(): + # Measure one configuration without extra search phases for correctness tests. + # Set this before kernel import: autotuners read these controls at import time. + # Subprocesses inherit it too, including distributed workers. + with mock.patch.dict( + os.environ, + { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "FINEGRAINED_AUTOTUNE_BASIN_PROBES": "0", + "FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS": "0", + "FINEGRAINED_AUTOTUNE_RUNOFF": "0", + }, + ): + yield diff --git a/tests/kernels/test_deepgemm.py b/tests/kernels/test_deepgemm.py index dbc618948bbb..2316a6b475c1 100644 --- a/tests/kernels/test_deepgemm.py +++ b/tests/kernels/test_deepgemm.py @@ -28,9 +28,9 @@ (packed int32 UE8M0 SFs, int32 grouped layout, `(qtensor, sf)` operand tuples, recipes, transformed Mega MoE weights, ...). -Arch-gated paths are exercised by faking the device capability (`is_sm100()` reads it): the -SF-packing / TMA-alignment / psum-layout code it selects is pure tensor arithmetic, so mocking the -capability to SM100 drives the Blackwell paths on any device. +Arch-gated paths are exercised by faking the device (`is_sm100()` reads availability, then the +capability): the SF-packing / TMA-alignment / psum-layout code it selects is pure tensor arithmetic, +so mocking both to a Blackwell device drives those paths anywhere, a CPU-only runner included. """ import contextlib @@ -43,21 +43,27 @@ import torch from parameterized import parameterized -from test_utils import make_experts, make_fp8_experts import transformers.integrations.deepgemm as dg -from transformers.distributed.utils import is_dtensor +import transformers.integrations.finegrained.layers as fg from transformers.integrations.deepgemm import ( deepgemm_bf16_experts_forward, deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_linear, deepgemm_fp8_fp4_megamoe_experts_forward, ) +from transformers.integrations.finegrained import FineGrainedLinear from transformers.testing_utils import ( require_torch, + require_torch_gpu, require_torch_greater_or_equal, + require_torch_multi_gpu, torch_device, ) +from transformers.utils.quantization_config import BlockFp8Config, Mxfp8Config + +from ..integrations.finegrained.utils import _Cfg, fake_finegrained_kernel, patch_finegrained_kernel +from .test_utils import make_experts, make_fp8_experts def _add_one(x, *args, **kwargs): @@ -90,6 +96,13 @@ class _DeepGemmKernelMissingSymbol(_FakeDeepGemmKernel): @require_torch class DeepGemmLoaderTest(unittest.TestCase): + def test_uint8_scale_containers_use_exponents_on_both_architectures(self): + scales = torch.tensor([[126, 127, 128, 129]], dtype=torch.uint8).repeat(2, 1) + hopper = dg._coerce_sf_for_kernel(scales, is_sm100=False) + torch.testing.assert_close(hopper, torch.tensor([[0.5, 1.0, 2.0, 4.0]]).repeat(2, 1)) + blackwell = dg._coerce_sf_for_kernel(scales, is_sm100=True) + torch.testing.assert_close(blackwell, scales.view(torch.int32)) + def setUp(self): dg._DEEPGEMM = None self.addCleanup(setattr, dg, "_DEEPGEMM", None) @@ -230,21 +243,6 @@ def run(x): run(torch.zeros(3, device=torch_device)) # a graph break / traced probe would raise here - def test_to_local_is_compile_safe(self): - # Regression guard: every experts forward here unwraps its weights through `to_local`, so it runs - # inside the traced region. It calls `is_dtensor`, whose torch-distributed availability check must - # therefore fold to a constant instead of being traced — its body reaches `importlib.metadata`, - # which dynamo cannot follow (and follows differently across Python versions). `@lru_cache` is no - # protection: dynamo steps past cache wrappers, so the marker has to sit on the wrapped function. - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return x.to_local() + 1 if is_dtensor(x) else x + 1 - - out = run(torch.zeros(3, device=torch_device)) # a graph break / traced probe would raise here - self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) - # ── Capturing fake DeepGEMM bundle ───────────────────────────────────────────── # @@ -361,18 +359,49 @@ def setUp(self): def _bundle(self, *, is_sm100): captured = {} bundle = _make_bundle(captured) - # The forwards read the arch via `is_sm100()` (which queries `get_device_capability`), so fake the - # device to the requested arch: lets SM100 dispatch/packing run on this SM80 box and drives the - # Hopper-rejection guards. `[0]` is all `is_sm100()` reads. + # The forwards read the arch via `is_sm100()`, so fake the device to the requested arch: lets + # SM100 dispatch/packing run on this SM80 box and drives the Hopper-rejection guards. `[0]` is + # all `is_sm100()` reads of the capability, and it reaches that only when CUDA is available — + # so a CPU-only runner needs both faked, or every arch-gated branch below goes untested. capability = (10, 0) if is_sm100 else (9, 0) with ( mock.patch.object(dg, "load_deepgemm_kernel", return_value=bundle), + mock.patch.object(torch.cuda, "is_available", return_value=True), mock.patch.object(torch.cuda, "get_device_capability", return_value=capability), ): yield captured # ── deepgemm_fp8_fp4_linear ──────────────────────────────────────────────── + def test_every_arm_refuses_hidden_states_that_are_not_bfloat16(self): + """Each arm builds its intermediates and its output in bfloat16, so anything else changes + the dtype the caller gets back. The check belongs to the guard, not to one forward: it + used to be copied into two of the three and missing from Mega MoE, which allocates its + output bfloat16 regardless.""" + import transformers.integrations.deepgemm as dg + + class _Experts(torch.nn.Module): + activation_scheme = "dynamic" + + arms = ( + dg.deepgemm_bf16_experts_forward, + dg.deepgemm_fp8_fp4_experts_forward, + dg.deepgemm_fp8_fp4_megamoe_experts_forward, + ) + for forward in arms: + with self.subTest(forward=forward.__name__), self.assertRaises(ValueError) as caught: + forward(_Experts(), torch.zeros(2, 4, dtype=torch.float16), None, None) + self.assertIn("bfloat16", str(caught.exception)) + + def test_megamoe_refuses_a_post_expert_norm(self): + """Mega MoE fuses the routing-weighted reduce, so a per-expert output norm has nowhere to go.""" + import transformers.integrations.deepgemm as dg + + module = torch.nn.Module() + module.has_post_expert_norm = True + with self.assertRaises(NotImplementedError): + dg._assert_no_post_expert_norm(module) + def test_linear_fp8_sm90_kernel_inputs(self): # FP8 weights + float32 block SF on SM90: recipe stays None, SFs are handed over row-major # float32 (SM90 dispatch transforms SFA itself and only checks SFB — see `_coerce_sf_for_kernel`). @@ -431,11 +460,12 @@ def test_assert_sm100_requirements(self): int8_w = torch.zeros(1, 1, dtype=torch.int8, device=torch_device) f32_sf = torch.ones(1, 1, dtype=torch.float32, device=torch_device) ue8m0_sf = f32_sf.to(torch.float8_e8m0fnu) - with mock.patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)): # SM90 + available = mock.patch.object(torch.cuda, "is_available", return_value=True) + with available, mock.patch.object(torch.cuda, "get_device_capability", return_value=(9, 0)): # SM90 with self.assertRaisesRegex(NotImplementedError, "Blackwell"): dg._assert_sm100_requirements(int8_w, ue8m0_sf) # FP4 has no Hopper kernel dg._assert_sm100_requirements(fp8_w, f32_sf) # float32 SF is fine on SM90 -> no raise - with mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 0)): # SM100 + with available, mock.patch.object(torch.cuda, "get_device_capability", return_value=(10, 0)): # SM100 with self.assertRaisesRegex(NotImplementedError, "float32 scale-factor path"): dg._assert_sm100_requirements(fp8_w, f32_sf) # no float32 SF path on Blackwell dg._assert_sm100_requirements(int8_w, ue8m0_sf) # UE8M0 on SM100 -> no raise @@ -580,6 +610,7 @@ def _make_megamoe(self, *, num_experts=4, hidden_dim=64, inter=32, swiglu_limit= requires_grad=False, ), config=types.SimpleNamespace(swiglu_limit=swiglu_limit), + has_post_expert_norm=False, ) @require_torch_greater_or_equal("2.7") # torch.float8_e8m0fnu (UE8M0) landed in 2.7 @@ -658,3 +689,150 @@ def test_megamoe_requires_fp4_weights_and_process_group(self): @staticmethod def _megamoe_flag_set(module): return getattr(module, "_megamoe_transformed", False) + + +@require_torch +class DeepGemmFineGrainedDispatchTest(unittest.TestCase): + """`prefers_deepgemm_linear` carries two independent gates: a correctness one (a pre-swizzled scale is not + readable as row-major) and an SM100 perf one. Block-FP8 scales have no swizzled layout, so the first never + fires for the shape the second catches.""" + + def _routed_to(self, *, sm100, scale_ndim, trains=False): + """Which backend a block-FP8 linear reaches, given arch, scale layout, and whether the call needs a + gradient.""" + kernel, _ = fake_finegrained_kernel() + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + s = torch.randn(*([1, 1, 1, 1, 1][:scale_ndim] if scale_ndim > 2 else [1, 1])) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with ( + p1, + p2, + p3, + mock.patch.dict(os.environ, {"TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR": "0"}), + mock.patch.object(dg, "is_deepgemm_loadable", return_value=True), + mock.patch.object(dg, "is_sm100", return_value=sm100), + mock.patch.object(fg, "deepgemm_fp8_fp4_linear") as deepgemm_linear, + ): + x = torch.randn(4, 64, dtype=torch.bfloat16, requires_grad=trains) + fg.finegrained_linear(x, w, s, block_size=[128, 128]) + return "deepgemm" if deepgemm_linear.called else "triton" + + def test_sm100_never_prefers_deepgemm(self): + self.assertEqual(self._routed_to(sm100=True, scale_ndim=2), "triton") + + def test_pre_sm100_still_uses_deepgemm(self): + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2), "deepgemm") + + def test_a_linear_that_needs_a_gradient_prefers_triton(self): + self.assertEqual(self._routed_to(sm100=False, scale_ndim=2, trains=True), "triton") + + def test_swizzled_scales_never_reach_deepgemm(self): + """The correctness gate holds on any arch, not just where the perf gate does.""" + self.assertEqual(self._routed_to(sm100=False, scale_ndim=5), "triton") + + def test_deepgemm_refuses_a_call_that_needs_a_gradient(self): + """DeepGEMM has no backward pass: every entry point fails a call that needs a gradient, before touching + the kernel.""" + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size = 256, 128 + experts = fg.FineGrainedExperts(cfg, quantization_config=BlockFp8Config(weight_block_size=(128, 128))) + experts._deepgemm_disabled = False + hs = torch.randn(2, 256, dtype=torch.bfloat16, requires_grad=True) + routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + w, s = torch.randn(32, 256).to(torch.float8_e4m3fn), torch.ones(1, 2) + calls = [ + lambda: deepgemm_fp8_fp4_linear(hs, w, s, block_size=(128, 128)), + lambda: deepgemm_fp8_fp4_experts_forward(experts, hs, *routing), + lambda: deepgemm_fp8_fp4_megamoe_experts_forward(experts, hs, *routing), + ] + with mock.patch.object(dg, "load_deepgemm_kernel") as load: + for call in calls: + with self.assertRaises(NotImplementedError): + call() + load.assert_not_called() + + def test_deepgemm_experts_refuse_swizzled_scales(self): + """A module loaded for a triton backend holds swizzled scales; switching it to a DeepGEMM experts backend + afterwards fails loudly rather than reading the permuted buffer as affine.""" + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, "grouped_mm" + with ( + mock.patch.object(fg, "is_sm100", return_value=True), + mock.patch("torch.cuda.is_available", return_value=True), + ): + experts = fg.FineGrainedExperts(cfg, quantization_config=Mxfp8Config()) + self.assertEqual(experts.gate_up_proj_scale_inv.ndim, 5) + hs = torch.randn(2, 256, dtype=torch.bfloat16) + idx, wts = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): + with self.assertRaises(RuntimeError): + forward(experts, hs, idx, wts) + + +class DeepGemmMultiDeviceTest(unittest.TestCase): + """`disable_deepgemm_on_multi_device` flags the quantized modules by the devices they occupy: DeepGEMM's + kernels are bound to one CUDA context and corrupt across devices, but a model on one device keeps DeepGEMM + even with other GPUs visible.""" + + @staticmethod + def _linear(device): + return FineGrainedLinear(256, 256, quantization_config=BlockFp8Config()).to(device) + + @require_torch_multi_gpu + def test_multi_device_disables_deepgemm(self): + model = torch.nn.Module() + model.a = self._linear("cuda:0") + model.b = self._linear("cuda:1") + fg.disable_deepgemm_on_multi_device(model) + self.assertTrue(model.a._deepgemm_disabled) + self.assertTrue(model.b._deepgemm_disabled) + + @require_torch_gpu + def test_single_device_keeps_deepgemm(self): + model = torch.nn.Module() + model.a = self._linear("cuda:0") + model.b = self._linear("cuda:0") + fg.disable_deepgemm_on_multi_device(model) + self.assertFalse(model.a._deepgemm_disabled) + self.assertFalse(model.b._deepgemm_disabled) + + +@require_torch +class DeepGemmQuantizedExpertsOperandsTest(unittest.TestCase): + """The quantized DeepGEMM arms read the stacked weights and their block scales and nothing else: whatever + else a module carries is refused before any kernel work rather than dropped.""" + + def _experts(self, quantization_config, impl="deepgemm", has_bias=False): + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg._experts_implementation = 256, 128, impl + experts = fg.FineGrainedExperts(cfg, quantization_config=quantization_config, has_bias=has_bias) + # `from_module` wraps the class's `__init__` to stamp the flags of the module it last replaced + experts.has_bias = has_bias + return experts + + def _assert_refused(self, experts, message): + hidden = torch.randn(2, 256, dtype=torch.bfloat16) + routing = torch.zeros(2, 2, dtype=torch.long), torch.ones(2, 2) + with mock.patch.object(dg, "load_deepgemm_kernel") as load: + for forward in (deepgemm_fp8_fp4_experts_forward, deepgemm_fp8_fp4_megamoe_experts_forward): + with self.subTest(forward=forward.__name__), self.assertRaisesRegex(NotImplementedError, message): + forward(experts, hidden, *routing) + load.assert_not_called() + + def test_an_interleaved_gate_up_is_refused(self): + """A module loaded for a triton backend holds gate|up interleaved, its block grid interleaved per block, + which a row-major read pairs with the wrong half.""" + experts = self._experts(BlockFp8Config(weight_block_size=(128, 128)), impl="grouped_mm") + self.assertTrue(experts.holds_interleaved_gate_up) + self._assert_refused(experts, "stacked") + + def test_operands_the_kernels_would_drop_are_refused(self): + from transformers.utils.quantization_config import Mxfp4Config, NVFP4Config + + for experts, message in ( + (self._experts(BlockFp8Config(weight_block_size=(128, 128)), has_bias=True), "bias"), + (self._experts(NVFP4Config()), "NVFP4"), + (self._experts(Mxfp4Config(activation_format="bf16")), "weight-only"), + ): + with self.subTest(message=message): + self._assert_refused(experts, message) diff --git a/tests/kernels/test_finegrained.py b/tests/kernels/test_finegrained.py new file mode 100644 index 000000000000..7a4bf22216fa --- /dev/null +++ b/tests/kernels/test_finegrained.py @@ -0,0 +1,340 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""What the `finegrained` integration hands `kernels-community/finegrained-kernels`: the kernel bundle is +mocked, so these pin the loader and the call contract (As-positional / no-block_size / b_global_scale) without +a GPU.""" + +import unittest +from unittest import mock + +import torch + +import transformers.integrations.deepgemm as deepgemm +import transformers.integrations.finegrained.layers as fg +from transformers.integrations.finegrained import ( + FineGrainedExperts, + FineGrainedLinear, + finegrained_linear, + load_finegrained_kernel, +) +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import NVFP4Config + +from ..integrations.finegrained.utils import _Cfg, _format_config, fake_finegrained_kernel, patch_finegrained_kernel + + +@require_torch +class FineGrainedLoaderTest(unittest.TestCase): + def test_loader_requires_every_symbol(self): + kernel, _ = fake_finegrained_kernel() + del kernel.matmul_grouped + # Mock auto-creates attributes; force the miss + kernel.matmul_grouped = None + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3, self.assertRaises(ImportError) as ctx: + load_finegrained_kernel() + self.assertIn("matmul_grouped", str(ctx.exception)) + + def test_loader_is_compile_safe_cold(self): + """Cold path: the compiled call is first to load, so the opaque loader node runs its full + body under compile and must return None, never the bundle (`torch.* op returned + non-Tensor`). The loader has no arch gate, so nothing here fakes a device.""" + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_is_compile_safe_when_warm(self): + """Warm path, which is the production order: eager warm-up, then compile. The loader hits + its short-circuit at trace time — the branch that must also return None.""" + kernel, _ = fake_finegrained_kernel() + kernel.matmul_2d = lambda x, *a, **k: x + 1 + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + load_finegrained_kernel() + torch.compiler.reset() + + @torch.compile(fullgraph=True) + def run(x): + return load_finegrained_kernel().matmul_2d(x) + + out = run(torch.zeros(3)) + self.assertTrue(torch.equal(out, torch.ones(3))) + + def test_loader_binds_all_symbols(self): + kernel, _ = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + bundle = load_finegrained_kernel() + self.assertIs(bundle.matmul_2d, kernel.matmul_2d) + self.assertIs(bundle.get_supported_act_fns, kernel.get_supported_act_fns) + + +@require_torch +class FineGrainedLinearMarshallingTest(unittest.TestCase): + def _run(self, **linear_kwargs): + kernel, rec = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + x = torch.randn(3, 5, 64, dtype=torch.bfloat16) + w = torch.randn(32, 64).to(torch.float8_e4m3fn) + ws = torch.randn(1, 1, dtype=torch.float32) + out = finegrained_linear(x, w, ws, **linear_kwargs) + return out, rec.calls["matmul_2d"][-1] + + def test_scale_rides_positionally_and_block_size_dies(self): + out, call = self._run(block_size=[128, 128]) + # A is flattened 2D; As slot (positional 3rd) is the activation scale = None here + self.assertEqual(call.args[0].shape, (15, 64)) + self.assertIsNone(call.args[2]) + self.assertEqual(call.args[3].dtype, torch.float32) + self.assertNotIn("block_size", call.kwargs) + self.assertEqual(out.shape, (3, 5, 32)) + + def test_static_activation_scale_is_As(self): + scale = torch.tensor(0.5) + _, call = self._run(activation_scale=scale) + self.assertIs(call.args[2], scale) + + def test_nvfp4_global_and_weight_only_format(self): + g = torch.tensor(2.0) + _, call = self._run(weight_global_scale=g, activation_format="bf16") + self.assertIs(call.kwargs["b_global_scale"], g) + self.assertEqual(call.kwargs["activation_format"], "bf16") + + def test_module_forward_threads_everything(self): + kernel, rec = fake_finegrained_kernel() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + m = FineGrainedLinear( + 64, 32, quantization_config=NVFP4Config(weight_block_size=None, activation_format="bf16") + ) + m.weight.data = torch.randint(-127, 127, (32, 32), dtype=torch.int8) + out = m(torch.randn(2, 64, dtype=torch.bfloat16)) + call = rec.calls["matmul_2d"][-1] + self.assertIs(call.kwargs["b_global_scale"], m.weight_global_scale) + self.assertEqual(call.kwargs["activation_format"], "bf16") + self.assertEqual(out.shape, (2, 32)) + + +@require_torch +class FineGrainedExpertsMarshallingTest(unittest.TestCase): + def _experts(self, **kw): + cfg = _Cfg() + weight_format = kw.pop("weight_format", "fp8") + block = {"weight_block_size": (4, 4)} if weight_format == "fp8" else {} + quantized_type = _format_config(weight_format, activation_format=kw.pop("activation_format", None), **block) + m = FineGrainedExperts(cfg, quantization_config=quantized_type, **kw) + for name, p in m.named_parameters(): + if p is None: + continue + if p.dtype == torch.float8_e4m3fn: + p.data = torch.randn(p.shape, dtype=torch.float32).to(torch.float8_e4m3fn) + elif p.dtype.is_floating_point: + p.data = torch.randn_like(p.data) if p.dim() else torch.ones_like(p.data) + return m + + def _route(self, tokens=6): + hs = torch.randn(tokens, 64, dtype=torch.bfloat16) + idx = torch.randint(0, 4, (tokens, 2)) + wts = torch.rand(tokens, 2) + return hs, idx, wts + + def test_eager_linear_threads_bias_global_and_activation_format(self): + """The eager per-expert loop hands `matmul_2d` the same operands the fused forwards do: + the expert's bias, its NVFP4 global and the module's activation format (W4A16 here).""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(weight_format="nvfp4", has_bias=True, activation_format="bf16") + p1, p2, p3 = patch_finegrained_kernel(kernel) + linear = mock.Mock(wraps=fg.finegrained_linear) + with ( + p1, + p2, + p3, + mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False), + mock.patch.object(fg, "finegrained_linear", linear), + ): + m(*self._route()) + self.assertTrue(linear.call_args_list) + for call in linear.call_args_list: + self.assertEqual(call.kwargs["bias"].ndim, 1) # this expert's bias (added after the matmul) + self.assertEqual(call.kwargs["weight_global_scale"].shape, ()) # this expert's NVFP4 weight global + self.assertEqual(call.kwargs["activation_format"], "bf16") + # weight-only: the activations stay bf16, so no activation global is quantized against + self.assertIsNone(call.kwargs["input_global_scale"]) + for call in rec.calls["matmul_2d"]: + self.assertIsNotNone(call.kwargs["b_global_scale"]) + self.assertEqual(call.kwargs["activation_format"], "bf16") + + def test_batched_marshalling(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + hidden, top_k_index, top_k_weights = call.args + self.assertEqual(hidden.shape[0], 6) # UNEXPANDED tokens — the kernel gathers + self.assertEqual(top_k_index.shape, (6, 2)) + self.assertIs(call.kwargs["gate_up_proj"], m.gate_up_proj) + # one scale tensor per projection: the affine Parameter, or the swizzled cache when the + # post-load hook built one (not here — the hook runs on SM100 for dot_scaled chains) + self.assertIs(call.kwargs["gate_up_proj_scale_inv"], m.gate_up_proj_scale_inv) + self.assertIsNone(call.kwargs["gate_up_proj_weight_global_scale"]) + self.assertEqual(call.kwargs["act_fn"], "silu") # fusable: passed by name + self.assertIs(call.kwargs["gate"], True) + self.assertIsNone(call.kwargs["activation_format"]) # None = the weight family's format + self.assertEqual(out.shape, (6, 64)) + + def test_grouped_marshalling(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIs(call.kwargs["down_proj"], m.down_proj) + self.assertIs(call.kwargs["down_proj_scale_inv"], m.down_proj_scale_inv) + self.assertNotIn("moe_fused_batched", rec.calls) + + def test_activation_format_passes_through(self): + kernel, rec = fake_finegrained_kernel() + # the format's own activations read as `None` + for activation_format, handed in ((None, None), ("bf16", "bf16"), ("nvfp4", None)): + m = self._experts(has_gate=True, weight_format="nvfp4", activation_format=activation_format) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_batched"][-1].kwargs["activation_format"], handed) + + def test_nvfp4_experts_thread_per_expert_globals(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, weight_format="nvfp4") + # the format table must resolve the ATTRIBUTE the forwards gate on, not just the + # param allocation — a None here silently drops the global at every forward + self.assertIsNotNone(m.global_scale_dtype) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + self.assertIs(call.kwargs["gate_up_proj_weight_global_scale"], m.gate_up_proj_weight_global_scale) + self.assertIs(call.kwargs["down_proj_weight_global_scale"], m.down_proj_weight_global_scale) + # one weight global per expert on both projections: a gate|up stack calibrated per half + # is merged to that at load, so nothing downstream carries the pair + self.assertEqual(tuple(m.gate_up_proj_weight_global_scale.shape), (m.num_experts,)) + self.assertEqual(tuple(m.down_proj_weight_global_scale.shape), (m.num_experts,)) + # the calibrated activation globals (the checkpoint's `input_scale`): one per expert on both + # projections, each routed row quantized against its own expert's + self.assertIs(call.kwargs["gate_up_proj_input_global_scale"], m.gate_up_proj_activation_scale) + self.assertIs(call.kwargs["down_proj_input_global_scale"], m.down_proj_activation_scale) + self.assertEqual(tuple(m.gate_up_proj_activation_scale.shape), (m.num_experts,)) + self.assertEqual(tuple(m.down_proj_activation_scale.shape), (m.num_experts,)) + + def test_weight_only_experts_hold_no_activation_global(self): + """W4A16 keeps the activations bf16 and requantizes nothing, so there is no activation + quant for a calibrated global to normalize: the module never allocates one, and the ops + would refuse a two-level path without it.""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, weight_format="nvfp4", activation_format="bf16") + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_grouped"] + self.assertIsNone(m.gate_up_proj_activation_scale) + self.assertIsNone(m.down_proj_activation_scale) + self.assertIsNone(call.kwargs["gate_up_proj_input_global_scale"]) + self.assertIsNone(call.kwargs["down_proj_input_global_scale"]) + # the WEIGHT globals are the format's own second level and stay + self.assertIsNotNone(call.kwargs["gate_up_proj_weight_global_scale"]) + + def test_post_expert_norm_rides_the_chain_and_the_eager_loop(self): + """A model whose experts norm the down output before the routing weights: + the swap carries the norm over, the kernel chain runs it on the routed rows, and the + eager loop applies it per expert application — the same place the reference forwards do. + A form the kernels do not implement rides as the module's own ``post_expert_norm``.""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.LayerNorm(m.hidden_dim), True + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + self.assertEqual(rec.calls["moe_fused_grouped"][-1].kwargs["post_expert_norm"], m.post_expert_norm) + + class _Recording(torch.nn.Module): + def __init__(self): + super().__init__() + self.rows = 0 + + def forward(self, x): + self.rows += x.shape[0] + return x + + m.post_expert_norm = _Recording() + with p1, p2, p3, mock.patch.object(deepgemm, "is_deepgemm_loadable", return_value=False): + m(*self._route()) + self.assertTrue(m.post_expert_norm.rows, "the eager loop never applied the post-expert norm") + + def test_post_expert_norm_rides_by_name_when_the_kernels_know_it(self): + """The norm follows `act_fn`'s shape: the name its class maps to goes to the kernels when + they implement it, with the weight and epsilon the fused form needs, so the chain folds it + into its reduce instead of calling the module.""" + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True) + m.post_expert_norm, m.has_post_expert_norm = torch.nn.RMSNorm(m.hidden_dim, eps=1e-4), True + p1, p2, p3 = patch_finegrained_kernel(kernel) + for name, fused in (("rms_norm", True), ("input_scaled_rms_norm", True), ("a_models_own_norm", False)): + with p1, p2, p3, mock.patch.dict(fg.SUPPORTED_FUSED_NORMS, {torch.nn.RMSNorm: name}): + fg.finegrained_grouped_mm_experts_forward(m, *self._route()) + call = rec.calls["moe_fused_grouped"][-1] + if fused: + self.assertEqual(call.kwargs["post_expert_norm"], name) + self.assertIs(call.kwargs["post_expert_norm_weight"], m.post_expert_norm.weight) + self.assertAlmostEqual(call.kwargs["post_expert_norm_eps"], 1e-4) + else: + self.assertEqual(call.kwargs["post_expert_norm"], m.post_expert_norm) + self.assertIsNone(call.kwargs["post_expert_norm_weight"]) + + def test_biases_ride_the_kernel_chain(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, has_bias=True) + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + out = fg.finegrained_batched_mm_experts_forward(m, *self._route()) + self.assertTrue(torch.isfinite(out).all()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels add gate_up's bias before the gated split and down's before the + # routing-weight multiply; nothing is added host-side + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) + self.assertIs(call.kwargs["down_proj_bias"], m.down_proj_bias) + + def test_unfusable_act_fn_is_passed_as_the_module_glu(self): + kernel, rec = fake_finegrained_kernel() + m = self._experts(has_gate=True, has_bias=True) + m.act_fn_name = "quick_gelu" # not in the kernels' get_supported_act_fns() + p1, p2, p3 = patch_finegrained_kernel(kernel) + with p1, p2, p3: + fg.finegrained_batched_mm_experts_forward(m, *self._route()) + (call,) = rec.calls["moe_fused_batched"] + # the kernels run the module's own GLU on the host between the two GEMMs — a new + # activation never waits for a kernel release; the bias still rides the GEMM + self.assertEqual(call.kwargs["act_fn"], m._apply_gate) # bound-method equality + self.assertIs(call.kwargs["gate_up_proj_bias"], m.gate_up_proj_bias) diff --git a/tests/kernels/test_finegrained_fp8.py b/tests/kernels/test_finegrained_fp8.py deleted file mode 100644 index 2e16b16bcead..000000000000 --- a/tests/kernels/test_finegrained_fp8.py +++ /dev/null @@ -1,348 +0,0 @@ -# Copyright 2026 The HuggingFace Team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# 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. -"""Integration tests for the finegrained-fp8 Triton kernel bindings. - -Two layers, both mocking only what needs the `kernels-community/finegrained-fp8` hub build: - -* `FinegrainedFp8LoaderTest` mocks `is_kernels_available` and `lazy_load_kernel` to drive - `load_finegrained_fp8_kernel` with no hub download — asserting it gates correctly (raises when a - precondition is unmet, returns the bundle otherwise) and stays torch-compile safe. -* `FinegrainedFp8ForwardTest` mocks only the loaded kernel bundle (`matmul` / `batched_matmul` / - `grouped_matmul`) with capturing fakes that return correctly-shaped tensors, and runs the real - `finegrained_fp8_linear`, `fp8_batched_mm_experts_forward` and `fp8_grouped_mm_experts_forward` so - their repeat-interleave / routing-flatten / sort-and-histogram / gating / sentinel-mask / reshape-sum - glue executes for real (device-agnostic, so it runs on CPU); it then asserts the tensors handed to the kernel are what the - kernel expects (flattened routing, unclamped sentinels, weight/scale pairing, ...) and - the value-exact output the surrounding marshalling produces. -""" - -import contextlib -import unittest -from unittest import mock - -import torch -from parameterized import parameterized -from test_utils import make_fp8_experts - -import transformers.integrations.finegrained_fp8 as fg -from transformers.integrations.finegrained_fp8 import ( - finegrained_fp8_linear, - fp8_batched_mm_experts_forward, - fp8_grouped_mm_experts_forward, -) -from transformers.testing_utils import require_torch, torch_device - - -def _add_one(x, *args, **kwargs): - return x + 1 - - -# ── Fake kernels (a "good" one exposing every symbol the loader resolves, plus a variant missing one) ── - - -class _FakeFinegrainedKernel: - matmul_2d = staticmethod(_add_one) - matmul_batched = staticmethod(_add_one) - matmul_grouped = staticmethod(_add_one) - - -class _FinegrainedKernelMissingSymbol: - matmul_2d = staticmethod(_add_one) # missing matmul_batched / matmul_grouped - - -@require_torch -class FinegrainedFp8LoaderTest(unittest.TestCase): - def setUp(self): - fg._FINEGRAINED_FP8 = None - self.addCleanup(setattr, fg, "_FINEGRAINED_FP8", None) - - def _env(self, *, kernels_available=True, kernel=_FakeFinegrainedKernel): - stack = contextlib.ExitStack() - stack.enter_context(mock.patch.object(fg, "is_kernels_available", return_value=kernels_available)) - stack.enter_context(mock.patch.object(fg, "lazy_load_kernel", return_value=kernel)) - return stack - - def test_loads_when_environment_is_valid(self): - with self._env(): - bundle = fg.load_finegrained_fp8_kernel() - self.assertIsInstance(bundle, fg.FineGrainedFP8) - self.assertIs(bundle.matmul, _add_one) - - @parameterized.expand( - [ - ("no_kernels", {"kernels_available": False}, "kernel unavailable"), - ("kernel_load_fails", {"kernel": None}, "Failed to load the finegrained-fp8 kernel"), - ("missing_symbols", {"kernel": _FinegrainedKernelMissingSymbol}, "missing required symbols"), - ] - ) - def test_loader_raises(self, _name, env_kwargs, pattern): - with self._env(**env_kwargs), self.assertRaisesRegex(ImportError, pattern): - fg.load_finegrained_fp8_kernel() - - def test_loader_is_compile_safe(self): - # Cold path: the compiled call is first to load, so the opaque loader node runs its full body - # under compile and must return None, never the bundle (`Unsupported: torch.* op returned - # non-Tensor`). Default (inductor) backend. The finegrained loader has no arch/CUDA gate, so - # nothing here fakes the device — torch.compile sees the real one; no GPU required. - with self._env(): - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return fg.load_finegrained_fp8_kernel().matmul(x) - - out = run(torch.zeros(3, device=torch_device)) - self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) - - def test_loader_is_compile_safe_when_warm(self): - # Warm path (production order: eager warmup, then compile). The loader hits its short-circuit at - # trace time — the branch that must also return None, not the already-loaded bundle. - with self._env(): - fg.load_finegrained_fp8_kernel() - torch.compiler.reset() - - @torch.compile(fullgraph=True) - def run(x): - return fg.load_finegrained_fp8_kernel().matmul(x) - - out = run(torch.zeros(3, device=torch_device)) - self.assertTrue(torch.equal(out, torch.ones(3, device=torch_device))) - - -@require_torch -class FinegrainedFp8ForwardTest(unittest.TestCase): - """Drives the real finegrained-fp8 forwards with only the loaded kernel bundle mocked.""" - - def setUp(self): - fg._FINEGRAINED_FP8 = None - self.addCleanup(setattr, fg, "_FINEGRAINED_FP8", None) - - @contextlib.contextmanager - def _mocked_kernel(self): - # Each fake records every call and returns a rightly-shaped, deterministic tensor: - # `matmul` fills 3.0 (so the bias add is value-checkable); the batched/grouped experts fakes - # return ones (so the routing-weight * mask * reshape-sum reduction is analytically exact and - # independent of the — non-stable on CUDA — expert sort). - calls = {"matmul": [], "batched_matmul": [], "grouped_matmul": []} - - def fake_matmul(input, weight, weight_scale_inv, block_size, output_dtype, *, activation_scale=None): - calls["matmul"].append( - { - "input": input, - "weight": weight, - "weight_scale_inv": weight_scale_inv, - "block_size": block_size, - "output_dtype": output_dtype, - "activation_scale": activation_scale, - } - ) - return torch.full((input.shape[0], weight.shape[0]), 3.0, dtype=output_dtype, device=input.device) - - def fake_batched(input, weight, weight_scale, *, block_size, expert_ids): - calls["batched_matmul"].append( - { - "input": input, - "weight": weight, - "weight_scale": weight_scale, - "block_size": block_size, - "expert_ids": expert_ids, - } - ) - return torch.ones(input.shape[0], weight.shape[1], dtype=torch.float32, device=input.device) - - def fake_grouped(input, weight, weight_scale, *, offsets, tokens_per_expert, block_size): - calls["grouped_matmul"].append( - { - "input": input, - "weight": weight, - "weight_scale": weight_scale, - "offsets": offsets, - "tokens_per_expert": tokens_per_expert, - "block_size": block_size, - } - ) - return torch.ones(input.shape[0], weight.shape[1], dtype=torch.float32, device=input.device) - - bundle = fg.FineGrainedFP8(matmul=fake_matmul, batched_matmul=fake_batched, grouped_matmul=fake_grouped) - with mock.patch.object(fg, "load_finegrained_fp8_kernel", return_value=bundle): - yield calls - - @staticmethod - def _expected_experts_output(hidden_states, top_k_index, top_k_weights, num_experts): - # Mirror the shared tail of both experts forwards with the ones-returning down projection: - # weighted_out = ones * routing_weight, sentinel rows zeroed, then per-token reduction. - num_tokens, top_k = top_k_index.shape - hidden = hidden_states.size(-1) - sw = top_k_weights.reshape(-1).to(torch.float32) - weighted = torch.ones(sw.shape[0], hidden, device=hidden_states.device) * sw.unsqueeze(-1) - sentinel = (top_k_index.reshape(-1) >= num_experts).unsqueeze(-1) - weighted = weighted.masked_fill(sentinel, 0.0) - return weighted.view(num_tokens, top_k, hidden).sum(dim=1).to(hidden_states.dtype) - - # ── finegrained_fp8_linear ──────────────────────────────────────────────────────────────────── - - def test_linear_marshals_args_and_defaults_output_dtype(self): - input = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - weight = torch.randn(16, 8, device=torch_device).to(torch.float8_e4m3fn) - weight_scale_inv = torch.randn(1, 1, dtype=torch.float32, device=torch_device) - block_size = [128, 128] - with self._mocked_kernel() as calls: - out = finegrained_fp8_linear(input, weight, weight_scale_inv, block_size) - call = calls["matmul"][0] - # input / weight / scale pass straight through, positionally. - self.assertIs(call["input"], input) - self.assertIs(call["weight"], weight) - self.assertIs(call["weight_scale_inv"], weight_scale_inv) - # output_dtype defaults to input.dtype when the caller leaves it None, and no activation scale. - self.assertEqual(call["output_dtype"], torch.bfloat16) - self.assertIsNone(call["activation_scale"]) - self.assertEqual(out.dtype, torch.bfloat16) - - def test_linear_adds_bias_in_place(self): - input = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - weight = torch.randn(16, 8, device=torch_device).to(torch.float8_e4m3fn) - weight_scale_inv = torch.randn(1, 1, dtype=torch.float32, device=torch_device) - bias = torch.randn(16, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel(): - out = finegrained_fp8_linear(input, weight, weight_scale_inv, [128, 128], bias=bias) - # The kernel returned a 3.0-filled (3, 16) tensor; bias is broadcast-added onto it. - self.assertTrue(torch.equal(out, torch.full((3, 16), 3.0, dtype=torch.bfloat16, device=torch_device) + bias)) - - # ── fp8_batched_mm_experts_forward ────────────────────────────────────────────────────────────── - - def test_batched_mm_kernel_inputs_and_output(self): - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - - up, down = calls["batched_matmul"] - # Up projection: each token replicated top_k times (S = 6), routing flattened row-major. - self.assertTrue(torch.equal(up["input"], hidden_states.repeat_interleave(2, dim=0))) - self.assertTrue(torch.equal(up["expert_ids"], top_k_index.reshape(-1))) - # Weight/scale pairing passes straight through (no copies). - self.assertIs(up["weight"], experts.gate_up_proj) - self.assertIs(up["weight_scale"], experts.gate_up_proj_scale_inv) - # Down projection sees the gated activation: 2*inter (32) collapsed to inter (16). - self.assertEqual(down["input"].shape, (6, 16)) - self.assertIs(down["weight"], experts.down_proj) - self.assertIs(down["weight_scale"], experts.down_proj_scale_inv) - self.assertTrue(torch.equal(down["expert_ids"], top_k_index.reshape(-1))) - - # Value-exact per-token reduction, recast to the input dtype. - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_batched_mm_non_gated_uses_up_proj_and_activation(self): - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, has_gate=False) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - up, down = calls["batched_matmul"] - self.assertIs(up["weight"], experts.up_proj) - self.assertIs(up["weight_scale"], experts.up_proj_scale_inv) - # Non-gated: act_fn keeps the inter dim (16), no 2*inter chunk. - self.assertEqual(down["input"].shape, (6, 16)) - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_batched_mm_passes_sentinel_expert_ids_unclamped(self): - # EP sentinels (expert_ids >= num_experts) reach the kernel unclamped; the post-mask zeroes the - # matching output rows before the per-token reduction (the kernel leaves them uninitialized). - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device) # 4 == num_experts -> sentinel - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - self.assertTrue(torch.equal(calls["batched_matmul"][0]["expert_ids"], top_k_index.reshape(-1))) - self.assertEqual(int(calls["batched_matmul"][0]["expert_ids"].max()), 4) - # Sentinel token-expert pairs contribute 0 to the reduction -> each token keeps only its non-sentinel weight. - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - self.assertTrue(torch.equal(out, top_k_weights[:, :1].to(torch.float32).expand(3, 8).to(torch.bfloat16))) - - def test_batched_mm_rejects_static_activation_scheme(self): - # Static activation quant needs a per-tensor activation scale the batched kernel can't consume; - # this guards a genuine unsupported-config path, not a trivial type/device check. - experts = make_fp8_experts(activation_scheme="static") - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel(), self.assertRaisesRegex(NotImplementedError, "activation_scheme='static'"): - fp8_batched_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - - # ── fp8_grouped_mm_experts_forward ──────────────────────────────────────────────────────────── - - def test_grouped_mm_kernel_inputs_and_output(self): - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_grouped_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - - up, down = calls["grouped_matmul"] - # S = num_tokens * top_k selected pairs, gathered by the expert-sort permutation. - self.assertEqual(up["input"].shape, (6, 8)) - self.assertIs(up["weight"], experts.gate_up_proj) - self.assertIs(up["weight_scale"], experts.gate_up_proj_scale_inv) - # Down projection sees the gated activation: 2*inter (32) collapsed to inter (16). - self.assertEqual(down["input"].shape, (6, 16)) - self.assertIs(down["weight"], experts.down_proj) - self.assertIs(down["weight_scale"], experts.down_proj_scale_inv) - - # offsets / tokens_per_expert are the per-expert histogram over the sorted expert ids. - # CPU histc requires float input, CUDA requires int (matches the source's dispatch). - expert_ids_g, _ = torch.sort(top_k_index.reshape(-1)) - histc_input = expert_ids_g.float() if expert_ids_g.device.type == "cpu" else expert_ids_g.int() - expected_tpe = torch.histc(histc_input, bins=experts.num_experts, min=0, max=experts.num_experts - 1) - expected_offsets = torch.cumsum(expected_tpe, dim=0, dtype=torch.int32) - self.assertTrue(torch.equal(up["tokens_per_expert"], expected_tpe)) - self.assertTrue(torch.equal(up["offsets"], expected_offsets)) - # Both projections share the same offsets/tokens_per_expert tensors. - self.assertIs(down["offsets"], up["offsets"]) - self.assertIs(down["tokens_per_expert"], up["tokens_per_expert"]) - - # Output restored to original token order (inv_perm) then reduced; value-exact. - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_grouped_mm_sentinels_dropped_from_histogram(self): - # Sentinels are left unclamped so the sort pushes them to the tail and histc(max=num_experts-1) - # drops them from tokens_per_expert -> no wasted GEMM rows; the post-mask zeroes their output. - experts = make_fp8_experts(num_experts=4, hidden=8, inter=16, is_expert_parallel=True) - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.tensor([[0, 4], [1, 4], [2, 4]], device=torch_device) # three sentinels (== num_experts) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel() as calls: - out = fp8_grouped_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) - tpe = calls["grouped_matmul"][0]["tokens_per_expert"] - # Only experts 0,1,2 got one token each; the 3 sentinels are absent from the histogram. - self.assertTrue(torch.equal(tpe, torch.tensor([1.0, 1.0, 1.0, 0.0], device=torch_device))) - self.assertEqual(int(tpe.sum()), 3) - expected = self._expected_experts_output(hidden_states, top_k_index, top_k_weights, experts.num_experts) - self.assertTrue(torch.equal(out, expected)) - - def test_grouped_mm_rejects_static_activation_scheme(self): - experts = make_fp8_experts(activation_scheme="static") - hidden_states = torch.randn(3, 8, dtype=torch.bfloat16, device=torch_device) - top_k_index = torch.randint(0, 4, (3, 2), device=torch_device) - top_k_weights = torch.rand(3, 2, dtype=torch.bfloat16, device=torch_device) - with self._mocked_kernel(), self.assertRaisesRegex(NotImplementedError, "activation_scheme='static'"): - fp8_grouped_mm_experts_forward(experts, hidden_states, top_k_index, top_k_weights) diff --git a/tests/kernels/test_kernels.py b/tests/kernels/test_kernels.py index 8dbe08a7dbce..b3a29ac4523e 100644 --- a/tests/kernels/test_kernels.py +++ b/tests/kernels/test_kernels.py @@ -489,7 +489,7 @@ def test_is_kernel_regex(self): def test_lazy_load_kernel_success_and_cache(self): sentinel = types.ModuleType("sentinel_kernel_module") - def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False): + def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False, check_arch=True): self.assertIn(repo_id, {"kernels-community/causal-conv1d"}) self.assertFalse(allow_all_kernels) return sentinel @@ -534,7 +534,7 @@ def test_lazy_load_kernel_version(self): sentinel_mod = types.ModuleType("sentinel_kernel_module") call_count = {"n": 0} - def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False): + def fake_get_kernel(repo_id, revision=None, version=None, allow_all_kernels=False, check_arch=True): call_count["n"] += 1 self.assertEqual(repo_id, "kernels-community/causal-conv1d") self.assertIsNone(revision) diff --git a/tests/kernels/test_sonicmoe.py b/tests/kernels/test_sonicmoe.py index b4eb58db6dcc..3c6ce2567aee 100644 --- a/tests/kernels/test_sonicmoe.py +++ b/tests/kernels/test_sonicmoe.py @@ -32,12 +32,13 @@ import torch from parameterized import parameterized -from test_utils import make_experts import transformers.integrations.sonicmoe as sm from transformers.integrations.sonicmoe import sonicmoe_experts_forward from transformers.testing_utils import require_torch, torch_device +from .test_utils import make_experts + class _FakeActivationType: SWIGLU = "swiglu" diff --git a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py index e5d3df31d7bc..57e38bcbc77a 100644 --- a/tests/models/deepseek_v4/test_modeling_deepseek_v4.py +++ b/tests/models/deepseek_v4/test_modeling_deepseek_v4.py @@ -38,7 +38,7 @@ AutoModelForCausalLM, AutoTokenizer, DeepseekV4Model, - FineGrainedFP8Config, + FineGrainedConfig, ) from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester @@ -179,22 +179,9 @@ def _check_past_key_values_for_generate(self, batch_size, past_key_values, seq_l self.assertEqual(keys.shape[3], head_dim) self.assertEqual(keys.shape, values.shape) - @unittest.skip( - reason=( - "V4's conversion mapping is two-pass: a structural prefix rename " - "(``layers.X.attn.`` → ``model.layers.X.self_attn.``) runs first, then specific in-prefix " - "renames operate on the already-prefixed HF-form keys (``model.layers.X.self_attn.compressor.norm.`` " - "→ ``...compressor.kv_norm.``). This split is load-bearing for save / load round-tripping — " - "any single-pass ordering loses information in either direction (the general prefix rule " - "and a specific in-prefix rule both want to match the same upstream key, and one of the " - "two directions ends up with the general rule stealing the match). The base " - "``test_reverse_loading_mapping`` checks every source pattern against the *upstream-form* " - "serialized keys, so the Pass 2 patterns (written in HF form) inherently can't satisfy " - "that invariant. The actual round-trip is exercised by ``test_save_load``." - ) - ) def test_reverse_loading_mapping(self): - pass + # the base model has no `model.` prefix for the checkpoint's prefix renames to reverse onto + super().test_reverse_loading_mapping(skip_base_model=True) @unittest.skip( reason=( @@ -233,7 +220,7 @@ class DeepseekV4IntegrationTest(unittest.TestCase): """End-to-end check on the published DeepSeek-V4-Flash checkpoint. Loads the real 43-layer FP8 weights, dequantizes on the fly via - :class:`FineGrainedFP8Config`, and greedy-generates a continuation of a fixed + :class:`FineGrainedConfig`, and greedy-generates a continuation of a fixed prompt. The forward path that this test covers is everything past the typical tiny-config tests can reach: the per-layer FP8 dequant in ``update_weight_conversions``, the ``compress_ratios → layer_types`` config @@ -264,7 +251,7 @@ class DeepseekV4IntegrationTest(unittest.TestCase): ) ) def test_v4_flash_dequantized_generation(self): - quantization_config = FineGrainedFP8Config(dequantize=True) + quantization_config = FineGrainedConfig(dequantize=True) config = AutoConfig.from_pretrained(self.model_id) tokenizer = AutoTokenizer.from_pretrained(self.model_id) model = AutoModelForCausalLM.from_pretrained( @@ -382,7 +369,7 @@ def test_v4_flash_dequantized_chat_seven_prompts(self): ), ] - quantization_config = FineGrainedFP8Config(dequantize=True) + quantization_config = FineGrainedConfig(dequantize=True) config = AutoConfig.from_pretrained(self.model_id) tokenizer = AutoTokenizer.from_pretrained(self.model_id) model = AutoModelForCausalLM.from_pretrained( @@ -426,7 +413,7 @@ def test_v4_flash_dequantized_chat_seven_prompts(self): from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed import DistributedConfig -from transformers.utils.quantization_config import FineGrainedFP8Config +from transformers.utils.quantization_config import FineGrainedConfig LOADTIME_DISPATCH = {loadtime_dispatch!r} @@ -455,12 +442,12 @@ def main() -> int: tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) inputs = tokenizer(PROMPT, return_tensors="pt", add_special_tokens=ADD_SPECIAL_TOKENS).to(model.device) - # FP8 only: Triton multi-expert kernels generate wrong tokens after DeepGEMM has run on + # FP8 experts only: Triton multi-expert kernels generate wrong tokens after DeepGEMM has run on # the same module (root cause unknown — per-row kernel outputs measure bit-perfect, yet # end-to-end generation degrades; not reproducible with DeepGEMM disabled). Force the - # Triton fallback in `fp8_linear` for batched_mm / grouped_mm so attention/MLP linears - # stay Triton-only too. FP4 is unaffected and keeps DeepGEMM for performance. - is_fp8 = getattr(model.config, "expert_dtype", "fp8") != "fp4" + # Triton dense-linear dispatch for batched_mm / grouped_mm so attention/MLP linears + # stay Triton-only too. FP4 experts are unaffected and keep DeepGEMM for performance. + is_fp8 = next((m.weight_format for m in model.modules() if hasattr(m, "holds_interleaved_gate_up")), None) == "fp8" failed = [] for dispatch in RUNTIME_DISPATCHES: @@ -535,7 +522,7 @@ def main() -> int: inputs = tokenizer(PROMPT, return_tensors="pt", add_special_tokens=ADD_SPECIAL_TOKENS).to(model.device) expected_logits_shape = inputs.input_ids.shape - is_fp8 = getattr(model.config, "expert_dtype", "fp8") != "fp4" + is_fp8 = next((m.weight_format for m in model.modules() if hasattr(m, "holds_interleaved_gate_up")), None) == "fp8" failed = [] for dispatch in RUNTIME_DISPATCHES: diff --git a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py index a21a81b5f677..868fa6eb2d4a 100644 --- a/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py +++ b/tests/models/glm_moe_dsa/test_modeling_glm_moe_dsa.py @@ -22,7 +22,7 @@ from transformers import ( AutoModelForCausalLM, AutoTokenizer, - FineGrainedFP8Config, + BlockFp8Config, GlmMoeDsaConfig, is_torch_available, set_seed, @@ -164,7 +164,7 @@ def test_glm_moe_dsa_fp8_inference(self): set_seed(0) # different ranks need the same seed model_id = "zai-org/GLM-5-FP8" - quantization_config = FineGrainedFP8Config( + quantization_config = BlockFp8Config( modules_to_not_convert=[ "model.layers.*.mlp.gate$", "model.layers.*.self_attn.indexer.weights_proj$", diff --git a/tests/models/gpt_oss/test_modeling_gpt_oss.py b/tests/models/gpt_oss/test_modeling_gpt_oss.py index 839054f87625..a94db794d928 100644 --- a/tests/models/gpt_oss/test_modeling_gpt_oss.py +++ b/tests/models/gpt_oss/test_modeling_gpt_oss.py @@ -28,7 +28,6 @@ from transformers import ( AutoModelForCausalLM, AutoTokenizer, - Mxfp4Config, is_torch_available, ) from transformers.testing_utils import ( @@ -46,6 +45,7 @@ torch_device, ) from transformers.utils import is_rocm_platform, is_torch_greater_or_equal +from transformers.utils.quantization_config import Mxfp4Config from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester diff --git a/tests/models/llama4/test_modeling_llama4.py b/tests/models/llama4/test_modeling_llama4.py index cd07907fe892..1a7d4516bdf1 100644 --- a/tests/models/llama4/test_modeling_llama4.py +++ b/tests/models/llama4/test_modeling_llama4.py @@ -13,11 +13,15 @@ # limitations under the License. """Testing suite for the PyTorch Llama4 model.""" +import tempfile import unittest +from parameterized import parameterized + from transformers import is_torch_available from transformers.testing_utils import ( Expectations, + require_torch, require_torch_large_accelerator, slow, torch_device, @@ -29,13 +33,77 @@ if is_torch_available(): import torch + from safetensors.torch import load_file from transformers import ( + Llama4ForCausalLM, Llama4ForConditionalGeneration, Llama4Processor, + Llama4TextConfig, ) +@require_torch +class Llama4TextMoeTest(unittest.TestCase): + def get_config(self, top_k=1): + return Llama4TextConfig( + vocab_size=97, + hidden_size=32, + intermediate_size=48, + intermediate_size_mlp=48, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_local_experts=4, + num_experts_per_tok=top_k, + max_position_embeddings=64, + attention_chunk_size=16, + ) + + @parameterized.expand([1, 2]) + def test_input_router_weights(self, top_k): + config = self.get_config(top_k) + model = Llama4ForCausalLM(config).to(torch_device).eval() + moe = model.model.layers[0].feed_forward + hidden_states = torch.randn(6, config.hidden_size, device=torch_device) + with torch.no_grad(): + output, logits = moe(hidden_states) + # Independent reference: the original expert-major BMM computes every expert, + # with zero inputs for experts outside the top-k selection. + scores = torch.zeros_like(logits) + values, indices = logits.topk(top_k, dim=-1) + scores.scatter_(1, indices, values.sigmoid()) + routed = hidden_states.unsqueeze(0) * scores.t().unsqueeze(-1) + gate, up = torch.bmm(routed, moe.experts.gate_up_proj.transpose(1, 2)).chunk(2, dim=-1) + routed = torch.bmm(moe.experts.act_fn(gate) * up, moe.experts.down_proj.transpose(1, 2)) + expected = moe.shared_expert(hidden_states) + routed.sum(dim=0) + torch.testing.assert_close(output, expected) + + def test_legacy_expert_checkpoint_layout(self): + model = Llama4ForCausalLM(self.get_config(top_k=2)).eval() + input_ids = torch.tensor([[1, 4, 7]]) + with tempfile.TemporaryDirectory() as checkpoint: + model.save_pretrained(checkpoint) + state = load_file(f"{checkpoint}/model.safetensors") + prefix = "model.layers.0.feed_forward.experts." + torch.testing.assert_close( + state[prefix + "gate_up_proj"], model.state_dict()[prefix + "gate_up_proj"].transpose(1, 2) + ) + torch.testing.assert_close( + state[prefix + "down_proj"], model.state_dict()[prefix + "down_proj"].transpose(1, 2) + ) + loaded = Llama4ForCausalLM.from_pretrained(checkpoint).eval() + with torch.no_grad(): + expected = model(input_ids, output_router_logits=True) + actual = loaded(input_ids, output_router_logits=True) + torch.testing.assert_close(actual.logits, expected.logits) + self.assertEqual(len(actual.router_logits), model.config.num_hidden_layers) + for actual_router, expected_router in zip(actual.router_logits, expected.router_logits): + self.assertEqual(actual_router.shape, (input_ids.numel(), model.config.num_local_experts)) + torch.testing.assert_close(actual_router, expected_router) + + @slow @require_torch_large_accelerator class Llama4IntegrationTest(MemoryCleanupMixin, unittest.TestCase): diff --git a/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py b/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py index 482bea6078ee..688f58955ff3 100644 --- a/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py +++ b/tests/models/minimax_m3_vl/test_modeling_minimax_m3_vl.py @@ -176,9 +176,9 @@ def test_indexer_selects_blocks_per_head(self): heads_differ = (block_indices[0, 0] != block_indices[0, 1]).any() self.assertTrue(bool(heads_differ), "indexer heads selected identical blocks -- head axis was collapsed") - @unittest.skip(reason="IDK exactly why, can be adressed later") def test_reverse_loading_mapping(self): - pass + # the base model has no `model.` prefix for the `^language_model\.model\.` rename to reverse onto, as in llava + super().test_reverse_loading_mapping(skip_base_model=True) @unittest.skip( reason=( diff --git a/tests/models/qwen4_exp/test_modeling_qwen4_exp.py b/tests/models/qwen4_exp/test_modeling_qwen4_exp.py index 37ddc4a960bb..9726fa0d4aec 100644 --- a/tests/models/qwen4_exp/test_modeling_qwen4_exp.py +++ b/tests/models/qwen4_exp/test_modeling_qwen4_exp.py @@ -22,6 +22,7 @@ from transformers import is_torch_available from transformers.testing_utils import require_torch, torch_device +from transformers.utils.quantization_config import NVFP4Config from ...causal_lm_tester import CausalLMModelTest, CausalLMModelTester from ...test_modeling_common import floats_tensor @@ -35,7 +36,6 @@ from transformers import ( AutoModelForCausalLM, DynamicCache, - FineGrainedFP8Config, Qwen4ExpConfig, Qwen4ExpForCausalLM, Qwen4ExpForConditionalGeneration, @@ -46,8 +46,10 @@ StaticCache, ) from transformers.distributed.fsdp import verify_fsdp_plan - from transformers.integrations import FP8Embedding - from transformers.quantizers.quantizer_finegrained_fp8 import FineGrainedFP8HfQuantizer + from transformers.integrations.finegrained import FineGrainedEmbedding + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.utils.quantization_config import BlockFp8Config, FineGrainedConfig class Qwen4ExpTextModelTester(CausalLMModelTester): @@ -192,17 +194,42 @@ def test_ple_layers_must_use_linear_attention(self): layer_types=["linear_attention", "indexed_attention"], ) - def test_finegrained_fp8_embedding_conversion(self): + def test_finegrained_embedding_conversion(self): + """The n-gram table through the CURRENT integration, not the frozen one: a table too big + to dequantize at load is held in FP8 with one per-tensor scale, and the rescale runs on + the rows a lookup gathers.""" config = self.model_tester.get_config() with torch.device("meta"): model = Qwen4ExpForCausalLM(config) - quantization_config = FineGrainedFP8Config( - modules_to_convert=["ple.ple_embedding.ngram_embedding"], modules_to_not_convert=[] + quantizer = BlockFp8HfQuantizer( + BlockFp8Config(modules_to_convert=["ple.ple_embedding.ngram_embedding"], modules_to_not_convert=[]) ) - FineGrainedFP8HfQuantizer(quantization_config).preprocess_model(model) + quantizer.pre_quantized = True + quantizer.preprocess_model(model) - self.assertIsInstance(model.model.layers[0].ple.ple_embedding.ngram_embedding, FP8Embedding) + table = model.model.layers[0].ple.ple_embedding.ngram_embedding + self.assertIsInstance(table, FineGrainedEmbedding) + self.assertEqual(table.weight.dtype, torch.float8_e4m3fn) + self.assertEqual(table.weight_scale.shape, torch.Size([1])) + + def test_a_table_beside_another_format_is_held_in_fp8(self): + """An embedding table is FP8 whatever its model's other modules are: beside NVFP4 linears, the plan's + `modules_to_convert` gives it a block-FP8 entry of its own.""" + config = self.model_tester.get_config() + with torch.device("meta"): + model = Qwen4ExpForCausalLM(config) + + quantizer = FineGrainedHfQuantizer( + FineGrainedConfig( + dtype_plan={".*": NVFP4Config()}, + modules_to_convert=["ple.ple_embedding.ngram_embedding"], + modules_to_not_convert=[], + ) + ) + quantizer.pre_quantized = True + quantizer.preprocess_model(model) + self.assertEqual(model.model.layers[0].ple.ple_embedding.ngram_embedding.weight.dtype, torch.float8_e4m3fn) def test_ple_padding_and_static_cache_match_unpadded_sequence(self): torch.manual_seed(0) diff --git a/tests/quantization/block_scaled_checkpoint_tester.py b/tests/quantization/block_scaled_checkpoint_tester.py new file mode 100644 index 000000000000..c6626db88693 --- /dev/null +++ b/tests/quantization/block_scaled_checkpoint_tester.py @@ -0,0 +1,293 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Accuracy and serialization checks against independently generated tiny checkpoints.""" + +import os +import re + +import torch + +from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText +from transformers.testing_utils import ( + Expectations, + cleanup, + get_device_properties, + require_kernels, + require_torch_accelerator, + slow, + torch_device, +) + +from .block_scaled_tester import FormatSpec, _cheap_tuning, _require_kernel_capability + + +def _checkpoint_bytes(path): + """`{key: (dtype, shape, raw bytes)}` of a saved checkpoint, for comparing one save against another.""" + import glob + + from safetensors import safe_open + + tensors = {} + for shard in sorted(glob.glob(os.path.join(path, "*.safetensors"))): + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + tensor = handle.get_tensor(key) + raw = tensor.reshape(-1).view(torch.uint8).numpy().tobytes() + tensors[key] = (tensor.dtype, tuple(tensor.shape), raw) + return tensors + + +# The logits `[0, :3, :8]` each checkpoint gives on `torch.arange(3, 35)`, per device: bit-stable across tuner +# caches, so a tight tolerance holds. +EXPECTED_SLICES = { + "IlyasMoutawwakil/tiny-random-Qwen3-FP8": Expectations( + { + ("cuda", (10, 0)): [ + [0.6094, -0.6055, -0.3262, 0.1396, -0.1270, -0.0420, 0.3750, -0.1592], + [0.1689, -0.4961, 0.1621, 0.0752, -0.1895, 0.0688, 0.3320, -0.0234], + [0.1069, -0.3027, 0.1050, 0.1553, 0.1699, -0.1572, 0.1699, 0.2598], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Mistral3-FP8-static": Expectations( + { + ("cuda", (10, 0)): [ + [-0.1992, -0.4141, -0.1650, -0.1826, 0.1309, 0.4375, -0.2832, 0.2480], + [-0.1367, -0.2422, -0.1138, -0.3418, 0.2910, 0.3965, -0.0293, -0.0330], + [-0.1436, -0.1055, -0.3320, -0.3203, 0.3633, 0.2480, 0.1523, -0.1621], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Mistral4-FP8-static": Expectations( + { + ("cuda", (10, 0)): [ + [-0.0850, 0.1816, -0.2344, 0.5547, -0.0747, 0.6875, -0.2949, 0.1982], + [-0.0056, -0.0065, 0.0530, 0.0996, 0.1436, 0.2217, -0.3281, 0.1846], + [0.0055, 0.2559, -0.1045, -0.0031, 0.0608, 0.0449, 0.0962, 0.1416], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Llama-NVFP4": Expectations( + { + ("cuda", (10, 0)): [ + [0.7617, -0.5938, -0.2832, 0.2559, -0.1021, 0.0082, 0.3418, -0.1699], + [0.6523, -0.7148, 0.2930, 0.3750, -0.5039, 0.0835, 0.2793, -0.0153], + [0.3457, -0.4785, 0.1924, 0.4102, -0.2715, -0.0835, 0.2070, 0.1602], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-DeepseekV3-FP8": Expectations( + { + ("cuda", (10, 0)): [ + [-0.1260, -0.2051, -0.1836, -0.6523, 0.4395, -0.0771, -0.7461, 0.3613], + [-0.1025, -0.0640, -0.0132, -0.6602, 0.3672, -0.1270, -0.9023, -0.0786], + [-0.0874, 0.0552, 0.0442, -0.4824, 0.4375, -0.4121, -0.4980, -0.1836], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-DeepseekV4-Flash": Expectations( + { + ("cuda", (10, 0)): [ + [-0.3984, 0.3633, 0.1348, -0.0476, -0.1118, -0.0518, -0.0952, -0.4590], + [-0.4336, 0.1836, 0.0049, -0.1514, -0.0449, -0.2383, -0.1113, -0.4766], + [-0.2559, 0.0693, -0.0664, -0.2207, -0.2500, -0.1177, -0.1562, -0.2949], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-GptOss-MXFP4": Expectations( + { + ("cuda", (10, 0)): [ + [-0.1309, 0.2061, 0.2559, 0.5156, -0.2715, 0.5469, 0.4004, -0.0854], + [-0.2490, 0.6406, 0.1289, 0.4277, -0.3809, 0.3125, 0.3516, -0.0811], + [-0.2334, 0.1367, 0.2930, 0.6875, -0.1836, 0.4844, 0.1953, -0.2363], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": Expectations( + { + ("cuda", (10, 0)): [ + [0.1709, -0.4199, -0.4199, -0.2402, -0.3125, 0.6016, 0.3008, 0.1914], + [-0.2051, -0.6328, -0.5391, 0.3184, 0.1299, 0.2773, 0.4824, 0.0938], + [0.0447, -0.5078, -0.3672, -0.1992, 0.1553, 0.6523, 0.4141, -0.0038], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": Expectations( + { + ("cuda", (10, 0)): [ + [0.3086, -0.2773, -0.4883, -0.1797, -0.0073, 0.4531, 0.2158, 0.2754], + [0.0295, -0.2598, -0.5586, 0.2432, -0.0106, 0.3496, 0.4805, 0.1279], + [0.1533, -0.4551, -0.5742, -0.1973, 0.2139, 0.4141, 0.3574, 0.1338], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-Muse-Spark-NVFP4": Expectations( + { + ("cuda", (10, 0)): [ + [0.0430, -0.0131, -0.0304, 0.0415, -0.0082, -0.0029, 0.0042, -0.0498], + [0.0105, 0.0723, 0.0615, -0.0175, 0.0522, 0.0015, 0.0023, 0.0304], + [0.0089, -0.0251, 0.0131, 0.0234, 0.0243, -0.0283, -0.0703, 0.0664], + ], + } + ), # fmt: skip + "IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": Expectations( + { + ("cuda", (10, 0)): [ + [0.2637, -0.0815, -0.3418, 0.1553, -0.2715, 0.1562, 0.4590, 0.0635], + [0.4355, -0.1016, -0.2002, -0.2412, -0.2930, 0.1187, 0.5664, 0.1104], + [0.1865, -0.0996, -0.6875, 0.2988, 0.0635, 0.1035, 0.5586, -0.0703], + ], + } + ), # fmt: skip +} + + +class TinyCheckpointTesterMixin: + """A format's pushed tiny checkpoints, random copies of the shipped ones written in their producers' + layouts (modelopt, OpenAI's and DeepSeek's quantizers): each loads with every key accounted for, computes + what its dequantized weights compute within the format's rounding, matches its recorded logits, and saves + back to its own layout. So a converter is checked against what the format really is, not against what this + code believes it to be. A format's test file sets `spec`.""" + + spec: FormatSpec + + def setUp(self): + super().setUp() + self.enterContext(_cheap_tuning()) + _require_kernel_capability(self) + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + @staticmethod + def _load(repo, **kwargs): + config = AutoConfig.from_pretrained(repo, **kwargs) + model_cls = AutoModelForImageTextToText if "vision_config" in config.to_dict() else AutoModelForCausalLM + return model_cls.from_pretrained(repo, device_map=torch_device, output_loading_info=True, **kwargs) + + @staticmethod + def _unbuilt_mtp_keys(model, keys): + """The multi-token-prediction layer the checkpoint ships and the model does not build: it + sits at index `num_hidden_layers`, which the models' own ignore rule only names at the + full-size index.""" + mtp = model.config.get_text_config().num_hidden_layers + return {key for key in keys if re.search(rf"(^|\.)layers\.{mtp}\.", key)} + + def _not_loaded(self, model, keys): + """Checkpoint keys the model does not load, so a save does not write them back: the unbuilt MTP + layer, and what the load declares it ignores (an FP8 KV cache's scales).""" + ignored = model._keys_to_ignore_on_load_unexpected or [] + return self._unbuilt_mtp_keys(model, keys) | {k for k in keys if any(re.search(p, k) for p in ignored)} + + def _reference_error(self, model, repo): + """The logits' error against the dequantized reference, and the logits.""" + reference, _ = self._load(repo, subfolder="reference", dtype=torch.bfloat16) + ids = torch.arange(3, 35, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + got, want = model(ids).logits.float(), reference(ids).logits.float() + return ((got - want).norm() / want.norm()).item(), got + + @staticmethod + def _recorded_slice(repo): + """The slice recorded on this exact device, if any: another device's kernels round differently.""" + return dict(EXPECTED_SLICES[repo].unpacked()).get(get_device_properties()) + + @slow + @require_kernels + @require_torch_accelerator + def test_logits_match_the_dequantized_reference_and_the_recorded_slice(self): + for repo, budget in self.spec.tiny_checkpoints.items(): + with self.subTest(repo=repo): + model, info = self._load(repo) + unexpected = set(info["unexpected_keys"]) - self._unbuilt_mtp_keys(model, info["unexpected_keys"]) + self.assertEqual( + (set(info["missing_keys"]), unexpected, set(info["mismatched_keys"])), (set(), set(), set()) + ) + error, logits = self._reference_error(model, repo) + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + expected = self._recorded_slice(repo) + if expected is not None: + torch.testing.assert_close(logits[0, :3, :8].cpu(), torch.tensor(expected), rtol=1e-2, atol=5e-3) + + @slow + @require_kernels + @require_torch_accelerator + def test_dequantized_load_matches_the_reference(self): + """`dequantize=True` folds each weight against its scales (and its global, for a two-level format).""" + for repo, budget in self.spec.tiny_checkpoints.items(): + with self.subTest(repo=repo): + model, info = self._load(repo, quantization_config=self.spec.config_cls(dequantize=True)) + unexpected = set(info["unexpected_keys"]) - self._unbuilt_mtp_keys(model, info["unexpected_keys"]) + self.assertEqual((set(info["missing_keys"]), unexpected), (set(), set())) + # nothing stays quantized; a model keeps its own fp32 parameters (DeepSeek-V4's hyper-connections) + quantized = {torch.float8_e4m3fn, torch.float8_e8m0fnu, torch.int8, torch.uint8} + self.assertEqual([name for name, param in model.named_parameters() if param.dtype in quantized], []) + error, _ = self._reference_error(model, repo) + self.assertLessEqual(error, budget, f"{repo}: logits {error:.4f} away from the dequantized reference") + + @slow + @require_kernels + @require_torch_accelerator + def test_save_writes_the_checkpoint_layout_back(self): + """A pre-quantized model saved again lands on its checkpoint's keys, dtypes and bytes.""" + from huggingface_hub import snapshot_download + + for repo in self.spec.tiny_checkpoints: + with self.subTest(repo=repo): + model, _ = self._load(repo) + saved = os.path.join(self.get_auto_remove_tmp_dir(), "saved") + model.save_pretrained(saved) + source = _checkpoint_bytes(snapshot_download(repo, allow_patterns=["*.safetensors", "*.json"])) + source = {key: value for key, value in source.items() if key not in self._not_loaded(model, source)} + written = _checkpoint_bytes(saved) + self.assertEqual(sorted(written), sorted(source)) + for key, (dtype, shape, raw) in written.items(): + with self.subTest(repo=repo, key=key): + source_dtype, source_shape, source_raw = source[key] + self.assertEqual(shape, source_shape) + if dtype != source_dtype: + # Loading honors initialized dtypes, including FP32 norms and BF16 + # MiniMax routers/patch embeddings. FP8 scales normalize to FP32; + # float32 UE8M0 containers save natively. + self.assertTrue( + ( + dtype == torch.float32 + and key.endswith("norm.weight") + and source_dtype == torch.bfloat16 + ) + or ( + dtype == torch.bfloat16 + and source_dtype == torch.float32 + and key.endswith( + ("gate.weight", "e_score_correction_bias", "patch_embedding.weight") + ) + ) + or ( + dtype == torch.float32 + and source_dtype == torch.bfloat16 + and key.endswith("_scale_inv") + and written.get(key.removesuffix("_scale_inv"), (None,))[0] == torch.float8_e4m3fn + ) + or ( + dtype == torch.float8_e8m0fnu + and source_dtype == torch.float32 + and key.endswith((".scale", "_scale_inv")) + ), + f"{repo}: {key} unexpectedly changed dtype from {source_dtype} to {dtype}", + ) + expected = torch.frombuffer(bytearray(source_raw), dtype=source_dtype).to(dtype) + self.assertEqual(raw, expected.view(torch.uint8).numpy().tobytes()) + else: + self.assertEqual(raw, source_raw) diff --git a/tests/quantization/block_scaled_parallel_tester.py b/tests/quantization/block_scaled_parallel_tester.py new file mode 100644 index 000000000000..08c46f098fec --- /dev/null +++ b/tests/quantization/block_scaled_parallel_tester.py @@ -0,0 +1,513 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Distributed loading, sharding, and compiled static-cache decode checks.""" + +import os +import re +import tempfile +from unittest import mock + +import torch + +from transformers import CompileConfig +from transformers.testing_utils import ( + cleanup, + require_kernels, + require_torch_multi_accelerator, + set_model_for_less_flaky_test, + slow, + torch_device, +) +from transformers.utils.quantization_config import BlockFp8Config, Mxfp4Config, Mxfp8Config + +from ..test_tensor_parallel_mixin import _init_distributed +from .block_scaled_checkpoint_tester import _checkpoint_bytes +from .block_scaled_tester import NVFP4, FormatSpec, _cheap_tuning + + +def _check_static_cache_decode(model, ids): + from torch._dynamo.backends.registry import lookup_backend + + generation_kwargs = { + "input_ids": ids.to(model.device), + "attention_mask": torch.ones_like(ids, device=model.device), + "max_new_tokens": 3, + "do_sample": False, + "eos_token_id": None, + "pad_token_id": 0, + "cache_implementation": "static", + "output_logits": True, + "return_dict_in_generate": True, + } + # Match the generation compile tests' norm stabilization for tiny random models. + set_model_for_less_flaky_test(model) + eager = model.generate(**generation_kwargs, disable_compile=True) + graphs = [] + + def compile_backend(graph, inputs, **kwargs): + graphs.append(graph) + return lookup_backend("inductor")(graph, inputs, **kwargs) + + compiled = model.generate( + **generation_kwargs, + compile_config=CompileConfig(backend=compile_backend, fullgraph=True, dynamic=False, mode="default"), + ) + assert graphs, "Static-cache generation did not compile the decode forward" + torch.testing.assert_close(compiled.sequences, eager.sequences) + for actual, expected in zip(compiled.logits, eager.logits): + torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3) + + +def _check_sharded_load_paths(rank, model_dir, model_cls, modes, out_dir, compile_decode): + from transformers.distributed import DistributedConfig + + world = torch.distributed.get_world_size() + ids = torch.arange(16, dtype=torch.long).unsqueeze(0) + for mode in modes: + model = model_cls.from_pretrained( + model_dir, + dtype="auto", + attn_implementation="eager", + distributed_config=DistributedConfig(tp_size=world, ep_size=world if mode == "ep" else None), + ).eval() + experts = next(m for n, m in model.named_modules() if n.endswith("mlp.experts")) + weight = getattr(experts, experts.gate_up_name) + local = weight.to_local() if hasattr(weight, "to_local") else weight + payload = {"expert_local": tuple(local.shape)} + with torch.no_grad(): + if compile_decode: + _check_static_cache_decode(model, ids) + else: + payload["logits"] = model(ids.to(model.device)).logits.float().cpu() + if rank == 0: + torch.save(payload, os.path.join(out_dir, mode + ".pt")) + del model + torch.accelerator.empty_cache() + + +def _checkpoint_expert_shape(model_dir): + """The unsharded `(experts, rows)` of the stacked gate|up projection the module holds, read from the + checkpoint, which ships the experts either fused (`experts.gate_up_proj`) or one per expert + (`experts.0.gate_proj.weight`, or DeepSeek's `w1`/`w3`).""" + import glob + import json + + from safetensors import safe_open + + shards = sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))) + index = os.path.join(model_dir, "model.safetensors.index.json") + if os.path.exists(index): + with open(index, encoding="utf-8") as fh: + weight_map = json.load(fh)["weight_map"] + shards = sorted({os.path.join(model_dir, f) for f in weight_map.values()}) + + experts: set[int] = set() + rows = 0 + for shard in shards: + with safe_open(shard, framework="pt") as handle: + for key in handle.keys(): + shape = None + if re.search(r"\.experts\.gate_up_proj$", key): + shape = handle.get_slice(key).get_shape() + if len(shape) == 3: + return (shape[0], shape[1]) + if re.search(r"\.experts\.(\d+)\.(gate_proj|up_proj|w1|w3)\.weight$", key): + experts.add(int(key.rsplit(".experts.", 1)[1].split(".", 1)[0])) + rows = max(rows, handle.get_slice(key).get_shape()[0]) + if not experts or not rows: + raise AssertionError(f"no expert projection in {model_dir} to size the shard against") + return (len(experts), 2 * rows) # the module stacks gate and up into one row extent + + +class LoadPathEquivalenceTesterMixin: + """Expert parallelism and intra-expert tensor parallelism must both reproduce what + `device_map` gives, which is the baseline because it places modules across devices without + splitting a single tensor — no process group, no DTensor, no collectives. So it computes the + unsharded answer, on the multi-GPU path people actually deploy. + + Each model is built in the format it actually ships, because the conversion path differs by + format and by checkpoint layout: a per-expert FP8 checkpoint and a fused NVFP4 one reach the + experts through different converters. + + Three properties make this able to fail: + * the fixture is PRE-QUANTIZED. Quantizing on the fly derives each rank's scales from the + shard it already holds, so they come out correctly sized whatever the plan says and a + plan that shards no scale at all still produces the right answer. + * the model's plan must carry expert entries. A model whose `base_model_tp_plan` is empty + (GPT-OSS, DeepSeek-V4) shards nothing under TP, so that leg is not run. + * every leg reports the local expert shard, so a leg that placed nothing fails loudly. + + A format's test file sets `spec`; the fixtures are its `distributed_models`.""" + + spec: FormatSpec + + @staticmethod + def _model_table(): + """`{label: (config_cls, model_cls, config_kwargs, quantization_config)}`: each model in the format and + under the quantization config it ships with. Dims are multiples of the 128 block so a 2-way split stays + block-aligned, and the expert count divides the mesh.""" + from transformers import ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + DeepseekV4Config, + DeepseekV4ForCausalLM, + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + Glm4vMoeTextConfig, + Glm4vMoeVisionConfig, + GptOssConfig, + GptOssForCausalLM, + MiniMaxM3SparseForConditionalGeneration, + MiniMaxM3VLConfig, + Mistral4Config, + Mistral4ForCausalLM, + ) + + return { + "deepseek_v3-fp8": ( + DeepseekV3Config, + DeepseekV3ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "n_group": 1, + "topk_group": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "q_lora_rank": None, + "kv_lora_rank": 32, + "qk_nope_head_dim": 32, + "qk_rope_head_dim": 16, + "v_head_dim": 32, + }, + BlockFp8Config(weight_block_size=(128, 128)), + ), + # the only shipped STATIC scheme: per-TENSOR weights (no `weight_block_size`) and a + # calibrated activation scale per quantized module, which for a MoE is one per expert. + # Its conversion script asserts `qscheme_act == "TENSOR"`; Ministral-3 is the dense + # counterpart of the same export. + "mistral4-fp8_tensor_static": ( + Mistral4Config, + Mistral4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + }, + BlockFp8Config(weight_block_size=None, activation_scheme="static"), + ), + # multimodal: the experts' plans live on `text_config`, not on the config the quantizer is handed + "glm4v_moe-nvfp4": ( + Glm4vMoeConfig, + Glm4vMoeForConditionalGeneration, + { + "text_config": Glm4vMoeTextConfig( + vocab_size=64, + hidden_size=256, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=32, + rope_parameters={"type": "default", "mrope_section": [16, 8, 8], "partial_rotary_factor": 1.0}, + rope_theta=10000, + tie_word_embeddings=True, + bos_token_id=0, + eos_token_id=0, + pad_token_id=0, + n_routed_experts=4, + n_shared_experts=1, + n_group=1, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + ), + "vision_config": Glm4vMoeVisionConfig( + depth=2, + num_heads=4, + hidden_size=64, + out_hidden_size=256, + intermediate_size=64, + patch_size=14, + spatial_merge_size=1, + temporal_patch_size=2, + ), + }, + NVFP4.checkpoint_config(), + ), + # interleaved rows (`is_concatenated=False`), transposed, with expert biases its EP plan already names + "gpt_oss-mxfp4": ( + GptOssConfig, + GptOssForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "max_position_embeddings": 32, + }, + Mxfp4Config(activation_format="bf16"), + ), + # mixed: `expert_dtype` makes the experts MXFP4 over block-FP8 with UE8M0 scales + "deepseek_v4-fp4_experts+fp8_dense": ( + DeepseekV4Config, + DeepseekV4ForCausalLM, + { + "vocab_size": 64, + "hidden_size": 256, + "intermediate_size": 256, + "moe_intermediate_size": 256, + "num_hidden_layers": 1, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "n_routed_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "first_k_dense_replace": 0, + "max_position_embeddings": 32, + "expert_dtype": "fp4", + }, + BlockFp8Config(weight_block_size=(128, 128), scale_fmt="ue8m0"), + ), + # a second multimodal nesting; the MoE dims are raised so a 2-way split stays 128-aligned + "minimax_m3_vl-mxfp8": ( + MiniMaxM3VLConfig, + MiniMaxM3SparseForConditionalGeneration, + { + "text_config": { + "hidden_size": 256, + # the 2-way split leaves 256: a swizzled tile is 128 wide along the down's K + "intermediate_size": 512, + "dense_intermediate_size": 256, + "shared_intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 64, + "rotary_dim": 32, + "vocab_size": 64, + "max_position_embeddings": 32, + "bos_token_id": 0, + "eos_token_id": 1, + "pad_token_id": 2, + "num_local_experts": 4, + "num_experts_per_tok": 2, + "n_shared_experts": 1, + "moe_layer_freq": [0, 1], + "layer_types": ["full_attention", "minimax_m3_sparse"], + "tie_word_embeddings": False, + "index_n_heads": 2, + "index_head_dim": 16, + "index_block_size": 8, + "index_topk_blocks": 4, + "index_local_blocks": 1, + }, + "vision_config": { + # the 2-way split leaves 128, one swizzled K tile + "hidden_size": 256, + "intermediate_size": 256, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_channels": 3, + "image_size": 14, + "patch_size": 14, + "temporal_patch_size": 2, + "spatial_merge_size": 1, + }, + "image_token_index": 4, + "video_token_index": 5, + "projector_hidden_size": 256, + "pad_token_id": 2, + }, + Mxfp8Config(), + ), + } + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls._tmp = cls.enterClassContext(tempfile.TemporaryDirectory()) + cls._fixtures = {} + cls.enterClassContext(_cheap_tuning()) + cls.enterClassContext( + mock.patch.dict( + os.environ, {"TRITON_CACHE_DIR": os.environ.get("TRITON_CACHE_DIR", os.path.join(cls._tmp, "triton"))} + ) + ) + + def tearDown(self): + super().tearDown() + cleanup(torch_device, gc_collect=True) + + def _fixture(self, label): + """A pre-quantized checkpoint of one tiny model, in the format that model ships.""" + if label in self._fixtures: + return self._fixtures[label] + config_cls, model_cls, kwargs, quantization_config = self._model_table()[label] + config = config_cls(**kwargs) + + bf16_dir = os.path.join(self._tmp, label, "bf16") + quant_dir = os.path.join(self._tmp, label, "quantized") + torch.manual_seed(0) + # saved in bf16, the dtype every model here ships, so `dtype="auto"` reloads bf16 + model_cls(config).to(torch.bfloat16).save_pretrained(bf16_dir, safe_serialization=True) + quantized = model_cls.from_pretrained( + bf16_dir, + dtype="auto", + attn_implementation="eager", + quantization_config=quantization_config, + device_map=f"{torch_device}:0", + ) + quantized.save_pretrained(quant_dir, safe_serialization=True) + + # a reload saved again lands on the same bytes: every layout op's reverse restores the checkpoint + reloaded = model_cls.from_pretrained(quant_dir, dtype="auto", device_map=f"{torch_device}:0") + round_trip = os.path.join(self._tmp, label, "round_trip") + reloaded.save_pretrained(round_trip, safe_serialization=True) + saved, resaved = _checkpoint_bytes(quant_dir), _checkpoint_bytes(round_trip) + mismatched = sorted(key for key in saved.keys() | resaved.keys() if saved.get(key) != resaved.get(key)) + self.assertEqual(mismatched, [], f"{label}: save did not round-trip") + self._fixtures[label] = quant_dir, config, model_cls + return self._fixtures[label] + + def _logits(self, model_dir, model_cls, **load_kwargs): + model = model_cls.from_pretrained(model_dir, dtype="auto", attn_implementation="eager", **load_kwargs).eval() + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + logits = model(ids).logits.float().cpu() + return logits + + def _rounding_floor(self, model_dir, model_cls, reference): + """How far this model's logits move when every sharded block is perturbed by BF16 rounding. + + Sharding cannot be bit-exact: a rowwise all-reduce sums the same terms in a different order. How far + that travels is a property of the model (a chain of MoE layers can amplify it a hundredfold), so the + same magnitude injected on one device gives each model the budget its own conditioning earns. + """ + model = model_cls.from_pretrained( + model_dir, dtype="auto", attn_implementation="eager", device_map=f"{torch_device}:0" + ).eval() + # every block a sharded reduction passes through: the rounding accumulates down the stack + blocks = [ + m + for n, m in model.named_modules() + if n.endswith((".self_attn", ".mlp", ".block_sparse_moe")) and n.count(".layers.") == 1 + ] + generator = torch.Generator(device=model.device).manual_seed(0) + + def perturb(module, args, output): + tensor = output[0] if isinstance(output, tuple) else output + if not torch.is_tensor(tensor): + return output + ulp = torch.finfo(tensor.dtype).eps * tensor.abs().max() + noise = torch.randn(tensor.shape, generator=generator, device=tensor.device, dtype=tensor.dtype) * ulp + tensor = tensor + noise + return (tensor,) + output[1:] if isinstance(output, tuple) else tensor + + handles = [b.register_forward_hook(perturb) for b in blocks] + ids = torch.arange(16, dtype=torch.long, device=model.device).unsqueeze(0) + with torch.no_grad(): + perturbed = model(ids).logits.float().cpu() + for handle in handles: + handle.remove() + return (perturbed - reference).abs().max().item() + + def _sharded(self, model_dir, model_cls, modes, compile_decode=False): + """`{mode: payload}` from two ranks, through the real EP / TP load path.""" + out = os.path.join(self._tmp, "out") + os.makedirs(out, exist_ok=True) + _init_distributed(tp=2)(_check_sharded_load_paths)(model_dir, model_cls, modes, out, compile_decode) + return {m: torch.load(os.path.join(out, m + ".pt")) for m in modes} + + @slow + @require_kernels + @require_torch_multi_accelerator + def test_every_load_path_agrees(self): + for label in self.spec.distributed_models: + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + # `max_memory` forces the baseline across both devices; `auto` alone fits it on one + reference = self._logits(model_dir, model_cls, device_map="auto", max_memory={0: "120MiB", 1: "40GiB"}) + noise_floor = self._rounding_floor(model_dir, model_cls, reference) + + modes = self._parallel_modes(config) + self.assertTrue(modes, f"{label}: no plan shards experts") + whole = _checkpoint_expert_shape(model_dir) + for mode, payload in self._sharded(model_dir, model_cls, modes).items(): + with self.subTest(model=label, mode=mode): + expected = (whole[0] // 2, whole[1]) if mode == "ep" else (whole[0], whole[1] // 2) + self.assertEqual(payload["expert_local"][:2], expected, f"{label}/{mode}: wrong shard axis") + # a wrong shard axis lands orders of magnitude above the model's rounding floor + budget = max(2e-2, noise_floor) + gap = (payload["logits"] - reference).abs().max().item() + self.assertLessEqual( + gap, + budget, + f"{label}/{mode}: sharded logits differ by {gap:.4g}, beyond this " + f"model's own rounding budget {budget:.4g} — that is a sharding bug, " + f"not reduction order", + ) + if noise_floor <= 2e-2: # rounding alone does not move the argmax + self.assertTrue( + torch.equal(payload["logits"].argmax(-1), reference.argmax(-1)), + f"{label}/{mode}: argmax diverged from the unsharded model", + ) + + @staticmethod + def _parallel_modes(config): + owners = [config] + [ + c for name in getattr(type(config), "sub_configs", {}) if (c := getattr(config, name, None)) + ] + return [ + mode + for mode, attr in (("ep", "base_model_ep_plan"), ("tp", "base_model_tp_plan")) + if any(".experts." in key for owner in owners for key in (getattr(owner, attr, None) or {})) + ] + + @slow + @require_kernels + @require_torch_multi_accelerator + def test_static_cache_decode_compiles(self): + supported = [ + label for label in self.spec.distributed_models if self._model_table()[label][1]._can_compile_fullgraph + ] + if not supported: + self.skipTest("These model architectures do not support fullgraph compilation") + for label in supported: + with self.subTest(model=label): + model_dir, config, model_cls = self._fixture(label) + modes = self._parallel_modes(config) + self.assertTrue(modes, f"{label}: no plan shards experts") + self._sharded(model_dir, model_cls, modes, compile_decode=True) diff --git a/tests/quantization/block_scaled_tester.py b/tests/quantization/block_scaled_tester.py new file mode 100644 index 000000000000..79c099b5140d --- /dev/null +++ b/tests/quantization/block_scaled_tester.py @@ -0,0 +1,365 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Shared format specifications and fast quantizer contract tests. + +Real-checkpoint tests live in `block_scaled_checkpoint_tester.py`; distributed +loading and compiled decode tests live in `block_scaled_parallel_tester.py`. +""" + +import os +from collections.abc import Callable +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, field +from unittest import mock + +import torch + +from transformers import AutoModelForCausalLM +from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key +from transformers.integrations.finegrained import FineGrainedExperts, FineGrainedLinear +from transformers.quantizers.auto import AutoHfQuantizer, AutoQuantizationConfig +from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer +from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer +from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer +from transformers.quantizers.quantizer_mxfp8 import Mxfp8HfQuantizer +from transformers.quantizers.quantizer_nvfp4 import NVFP4HfQuantizer +from transformers.testing_utils import ( + require_accelerate, + require_kernels, + require_torch_accelerator, + torch_device, +) +from transformers.utils.quantization_config import ( + BlockFp8Config, + FineGrainedConfig, + Mxfp4Config, + Mxfp8Config, + NVFP4Config, +) + +from ..integrations.finegrained.utils import _Cfg, _conversion_model, _model_mapping, _quantizer_for + + +def _cheap_tuning(): + """Measure one configuration without extra search phases for correctness tests.""" + return mock.patch.dict( + os.environ, + { + "FINEGRAINED_AUTOTUNE_TRIALS": "1", + "FINEGRAINED_AUTOTUNE_BASIN_PROBES": "0", + "FINEGRAINED_AUTOTUNE_MAX_REFINE_ITERS": "0", + "FINEGRAINED_AUTOTUNE_RUNOFF": "0", + }, + ) + + +@dataclass(frozen=True) +class FormatSpec: + """What the shared tests need to know about one format; no test hard-codes it. + + `quantize_kwargs` configure the on-the-fly quantization whose layout is checked (`None`: the format's + quantize op is tested on its own). `checkpoint_keys` are keys of the format's shipped layouts, each of which + at most one converter may claim; `dequantize_keys` are a quantized linear's keys, each of which a + dequantizing load must route to a converter. `tiny_checkpoints` maps each pushed tiny + checkpoint to its logits budget against its `reference/` bf16 copy; `distributed_models` names the + `LoadPathEquivalenceTesterMixin` fixtures built in this format.""" + + name: str + config_cls: type + quantizer_cls: type + checkpoint_config: Callable + activation_formats: tuple = () + weight_bits: int = 16 + two_level: bool = False + quantize_kwargs: dict | None = None + checkpoint_keys: tuple = () + dequantize_keys: tuple = () + tiny_checkpoints: dict = field(default_factory=dict) + distributed_models: tuple = () + + +BLOCK_FP8 = FormatSpec( + name="fp8", + config_cls=BlockFp8Config, + quantizer_cls=BlockFp8HfQuantizer, + activation_formats=(None, "fp8", "bf16"), + checkpoint_config=BlockFp8Config, + weight_bits=8, + checkpoint_keys=( + "model.layers.3.mlp.experts.7.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.gate_up_proj_scale_inv", + "model.layers.3.self_attn.q_proj.weight_scale_inv", + ), + dequantize_keys=("model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.weight_scale_inv"), + tiny_checkpoints={ + "IlyasMoutawwakil/tiny-random-Qwen3-FP8": 0.1, + "IlyasMoutawwakil/tiny-random-Mistral3-FP8-static": 0.1, + "IlyasMoutawwakil/tiny-random-Mistral4-FP8-static": 0.15, + "IlyasMoutawwakil/tiny-random-DeepseekV3-FP8": 0.25, + }, + distributed_models=("deepseek_v3-fp8", "mistral4-fp8_tensor_static"), +) +MXFP8 = FormatSpec( + name="mxfp8", + config_cls=Mxfp8Config, + quantizer_cls=Mxfp8HfQuantizer, + activation_formats=(None, "mxfp8", "mxfp4", "bf16"), + checkpoint_config=Mxfp8Config, + weight_bits=8, + quantize_kwargs={}, + checkpoint_keys=( + "model.layers.3.mlp.experts.7.gate_proj.weight_scale_inv", + "model.layers.3.mlp.experts.down_proj_scale_inv", + ), + dequantize_keys=("model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.down_proj.weight_scale_inv"), + tiny_checkpoints={"IlyasMoutawwakil/tiny-random-MiniMaxM3-MXFP8": 0.3}, + distributed_models=("minimax_m3_vl-mxfp8",), +) +MXFP4 = FormatSpec( + name="mxfp4", + config_cls=Mxfp4Config, + quantizer_cls=Mxfp4HfQuantizer, + activation_formats=(None, "mxfp8", "mxfp4", "bf16"), + checkpoint_config=Mxfp4Config, + weight_bits=4, + quantize_kwargs={"activation_format": None}, + checkpoint_keys=( + "model.layers.3.mlp.experts.gate_up_proj_blocks", + "model.layers.3.mlp.experts.gate_up_proj_scales", + "model.layers.3.mlp.experts.7.gate_proj.weight_scale_inv", + ), + # GPT-OSS's blocks + scales pair + dequantize_keys=( + "model.layers.0.mlp.experts.gate_up_proj_blocks", + "model.layers.0.mlp.experts.gate_up_proj_scales", + ), + tiny_checkpoints={"IlyasMoutawwakil/tiny-random-GptOss-MXFP4": 0.05}, + distributed_models=("gpt_oss-mxfp4",), +) +NVFP4 = FormatSpec( + name="nvfp4", + config_cls=NVFP4Config, + quantizer_cls=NVFP4HfQuantizer, + activation_formats=(None, "nvfp4", "bf16"), + # what a modelopt export reads as + checkpoint_config=lambda: AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "NVFP4"}), + weight_bits=4, + two_level=True, + quantize_kwargs={}, + checkpoint_keys=( + "model.layers.3.mlp.experts.gate_up_proj", + "model.layers.3.mlp.experts.down_proj", + "model.layers.3.mlp.experts.gate_up_proj_weight_scale_2", + "model.layers.3.mlp.experts.7.up_proj.weight_scale_2", + ), + dequantize_keys=( + "model.layers.0.mlp.down_proj.weight", + "model.layers.0.mlp.down_proj.weight_scale", + "model.layers.0.mlp.down_proj.weight_scale_2", + ), + tiny_checkpoints={ + "IlyasMoutawwakil/tiny-random-Llama-NVFP4": 0.25, + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4": 0.25, + "IlyasMoutawwakil/tiny-random-GlmMoeDsa-NVFP4-all-linears": 0.6, + "IlyasMoutawwakil/tiny-random-Muse-Spark-NVFP4": 0.3, + }, + distributed_models=("glm4v_moe-nvfp4",), +) +FORMATS = (BLOCK_FP8, MXFP8, MXFP4, NVFP4) + +# Plans mixing the formats, as DeepSeek-V4 ships one (block-FP8 over MXFP4 experts). +FINEGRAINED = FormatSpec( + name="finegrained", + config_cls=FineGrainedConfig, + quantizer_cls=FineGrainedHfQuantizer, + checkpoint_config=FineGrainedConfig, + tiny_checkpoints={"IlyasMoutawwakil/tiny-random-DeepseekV4-Flash": 0.25}, + distributed_models=("deepseek_v4-fp4_experts+fp8_dense",), +) + + +@contextmanager +def no_accelerator(): + with ExitStack() as stack: + stack.enter_context(mock.patch("torch.cuda.is_available", return_value=False)) + if hasattr(torch, "xpu"): + stack.enter_context(mock.patch("torch.xpu.is_available", return_value=False)) + stack.enter_context( + mock.patch("transformers.quantizers.block_scaled_utils.is_torch_xpu_available", return_value=False) + ) + yield + + +def _require_kernel_capability(test): + # below it the quantizer dequantizes instead, which is not what is under test + if torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9): + test.skipTest("fine-grained quantized compute needs compute capability >= 8.9") + + +class BlockScaledFormatTesterMixin: + """The contract every format keeps. A format's test file sets `spec` and inherits these.""" + + spec: FormatSpec + + def test_config_round_trips_through_a_dict(self): + config = self.spec.config_cls(modules_to_not_convert=["lm_head"]) + read = self.spec.config_cls.from_dict(config.to_dict()) + self.assertEqual(read.to_dict(), config.to_dict()) + self.assertEqual(read.quant_method.value, self.spec.name) + + def test_from_pretrained_picks_the_format_quantizer(self): + quantizer = AutoHfQuantizer.from_config(self.spec.config_cls()) + self.assertIsInstance(quantizer, self.spec.quantizer_cls) + + def test_the_config_takes_only_the_activation_formats_its_kernels_serve(self): + """The kernels quantize activations only to what the weight format's MMA takes, and refuse the rest at + the first forward; the config refuses them up front.""" + for activation_format in {None} | {f for spec in FORMATS for f in spec.activation_formats}: + with self.subTest(activation_format=activation_format): + if activation_format in self.spec.activation_formats: + self.spec.config_cls(activation_format=activation_format) + else: + with self.assertRaises(ValueError): + self.spec.config_cls(activation_format=activation_format) + + @require_accelerate + def test_no_accelerator_refuses_to_quantize(self): + # quantizing a bf16 checkpoint needs the hardware; there is nothing to fall back to + quantizer = _quantizer_for(self.spec.config_cls()) + quantizer.pre_quantized = False + with no_accelerator(), self.assertRaises(RuntimeError): + quantizer.validate_environment() + + @require_accelerate + def test_no_accelerator_dequantizes_a_quantized_checkpoint(self): + quantizer = _quantizer_for(self.spec.checkpoint_config()) + quantizer.pre_quantized = True + with no_accelerator(): + quantizer.validate_environment() + self.assertTrue(quantizer.quantization_config.dequantize) + + @require_accelerate + def test_unsupported_gpu_dequantizes_a_quantized_checkpoint(self): + quantizer = _quantizer_for(self.spec.checkpoint_config()) + quantizer.pre_quantized = True + with ( + mock.patch("torch.cuda.is_available", return_value=True), + mock.patch("torch.cuda.get_device_capability", return_value=(6, 0)), + ): + quantizer.validate_environment() + self.assertTrue(quantizer.quantization_config.dequantize) + + def test_a_checkpoint_key_is_claimed_by_at_most_one_converter(self): + """Two converters claiming one key is not additive: the later one wins and drops the ops the first carried. + Keys are routed as the loader routes them, renamed first.""" + quantizer = _quantizer_for(self.spec.checkpoint_config()) + quantizer.pre_quantized = True + conversions = quantizer.update_weight_conversions(_model_mapping(), _conversion_model(quantizer)) + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + for key in self.spec.checkpoint_keys: + renamed, _ = rename_source_key(key, renamings, converters) + claimed = [c for c in converters if c.compiled_sources.search(renamed)] + with self.subTest(key=key): + self.assertLessEqual( + len(claimed), + 1, + f"{len(claimed)} converters claim this key; the last one wins and drops " + f"the others' ops: {[[type(o).__name__ for o in c.operations] for c in claimed]}", + ) + + def test_dequantize_claims_every_key_of_a_quantized_linear(self): + """A dequantizing load swaps in no fine-grained module, so no scale slot exists: every key of the pair + (and the global, for a two-level format) must reach a converter.""" + quantizer = _quantizer_for(self.spec.config_cls(dequantize=True)) + quantizer.pre_quantized = True + conversions = quantizer.get_weight_conversions() + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + for key in self.spec.dequantize_keys: + renamed, _ = rename_source_key(key, renamings, converters) + with self.subTest(key=key): + self.assertTrue( + any(c.compiled_sources.search(renamed) for c in converters), f"nothing dequantizes {key}" + ) + + def test_module_sizes_count_the_quantized_weights(self): + """The device map is sized before any weight loads, from the replaced modules' dtypes and shapes.""" + from transformers import Qwen3MoeConfig + from transformers.integrations.accelerate import compute_module_sizes + from transformers.modeling_utils import expand_device_map, get_total_byte_count + + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=256, + intermediate_size=256, + moe_intermediate_size=256, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + num_experts=4, + num_experts_per_tok=2, + ) + with torch.device("meta"): + model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16) + sizes, _ = compute_module_sizes(model, only_modules=False) + quantizer = AutoHfQuantizer.from_config(self.spec.config_cls(), pre_quantized=False) + quantizer.preprocess_model(model=model, config=model.config) + quantized_sizes, _ = compute_module_sizes(model, quantizer, only_modules=False) + keys = [name for name, _ in model.named_parameters()] + [name for name, _ in model.named_buffers()] + (total,) = get_total_byte_count(model, expand_device_map({"": torch_device}, keys), quantizer).values() + + weights = [ + f"{name}.{param}" + for name, module in model.named_modules() + if isinstance(module, (FineGrainedLinear, FineGrainedExperts)) + for param in (("weight",) if isinstance(module, FineGrainedLinear) else ("gate_up_proj", "down_proj")) + ] + self.assertTrue(weights) + for weight in weights: + self.assertEqual(quantized_sizes[weight], sizes[weight] * self.spec.weight_bits // 16, weight) + self.assertEqual(total, quantized_sizes[""]) + + @require_kernels + @require_torch_accelerator + def test_on_the_fly_quantization_lands_in_the_module_layout(self): + """The kernels' quantizers then the module's layout ops fill the held layout (swizzled where it is held + so), with the second-level global only for a two-level format.""" + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + if self.spec.quantize_kwargs is None: + self.skipTest(f"{self.spec.name}'s quantize op is tested on its own") + _require_kernel_capability(self) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 512, 256, 4 + cfg._experts_implementation = "grouped_mm" + experts = FineGrainedExperts(cfg, quantization_config=self.spec.config_cls(**self.spec.quantize_kwargs)) + experts = experts.to(torch_device) + model = torch.nn.Module() + model.experts = experts + original = torch.randn(4, 512, 512, device=torch_device, dtype=torch.bfloat16) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + for name in ("gate_up_proj", "gate_up_proj_scale_inv"): + for op in _quantizer_for(self.spec.config_cls()).get_param_ops(model, f"experts.{name}"): + key = f"experts.{name}" + out.update(op.convert({key: out[key]}, model=model, full_layer_name=key)) + for name in ("gate_up_proj", "gate_up_proj_scale_inv"): + held = getattr(experts, name) + self.assertEqual((out[f"experts.{name}"].dtype, out[f"experts.{name}"].shape), (held.dtype, held.shape)) + global_key = "experts.gate_up_proj_weight_global_scale" + if self.spec.two_level: + self.assertEqual(out[global_key].shape, experts.gate_up_proj_weight_global_scale.shape) + else: + self.assertNotIn(global_key, out) diff --git a/tests/quantization/blockfp8/__init__.py b/tests/quantization/blockfp8/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/blockfp8/test_blockfp8.py b/tests/quantization/blockfp8/test_blockfp8.py new file mode 100644 index 000000000000..1254003114c0 --- /dev/null +++ b/tests/quantization/blockfp8/test_blockfp8.py @@ -0,0 +1,557 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Block-FP8 (`BlockFp8Config`, `quant_method="fp8"`): the config, the quantize op and the checkpoint +layouts it loads and saves back, and model-level loading.""" + +import tempfile +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch +from parameterized import parameterized + +import transformers.integrations.finegrained.layers as fg +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.integrations.finegrained import FineGrainedExperts +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import ( + TestCasePlus, + get_device_properties, + require_accelerate, + require_torch, + require_torch_accelerator, + require_torch_multi_accelerator, + slow, + torch_device, +) +from transformers.utils.quantization_config import ( + BlockFp8Config, +) + +from ...integrations.finegrained.utils import ( + E, + H, + I, + _Cfg, + _ChainAssertions, + _compressed_tensors, + _conversion_model, + _load_and_save, + _per_expert, + _quantizer_for, +) +from ...test_memory_cleanup_mixin import MemoryCleanupMixin +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin +from ..block_scaled_tester import ( + BLOCK_FP8, + BlockScaledFormatTesterMixin, + no_accelerator, +) + + +@require_torch +class BlockFp8ConfigTest(unittest.TestCase): + def test_a_passed_config_overrides_the_checkpoints_loading_attributes(self): + from transformers.quantizers.auto import AutoHfQuantizer + + passed = BlockFp8Config(dequantize=True, modules_to_not_convert=["lm_head", "model.layers.0"]) + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "fp8", "modules_to_not_convert": ["lm_head"]}, passed + ) + self.assertEqual((merged.dequantize, merged.modules_to_not_convert), (True, ["lm_head", "model.layers.0"])) + + def test_a_compressed_tensors_block_fp8_checkpoint_reads_as_block_fp8(self): + block = {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]} + fp8 = {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 128, "dynamic": True} + config = AutoQuantizationConfig.from_dict(_compressed_tensors(block, fp8, fmt="float-quantized")) + self.assertEqual((type(config), config.weight_block_size), (BlockFp8Config, (128, 128))) + + +@require_torch +class BlockFp8QuantizeOpTest(unittest.TestCase): + """`FineGrainedQuantize` reads the module's format: block-FP8 in torch (the group formats run + the kernels' own quantizers, which need a GPU), emitted in the module's layout.""" + + def test_a_static_scheme_gets_its_activation_scales_written(self): + """A calibration-fed slot the checkpoint does not supply is still WRITTEN, as the identity. + + The loader materializes a missing key with `torch.empty_like` and `_init_weights` has no + branch for a scale, so a slot this op leaves out reaches the kernels as uninitialized + memory — a zero divides the activations by zero, a negative flips their sign. + """ + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts( + cfg, quantization_config=BlockFp8Config(weight_block_size=(128, 128), activation_scheme="static") + ) + model.proj = fg.FineGrainedLinear( + quantization_config=BlockFp8Config(weight_block_size=(128, 128), activation_scheme="static"), + in_features=256, + out_features=256, + ) + op = FineGrainedQuantize(hf_quantizer=None) + for key, tensor, slot in ( + ("experts.gate_up_proj", torch.randn(4, 256, 256), "experts.gate_up_proj_activation_scale"), + ("proj.weight", torch.randn(256, 256), "proj.activation_scale"), + ): + with self.subTest(key=key): + out = op.convert({key: tensor}, model=model) + self.assertIn(slot, out, f"{key}: the static activation scale was not written") + held = model.get_parameter(slot) + self.assertEqual((out[slot].shape, out[slot].dtype), (held.shape, held.dtype)) + torch.testing.assert_close(out[slot], torch.ones_like(out[slot]), rtol=0, atol=0) + + def test_an_expert_bias_is_passed_through_not_quantized(self): + """A GPT-OSS expert bias is `(E, rows)` — 2-D, like a dense weight — so the rank guard + alone lets it through. It has no scale slot, and emitting one gives the loader a + `_bias_scale_inv` no module holds.""" + from transformers.integrations.finegrained.conversions import FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + model = torch.nn.Module() + model.experts = fg.FineGrainedExperts( + cfg, quantization_config=BlockFp8Config(weight_block_size=(128, 128)), has_bias=True + ) + bias = torch.randn(4, 256) + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.down_proj_bias": bias}, model=model) + self.assertEqual(list(out), ["experts.down_proj_bias"]) + self.assertIs(out["experts.down_proj_bias"], bias) + + def test_a_partial_trailing_block_round_trips_against_the_configured_block(self): + """DeepSeek-V3 ships `kv_a_proj_with_mqa` as `(576, 7168)` against a 128x128 block, so the + grid rounds UP to `(5, 56)` and the short block quantizes on its own values. + + The grid cannot be inverted to recover the block: 576 // 5 is 116, and `(192, 256)` divides + evenly at 96 while the real block is still 128. A dequantize that guesses then reads the + wrong rows per scale, which only lifts the error to 4.4e-2 against a 2.6e-2 floor — hence + the tight bound, since a loose one passes the broken arm too.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + quantizer = _quantizer_for(BlockFp8Config(weight_block_size=(128, 128))) + for rows, cols in ((576, 256), (192, 256), (256, 256)): + with self.subTest(shape=(rows, cols)): + original = torch.randn(rows, cols) + weight, scale = FineGrainedQuantize._quantize_block_fp8(original, (128, 128), ue8m0=False) + self.assertEqual(weight.shape, (rows, cols)) + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(scale.shape, (-(-rows // 128), -(-cols // 128))) + + recovered = FineGrainedDequantize(quantizer)._dequantize_one( + weight, scale.float(), output_dtype=torch.float32 + ) + rel = ((recovered - original).norm() / original.norm()).item() + self.assertLess(rel, 3e-2, f"{(rows, cols)}: {rel:.4f}") + + def test_block_fp8_round_trips_within_its_floor(self): + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedQuantize + + torch.manual_seed(0) + cfg = _Cfg() + cfg.hidden_size, cfg.intermediate_size, cfg.num_local_experts = 256, 128, 4 + for scale_fmt in ("float", "ue8m0"): + with self.subTest(scale_fmt=scale_fmt), mock.patch.object(fg, "is_sm100", return_value=False): + experts = fg.FineGrainedExperts( + cfg, quantization_config=BlockFp8Config(weight_block_size=(128, 128), scale_fmt=scale_fmt) + ) + model = torch.nn.Module() + model.experts = experts + original = torch.randn(4, 256, 256) * 0.02 + out = FineGrainedQuantize(hf_quantizer=None).convert({"experts.gate_up_proj": original}, model=model) + weight, scale = out["experts.gate_up_proj"], out["experts.gate_up_proj_scale_inv"] + self.assertEqual( + (weight.dtype, weight.shape), (experts.gate_up_proj.dtype, experts.gate_up_proj.shape) + ) + self.assertEqual( + (scale.dtype, scale.shape), + (experts.gate_up_proj_scale_inv.dtype, experts.gate_up_proj_scale_inv.shape), + ) + deq = FineGrainedDequantize(None)._dequantize_one(weight, scale.float(), output_dtype=torch.float32) + rel = ((deq - original).norm() / original.norm()).item() + self.assertLess(rel, 4e-2, f"{scale_fmt}: {rel:.3f}") + + def test_float_scale_fmt_stays_bit_identical_to_its_formula(self): + """The inverse scale is `amax / MAX` and the weight is DIVIDED by it. The algebraically + equal `weight * (MAX / amax)` is not bit-equal — it moves 29 of these 100 blocks' scales — + so the divergence only shows across many blocks, on particular fp32 maxima.""" + from transformers.integrations.finegrained.conversions import ( + _FP8_DTYPE, + _FP8_MAX, + _FP8_MIN, + FineGrainedQuantize, + ) + + quantizer = FineGrainedQuantize( + SimpleNamespace(quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="float")) + ) + for seed in range(100): + torch.manual_seed(seed) + weight = torch.randn(128, 128, dtype=torch.float32) + out = quantizer._quantize_one("layer.weight", weight, None) + ref_inv = (weight.abs().amax() / _FP8_MAX).to(torch.float32).reshape(1, 1) + ref_q = torch.clamp(weight / ref_inv, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) + self.assertTrue(torch.equal(out["layer.weight"], ref_q), f"float weight diverged at seed {seed}") + self.assertTrue( + torch.equal(out["layer.weight_scale_inv"], ref_inv), f"float scale diverged at seed {seed}" + ) + + +@require_torch +class BlockFp8ChainTest(_ChainAssertions, unittest.TestCase): + """Block-FP8 and per-tensor FP8, as DeepSeek-V3, DeepSeek-V4 and a static (Mistral) checkpoint ship them.""" + + def test_deepseek_v3_block_fp8(self): + """Per-expert e4m3 weights with fp32 `weight_scale_inv` per block, merged onto the stacked experts.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "weight_scale_inv": torch.rand(-(-n // 16), -(-k // 16)) + 0.5, + } + ) + model, info, saved = _load_and_save(checkpoint, BlockFp8Config(weight_block_size=(16, 16))) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + self.assertHoldsInterleaved(experts, checkpoint, "weight", experts.gate_up_proj) + self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) + self.assertSavesBack(saved, checkpoint) + + def test_ue8m0_block_scales(self): + """Block scales shipped as E8M0 load as `float8_e8m0fnu` and save back unchanged. (DeepSeek-V4's ship under + `.scale`, which its own conversion mapping renames.)""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "weight_scale_inv": torch.randint(120, 130, (-(-n // 16), -(-k // 16)), dtype=torch.uint8).view( + torch.float8_e8m0fnu + ), + } + ) + config = BlockFp8Config(weight_block_size=(16, 16), scale_fmt="ue8m0") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + self.assertEqual(model.model.layers[0].self_attn.q_proj.weight_scale_inv.dtype, torch.float8_e8m0fnu) + self.assertSavesBack(saved, checkpoint) + + def test_float32_ue8m0_containers_save_as_native_ue8m0(self): + """DeepSeek-V4-Base's power-of-two float32 scale containers normalize to native UE8M0.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "weight_scale_inv": torch.pow(2.0, torch.randint(-8, 8, (-(-n // 16), -(-k // 16))).float()), + } + ) + config = BlockFp8Config(weight_block_size=(16, 16), scale_fmt="ue8m0") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + self.assertEqual(set(saved), set(checkpoint)) + for name, tensor in checkpoint.items(): + with self.subTest(name=name): + self.assertEqual( + saved[name].dtype, torch.float8_e8m0fnu if name.endswith("scale_inv") else tensor.dtype + ) + torch.testing.assert_close(saved[name].float(), tensor.float(), rtol=0, atol=0) + _, info, resaved = _load_and_save(saved, config) + self.assertLoadsCleanly(info) + self.assertSavesBack(resaved, saved) + + def test_mistral4_static_per_tensor(self): + """Mistral-Small-4 ships its experts already stacked under the module's own names: e4m3, one + BF16 weight scale per expert `(E, 1, 1)` and a calibrated BF16 activation scale per expert; + its dense linears hold one 0-dim BF16 scale of each.""" + torch.manual_seed(0) + checkpoint = { + "model.layers.0.self_attn.q_proj.weight": torch.randn(H, H).to(torch.float8_e4m3fn), + "model.layers.0.self_attn.q_proj.weight_scale_inv": (torch.rand(()) + 0.5).to(torch.bfloat16), + "model.layers.0.self_attn.q_proj.activation_scale": (torch.rand(()) + 0.5).to(torch.bfloat16), + } + for proj, (rows, cols) in (("gate_up_proj", (2 * I, H)), ("down_proj", (H, I))): + prefix = f"model.layers.0.mlp.experts.{proj}" + checkpoint[prefix] = torch.randn(E, rows, cols).to(torch.float8_e4m3fn) + checkpoint[f"{prefix}_scale_inv"] = (torch.rand(E, 1, 1) + 0.5).to(torch.bfloat16) + checkpoint[f"{prefix}_activation_scale"] = (torch.rand(E) + 0.5).to(torch.bfloat16) + config = BlockFp8Config(weight_block_size=None, activation_scheme="static") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + torch.testing.assert_close( + experts.gate_up_proj_activation_scale.float(), + checkpoint["model.layers.0.mlp.experts.gate_up_proj_activation_scale"].float(), + ) + # The loader uses the initialized FP32 dtype for a dense block scale. + expected = dict(checkpoint) + key = "model.layers.0.self_attn.q_proj.weight_scale_inv" + expected[key] = expected[key].float() + self.assertSavesBack(saved, expected) + + def test_a_calibrated_checkpoint_s_input_scale_reaches_the_module_s_slot(self): + """A static checkpoint calibrates one `input_scale` per quantized module, and each key + must reach the slot its module holds, in that module's shape: one value on a dense linear, + one per expert on the stacked experts, and the gate|up pair reduced to one per expert + rather than flattened to 2E. Unrouted, the module keeps its registered default of 1.0 and + every activation quantizes against the wrong scale, silently.""" + + from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key + + quantizer = _quantizer_for(BlockFp8Config(activation_scheme="static")) + quantizer.pre_quantized = True + conversions = quantizer.update_weight_conversions([], _conversion_model(quantizer)) + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + + for key, expected in ( + ("model.layers.3.self_attn.q_proj.input_scale", "model.layers.3.self_attn.q_proj.activation_scale"), + ("model.layers.3.mlp.experts.7.gate_proj.input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.up_proj.input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.7.down_proj.input_scale", "experts.down_proj_activation_scale"), + ("model.layers.3.mlp.experts.gate_up_proj_input_scale", "experts.gate_up_proj_activation_scale"), + ("model.layers.3.mlp.experts.down_proj_input_scale", "experts.down_proj_activation_scale"), + ): + with self.subTest(key=key.rsplit(".", 2)[-2]): + # as the loader routes it: every renaming, then the one converter that matches + renamed, pattern = rename_source_key(key, renamings, converters) + targets = ( + [renamed] + if pattern is None + else [t for c in converters if pattern in c.source_patterns for t in c.target_patterns] + ) + self.assertTrue(any(expected in t for t in targets), f"{key} -> {targets}") + # a dynamic checkpoint has no such key to route, and a weight-only one ignores it + dynamic = _quantizer_for(BlockFp8Config()) + dynamic.pre_quantized = True + self.assertEqual(dynamic.get_weight_conversions(), []) + + +@slow +@require_accelerate +@require_torch_accelerator +@unittest.skipIf( + get_device_properties()[0] == "cuda" + and (get_device_properties()[1] < 8 or (get_device_properties()[1] == 8 and get_device_properties()[2] < 9)), + "Skipping BlockFp8QuantizerTest because it is not supported on GPU with capability < 8.9", +) +class BlockFp8QuantizerTest(MemoryCleanupMixin, unittest.TestCase): + model_name = "meta-llama/Llama-3.2-1B" + quantized_model_name = "hf-internal-testing/Llama-3.2-1B-Instruct-fp8" + input_text = "Once upon a time" + max_new_tokens = 10 + EXPECTED_OUTPUTS = { + "Once upon a time, there was a little girl who loved to play", + "Once upon a time, there was a man who was very rich.", + } + EXPECTED_DEQUANTIZED_OUTPUT = "Once upon a time, in a small village nestled in the rolling hills" + device_map = torch_device + offload_device_map = { + "model.embed_tokens": 0, + "model.layers.0": 0, + "model.layers.1": 0, + "model.layers.2": 0, + "model.layers.3": 0, + "model.layers.4": 0, + "model.layers.5": 0, + "model.layers.6": 0, + "model.layers.7": "cpu", + "model.layers.8": "cpu", + "model.layers.9": "cpu", + "model.layers.10": "cpu", + "model.layers.11": "cpu", + "model.layers.12": "cpu", + "model.layers.13": "cpu", + "model.layers.14": "cpu", + "model.layers.15": "cpu", + "model.rotary_emb": "cpu", + "model.norm": "cpu", + "lm_head": 0, + } + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.quantization_config = BlockFp8Config() + cls.tokenizer = AutoTokenizer.from_pretrained(cls.model_name) + cls.quantized_model = AutoModelForCausalLM.from_pretrained( + cls.model_name, device_map=cls.device_map, quantization_config=cls.quantization_config + ) + + @parameterized.expand( + [ + "hf-internal-testing/tiny-random-Qwen3MoeForCausalLM", + "hf-internal-testing/tiny-random-MixtralForCausalLM", + ] + ) + def test_moe_conversion(self, model_id): + quantization_config = BlockFp8Config(weight_block_size=(32, 32)) + model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quantization_config) + self.assertTrue(any(isinstance(module, FineGrainedExperts) for module in model.modules())) + + def test_quantized_model(self): + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + output = self.quantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + output_tokens = self.tokenizer.decode(output[0], skip_special_tokens=True) + self.assertIn(output_tokens, self.EXPECTED_OUTPUTS) + + def test_dequantized_model(self): + quantization_config = BlockFp8Config(dequantize=True) + dequantized_model = AutoModelForCausalLM.from_pretrained( + self.quantized_model_name, device_map=self.device_map, quantization_config=quantization_config + ) + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + output = dequantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + output_tokens = self.tokenizer.decode(output[0], skip_special_tokens=True) + self.assertEqual(output_tokens, self.EXPECTED_DEQUANTIZED_OUTPUT) + del dequantized_model + + def test_dequantize_when_no_accelerator(self): + with no_accelerator(): + dequantized_model = AutoModelForCausalLM.from_pretrained(self.quantized_model_name, device_map="cpu") + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to("cpu") + output = dequantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + output_tokens = self.tokenizer.decode(output[0], skip_special_tokens=True) + self.assertEqual(output_tokens, self.EXPECTED_DEQUANTIZED_OUTPUT) + del dequantized_model + + def test_save_pretrained(self): + with tempfile.TemporaryDirectory() as tmpdirname: + self.quantized_model.save_pretrained(tmpdirname) + + model = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=self.device_map) + + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + output = model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + def test_weight_and_weight_scale_inv(self): + weight = self.quantized_model.model.layers[0].self_attn.q_proj.weight + weight_scale_inv = self.quantized_model.model.layers[0].self_attn.q_proj.weight_scale_inv + self.assertEqual(weight.dtype, torch.float8_e4m3fn) + self.assertEqual(weight_scale_inv.dtype, torch.float32) + self.assertEqual(weight.shape, (weight_scale_inv.shape[0] * 128, weight_scale_inv.shape[1] * 128)) + + def test_block_size(self): + self.assertEqual(self.quantized_model.config.quantization_config.weight_block_size, (128, 128)) + quantization_config = BlockFp8Config(weight_block_size=(32, 32)) + quantized_model = AutoModelForCausalLM.from_pretrained( + self.model_name, device_map=self.device_map, quantization_config=quantization_config + ) + self.assertEqual(quantized_model.config.quantization_config.weight_block_size, (32, 32)) + + @require_torch_multi_accelerator + def test_quantized_model_multi_accelerators(self): + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + quantization_config = BlockFp8Config() + # need to empty cache or set max_memory, otherwise we will use the reserved memory that was not allocated when computing max-memory + # this will lead to put the entire model to device 0. + quantized_model = AutoModelForCausalLM.from_pretrained( + self.model_name, + device_map="auto", + quantization_config=quantization_config, + max_memory={0: "1GB", 1: "10GB"}, + ) + self.assertTrue(set(quantized_model.hf_device_map.values()) == {0, 1}) + + output = quantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + @require_torch_multi_accelerator + def test_save_pretrained_multi_accelerators(self): + with tempfile.TemporaryDirectory() as tmpdirname: + self.quantized_model.save_pretrained(tmpdirname) + # need to empty cache or set max_memory, otherwise we will use the reserved memory that was not allocated when computing max-memory + # this will lead to put the entire model to device 0. + model = AutoModelForCausalLM.from_pretrained( + tmpdirname, device_map="auto", max_memory={0: "1GB", 1: "10GB"} + ) + self.assertTrue(set(model.hf_device_map.values()) == {0, 1}) + + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + output = model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + def test_quantized_model_offload(self): + with self.assertRaisesRegex( + ValueError, + "You are attempting to load a fine-grained quantized model with a device_map that contains a cpu/disk", + ): + AutoModelForCausalLM.from_pretrained( + self.model_name, device_map=self.offload_device_map, quantization_config=self.quantization_config + ) + + def test_save_pretrained_offload(self): + with tempfile.TemporaryDirectory() as tmpdirname: + self.quantized_model.save_pretrained(tmpdirname) + + input_ids = self.tokenizer(self.input_text, return_tensors="pt").to(self.device_map) + + quantized_model = AutoModelForCausalLM.from_pretrained(tmpdirname, device_map=self.offload_device_map) + output = quantized_model.generate(**input_ids, max_new_tokens=self.max_new_tokens, do_sample=False) + self.assertIn(self.tokenizer.decode(output[0], skip_special_tokens=True), self.EXPECTED_OUTPUTS) + + @parameterized.expand(["eager", "batched_mm", "grouped_mm", "deepgemm"]) + def test_quantized_moe_forward(self, experts_implementation): + """MiniMax-M2 routes indices past `top_k`; every experts implementation has to run them.""" + if experts_implementation == "deepgemm" and torch_device != "cuda": + self.skipTest("deepgemm is only supported on CUDA") + + model = AutoModelForCausalLM.from_pretrained( + "hf-internal-testing/MiniMax-M2-Tiny-FP8", # single layer version + experts_implementation=experts_implementation, + device_map=self.device_map, + ) + assert model.config._experts_implementation == experts_implementation + + tokenizer = AutoTokenizer.from_pretrained("MiniMaxAI/MiniMax-M2") + messages = [ + {"role": "user", "content": [{"type": "text", "text": "What is your favourite condiment?"}]}, + { + "role": "assistant", + "content": [ + { + "type": "text", + "text": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!", + } + ], + }, + {"role": "user", "content": [{"type": "text", "text": "Do you have mayonnaise recipes?"}]}, + ] + model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to( + self.device_map + ) + + # Only caring about this not crashing + _ = model.generate(**model_inputs, max_new_tokens=24) + + +class BlockFp8Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = BLOCK_FP8 + + +class BlockFp8TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = BLOCK_FP8 + + +class BlockFp8LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = BLOCK_FP8 diff --git a/tests/quantization/finegrained/__init__.py b/tests/quantization/finegrained/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/finegrained/test_config_readers.py b/tests/quantization/finegrained/test_config_readers.py new file mode 100644 index 000000000000..879cf1995207 --- /dev/null +++ b/tests/quantization/finegrained/test_config_readers.py @@ -0,0 +1,375 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""Normalize producer metadata without silently changing module targets or quantization formats.""" + +import unittest + +import torch + +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import require_torch +from transformers.utils.quantization_config import ( + BlockFp8Config, + FbgemmFp8Config, + NVFP4Config, + read_compressed_tensors, + read_modelopt, +) + +from ...integrations.finegrained.utils import _compressed_tensors, _modelopt, _quantizer_for + + +@require_torch +class BlockScaledCheckpointReaderTest(unittest.TestCase): + """modelopt and compressed-tensors NVFP4 exports read as `NVFP4Config`, skip-lists bounded to the modules they name.""" + + def test_rowwise_fp8_producers_share_the_finegrained_runtime(self): + from transformers.quantizers.auto import AutoHfQuantizer + from transformers.quantizers.quantizer_blockfp8 import BlockFp8HfQuantizer + + channel = {"num_bits": 8, "type": "float", "strategy": "channel"} + token = {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True} + tensor = {"num_bits": 8, "type": "float", "strategy": "tensor", "dynamic": False} + for checkpoint, activation_format, scheme, upper_bound in ( + ({"quant_method": "fbgemm_fp8", "activation_scale_ub": 1200.0}, None, "dynamic", 1200.0), + ({"quant_method": "modelopt", "quant_algo": "FP8_PER_CHANNEL_PER_TOKEN"}, None, "dynamic", None), + (_compressed_tensors(channel, token, fmt="float-quantized"), None, "dynamic", None), + (_compressed_tensors(channel, tensor, fmt="float-quantized"), None, "static", None), + (_compressed_tensors(channel, fmt="float-quantized"), "bf16", "dynamic", None), + ): + with self.subTest(checkpoint=checkpoint): + config = AutoQuantizationConfig.from_dict(checkpoint) + self.assertIsInstance(config, BlockFp8Config) + self.assertEqual(config.weight_block_size, "row") + self.assertEqual(config.activation_format, activation_format) + self.assertEqual(config.activation_scheme, scheme) + self.assertEqual(config.activation_scale_ub, upper_bound) + self.assertIsInstance(AutoHfQuantizer.from_config(config), BlockFp8HfQuantizer) + self.assertEqual(AutoQuantizationConfig.from_dict(config.to_dict()).to_dict(), config.to_dict()) + self.assertIsInstance(AutoHfQuantizer.from_config(FbgemmFp8Config()), BlockFp8HfQuantizer) + checkpoint = AutoQuantizationConfig.from_dict(FbgemmFp8Config().to_dict()) + loading = AutoQuantizationConfig.from_dict(FbgemmFp8Config(activation_scale_ub=600.0).to_dict()) + self.assertEqual(AutoHfQuantizer.merge_quantization_configs(checkpoint, loading).activation_scale_ub, 600.0) + weight_only = AutoHfQuantizer.merge_quantization_configs(checkpoint, BlockFp8Config(activation_format="bf16")) + self.assertEqual(weight_only.activation_format, "bf16") + self.assertIsNone(weight_only.activation_scale_ub) + + def test_a_producers_config_groups_are_normalized(self): + """modelopt describes a format by its parameters; we name it. GLM-5.2-NVFP4 ships exactly + this, and `targets: ["Linear"]` is its spelling of the catch-all, so the export is one format.""" + config = _modelopt( + config_groups={ + "group_0": { + "weights": {"num_bits": 4, "type": "float", "group_size": 16}, + "input_activations": {"num_bits": 4, "type": "float", "group_size": 16, "dynamic": False}, + "targets": ["Linear"], + } + }, + ) + self.assertIsInstance(config, NVFP4Config) + self.assertIsNone(config.activation_format) + # NVFP4's `dynamic: False` is its calibrated activation global, not a static activation scale + self.assertEqual(config.activation_scheme, "dynamic") + + def test_an_unsupported_modelopt_group_is_rejected(self): + with self.assertRaisesRegex(ValueError, "Unsupported modelopt config_groups entry 'group_0'"): + _modelopt( + config_groups={ + "group_0": {"weights": {"num_bits": 3, "type": "int", "group_size": 64}, "targets": ["Linear"]} + }, + ) + + def test_mixed_precision_maps_explicit_modules_and_leaves_others_unquantized(self): + from transformers.quantizers.quantizers_utils import plan_entry_for + from transformers.utils.quantization_config import BlockFp8Config, FineGrainedConfig + + spec = { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "layers.8.attention.wq": {"quant_algo": "FP8_PB_WO"}, + "layers.9.experts": {"quant_algo": "NVFP4", "group_size": 16}, + }, + } + for fields in (spec, {"quantization": spec}): + with self.subTest(nested="quantization" in fields): + config = AutoQuantizationConfig.from_dict({"quant_method": "modelopt", **fields}) + self.assertIsInstance(config, FineGrainedConfig) + dense = plan_entry_for(config.dtype_plan, "model.layers.8.attention.wq") + self.assertIsInstance(dense, BlockFp8Config) + self.assertEqual(dense.activation_format, "bf16") + self.assertEqual(dense.weight_block_size, (128, 128)) + self.assertIsInstance(plan_entry_for(config.dtype_plan, "layers.9.experts.0.gate_proj"), NVFP4Config) + self.assertIsNone(plan_entry_for(config.dtype_plan, "layers.90.experts")) + self.assertIsNone(plan_entry_for(config.dtype_plan, "layers.8.attention.wk")) + restored = AutoQuantizationConfig.from_dict(config.to_dict()) + self.assertEqual(restored.to_dict(), config.to_dict()) + + def test_mixed_precision_rejects_incomplete_or_unsupported_assignments(self): + for fields, message in ( + ({}, "non-empty quantized_layers"), + ({"quantized_layers": {}}, "non-empty quantized_layers"), + ({"quantized_layers": {"layers.0": {"quant_algo": "W4A16_AWQ"}}}, "layers.0"), + ({"quantized_layers": {"layers.0": {"quant_algo": "NVFP4", "group_size": 32}}}, "group_size=16"), + ): + with self.subTest(fields=fields), self.assertRaisesRegex(ValueError, message): + AutoQuantizationConfig.from_dict( + {"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION", **fields} + ) + + def test_internal_globs_match_skipped_subtrees(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + skip = _modelopt(exclude_modules=["model.layers.*.self_attn", "*.shared_experts.*"]).modules_to_not_convert + for name in ( + "model.layers.0.self_attn.q_proj", + "model.layers.12.self_attn", + "model.layers.0.mlp.shared_experts.up_proj", + ): + self.assertFalse(should_convert_module(name, skip), name) + for name in ("model.layers.0.mlp.experts", "model.layers.0.my_self_attn", "modelXlayers.0.self_attn"): + self.assertTrue(should_convert_module(name, skip), name) + + def test_modelopt_runtime_transforms_are_not_dropped(self): + from transformers.utils.quantization_config import read_compressed_tensors + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + for field, value in ( + ("pre_quant_scale", True), + ("lora_rank", 16), + ("has_zero_point", True), + ("weight_scale_encoding", "csf"), + ): + for placement in ("top", "group", "layer"): + with self.subTest(field=field, placement=placement), self.assertRaisesRegex(ValueError, field): + if placement == "top": + _modelopt(**{field: value}) + elif placement == "group": + _modelopt(config_groups={"g": {"weights": nvfp4, field: value}}) + else: + AutoQuantizationConfig.from_dict( + { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": {"layers.0": {"quant_algo": "NVFP4", field: value}}, + } + ) + # The compressed-tensors reader likewise must leave unsupported transforms to its own backend. + self.assertIsNone(read_compressed_tensors(_compressed_tensors(nvfp4, nvfp4, pre_quant_scale=True))) + + def test_modelopt_groups_accept_exported_module_names(self): + from transformers.quantizers.quantizers_utils import plan_entry_for + + config = _modelopt( + config_groups={ + "g": { + "weights": {"num_bits": 4, "type": "float", "group_size": 16}, + "targets": ["model.layers.0.mlp.experts"], + } + } + ) + self.assertIsInstance(plan_entry_for(config.dtype_plan, "model.layers.0.mlp.experts"), NVFP4Config) + self.assertIsNone(plan_entry_for(config.dtype_plan, "model.layers.1.mlp.experts")) + + def test_a_compressed_tensors_ignore_entry_is_that_module_not_a_prefix(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + skip = AutoQuantizationConfig.from_dict( + _compressed_tensors(nvfp4, ignore=["model.layers.0.mlp.gate", "re:.*router$"]) + ).modules_to_not_convert + self.assertFalse(should_convert_module("model.layers.0.mlp.gate", skip)) + self.assertFalse(should_convert_module("model.layers.3.mlp.router", skip)) + self.assertTrue(should_convert_module("model.layers.0.mlp.gate_proj", skip)) + + def test_a_checkpoints_glob_skip_list_is_bounded(self): + """modelopt names skipped subtrees with globs. Passed through as regexes the dots match any + character and the star is greedy, so "model.layers.1.*" also takes layers 10-19.""" + from transformers.quantizers.quantizers_utils import should_convert_module + + # as a modelopt checkpoint ships them, through the config that normalizes them + skip = AutoQuantizationConfig.from_dict( + { + "quant_method": "modelopt", + "quant_algo": "NVFP4", + "ignore": ["model.layers.0*", "model.layers.1.*", "self_attn"], + } + ).modules_to_not_convert + for name in ( + "model.layers.0.mlp", + "model.layers.01.mlp", + "model.layers.1.mlp.down_proj", + "model.layers.2.self_attn.q_proj", + ): + self.assertFalse(should_convert_module(name, skip), name) + for name in ("model.layers.16.mlp.down_proj", "model.layers.2.mlp.gate_proj", "model.my_self_attn.q_proj"): + self.assertTrue(should_convert_module(name, skip), name) + + def test_a_compressed_tensors_nvfp4_checkpoint_reads_as_nvfp4(self): + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} + config = AutoQuantizationConfig.from_dict(_compressed_tensors(nvfp4, nvfp4, ignore=["lm_head"])) + self.assertIsInstance(config, NVFP4Config) + self.assertEqual(config.checkpoint_format, "compressed-tensors") + + def test_compressed_tensors_checkpoints_we_do_not_serve_stay_with_it(self): + """Integer, unsupported activation recipes and sparse checkpoints stay compressed-tensors' own.""" + from transformers.utils.quantization_config import read_compressed_tensors + + int4 = {"num_bits": 4, "type": "int", "group_size": 128, "strategy": "group"} + channel = {"num_bits": 8, "type": "float", "strategy": "channel"} + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16} + for checkpoint in ( + _compressed_tensors(int4, fmt="pack-quantized"), + _compressed_tensors( + channel, {"num_bits": 8, "type": "float", "strategy": "tensor", "dynamic": True}, fmt="float-quantized" + ), + _compressed_tensors(nvfp4, sparsity_config={"format": "sparse-24-bitmask"}), + ): + self.assertIsNone(read_compressed_tensors(checkpoint)) + + def test_activations_the_weights_kernels_do_not_serve_stay_with_compressed_tensors(self): + """An activation scheme is read only where the weights' kernels serve it: NVFP4 weights with FP8 per-token + activations and FP8 weights with int8 activations stay compressed-tensors' own, as do + group-32 FP8 weights whose scales are not E8M0 (MXFP8's).""" + from transformers.utils.quantization_config import Mxfp8Config, read_compressed_tensors + + nvfp4 = {"num_bits": 4, "type": "float", "group_size": 16, "strategy": "tensor_group"} + block = {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]} + fp8_token = {"num_bits": 8, "type": "float", "strategy": "token", "dynamic": True} + int8_token = {"num_bits": 8, "type": "int", "strategy": "token", "dynamic": True} + group_32 = {"num_bits": 8, "type": "float", "strategy": "group", "group_size": 32} + for name, checkpoint in ( + ("nvfp4 x fp8", _compressed_tensors(nvfp4, fp8_token)), + ("block-fp8 x int8", _compressed_tensors(block, int8_token, fmt="float-quantized")), + ( + "group-32 fp32 scales", + _compressed_tensors({**group_32, "scale_dtype": "torch.float32"}, group_32, fmt="float-quantized"), + ), + ): + with self.subTest(name): + self.assertIsNone(read_compressed_tensors(checkpoint)) + e8m0 = {**group_32, "scale_dtype": "torch.float8_e8m0fnu"} + self.assertIsInstance( + read_compressed_tensors(_compressed_tensors(e8m0, group_32, fmt="float-quantized")), Mxfp8Config + ) + + def test_unrepresentable_modelopt_targets_raise_before_loading(self): + model = torch.nn.Module() + model.linear = torch.nn.Linear(16, 16) + for target in ("unknown", "experts.0.gate_proj"): + config = AutoQuantizationConfig.from_dict( + { + "quant_method": "modelopt", + "quant_algo": "MIXED_PRECISION", + "quantized_layers": {target: {"quant_algo": "NVFP4"}}, + } + ) + with self.subTest(target=target), self.assertRaisesRegex(ValueError, "matches no supported"): + _quantizer_for(config)._process_model_before_weight_loading(model) + + +class FineGrainedConfigGroupReaderTest(unittest.TestCase): + def _groups(self, reverse=False, identical=False): + nvfp4 = {"targets": ["Linear"], "weights": {"num_bits": 4, "type": "float", "group_size": 16}} + fp8 = { + "targets": ["Linear"], + "weights": {"num_bits": 8, "type": "float", "strategy": "block", "block_structure": [128, 128]}, + } + groups = [("nvfp4", nvfp4), ("other", nvfp4 if identical else fp8)] + return dict(reversed(groups) if reverse else groups) + + def test_modelopt_rejects_conflicting_group_targets_in_either_order(self): + for reverse in (False, True): + with ( + self.subTest(reverse=reverse), + self.assertRaisesRegex(ValueError, "Conflicting quantization configs for target"), + ): + read_modelopt({"quant_algo": "NVFP4", "config_groups": self._groups(reverse=reverse)}) + + def test_compressed_tensors_keeps_its_backend_for_conflicting_group_targets(self): + for reverse in (False, True): + with self.subTest(reverse=reverse): + self.assertIsNone( + read_compressed_tensors( + {"format": "float-quantized", "config_groups": self._groups(reverse=reverse)} + ) + ) + + def test_identical_group_targets_can_share_a_config(self): + for reverse in (False, True): + for reader, fields, checkpoint_format in ( + (read_modelopt, {"quant_algo": "NVFP4"}, "modelopt"), + (read_compressed_tensors, {"format": "float-quantized"}, "compressed-tensors"), + ): + with self.subTest(reverse=reverse, producer=checkpoint_format): + config = reader({**fields, "config_groups": self._groups(reverse=reverse, identical=True)}) + self.assertIsInstance(config, NVFP4Config) + self.assertEqual(config.checkpoint_format, checkpoint_format) + + def test_modelopt_trailing_separator_exclusions_match_only_the_named_layers(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + config = read_modelopt( + { + "quant_algo": "NVFP4", + "exclude_modules": ["layers.0.", "layers.61.", "self_attn", "lm_head"], + } + ) + for name in ( + "model.layers.0.mlp.experts", + "model.layers.0", + "model.layers.61.mlp.up_proj", + "model.layers.61", + "model.layers.3.self_attn.q_proj", + "lm_head", + ): + with self.subTest(skipped=name): + self.assertFalse(should_convert_module(name, config.modules_to_not_convert)) + for name in ( + "model.layers.00.mlp.up_proj", + "model.layers.610.mlp.up_proj", + "model.layers.6.mlp.up_proj", + "model.layers.1.mlp.up_proj", + ): + with self.subTest(converted=name): + self.assertTrue(should_convert_module(name, config.modules_to_not_convert)) + + def test_modelopt_merges_native_regex_skip_list_with_both_producer_exclusions(self): + from transformers.quantizers.quantizers_utils import should_convert_module + + native_pattern = r"^model\.audio_tower(?:\.|$)" + config = read_modelopt( + { + "quant_algo": "NVFP4", + "modules_to_not_convert": [native_pattern], + "ignore": ["model.visual_tower.*"], + "exclude_modules": ["layers.0."], + } + ) + self.assertIn(native_pattern, config.modules_to_not_convert) + for name in ( + "model.audio_tower.layers.1.proj", + "model.visual_tower.layers.1.proj", + "model.layers.0.mlp.up_proj", + ): + with self.subTest(skipped=name): + self.assertFalse(should_convert_module(name, config.modules_to_not_convert)) + for name in ( + "model.audio_tower_extra.proj", + "model.visual_tower_extra.proj", + "model.layers.00.mlp.up_proj", + ): + with self.subTest(converted=name): + self.assertTrue(should_convert_module(name, config.modules_to_not_convert)) + restored = AutoQuantizationConfig.from_dict(config.to_dict()) + self.assertEqual(restored.modules_to_not_convert, config.modules_to_not_convert) diff --git a/tests/quantization/finegrained/test_finegrained.py b/tests/quantization/finegrained/test_finegrained.py new file mode 100644 index 000000000000..a87221cf4064 --- /dev/null +++ b/tests/quantization/finegrained/test_finegrained.py @@ -0,0 +1,676 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""The fine-grained quantizer (`FineGrainedConfig`): a plan mixing the format quantizers, the checkpoints that +read into one, and the conversions it builds. Each format's own coverage lives in its folder.""" + +import os +import re +import tempfile +import unittest +import warnings +from types import SimpleNamespace +from unittest import mock + +import torch +from parameterized import parameterized +from torch import nn + +import transformers.integrations.finegrained.layers as fg +from transformers import AutoModelForCausalLM, PreTrainedConfig, PreTrainedModel +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import ( + TestCasePlus, + cleanup, + require_kernels, + require_torch, + require_torch_accelerator, + slow, + torch_device, +) +from transformers.utils.quantization_config import ( + BlockFp8Config, + BlockScaledConfig, + FineGrainedConfig, + Mxfp4Config, + Mxfp8Config, + NVFP4Config, +) + +from ...integrations.finegrained.utils import ( + _Cfg, + _conversion_model, + _model_config, + _modelopt, + _quantizer_for, + _targets, +) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin +from ..block_scaled_tester import ( + FINEGRAINED, + FORMATS, + _cheap_tuning, +) + + +@require_torch +class FineGrainedPlanTest(unittest.TestCase): + """`FineGrainedConfig`: a plan of formats per module group, how a checkpoint that mixes formats reads into one, + and what each plan entry loads.""" + + def test_a_legacy_expert_dtype_becomes_the_two_entries_it_meant(self): + """DeepSeek-V4 ships `quant_method="fp8"` and declares its mxfp4 experts on the MODEL + config, because a flat config cannot say "the experts differ". `shared_experts` is DENSE + despite the name, which is why the pattern matches a path segment and not a substring.""" + from transformers.quantizers.auto import finegrained_quantization_config + from transformers.quantizers.quantizers_utils import plan_entry_for + + config = finegrained_quantization_config(BlockFp8Config(scale_fmt="ue8m0"), _model_config(expert_dtype="fp4")) + self.assertIsInstance(config, FineGrainedConfig) + for module, expected in ( + ("model.layers.3.mlp.experts", "mxfp4"), + ("model.layers.3.mlp.experts.gate_up_proj", "mxfp4"), + ("model.layers.3.self_attn.q_proj", "fp8"), + ("model.layers.3.mlp.shared_experts.up_proj", "fp8"), + ): + self.assertEqual(plan_entry_for(config.dtype_plan, module).method.value, expected, module) + # without one, block-FP8 stays what it is + plain = BlockFp8Config() + self.assertIs(finegrained_quantization_config(plain, _model_config()), plain) + + def test_deepseek_v4_takes_one_config_whether_its_checkpoint_is_from_the_hub_or_a_save(self): + """Both sides are read for the model before they are compared, so the class a user passes does not depend on + whether the checkpoint still says `fp8` (the Hub) or already `finegrained` (a `save_pretrained`).""" + from transformers.quantizers.auto import AutoHfQuantizer, finegrained_quantization_config + + model_config = _model_config(expert_dtype="fp4") + hub = {"quant_method": "fp8", "scale_fmt": "ue8m0", "weight_block_size": [128, 128]} + saved = finegrained_quantization_config(AutoQuantizationConfig.from_dict(hub), model_config).to_dict() + for checkpoint in (hub, saved): + resolved = finegrained_quantization_config(AutoQuantizationConfig.from_dict(checkpoint), model_config) + passed = finegrained_quantization_config(BlockFp8Config(dequantize=True), model_config) + self.assertTrue(AutoHfQuantizer.merge_quantization_configs(resolved, passed).dequantize) + + def test_an_expert_dtype_in_the_quantization_config_reads_as_the_same_plan(self): + """DeepSeek-V4.1 declares its FP4 experts in the quantization config instead of the model config: the same + plan, block-FP8 at the checkpoint's block over W4A4 MXFP4 experts, whichever config the load is handed.""" + from transformers.quantizers.auto import finegrained_quantization_config + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + shipped = { + "quant_method": "fp8", + "activation_scheme": "dynamic", + "weight_block_size": [32, 32], + "scale_fmt": "ue8m0", + "expert_dtype": "fp4", + } + model_config = SimpleNamespace(quantization_config=shipped, get_text_config=lambda: SimpleNamespace()) + for passed in (AutoQuantizationConfig.from_dict(shipped), BlockFp8Config(dequantize=True)): + with self.subTest(passed=type(passed).__name__): + config = finegrained_quantization_config(passed, model_config) + self.assertIsInstance(config, FineGrainedConfig) + experts, rest = config.dtype_plan[EXPERTS_PATTERN], config.dtype_plan[".*"] + # W4A4: the weights' own activation format + self.assertEqual((type(experts), experts.activation_format), (Mxfp4Config, None)) + self.assertEqual( + (type(rest), rest.weight_block_size, rest.scale_fmt), + (BlockFp8Config, passed.weight_block_size, passed.scale_fmt), + ) + # the loading options are the plan's, not an entry's + self.assertEqual((config.dequantize, rest.dequantize), (passed.dequantize, False)) + + def test_a_mixed_plans_dequantize_reads_every_entrys_block(self): + """One dequantize chain serves the whole plan, so a grid that rounds up finds its block whichever entry names it.""" + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.quantizers_utils import FUSED_STACKS + + config = FineGrainedConfig( + dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config(weight_block_size=(128, 128))}, + modules_to_convert=["embed_tokens"], + dequantize=True, + ) + quantizer = FineGrainedHfQuantizer(config, pre_quantized=True) + dequantize = [ + op + for c in quantizer.get_weight_conversions() + for op in getattr(c, "operations", ()) + if isinstance(op, FineGrainedDequantize) + ] + # the dense weights' converter and one per fused expert stack + self.assertEqual(len(dequantize), 1 + len(FUSED_STACKS)) + self.assertTrue(all(op.hf_quantizer is quantizer for op in dequantize)) + + def test_resolution_survives_a_round_trip_through_config_json(self): + """`to_json_string` SORTS keys, so a config read back has lost the order it was written + in. Resolution must not depend on it: the catch-all sorts first here and would swallow + every module if targeted entries were not preferred outright.""" + import json + + from transformers.quantizers.quantizers_utils import plan_entry_for + + config = FineGrainedConfig( + dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config(weight_block_size=(128, 128))} + ) + back = FineGrainedConfig.from_dict(json.loads(config.to_json_string())) + self.assertIsInstance(plan_entry_for(back.dtype_plan, "model.layers.0.mlp.experts"), NVFP4Config) + self.assertIsInstance(plan_entry_for(back.dtype_plan, "model.layers.0.self_attn.q_proj"), BlockFp8Config) + # a list survives JSON where a tuple does not, and the block size is compared as a tuple + self.assertEqual(back.dtype_plan[".*"].weight_block_size, (128, 128)) + + def test_two_targeted_entries_claiming_one_module_is_an_error(self): + from transformers.quantizers.quantizers_utils import plan_entry_for + + plan = {r"\.experts": BlockFp8Config(weight_block_size=None), r"mlp\.": NVFP4Config()} + with self.assertRaises(ValueError): + plan_entry_for(plan, "model.layers.0.mlp.experts") + + def test_an_export_mixing_formats_is_a_finegrained_config(self): + config = _modelopt( + config_groups={ + "group_0": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}, "targets": ["re:.*mlp.*"]}, + "group_1": {"weights": {"num_bits": 8, "type": "float", "group_size": 32}, "targets": ["re:.*attn.*"]}, + }, + ) + self.assertIsInstance(config, FineGrainedConfig) + self.assertEqual( + {pattern: type(entry) for pattern, entry in config.dtype_plan.items()}, + {".*mlp.*": NVFP4Config, ".*attn.*": Mxfp8Config}, + ) + # no `input_activations` is weight-only, whatever the weights' own activation format + self.assertEqual(config.dtype_plan[".*mlp.*"].activation_format, "bf16") + + def test_expert_dtype_reaches_a_config_quantizing_on_the_fly(self): + """The checkpoint an on-the-fly quantization writes declares the model config's `expert_dtype`, and is read + back with it: the experts are quantized to MXFP4 then, so the save loads as what it holds.""" + from transformers.quantizers.auto import get_hf_quantizer + from transformers.quantizers.block_scaled_utils import BlockScaledQuantizer + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + model_config = PreTrainedConfig() + model_config.expert_dtype = "fp4" + with mock.patch.object(BlockScaledQuantizer, "validate_environment"): + quantizer, model_config, _ = get_hf_quantizer(model_config, BlockFp8Config(), None, True, {}) + self.assertIsInstance(quantizer.quantization_config, FineGrainedConfig) + self.assertIsInstance(quantizer.quantization_config.dtype_plan[EXPERTS_PATTERN], Mxfp4Config) + self.assertFalse(quantizer.pre_quantized) + + def test_a_passed_configs_activation_format_reaches_a_prequantized_checkpoint(self): + """The activation format picks the kernel arm the same weights run on, so a passed config sets it; the + checkpoint's format options stay its own, and a passed one that differs from them is warned about.""" + from transformers.quantizers.auto import AutoHfQuantizer, finegrained_quantization_config + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + with warnings.catch_warnings(): + warnings.simplefilter("error") + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "mxfp4"}, Mxfp4Config(activation_format="mxfp8") + ) + self.assertEqual((type(merged), merged.activation_format), (Mxfp4Config, "mxfp8")) + + # a plan's entries lend theirs to the checkpoint's own entries + model_config = _model_config(expert_dtype="fp4") + hub = {"quant_method": "fp8", "scale_fmt": "ue8m0", "weight_block_size": [128, 128]} + passed = FineGrainedConfig( + dtype_plan={ + EXPERTS_PATTERN: Mxfp4Config(activation_format="bf16"), + ".*": BlockFp8Config(weight_block_size=(64, 64)), + } + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + merged = AutoHfQuantizer.merge_quantization_configs( + finegrained_quantization_config(AutoQuantizationConfig.from_dict(hub), model_config), passed + ) + self.assertTrue(any(issubclass(w.category, UserWarning) for w in caught)) + experts, rest = merged.dtype_plan[EXPERTS_PATTERN], merged.dtype_plan[".*"] + self.assertEqual(experts.activation_format, "bf16") + self.assertEqual((rest.weight_block_size, rest.scale_fmt), ((128, 128), "ue8m0")) + + def test_a_passed_config_lends_only_the_activation_format_it_names(self): + """A config passed for its skip-list keeps the checkpoint's activation format; one naming a format NVFP4 + cannot run on a weight-only checkpoint (no calibrated activation globals) is refused.""" + from transformers.quantizers.auto import AutoHfQuantizer + + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "nvfp4", "activation_format": "bf16"}, NVFP4Config(modules_to_not_convert=["lm_head"]) + ) + self.assertEqual((merged.activation_format, merged.modules_to_not_convert), ("bf16", ["lm_head"])) + merged = AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "mxfp4", "activation_format": "mxfp8"}, Mxfp4Config(modules_to_not_convert=["lm_head"]) + ) + self.assertEqual(merged.activation_format, "mxfp8") + with self.assertRaises(ValueError): + AutoHfQuantizer.merge_quantization_configs( + {"quant_method": "nvfp4", "activation_format": "bf16"}, NVFP4Config(activation_format=None) + ) + self.assertFalse([name for name in NVFP4Config().to_dict() if name.startswith("_")]) + + def test_the_formats_own_activation_name_is_native(self): + """One spelling for a format's own activations: `None`.""" + for config_class, own in ( + (BlockFp8Config, "fp8"), + (Mxfp8Config, "mxfp8"), + (Mxfp4Config, "mxfp4"), + (NVFP4Config, "nvfp4"), + ): + with self.subTest(config_class=config_class.__name__): + self.assertIsNone(config_class(activation_format=own).activation_format) + self.assertIsNone(_modelopt().activation_format) + + def test_options_the_format_fixes_are_refused_and_not_written(self): + """MXFP8, MXFP4 and NVFP4 fix their block, scale dtype and activation scheme; only block-FP8 sets them. A + `config.json` spelling the format's own values (MiniMax-M3's `weight_block_size: [1, 32]`) or the old + defaults reads back.""" + for config_class in (Mxfp8Config, Mxfp4Config, NVFP4Config): + for option in ({"weight_block_size": (128, 128)}, {"activation_scheme": "static"}): + with self.subTest(config_class=config_class.__name__, option=option): + with self.assertRaises(ValueError): + config_class(**option) + for written in ( + {"weight_block_size": None, "scale_fmt": "float", "activation_scheme": "dynamic"}, + {"weight_block_size": list(config_class.format_options["weight_block_size"])}, + ): + config = config_class.from_dict({"quant_method": config_class.method.value, **written}) + self.assertFalse(set(written) & set(config.to_dict())) + self.assertEqual(BlockFp8Config(scale_fmt="ue8m0").to_dict()["scale_fmt"], "ue8m0") + + def test_a_format_configs_options_are_keyword_only_and_unknown_ones_reported(self): + with self.assertRaises(TypeError): + BlockFp8Config((64, 64)) + with self.assertLogs("transformers.utils.quantization_config", level="INFO") as logs: + Mxfp8Config(activaton_format="bf16") + self.assertIn("activaton_format", "".join(logs.output)) + + def test_the_base_config_is_not_a_format(self): + with self.assertRaises(TypeError): + BlockScaledConfig() + + def test_a_plan_entry_is_a_format_config_without_loading_options(self): + """The loading options hold for the whole model, so an entry naming one is refused rather than ignored.""" + for entry in ( + BlockFp8Config(dequantize=True), + NVFP4Config(modules_to_not_convert=["lm_head"]), + BlockFp8Config(modules_to_convert=["embed_tokens"]), + {"quant_method": "mxfp8", "dequantize": True}, + ): + with self.subTest(entry=entry): + with self.assertRaises(ValueError): + FineGrainedConfig(dtype_plan={".*": entry}) + with self.assertRaises(ValueError): + FineGrainedConfig(dtype_plan={".*": {"quant_method": "awq"}}) + with self.assertRaises(TypeError): + FineGrainedConfig(dtype_plan={".*": "fp8"}) + + def test_the_plans_members_load_as_the_host_does(self): + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + + config = FineGrainedConfig(dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config()}) + quantizer = FineGrainedHfQuantizer(config, pre_quantized=False) + self.assertEqual([member.pre_quantized for member in quantizer.members()], [False, False]) + + def test_a_loaded_model_is_finished_once_whatever_its_formats(self): + """Every format checks its modules and keeps DeepGEMM off a model spread over devices, once per load.""" + import transformers.integrations.finegrained as finegrained + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + + for config in ( + FineGrainedConfig(dtype_plan={r"\.experts($|\.)": NVFP4Config(), ".*": BlockFp8Config()}), + Mxfp8Config(), + ): + with self.subTest(config=type(config).__name__): + quantizer = ( + FineGrainedHfQuantizer(config) if isinstance(config, FineGrainedConfig) else _quantizer_for(config) + ) + with ( + mock.patch.object(finegrained, "assert_modules_are_quantized") as checked, + mock.patch.object(finegrained, "disable_deepgemm_on_multi_device") as deepgemm, + ): + quantizer._process_model_after_weight_loading(nn.Module()) + self.assertEqual((checked.call_count, deepgemm.call_count), (1, 1)) + + def test_the_tensor_parallel_plan_resolves_the_experts_entry_for_the_tp_glob(self): + """Under intra-expert TP, whether the expert scale grid splits with its weight is read off the experts' + plan entries for the modules the TP glob (`layers.*.mlp.experts`) names. The glob's one style serves every + layer, so a plan giving some of them a grid and others a per-tensor scale is refused.""" + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + + def tp_plan(dtype_plan): + config = SimpleNamespace( + base_model_tp_plan={ + "layers.*.mlp.experts": "moe_tp_experts", + "layers.*.mlp.experts.gate_up_proj": "colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + }, + num_hidden_layers=2, + _experts_implementation=None, + ) + quantizer = FineGrainedHfQuantizer(FineGrainedConfig(dtype_plan=dtype_plan)) + model = nn.Module() + model.config = config + model.layers = nn.ModuleList([nn.Module(), nn.Module()]) + with torch.device("meta"), mock.patch.object(fg, "is_sm100", return_value=False): + for i, layer in enumerate(model.layers): + layer.mlp = nn.Module() + entry = next( + value + for pattern, value in dtype_plan.items() + if re.fullmatch(pattern, f"layers.{i}.mlp.experts") + ) + layer.mlp.experts = fg.FineGrainedExperts(_Cfg(), quantization_config=entry) + plan = dict(config.base_model_tp_plan) + quantizer.update_parallel_plan(model, plan) + return plan + + per_tensor = BlockFp8Config(weight_block_size=None) + plan = tp_plan({r"layers\.\d+\.mlp\.experts": Mxfp8Config(), ".*": per_tensor}) + self.assertEqual(plan.get("layers.*.mlp.experts.gate_up_proj_scale_inv"), "moe_experts_colwise") + self.assertEqual(plan.get("layers.*.mlp.experts.down_proj_scale_inv"), "moe_experts_rowwise") + with self.assertRaises(ValueError): + tp_plan({r"layers\.1\.mlp\.experts": Mxfp8Config(), ".*": per_tensor}) + + def test_ep_plans_use_masked_routing_for_every_format(self): + configs = [spec.config_cls() for spec in FORMATS] + configs.append(FineGrainedConfig(dtype_plan={r"\.experts($|\.)": Mxfp4Config(), ".*": BlockFp8Config()})) + for config in configs: + for router in ("gate", "router"): + with self.subTest(config=type(config).__name__, router=router): + model = nn.Module() + model.config = SimpleNamespace(_experts_implementation=None) + model.layers = nn.ModuleList([nn.Module()]) + model.layers[0].mlp = nn.Module() + setattr(model.layers[0].mlp, router, nn.Module()) + model.layers[0].mlp.experts = nn.Module() + plan = { + "layers.*.mlp.experts": "ep_dispatch_experts", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + } + quantizer = _quantizer_for(config) + quantizer.update_parallel_plan(model, plan) + expected = { + "layers.*.mlp.experts": "moe_tp_experts", + f"layers.*.mlp.{router}": "ep_router", + "layers.*.mlp.experts.gate_up_proj": "grouped_gemm", + "layers.*.mlp.experts.down_proj": "grouped_gemm", + } + self.assertEqual(plan, expected) + quantizer.update_parallel_plan(model, plan) + self.assertEqual(plan, expected) + + def test_parallel_plans_replicate_the_per_tensor_scales(self): + """A per-tensor scale is 0-dim, with no axis to split: the quantizer names each one a plan would shard with + its module `replicate`, which the parameter's own entry makes take precedence. A scale grid follows its + module's style.""" + from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES + from transformers.integrations.finegrained import FineGrainedLinear + from transformers.quantizers.quantizer_nvfp4 import NVFP4HfQuantizer + + model = nn.Module() + model.q_proj = FineGrainedLinear(64, 64, quantization_config=NVFP4Config()) + model.o_proj = FineGrainedLinear(64, 64, quantization_config=NVFP4Config()) + model._tp_plan, model._ep_plan = {"q_proj": "colwise"}, {} + model.config = SimpleNamespace(_experts_implementation=None) + NVFP4HfQuantizer(NVFP4Config()).update_parallel_plan(model, model._tp_plan) + self.assertEqual( + model._tp_plan, + {"q_proj": "colwise", "q_proj.weight_global_scale": "replicate", "q_proj.activation_scale": "replicate"}, + ) + scale = model.q_proj.weight_global_scale + ALL_PARALLEL_STYLES["replicate"].shard_param(model.q_proj, "weight_global_scale", mesh=None) + self.assertIs(model.q_proj.weight_global_scale, scale) + + +@require_torch +class FineGrainedPlanConversionTest(unittest.TestCase): + """The conversions a plan of several formats builds.""" + + def test_input_scale_reaches_each_formats_own_slot(self): + """One checkpoint key, two formats: `input_scale` is a static FP8 module's activation scale and an + NVFP4 module's activation global, held in one slot either way. Only the format holding the experts + reads their keys, and two formats that both would refuse the plan.""" + from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key + from transformers.quantizers.quantizer_finegrained import FineGrainedHfQuantizer + from transformers.quantizers.quantizers_utils import EXPERTS_PATTERN + + config = FineGrainedConfig( + dtype_plan={ + EXPERTS_PATTERN: NVFP4Config(activation_format="nvfp4"), + ".*": BlockFp8Config(weight_block_size=None, activation_scheme="static"), + }, + ) + quantizer = FineGrainedHfQuantizer(config) + quantizer.pre_quantized = True + fp8 = next(member for member in quantizer.members() if member.quantization_config is config.dtype_plan[".*"]) + model = _conversion_model(quantizer) + self.assertFalse(fp8.holds_experts) + conversions = quantizer.update_weight_conversions([], model) + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + + renamed, _ = rename_source_key("model.layers.0.self_attn.q_proj.input_scale", renamings, converters) + self.assertEqual(renamed, "model.layers.0.self_attn.q_proj.activation_scale") + _, pattern = rename_source_key("model.layers.0.mlp.experts.3.down_proj.input_scale", renamings, converters) + (converter,) = [c for c in converters if pattern in c.source_patterns] + self.assertIn("experts.down_proj_activation_scale", _targets(converter)) + self.assertIn("experts.down_proj_weight_global_scale", _targets(converter)) + + fp8.holds_experts = True + with self.assertRaises(ValueError): + quantizer.update_weight_conversions([], model) + + @staticmethod + def _split(conversions): + from transformers.core_model_loading import WeightConverter, WeightRenaming + + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + return renamings, [c for c in conversions if isinstance(c, WeightConverter)] + + def test_a_rewritten_model_converter_keeps_its_placement(self): + """Qwen4-Exp concatenates its ~95 GiB n-gram table on CPU (`force_cpu`); the rewrite that adds the scales + keeps it there, both ways the quantizer rewrites a converter.""" + from transformers.core_model_loading import MergeModulelist, WeightConverter + + for dequantize in (False, True): + with self.subTest(dequantize=dequantize): + model_converter = WeightConverter( + "ngram_embeddings.*.weight", "ngram_embedding.weight", [MergeModulelist(dim=0)], force_cpu=True + ) + quantizer = _quantizer_for(BlockFp8Config(dequantize=dequantize)) + (rewritten, *_) = quantizer.update_weight_conversions([model_converter], _conversion_model(quantizer)) + self.assertTrue(rewritten.force_cpu) + + def test_the_fused_nvfp4_layout_reaches_the_dequantize_converter(self): + """vLLM's stacked layout names its scales `{proj}_weight_scale` / `{proj}_weight_scale_2`: under `dequantize` + they take the slot names the fused stack's dequantize converter collects, so the packed weight is folded.""" + from transformers.core_model_loading import rename_source_key + from transformers.integrations.finegrained.conversions import FineGrainedDequantize + + quantizer = _quantizer_for(NVFP4Config(dequantize=True)) + renamings, converters = self._split(quantizer.update_weight_conversions([], _conversion_model(quantizer))) + prefix = "model.layers.0.mlp.experts" + for proj in ("gate_up_proj", "down_proj"): + for shipped, slot in (("_weight_scale", "_scale_inv"), ("_weight_scale_2", "_weight_global_scale")): + with self.subTest(key=f"{proj}{shipped}"): + renamed, pattern = rename_source_key(f"{prefix}.{proj}{shipped}", renamings, converters) + self.assertEqual(renamed, f"{prefix}.{proj}") + (converter,) = [c for c in converters if pattern in c.source_patterns] + self.assertTrue(any(isinstance(op, FineGrainedDequantize) for op in converter.operations)) + self.assertIn(rf"(? quantized_model_size[""] * 1.5 - - @parameterized.expand(["eager", "batched_mm", "grouped_mm", "deepgemm"]) - def test_quantized_moe_forward(self, experts_implementation): - """ - Checks implicitly if the moe implementation is correct, i.e. it does not crash for cases - where the indices go over `top_k` as shown within the Minimax M2 model - """ - # deepgemm only has CUDA kernels, skip on other devices - if experts_implementation == "deepgemm" and torch_device != "cuda": - self.skipTest("deepgemm is only supported on CUDA") - - model = AutoModelForCausalLM.from_pretrained( - "hf-internal-testing/MiniMax-M2-Tiny-FP8", # single layer version - experts_implementation=experts_implementation, - device_map=self.device_map, - ) - assert model.config._experts_implementation == experts_implementation - - tokenizer = AutoTokenizer.from_pretrained("MiniMaxAI/MiniMax-M2") - messages = [ - {"role": "user", "content": [{"type": "text", "text": "What is your favourite condiment?"}]}, - { - "role": "assistant", - "content": [ - { - "type": "text", - "text": "Well, I'm quite partial to a good squeeze of fresh lemon juice. It adds just the right amount of zesty flavour to whatever I'm cooking up in the kitchen!", - } - ], - }, - {"role": "user", "content": [{"type": "text", "text": "Do you have mayonnaise recipes?"}]}, - ] - model_inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to( - self.device_map - ) - - # Only caring about this not crashing - _ = model.generate(**model_inputs, max_new_tokens=24) - - -@require_torch_accelerator -@unittest.skipIf( - get_device_properties()[0] == "cuda" - and (get_device_properties()[1] < 8 or (get_device_properties()[1] == 8 and get_device_properties()[2] < 9)), - "Skipping FP8LinearTest because it is not supported on GPU with capability < 8.9", -) -class FP8LinearTest(unittest.TestCase): - device = torch_device - - def test_linear_preserves_shape(self): - """ - Test that FP8Linear preserves shape when in_features == out_features. - """ - from transformers.integrations import FP8Linear - - linear = FP8Linear(256, 256, block_size=(128, 128)).to(self.device) - x = torch.rand((1, 5, 256)).to(self.device) - - x_ = linear(x) - self.assertEqual(x_.shape, x.shape) - - def test_linear_with_diff_feature_size_preserves_shape(self): - """ - Test that FP8Linear generates the correct shape when in_features != out_features. - """ - from transformers.integrations import FP8Linear - - linear = FP8Linear(128, 256, block_size=(128, 128)).to(self.device) - x = torch.rand((1, 5, 128)).to(self.device) - - x_ = linear(x) - self.assertEqual(x_.shape, (1, 5, 256)) - - -class FP8EmbeddingTest(unittest.TestCase): - """ - Some checkpoints quantize an embedding *table* rather than a linear (Qwen4-Exp's n-gram/PLE - table). - """ - - def test_lookup_applies_the_per_tensor_scale(self): - """Without the rescale the rows come back as raw FP8 and blow up the next matmul.""" - from transformers.integrations import FP8Embedding - - table = torch.randn(64, 16, dtype=torch.bfloat16) - scale = table.abs().max().float() / torch.finfo(torch.float8_e4m3fn).max - embedding = FP8Embedding(64, 16) - embedding.weight = torch.nn.Parameter((table.float() / scale).to(torch.float8_e4m3fn), requires_grad=False) - embedding.weight_scale = torch.nn.Parameter(scale.to(torch.bfloat16).reshape(1), requires_grad=False) - - ids = torch.randint(0, 64, (2, 5)) - rows, expected = embedding(ids), table[ids] - - self.assertEqual(rows.dtype, expected.dtype) - error = (rows.float() - expected.float()).abs().max() / expected.float().abs().max() - self.assertLess(error, 0.1) - - def test_patterns_match_by_suffix_and_honour_the_skip_list(self): - """Suffixes match whether or not the text model is nested under a multimodal wrapper.""" - from transformers.integrations import FP8Embedding, replace_with_fp8_embedding - - model = torch.nn.ModuleDict( - {n: torch.nn.ModuleDict({"table": torch.nn.Embedding(8, 4)}) for n in ("layers", "language_model")} - ) - model.other = torch.nn.Embedding(8, 4) - replace_with_fp8_embedding(model, ["table"], modules_to_not_convert=["language_model.table"]) - - self.assertIsInstance(model["layers"]["table"], FP8Embedding) # matched by suffix - self.assertIs(type(model["language_model"]["table"]), torch.nn.Embedding) # skip-listed - self.assertIs(type(model.other), torch.nn.Embedding) # no pattern match - - -class FP8DeepGEMMMultiDeviceTest(unittest.TestCase): - """`_disable_deepgemm_on_multi_device` must flag FP8 modules based on the devices they actually - occupy — DeepGEMM's kernels are bound to a single CUDA context and corrupt across devices, but a - model that fits on one device must keep DeepGEMM even when other GPUs are visible (no overshoot). - """ - - @staticmethod - def _fp8_module(device): - from transformers.integrations import FP8Linear - - return FP8Linear(256, 256, block_size=(128, 128)).to(device) - - @require_torch_multi_gpu - def test_multi_device_disables_deepgemm(self): - from transformers.integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device - - model = torch.nn.Module() - model.a = self._fp8_module("cuda:0") - model.b = self._fp8_module("cuda:1") - _disable_deepgemm_on_multi_device(model) - self.assertTrue(model.a._deepgemm_disabled) - self.assertTrue(model.b._deepgemm_disabled) - - @require_torch_gpu - def test_single_device_keeps_deepgemm(self): - from transformers.integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device - - model = torch.nn.Module() - model.a = self._fp8_module("cuda:0") - model.b = self._fp8_module("cuda:0") - _disable_deepgemm_on_multi_device(model) - self.assertFalse(model.a._deepgemm_disabled) - self.assertFalse(model.b._deepgemm_disabled) - - -class FP8Ue8m0LoadingTest(unittest.TestCase): - @unittest.skipUnless(hasattr(torch, "float8_e8m0fnu"), "needs torch.float8_e8m0fnu") - @patch.object(FineGrainedFP8HfQuantizer, "validate_environment", lambda *args, **kwargs: None) def test_float32_scale_container_loads_as_ue8m0(self): - # dsv4-flash-base ships power-of-two ue8m0 scales as float32 under `.scale`; they must load as e8m0 + import torch from safetensors.torch import load_file, save_file from transformers import MixtralConfig, MixtralForCausalLM + if not hasattr(torch, "float8_e8m0fnu"): + self.skipTest("needs torch.float8_e8m0fnu") config = MixtralConfig( - vocab_size=64, hidden_size=128, intermediate_size=128, num_hidden_layers=1, num_attention_heads=4, - num_key_value_heads=2, num_local_experts=4, num_experts_per_tok=2, + vocab_size=64, hidden_size=128, intermediate_size=128, num_hidden_layers=1, + num_attention_heads=4, num_key_value_heads=2, num_local_experts=4, num_experts_per_tok=2, ) # fmt: skip with tempfile.TemporaryDirectory() as tmp: MixtralForCausalLM(config).save_pretrained(f"{tmp}/bf16") - quantized = AutoModelForCausalLM.from_pretrained( - f"{tmp}/bf16", quantization_config=FineGrainedFP8Config(scale_fmt="ue8m0"), device_map="cpu" + quantized = MixtralForCausalLM.from_pretrained( + f"{tmp}/bf16", quantization_config=_legacy(scale_fmt="ue8m0"), device_map="cpu" ) - expected = {k: v.float() for k, v in quantized.state_dict().items() if k.endswith("scale_inv")} + expected = { + name: tensor.float() for name, tensor in quantized.state_dict().items() if name.endswith("scale_inv") + } quantized.save_pretrained(f"{tmp}/fp8") path = f"{tmp}/fp8/model.safetensors" state_dict = { - k.replace("proj.weight_scale_inv", "proj.scale"): v.float() if k.endswith("scale_inv") else v - for k, v in load_file(path).items() + name.replace("proj.weight_scale_inv", "proj.scale"): tensor.float() + if name.endswith("scale_inv") + else tensor + for name, tensor in load_file(path).items() } save_file(state_dict, path, metadata={"format": "pt"}) - - loaded = AutoModelForCausalLM.from_pretrained(f"{tmp}/fp8", device_map="cpu") - scales = {k: v for k, v in loaded.state_dict().items() if k.endswith("scale_inv")} + loaded = MixtralForCausalLM.from_pretrained( + f"{tmp}/fp8", quantization_config=_legacy(scale_fmt="ue8m0"), device_map="cpu" + ) + scales = {name: tensor for name, tensor in loaded.state_dict().items() if name.endswith("scale_inv")} self.assertEqual(scales.keys(), expected.keys()) for name, scale in scales.items(): self.assertEqual(scale.dtype, torch.float8_e8m0fnu, name) diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index da705d8a10f1..dd2b7d4330b0 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -12,394 +12,137 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc +"""MXFP4 (`Mxfp4Config`): GPT-OSS's blocks layout and the packed per-expert one it loads and saves back, the +generate path, and the published GPT-OSS checkpoints.""" + import tempfile import unittest -from contextlib import ExitStack, contextmanager from unittest.mock import patch -from transformers import AutoTokenizer, GptOssForCausalLM, Mxfp4Config +import torch + +from transformers import AutoTokenizer, GptOssConfig, GptOssForCausalLM, Mxfp4Config from transformers.testing_utils import ( + TestCasePlus, + cleanup, require_kernels, require_torch, + require_torch_accelerator, require_torch_accelerator_memory, require_torch_large_accelerator, - require_triton, slow, torch_device, ) -from transformers.utils import ( - is_torch_available, -) +from ...integrations.finegrained.utils import ( + _ChainAssertions, + _conversion_model, + _load_and_save, + _per_expert, + _quantizer_for, +) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin +from ..block_scaled_tester import ( + MXFP4, + BlockScaledFormatTesterMixin, + _cheap_tuning, +) -if is_torch_available(): - import torch - - -if torch.cuda.is_available(): - REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.4.0") -elif hasattr(torch, "xpu") and torch.xpu.is_available(): - REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0") -elif torch_device == "cpu": - REQUIRE_TRITON_MXFP4 = require_triton(min_version="3.5.0") -else: - REQUIRE_TRITON_MXFP4 = unittest.skip("test requires CUDA or XPU") - - -def _empty_accelerator_cache(): - if torch.cuda.is_available(): - torch.cuda.empty_cache() - elif hasattr(torch, "xpu") and torch.xpu.is_available(): - torch.xpu.empty_cache() - - -@contextmanager -def _patch_no_accelerator(): - with ExitStack() as stack: - stack.enter_context(patch("torch.cuda.is_available", return_value=False)) - if hasattr(torch, "xpu"): - stack.enter_context(patch("torch.xpu.is_available", return_value=False)) - stack.enter_context(patch("torch.accelerator.current_accelerator", return_value=None)) - yield - - -class Mxfp4ConfigTest(unittest.TestCase): - def test_basic_config_creation(self): - """Test basic configuration creation with default values""" - config = Mxfp4Config() - self.assertEqual(config.quant_method.value, "mxfp4") - self.assertIsNone(config.modules_to_not_convert) - self.assertFalse(config.dequantize) - - def test_config_with_modules_to_not_convert(self): - """Test configuration with modules to not convert""" - modules = ["model.layers.*.self_attn", "lm_head"] - config = Mxfp4Config(modules_to_not_convert=modules) - self.assertEqual(config.modules_to_not_convert, modules) - - def test_config_with_dequantize(self): - """Test configuration with dequantize enabled""" - config = Mxfp4Config(dequantize=True) - self.assertTrue(config.dequantize) - - def test_get_loading_attributes(self): - """Test get_loading_attributes method""" - config = Mxfp4Config(dequantize=True) - attrs = config.get_loading_attributes() - self.assertEqual(attrs["dequantize"], True) - - def test_to_dict(self): - """Test configuration serialization to dict""" - config = Mxfp4Config(modules_to_not_convert=["lm_head"], dequantize=True) - config_dict = config.to_dict() - self.assertEqual(config_dict["quant_method"], "mxfp4") - self.assertEqual(config_dict["modules_to_not_convert"], ["lm_head"]) - # we don't keep dequantize in config_dict - self.assertTrue("dequantize" not in config_dict) - - def test_from_dict(self): - """Test configuration creation from dict""" - config_dict = {"quant_method": "mxfp4", "modules_to_not_convert": ["lm_head"], "dequantize": True} - config = Mxfp4Config.from_dict(config_dict) - self.assertEqual(config.modules_to_not_convert, ["lm_head"]) - self.assertTrue(config.dequantize) - - -class Mxfp4QuantizerTest(unittest.TestCase): - """Test the Mxfp4HfQuantizer class""" +class Mxfp4EnvironmentTest(unittest.TestCase): + def test_weight_only_mxfp4_keeps_existing_cuda_architectures(self): + for capability in ((7, 5), (8, 0), (8, 6)): + with self.subTest(capability=capability): + config = Mxfp4Config() + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=capability), + patch("transformers.quantizers.block_scaled_utils.is_accelerate_available", return_value=True), + ): + _quantizer_for(config).validate_environment(device_map="cuda") + self.assertFalse(config.dequantize) + + def test_quantized_activations_keep_the_fp8_hardware_requirement(self): + for activation_format in ("mxfp4", "mxfp8"): + with self.subTest(activation_format=activation_format): + config = Mxfp4Config(activation_format=activation_format) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.get_device_capability", return_value=(8, 0)), + patch("transformers.quantizers.block_scaled_utils.is_accelerate_available", return_value=True), + ): + _quantizer_for(config).validate_environment(device_map="cuda") + self.assertTrue(config.dequantize) + + +@require_torch_accelerator +@require_kernels +class Mxfp4GenerateTest(unittest.TestCase): def setUp(self): - gc.collect() - _empty_accelerator_cache() - from transformers.utils.logging import warning_once - - warning_once.cache_clear() - - def test_quantizer_validation_no_torch(self): - """Test quantizer validation when torch is not available""" - with patch("transformers.quantizers.quantizer_mxfp4.is_torch_available", return_value=False): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - - with self.assertRaises(ImportError): - quantizer.validate_environment() - - def test_quantizer_validation_no_accelerator(self): - """Test quantizer validation when CUDA/XPU is not available""" - with ( - _patch_no_accelerator(), - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - # CPU already supported MXFP4 - quantizer.validate_environment() - - @unittest.skipUnless(torch_device == "cuda", "test requires CUDA") - def test_quantizer_validation_low_compute_capability(self): - """Test quantizer validation with CUDA low compute capability""" - with patch("torch.cuda.get_device_capability", return_value=(7, 0)): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - with self.assertRaises(ValueError): - quantizer.validate_environment() - - @unittest.skipUnless(torch_device == "cuda", "test requires CUDA") - def test_quantizer_validation_low_compute_capability_with_prequantized(self): - """Test pre-quantized validation with CUDA low compute capability""" - with patch("torch.cuda.get_device_capability", return_value=(7, 0)): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - - # Should automatically set dequantize=True and warn - quantizer.validate_environment() - self.assertTrue(quantizer.quantization_config.dequantize) - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_quantizer_validation_low_compute_capability_with_dequantize(self): - """Test quantizer validation with dequantize enabled""" - with patch("torch.cuda.get_device_capability", return_value=(7, 0)): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config(dequantize=True) - quantizer = Mxfp4HfQuantizer(config) - - # Should not raise error with dequantize=True - try: - quantizer.validate_environment() - except ValueError as e: - if "compute capability" in str(e): - self.fail("Should not raise compute capability error when dequantize=True") - - def test_quantizer_validation_order_dequantize_before_accelerator_check(self): - """Test that dequantize check happens before CUDA/XPU availability check""" - # Mock torch.cuda.is_available - with ( - _patch_no_accelerator(), - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - # Test with dequantize=True - should pass even without CUDA/XPU and accelerate - config = Mxfp4Config(dequantize=True) - quantizer = Mxfp4HfQuantizer(config) - - # This should not raise any error because dequantize check comes first - quantizer.validate_environment() - - # Test with dequantize=False - should still fail due to missing CUDA/XPU - config = Mxfp4Config(dequantize=False) - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - # CPU already supported MXFP4 - quantizer.validate_environment() - - def test_quantizer_validation_missing_triton(self): - """Test quantizer validation when triton is not available""" - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - with self.assertRaises(ValueError): - quantizer.validate_environment() - - def test_quantizer_validation_missing_triton_pre_quantized_no_dequantize(self): - """Test quantizer validation when triton is not available but model is pre-quantized and dequantize is False""" - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - ): - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = True - - # Should automatically set dequantize=True and warn - quantizer.validate_environment() - self.assertTrue(quantizer.quantization_config.dequantize) - - def test_is_trainable(self): - """Test trainability""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - - # MXFP4 is not trainable - self.assertFalse(quantizer.is_trainable) - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_warning_distinguishes_triton_from_kernels(self): - """When only one dependency is missing, warning should mention it specifically.""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - # Missing kernels only -> warning should mention kernels - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = True - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - self.assertLogs("transformers", level="WARNING") as cm, - ): - quantizer.validate_environment() - - warning_text = " ".join(cm.output) - self.assertIn("kernels", warning_text.lower()) - self.assertTrue(quantizer.quantization_config.dequantize) - - # Missing triton only -> warning should mention triton - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = True - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - self.assertLogs("transformers", level="WARNING") as cm, - ): - quantizer.validate_environment() - - warning_text = " ".join(cm.output) - self.assertIn("triton", warning_text.lower()) - self.assertTrue(quantizer.quantization_config.dequantize) - - @unittest.skipUnless(torch_device in {"cuda", "xpu"}, "test requires CUDA or XPU") - def test_error_distinguishes_triton_from_kernels(self): - """When quantizing without a dependency, ValueError should mention it specifically.""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - # Missing kernels only -> error should mention kernels - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=True), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=False), - ): - with self.assertRaises(ValueError) as ctx: - quantizer.validate_environment() - - self.assertIn("kernels", str(ctx.exception).lower()) - - # Missing triton only -> error should mention triton - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - with ( - patch("transformers.quantizers.quantizer_mxfp4.is_triton_available", return_value=False), - patch("transformers.quantizers.quantizer_mxfp4.is_kernels_available", return_value=True), - ): - with self.assertRaises(ValueError) as ctx: - quantizer.validate_environment() - - self.assertIn("triton", str(ctx.exception).lower()) - - -class Mxfp4IntegrationTest(unittest.TestCase): - """Test mxfp4 integration functions""" - - def test_should_convert_module(self): - """Test module conversion decision logic""" - from transformers.quantizers.quantizers_utils import should_convert_module - - # Should convert by default - self.assertTrue(should_convert_module("model", None)) - self.assertTrue(should_convert_module("model", [])) - - # Should not convert if in exclusion list - patterns = ["model.layers.*.self_attn", "lm_head"] - self.assertFalse(should_convert_module("lm_head", patterns)) - self.assertTrue(should_convert_module("experts", patterns)) - - @require_torch - def test_convert_moe_packed_tensors(self): - """Test unpacking of quantized tensors""" - from transformers.integrations.mxfp4 import convert_moe_packed_tensors - - # Create dummy packed tensors - blocks = torch.randint(0, 255, (2, 4, 8, 16), dtype=torch.uint8) - scales = torch.randint(100, 150, (2, 4, 8), dtype=torch.uint8) - - result = convert_moe_packed_tensors(blocks, scales, dtype=torch.bfloat16) - self.assertEqual(result.shape, (2, 8 * 16 * 2, 4)) - self.assertEqual(result.dtype, torch.bfloat16) - - @REQUIRE_TRITON_MXFP4 - @require_kernels - @require_torch - def test_quantize_to_mxfp4(self): - """Test quantization function""" - from transformers.integrations.hub_kernels import get_kernel - from transformers.integrations.mxfp4 import quantize_to_mxfp4 - - # Create dummy weight tensor - device = torch_device - w = torch.randn(32, 64, 128, dtype=torch.bfloat16, device=torch.device(device)) - - quantized_w, w_scale = quantize_to_mxfp4(w, get_kernel("kernels-community/gpt-oss-triton-kernels", version=1)) + self.enterContext(_cheap_tuning()) + + def test_generate_matches_forward_logits(self): + """`generate` switches the experts to the batched kernels, while a plain forward runs the grouped ones: on + a tiny GPT-OSS with biased experts, quantized on the fly (W4A16), every decoding step's logits must match + a forward over the same prefix. Large biases make differences between the two kernels' bias handling visible.""" + torch.manual_seed(0) + config = GptOssConfig( + vocab_size=512, + hidden_size=256, + intermediate_size=256, + head_dim=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + ) + model = GptOssForCausalLM(config).to(torch.bfloat16) + with torch.no_grad(): + for layer in model.model.layers: + layer.mlp.experts.gate_up_proj_bias.normal_() + layer.mlp.experts.down_proj_bias.normal_() + with tempfile.TemporaryDirectory() as tmp: + model.save_pretrained(tmp) + quantized = GptOssForCausalLM.from_pretrained( + tmp, quantization_config=Mxfp4Config(), dtype=torch.bfloat16, device_map=torch_device + ) - # Check that shapes are reasonable - self.assertEqual(quantized_w.dtype, torch.uint8) + prompt = torch.arange(1, 17, device=torch_device).unsqueeze(0) + for use_cache in (False, True): + with self.subTest(use_cache=use_cache), torch.no_grad(): + out = quantized.generate( + prompt, + max_new_tokens=3, + do_sample=False, + use_cache=use_cache, + output_logits=True, + return_dict_in_generate=True, + ) + for step, logits in enumerate(out.logits): + prefix = out.sequences[:, : prompt.shape[1] + step] + expected = quantized(prefix).logits[:, -1].float() + torch.testing.assert_close(logits.float(), expected, atol=5e-2, rtol=5e-2) -@require_torch @require_torch_large_accelerator -@REQUIRE_TRITON_MXFP4 @require_kernels @slow class Mxfp4ModelTest(unittest.TestCase): - """Test mxfp4 with actual models (requires specific model and hardware)""" + """The published GPT-OSS checkpoint, quantized and dequantized, and saved back.""" - # These should be paths to real OpenAI MoE models for proper testing model_name = "openai/gpt-oss-20b" - input_text = "Once upon a time" - - # Expected outputs for generation tests - EXPECTED_OUTPUTS = set() - EXPECTED_OUTPUTS.add("Once upon a time, in a small town, there lived a young") - - def setUp(self): - gc.collect() - _empty_accelerator_cache() + EXPECTED_OUTPUTS = {"Once upon a time, in a small town, there lived a young"} + # what the published checkpoint keeps in full precision + MODULES_TO_NOT_CONVERT = ["model.layers.*.self_attn", "model.layers.*.mlp.router", "model.embed_tokens", "lm_head"] def tearDown(self): - gc.collect() - _empty_accelerator_cache() + cleanup(torch_device, gc_collect=True) def check_inference_correctness_quantized(self, model, tokenizer): - # Check that inference pass works on the model encoded_input = tokenizer(self.input_text, return_tensors="pt").to(model.device) - - # Set pad token if not set - if tokenizer.pad_token is None: - tokenizer.pad_token = tokenizer.eos_token - with torch.no_grad(): output_sequences = model.generate( **encoded_input, @@ -414,8 +157,6 @@ def check_inference_correctness_quantized(self, model, tokenizer): self.assertIn(generated_text, self.EXPECTED_OUTPUTS) def test_gpt_oss_model_loading_quantized_with_device_map(self): - """Test loading OpenAI MoE model with mxfp4 quantization and device_map""" - model = GptOssForCausalLM.from_pretrained( self.model_name, dtype=torch.bfloat16, @@ -427,39 +168,18 @@ def test_gpt_oss_model_loading_quantized_with_device_map(self): # ~39 GiB for the dequantized bfloat16 weights. @require_torch_accelerator_memory(memory=48) def test_gpt_oss_model_loading_dequantized_with_device_map(self): - """Test loading OpenAI MoE model with mxfp4 dequantization and device_map""" - - quantization_config = Mxfp4Config(dequantize=True) - - # Test that config is properly set up - self.assertTrue(quantization_config.dequantize) - model = GptOssForCausalLM.from_pretrained( self.model_name, - quantization_config=quantization_config, + quantization_config=Mxfp4Config(dequantize=True), dtype=torch.bfloat16, device_map="auto", ) tokenizer = AutoTokenizer.from_pretrained(self.model_name) self.check_inference_correctness_quantized(model, tokenizer) - def test_model_device_map_validation(self): - """Test device map validation""" - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) - quantizer.pre_quantized = False - - # Test with CPU in device map (CPU already support mxfp4) - quantizer.validate_environment(device_map={"": "cpu"}) - # Holds the packed (~13 GiB) and the dequantized (~39 GiB) model at the same time, by design. @require_torch_accelerator_memory(memory=56) def test_memory_footprint_comparison(self): - """Test memory footprint differences between quantized and unquantized models""" - - # Expected: quantized < dequantized < unquantized memory usage quantization_config = Mxfp4Config(dequantize=True) quantized_model = GptOssForCausalLM.from_pretrained( self.model_name, @@ -479,32 +199,26 @@ def test_memory_footprint_comparison(self): # Keeps the packed model alive while reloading the saved copy, dequantized (~13 + ~39 GiB). @require_torch_accelerator_memory(memory=56) def test_save_mxfp4(self): - """Test saving quantized OpenAI MoE model with device_map""" - model = GptOssForCausalLM.from_pretrained( self.model_name, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) tokenizer = AutoTokenizer.from_pretrained(self.model_name) with tempfile.TemporaryDirectory() as tmp: - # Save the model in mxfp4 format model.save_pretrained(tmp) - _empty_accelerator_cache() - gc.collect() - # test quantized model + cleanup(torch_device, gc_collect=True) loaded_model = GptOssForCausalLM.from_pretrained( tmp, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) - # test dequantized model loaded_model = GptOssForCausalLM.from_pretrained( tmp, quantization_config=Mxfp4Config(dequantize=True), - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) @@ -512,83 +226,148 @@ def test_save_mxfp4(self): # Quantizes from a bfloat16 checkpoint, so it holds the bfloat16 source and the mxfp4 result at once. @require_torch_accelerator_memory(memory=56) def test_save_mxfp4_non_quantized(self): - """Test saving dequantized OpenAI MoE model with mxfp4 quantization and device_map""" non_quantized_model_name = "hf-internal-testing/gpt-oss-20b-bf16" tokenizer = AutoTokenizer.from_pretrained(non_quantized_model_name) loaded_model = GptOssForCausalLM.from_pretrained( non_quantized_model_name, - quantization_config=Mxfp4Config(), - torch_dtype=torch.bfloat16, + quantization_config=Mxfp4Config(modules_to_not_convert=self.MODULES_TO_NOT_CONVERT), + dtype=torch.bfloat16, device_map="auto", ) - # save the quantized model with tempfile.TemporaryDirectory() as tmp: loaded_model.save_pretrained(tmp) - _empty_accelerator_cache() - gc.collect() - # load it back to check with everything works as expected + cleanup(torch_device, gc_collect=True) loaded_model = GptOssForCausalLM.from_pretrained( tmp, - torch_dtype=torch.bfloat16, + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) loaded_model = GptOssForCausalLM.from_pretrained( tmp, - quantization_config=Mxfp4Config(dequantized=True), - torch_dtype=torch.bfloat16, + quantization_config=Mxfp4Config(dequantize=True), + dtype=torch.bfloat16, device_map="auto", ) self.check_inference_correctness_quantized(loaded_model, tokenizer) - def test_compute_module_sizes(self): - r""" - Test if we compute the right module sizes needed to generate the device map. - Also test if we get the right values for `total_byte_count` in `caching_allocator_warmup`. - """ - from transformers import AutoConfig, AutoModelForCausalLM - from transformers.integrations import Mxfp4GptOssExperts - from transformers.integrations.accelerate import compute_module_sizes - from transformers.modeling_utils import expand_device_map, get_total_byte_count - from transformers.quantizers import AutoHfQuantizer - - # we need to preprocess the model like that because device_map calculation happens before we load the weights inside the model. - # For normal wieghts, it's fine but for quantized weights, the tensors dtype might change during loading. - with torch.device("meta"): - config = AutoConfig.from_pretrained(self.model_name) - model = AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16) - model_size, _ = compute_module_sizes(model, only_modules=False) - - expected_keys = [name for name, _ in model.named_parameters()] + [ - name for name, _ in model.named_buffers() - ] - expanded_device_map = expand_device_map({"": torch_device}, expected_keys) - total_byte_count = list(get_total_byte_count(model, expanded_device_map).values())[0] - - # testing prequantized = False should be enough, the shape should be the same whether it is pre-quantized or not - hf_quantizer = AutoHfQuantizer.from_config(Mxfp4Config(), pre_quantized=False) - hf_quantizer.preprocess_model(model=model, config=model.config) - quantized_model_size, _ = compute_module_sizes(model, hf_quantizer, only_modules=False) - - expected_keys = [name for name, _ in model.named_parameters()] + [ - name for name, _ in model.named_buffers() - ] - expanded_device_map = expand_device_map({"": torch_device}, expected_keys) - quantized_total_byte_count = list(get_total_byte_count(model, expanded_device_map, hf_quantizer).values())[ - 0 - ] - for name, module in model.named_modules(): - if isinstance(module, Mxfp4GptOssExperts): - # from 16 bits to 4 bits - assert int(model_size[f"{name}.gate_up_proj"] // 4) == int( - quantized_model_size[f"{name}.gate_up_proj"] - ) - assert int(model_size[f"{name}.down_proj"] // 4) == int(quantized_model_size[f"{name}.down_proj"]) - # check that we get the same value, as we use `compute_module_sizes` in `get_total_byte_count` - assert total_byte_count == model_size[""] - assert quantized_total_byte_count == quantized_model_size[""] +@require_torch +class Mxfp4ChainTest(_ChainAssertions, unittest.TestCase): + """MXFP4 as GPT-OSS ships it: blocks (E, N, K/32, 16) uint8 low-nibble-first E2M1 + biased-127 + exponent scales, gate|up rows INTERLEAVED.""" + + FP4_VALUES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0] + + def _reference_dequant(self, blocks, scales): + # the mxfp4.py reference semantics: lut[lo], lut[hi] interleaved along K, x 2^(scale-127) + lut = torch.tensor(self.FP4_VALUES) + lo = lut[(blocks & 0xF).long()] + hi = lut[(blocks >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], -1) # (E, N, K/32, 32) + exp = (scales.long() - 127).unsqueeze(-1) + return (vals * torch.pow(torch.tensor(2.0), exp)).reshape(*blocks.shape[:2], -1) + + def test_packed_experts(self): + """Packed E2M1 per expert (two values per `int8` byte along K) with an E8M0 scale per 32, as + DeepSeek-V4 ships its experts, merged onto the stacked experts.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randint(-128, 128, (n, k // 2), dtype=torch.int8), + "weight_scale_inv": torch.randint(120, 130, (n, k // 32), dtype=torch.uint8).view( + torch.float8_e8m0fnu + ), + } + ) + model, info, saved = _load_and_save(checkpoint, Mxfp4Config()) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + self.assertHoldsInterleaved(experts, checkpoint, "weight", experts.gate_up_proj) + self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) + self.assertSavesBack(saved, checkpoint) + + def test_blocks_and_scales_convert_to_the_kernel_pair_and_back(self): + from transformers.integrations.finegrained.conversions import ( + FineGrainedPackedBlocks, + ) + + torch.manual_seed(0) + E, N, K = 2, 8, 64 # N = 2I interleaved gate|up rows + blocks = torch.randint(0, 256, (E, N, K // 32, 16), dtype=torch.uint8) + scales = torch.randint(110, 140, (E, N, K // 32), dtype=torch.uint8) + + weight = FineGrainedPackedBlocks().convert({"gate_up_proj_blocks$": blocks}, target_patterns=["gate_up_proj"])[ + "gate_up_proj" + ] + self.assertEqual((weight.dtype, weight.shape), (torch.int8, (E, N, K // 2))) + back = FineGrainedPackedBlocks().reverse_op.convert({"gate_up_proj": weight})["gate_up_proj"] + self.assertTrue(torch.equal(back, blocks)) + + # Exponent bytes remain uint8; only the packed weights need a layout conversion. + scale_inv = scales + + ref = self._reference_dequant(blocks, scales) # rows pass through interleaved + # dequantize the converted pair: packed E2M1 low-nibble-first x 2^(e8m0) + lut = torch.tensor(self.FP4_VALUES) + w_u8 = weight.view(torch.uint8) + lo = lut[(w_u8 & 0xF).long()] + hi = lut[(w_u8 >> 4).long()] + vals = torch.stack([lo, hi], dim=-1).reshape(E, N, -1) + exp = (scale_inv.view(torch.uint8).long() - 127).repeat_interleave(32, dim=-1) + got = vals * torch.pow(torch.tensor(2.0), exp) + torch.testing.assert_close(got, ref, rtol=0, atol=0) + + def test_blocks_dequantize_to_the_unquantized_experts_orientation(self): + """`dequantize=True` lands the blocks as bf16 in the `(E, hidden, 2I)` orientation the unquantized + GPT-OSS experts hold.""" + from transformers.core_model_loading import Transpose + from transformers.integrations.finegrained.conversions import FineGrainedDequantize, FineGrainedPackedBlocks + + torch.manual_seed(0) + num_experts, rows, hidden = 2, 8, 64 + blocks = torch.randint(0, 256, (num_experts, rows, hidden // 32, 16), dtype=torch.uint8) + scales = torch.randint(120, 134, (num_experts, rows, hidden // 32), dtype=torch.uint8) + reference = self._reference_dequant(blocks, scales).transpose(1, 2) + + sources = ["gate_up_proj_blocks$", "gate_up_proj_scales$"] + tensors = {sources[0]: blocks.clone(), sources[1]: scales.clone()} + for op in (FineGrainedPackedBlocks(None), FineGrainedDequantize(None), Transpose(1, 2)): + tensors = op.convert(tensors, source_patterns=sources, target_patterns=["gate_up_proj"]) + self.assertEqual(len(tensors), 1) # the scale is consumed, not passed down the chain + torch.testing.assert_close(next(iter(tensors.values())).float(), reference, rtol=0, atol=0) + + def test_blocks_and_scales_take_only_the_ops_their_layout_needs(self): + """GPT-OSS's blocks unpack to the packed weight and its scales only rename; the module's own layout ops + follow.""" + from transformers.core_model_loading import WeightConverter, WeightRenaming + from transformers.integrations.finegrained.conversions import FineGrainedDeserialize, FineGrainedPackedBlocks + + quantizer = _quantizer_for(Mxfp4Config()) + quantizer.pre_quantized = True + conversions = quantizer.update_weight_conversions([], _conversion_model(quantizer)) + blocks = next( + c + for c in conversions + if isinstance(c, WeightConverter) and c.compiled_sources.search("gate_up_proj_blocks") + ) + self.assertEqual([type(op) for op in blocks.operations], [FineGrainedPackedBlocks, FineGrainedDeserialize]) + self.assertTrue( + any( + isinstance(c, WeightRenaming) and c.compiled_sources.search("experts.gate_up_proj_scales") + for c in conversions + ) + ) + + +class Mxfp4Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = MXFP4 + + +class Mxfp4TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = MXFP4 + - # we should at least have 3 times memory reduction in total for this model - assert model_size[""] > quantized_model_size[""] * 3 +class Mxfp4LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = MXFP4 diff --git a/tests/quantization/mxfp8/__init__.py b/tests/quantization/mxfp8/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/mxfp8/test_mxfp8.py b/tests/quantization/mxfp8/test_mxfp8.py new file mode 100644 index 000000000000..b061fda2f63e --- /dev/null +++ b/tests/quantization/mxfp8/test_mxfp8.py @@ -0,0 +1,70 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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. +"""MXFP8 (`Mxfp8Config`): the config and the checkpoint layouts it loads and saves back.""" + +import unittest + +import torch + +from transformers.testing_utils import ( + TestCasePlus, + require_torch, +) +from transformers.utils.quantization_config import Mxfp8Config + +from ...integrations.finegrained.utils import ( + _ChainAssertions, + _load_and_save, + _per_expert, +) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin +from ..block_scaled_tester import ( + MXFP8, + BlockScaledFormatTesterMixin, +) + + +@require_torch +class Mxfp8ChainTest(_ChainAssertions, unittest.TestCase): + """MXFP8, as MiniMax-M3 ships it.""" + + def test_minimax_m3_e8m0_scale_bytes(self): + """MiniMax-M3-MXFP8 ships its E8M0 scales as raw `uint8` bytes: their bytes reach the kernels + unchanged, and saving preserves the runtime representation.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randn(n, k).to(torch.float8_e4m3fn), + "weight_scale_inv": torch.randint(120, 130, (n, k // 32), dtype=torch.uint8), + } + ) + model, info, saved = _load_and_save(checkpoint, Mxfp8Config()) + self.assertLoadsCleanly(info) + experts = model.model.layers[0].mlp.experts + self.assertEqual(experts.gate_up_proj_scale_inv.dtype, torch.uint8) + self.assertHoldsInterleaved(experts, checkpoint, "weight_scale_inv", experts.gate_up_proj_scale_inv) + self.assertSavesBack(saved, checkpoint) + + +class Mxfp8Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = MXFP8 + + +class Mxfp8TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = MXFP8 + + +class Mxfp8LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = MXFP8 diff --git a/tests/quantization/nvfp4/test_nvfp4.py b/tests/quantization/nvfp4/test_nvfp4.py index b903a68257dc..6c75f4e83e25 100644 --- a/tests/quantization/nvfp4/test_nvfp4.py +++ b/tests/quantization/nvfp4/test_nvfp4.py @@ -11,157 +11,358 @@ # 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. -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -from transformers import NVFP4Config -from transformers.integrations.nvfp4 import NVFP4Linear, NVFP4Quantize, replace_with_nvfp4_linear -from transformers.quantizers.quantizer_nvfp4 import NVFP4HfQuantizer -from transformers.utils import is_torch_available -from transformers.utils.quantization_config import QuantizationMethod - - -if is_torch_available(): - import torch - from torch import nn - - -class _PackedWeight: - def __init__(self, **kwargs): - self.__dict__.update(kwargs) +"""NVFP4 (`NVFP4Config`): the two-level global-scale ops, the modelopt and compressed-tensors checkpoint layouts +it loads and saves back, and how those exports read.""" - -def _fake_kernel(): - def pack(weight, device): - out_features, in_features = weight.shape - return SimpleNamespace( - qweight=torch.empty(out_features, in_features // 2, dtype=torch.uint8, device=device), - sf=torch.empty(out_features, in_features // 64, dtype=torch.int32, device=device), - sf_rowmajor=torch.empty(out_features, in_features // 16, dtype=torch.uint8, device=device), - global_scale=torch.ones(1, dtype=torch.float32, device=device), +import os +import unittest +from unittest import mock + +import torch + +from transformers.integrations.finegrained import FineGrainedExperts +from transformers.quantizers.auto import AutoQuantizationConfig +from transformers.testing_utils import ( + TestCasePlus, + require_torch, +) +from transformers.utils.quantization_config import NVFP4Config + +from ...integrations.finegrained.utils import ( + E, + _Cfg, + _ChainAssertions, + _load_and_save, + _per_expert, + _quantizer_for, + _targets, +) +from ..block_scaled_checkpoint_tester import TinyCheckpointTesterMixin +from ..block_scaled_parallel_tester import LoadPathEquivalenceTesterMixin +from ..block_scaled_tester import ( + NVFP4, + BlockScaledFormatTesterMixin, +) + + +@require_torch +class Nvfp4GlobalScaleOpsTest(unittest.TestCase): + """The NVFP4 second-level globals and calibrated input scales, in the layout the kernels index: one per + expert per projection, the gate|up pair folded.""" + + def test_the_gate_up_fold_survives_the_loader_s_list_wrapping(self): + """The loader hands a converter its tensors in a list. Unwrapped, a fused `(E, 2)` pair + reads as `(1, 2E)`, the per-half fold does not fire, and the module gets `2E` globals + where the kernels assert one per expert.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_activation_scale", + ] + experts = 8 + for wrap in (False, True): + pair = torch.rand(experts, 2) + 1 + src = { + "experts.gate_up_proj_weight_scale_2": [pair] if wrap else pair, + "experts.down_proj_weight_scale_2": torch.rand(experts) + 1, + "experts.down_proj_input_scale": torch.rand(experts) + 1, + } + out = FineGrainedWeightGlobals(None).convert(src, target_patterns=targets, model=None) + with self.subTest(list_wrapped=wrap): + self.assertEqual( + out["experts.gate_up_proj_weight_global_scale"].shape, + torch.Size([experts]), + "the gate|up halves were not folded to one global per expert", + ) + + def test_one_converter_merges_a_layer_s_globals(self): + """modelopt calibrates the gate|up halves separately. One converter owns every global of + a layer, because merging them to the one-per-expert the kernels take moves the up half's + onto the down projection: SwiGLU is linear in the up half, so the stack keeps the gate's + global and the down's weight global scales the expert output back. The down's calibrated + input scale moves the other way, keeping the requantized intermediate on the range the + checkpoint calibrated.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + gate, up = torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0]) + down, down_input = torch.tensor([5.0, 6.0]), torch.tensor([0.5, 2.0]) + targets = [ + "mlp.experts.gate_up_proj_weight_global_scale", + "mlp.experts.down_proj_weight_global_scale", + "mlp.experts.down_proj_activation_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "mlp.experts.*.up_proj.weight_scale_2": up, + "mlp.experts.*.gate_proj.weight_scale_2": gate, + "mlp.experts.*.down_proj.weight_scale_2": down, + "mlp.experts.*.down_proj.input_scale": down_input, + }, + target_patterns=targets, ) - - return SimpleNamespace( - PackedWeight=_PackedWeight, - gemm=lambda packed_weight, input: input.new_zeros((*input.shape[:-1], packed_weight.n)), - pack=pack, - swizzled_sf_shape=lambda rows, columns: (rows, columns // 64), - ) - - -@unittest.skipUnless(is_torch_available(), "test requires PyTorch") -class NVFP4ConfigTest(unittest.TestCase): - def test_config_uses_quantization_method_enum(self): - config = NVFP4Config(modules_to_not_convert=["lm_head"]) - - self.assertEqual(config.quant_method, QuantizationMethod.NVFP4) - self.assertEqual(config.to_dict()["modules_to_not_convert"], ["lm_head"]) - - -@unittest.skipUnless(is_torch_available(), "test requires PyTorch") -class NVFP4IntegrationTest(unittest.TestCase): - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_replaces_only_eligible_linears(self, _): - model = nn.ModuleDict( + ratio = up / gate + torch.testing.assert_close(out[targets[0]], gate) + torch.testing.assert_close(out[targets[1]], down * ratio) + torch.testing.assert_close(out[targets[2]], down_input / ratio) + + def test_globals_with_one_calibrated_projection_pass_through(self): + """A checkpoint that calibrates the stack as one matrix (its halves agree, or it ships a + single global) has nothing to merge, so every global reaches its module unchanged.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + "experts.down_proj_activation_scale", + ] + out = FineGrainedWeightGlobals().convert( { - "convert": nn.Linear(64, 32, bias=False), - "skip": nn.Linear(64, 32, bias=False), - "biased": nn.Linear(64, 32, bias=True), - "unaligned": nn.Linear(63, 32, bias=False), - } + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0], [2.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([3.0, 4.0]), + "experts.down_proj_input_scale": torch.tensor([0.5, 0.25]), + }, + target_patterns=targets, ) - - replace_with_nvfp4_linear(model, modules_to_not_convert=["skip"]) - - self.assertIsInstance(model["convert"], NVFP4Linear) - self.assertIsInstance(model["skip"], nn.Linear) - self.assertIsInstance(model["biased"], nn.Linear) - self.assertIsInstance(model["unaligned"], nn.Linear) - self.assertEqual(model["convert"].weight_sf.dtype, torch.int32) - self.assertEqual(model["convert"].weight_sf_rowmajor.dtype, torch.uint8) - self.assertEqual(model["convert"].weight_global_scale.dtype, torch.float32) - - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_linear_forward_uses_normal_function_entry_point(self, _): - layer = NVFP4Linear(64, 32) - output = layer(torch.randn(2, 3, 64)) - - self.assertEqual(output.shape, (2, 3, 32)) - - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_packed_tensors_are_registered_in_state_dict(self, _): - layer = NVFP4Linear(64, 32) - - for buffer_name in ("weight", "weight_sf", "weight_sf_rowmajor", "weight_global_scale"): - self.assertIsInstance(getattr(layer, buffer_name), nn.Buffer) - self.assertEqual( - set(layer.state_dict()), - {"weight", "weight_sf", "weight_sf_rowmajor", "weight_global_scale"}, + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([3.0, 4.0])) + torch.testing.assert_close(out[targets[2]], torch.tensor([0.5, 0.25])) + + def test_the_fused_layout_ships_both_halves_in_one_tensor(self): + """A stacked `(E, 2)` `weight_scale_2` is the same pair as the per-expert layout's two + keys, and merges the same way.""" + from transformers.integrations.finegrained.conversions import FineGrainedWeightGlobals + + targets = [ + "experts.gate_up_proj_weight_global_scale", + "experts.down_proj_weight_global_scale", + ] + out = FineGrainedWeightGlobals().convert( + { + "experts.gate_up_proj_weight_scale_2": torch.tensor([[1.0, 3.0], [2.0, 4.0]]), + "experts.down_proj_weight_scale_2": torch.tensor([5.0, 6.0]), + }, + target_patterns=targets, ) - self.assertNotIn("nvfp4_kernel", layer.state_dict()) - - @patch("transformers.integrations.nvfp4.load_nvfp4_kernel", side_effect=_fake_kernel) - def test_conversion_produces_all_registered_buffers(self, _): - conversion = NVFP4Quantize(torch.device("cpu")) + torch.testing.assert_close(out[targets[0]], torch.tensor([1.0, 2.0])) + torch.testing.assert_close(out[targets[1]], torch.tensor([5.0 * 3.0, 6.0 * 2.0])) - output = conversion.convert({"model.proj.weight": torch.randn(32, 64)}) + def test_activation_global_is_per_expert_for_both_projections(self): + """Each expert keeps its calibrated `input_scale`, the gate|up pair reducing to their max.""" + from transformers.integrations.finegrained.conversions import FineGrainedInputScales - self.assertEqual( - set(output), + up = FineGrainedInputScales().convert( { - "model.proj.weight", - "model.proj.weight_sf", - "model.proj.weight_sf_rowmajor", - "model.proj.weight_global_scale", + "mlp.experts.*.gate_proj.input_scale": torch.tensor([1.0, 3.0]), + "mlp.experts.*.up_proj.input_scale": torch.tensor([2.0, 1.0]), }, + full_layer_name="mlp.experts.gate_up_proj_activation_scale", + )["mlp.experts.gate_up_proj_activation_scale"] + torch.testing.assert_close(up, torch.tensor([2.0, 3.0])) + down = FineGrainedInputScales().convert( + {"mlp.experts.*.down_proj.input_scale": torch.tensor([1.0, 4.0])}, + full_layer_name="mlp.experts.down_proj_activation_scale", + )["mlp.experts.down_proj_activation_scale"] + torch.testing.assert_close(down, torch.tensor([1.0, 4.0])) + + def test_shared_gate_up_input_global_follows_the_env_var(self): + """`TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL=global` shares the gate_up's activation + global as the max over experts; the down and the per-expert parameters keep the checkpoint's values.""" + from transformers.integrations.finegrained.layers import register_shared_gate_up_input_globals + + variable = "TRANSFORMERS_FINEGRAINED_SHARED_GATE_UP_ACTIVATION_GLOBAL" + for value, expected in ((None, None), ("global", 4.0)): + model = torch.nn.Module() + model.experts = experts = FineGrainedExperts(_Cfg(), quantization_config=NVFP4Config()) + with torch.no_grad(): + experts.gate_up_proj_activation_scale.copy_(torch.tensor([1.0, 4.0, 2.0, 3.0])) + with mock.patch.dict(os.environ): + os.environ.pop(variable, None) + if value is not None: + os.environ[variable] = value + register_shared_gate_up_input_globals(model) + gate_up_global = experts._projection_operands("gate_up_proj")["_input_global_scale"] + if expected is None: + self.assertIsNone(experts.shared_gate_up_input_global) + self.assertIs(gate_up_global, experts.gate_up_proj_activation_scale) + else: + torch.testing.assert_close(gate_up_global, torch.tensor([expected])) + down_global = experts._projection_operands("down_proj")["_input_global_scale"] + self.assertIs(down_global, experts.down_proj_activation_scale) + self.assertEqual(experts.gate_up_proj_activation_scale.numel(), 4) + with mock.patch.dict(os.environ, {variable: "node"}), self.assertRaises(ValueError): + register_shared_gate_up_input_globals(model) + + +@require_torch +class Nvfp4ChainTest(_ChainAssertions, unittest.TestCase): + """NVFP4 as modelopt ships it, per expert (GLM-5.2-NVFP4), stacked per layer (vLLM) and on dense + linears (modelopt's default export). `MergeModulelist` is load-bearing beyond the merge itself: + `core_model_loading` stamps a source with its expert index ONLY when one is present in the chain, + and expert parallelism selects experts by that index.""" + + def _modelopt_checkpoint(self): + """modelopt's per-expert layout, dense linear included: packed E2M1 bytes, E4M3 block scales + per 16, fp32 second-level `weight_scale_2` and calibrated `input_scale`. Gate and up share + their globals, as modelopt's linear fusion exports them, so a save writes each back exactly.""" + torch.manual_seed(0) + checkpoint = _per_expert( + lambda n, k: { + "weight": torch.randint(0, 256, (n, k // 2), dtype=torch.uint8), + "weight_scale": (torch.rand(n, k // 16) + 0.5).to(torch.float8_e4m3fn), + "weight_scale_2": torch.rand(()) * 1e-3 + 1e-4, + "input_scale": torch.rand(()) * 1e-2 + 1e-3, + } ) - self.assertEqual(output["model.proj.weight"].shape, (32, 32)) - - -@unittest.skipUnless(is_torch_available(), "test requires PyTorch") -class NVFP4QuantizerValidationTest(unittest.TestCase): - def setUp(self): - self.quantizer = NVFP4HfQuantizer(NVFP4Config(), pre_quantized=False) - - def test_rejects_pre_quantized_checkpoints(self): - quantizer = NVFP4HfQuantizer(NVFP4Config(), pre_quantized=True) - - with self.assertRaisesRegex(ValueError, "pre-quantized NVFP4"): - quantizer.validate_environment(device_map={"": "cuda"}) - - @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) - @patch("torch.cuda.is_available", return_value=True) - @patch("torch.cuda.get_device_capability", return_value=(10, 0)) - def test_accepts_one_blackwell_cuda_device(self, _, __, ___): - self.quantizer.validate_environment(device_map={"": "cuda:1"}) - - self.assertEqual(self.quantizer.quantization_device, torch.device("cuda:1")) - - @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) - @patch("torch.cuda.is_available", return_value=True) - def test_rejects_multi_device_map_until_tp_metadata_is_supported(self, _, __): - with self.assertRaisesRegex(ValueError, "Tensor parallelism"): - self.quantizer.validate_environment(device_map={"model.layers.0": 0, "model.layers.1": 1}) - - def test_rejects_native_tp_until_scale_metadata_is_supported(self): - model = SimpleNamespace(config=SimpleNamespace(distributed_config=SimpleNamespace(tp_size=2))) - - with self.assertRaisesRegex(ValueError, "scale metadata does not have a sharding plan"): - self.quantizer.update_parallel_plan(model, {}) - - def test_update_parallel_plan_keeps_non_tp_configs(self): - configs = ( - SimpleNamespace(), - SimpleNamespace(distributed_config=SimpleNamespace(tp_size=1)), + for expert in range(E): + prefix = f"model.layers.0.mlp.experts.{expert}" + for suffix in ("weight_scale_2", "input_scale"): + checkpoint[f"{prefix}.up_proj.{suffix}"] = checkpoint[f"{prefix}.gate_proj.{suffix}"].clone() + return checkpoint + + def test_modelopt_per_expert_and_dense(self): + checkpoint = self._modelopt_checkpoint() + model, info, saved = _load_and_save( + checkpoint, AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "NVFP4"}) ) - - for config in configs: - with self.subTest(config=config): - self.quantizer.update_parallel_plan(SimpleNamespace(config=config), {}) - - @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) - @patch("torch.cuda.is_available", return_value=True) - def test_rejects_cpu_offload(self, _, __): - with self.assertRaisesRegex(ValueError, "CPU or disk offload"): - self.quantizer.validate_environment(device_map={"model": "cuda", "lm_head": "cpu"}) + self.assertLoadsCleanly(info) + experts, q_proj = model.model.layers[0].mlp.experts, model.model.layers[0].self_attn.q_proj + self.assertEqual((experts.gate_up_proj.dtype, q_proj.weight.dtype), (torch.int8, torch.int8)) + self.assertHoldsInterleaved(experts, checkpoint, "weight", experts.gate_up_proj) + torch.testing.assert_close( + experts.gate_up_proj_weight_global_scale, + torch.stack([checkpoint[f"model.layers.0.mlp.experts.{e}.gate_proj.weight_scale_2"] for e in range(E)]), + ) + for slot, key in ( + ("weight_global_scale", "weight_scale_2"), + ("activation_scale", "input_scale"), + ): + torch.testing.assert_close(getattr(q_proj, slot), checkpoint[f"model.layers.0.self_attn.q_proj.{key}"]) + self.assertSavesBack(saved, checkpoint) + + def test_compressed_tensors_per_expert_and_dense(self): + """compressed-tensors names the packed values `weight_packed` and the activation global `input_global_scale`, + and stores both globals as the reciprocal of modelopt's: it loads into the same slots, and saves back as it + shipped. Globals that are powers of two keep the reciprocal exact.""" + modelopt = self._modelopt_checkpoint() + checkpoint = {} + for key, tensor in modelopt.items(): + stem, _, suffix = key.rpartition(".") + if suffix in ("weight_scale_2", "input_scale"): + modelopt[key] = tensor = torch.exp2(torch.round(torch.log2(tensor))) + name = "weight_global_scale" if suffix == "weight_scale_2" else "input_global_scale" + checkpoint[f"{stem}.{name}"] = 1 / tensor + else: + checkpoint[f"{stem}.weight_packed" if suffix == "weight" and tensor.dtype == torch.uint8 else key] = ( + tensor + ) + config = NVFP4Config(activation_format="nvfp4", checkpoint_format="compressed-tensors") + model, info, saved = _load_and_save(checkpoint, config) + self.assertLoadsCleanly(info) + experts, q_proj = model.model.layers[0].mlp.experts, model.model.layers[0].self_attn.q_proj + torch.testing.assert_close( + experts.gate_up_proj_weight_global_scale, + torch.stack([modelopt[f"model.layers.0.mlp.experts.{e}.gate_proj.weight_scale_2"] for e in range(E)]), + ) + torch.testing.assert_close( + q_proj.weight_global_scale, modelopt["model.layers.0.self_attn.q_proj.weight_scale_2"] + ) + torch.testing.assert_close(q_proj.activation_scale, modelopt["model.layers.0.self_attn.q_proj.input_scale"]) + self.assertSavesBack(saved, checkpoint) + + def test_modelopt_fused_layout(self): + """The vLLM layout: the per-expert modelopt tensors already stacked per layer.""" + per_expert = self._modelopt_checkpoint() + checkpoint = {k: v for k, v in per_expert.items() if ".experts." not in k} + for proj, halves in (("gate_up_proj", ("gate_proj", "up_proj")), ("down_proj", ("down_proj",))): + for suffix, stacked in (("", "weight"), ("_weight_scale", "weight_scale")): + checkpoint[f"model.layers.0.mlp.experts.{proj}{suffix}"] = torch.stack( + [ + torch.cat([per_expert[f"model.layers.0.mlp.experts.{e}.{h}.{stacked}"] for h in halves]) + for e in range(E) + ] + ) + # the globals stack one per expert; the gate|up pair agrees here, so its one value covers both + # halves (a pair calibrated apart folds, which the global-scale op tests cover) + for suffix in ("weight_scale_2", "input_scale"): + checkpoint[f"model.layers.0.mlp.experts.{proj}_{suffix}"] = torch.stack( + [per_expert[f"model.layers.0.mlp.experts.{e}.{halves[0]}.{suffix}"] for e in range(E)] + ) + model, info, saved = _load_and_save( + checkpoint, AutoQuantizationConfig.from_dict({"quant_method": "modelopt", "quant_algo": "NVFP4"}) + ) + self.assertLoadsCleanly(info) + self.assertEqual(model.model.layers[0].mlp.experts.gate_up_proj.dtype, torch.int8) + self.assertSavesBack(saved, checkpoint) + + def _modelopt_conversions(self, **cfg_kwargs): + # a pure modelopt export reads as this config + config = NVFP4Config(**{"activation_format": "nvfp4", **cfg_kwargs}) + return _quantizer_for(config).format_conversions(dequantize=False) + + def test_every_declared_target_is_a_slot_the_module_holds(self): + """A converter target that no module holds is a load failure, and the activation global is + exactly such a target under a weight-only run: the experts allocate one only when they + quantize activations. Both runs are checked against a real module's parameter names, since + the declarations are the only thing standing between the two.""" + from transformers.integrations.finegrained import FineGrainedLinear + + cfg = _Cfg() + for activation_format, holds_global in ((None, True), ("bf16", False)): + experts = FineGrainedExperts( + cfg, + quantization_config=NVFP4Config(activation_format=activation_format), + ) + linear = FineGrainedLinear(32, 32, quantization_config=NVFP4Config(activation_format=activation_format)) + for module in (experts, linear): + slots = {name for name, _ in module.named_parameters()} + self.assertEqual( + any("activation_scale" in slot for slot in slots), + holds_global, + f"{type(module).__name__} with activation_format={activation_format!r} disagrees about the slot", + ) + for c in self._modelopt_conversions(activation_format=activation_format): + for target in _targets(c): + # an expert target names the experts' slot; a dense rename's, the linear's + module, slot = ( + (experts, target.split("experts.")[-1]) + if "experts." in target + else (linear, target.removeprefix("\\1.")) + ) + self.assertIn( + slot.rstrip("$"), + {name for name, _ in module.named_parameters()}, + f"activation_format={activation_format!r} converts onto a missing slot", + ) + + def test_global_scale_converters_carry_an_expert_index(self): + # the per-expert layout only (`experts.*.`): a checkpoint that ships one already-stacked + # tensor per layer has no per-expert sources to stamp — it is sharded like the fused + # expert weight next to it + from transformers.core_model_loading import MergeModulelist, WeightConverter, WeightRenaming, rename_source_key + + conversions = self._modelopt_conversions() + renamings = [c for c in conversions if isinstance(c, WeightRenaming)] + converters = [c for c in conversions if isinstance(c, WeightConverter)] + _, pattern = rename_source_key("model.layers.0.mlp.experts.7.gate_proj.weight_scale_2", renamings, converters) + globals_converters = [c for c in converters if pattern in c.source_patterns] + self.assertTrue(globals_converters, "no global-scale converter found") + for conv in globals_converters: + self.assertTrue( + any(isinstance(op, MergeModulelist) for op in conv.operations), + f"{_targets(conv)} has no MergeModulelist, so expert parallelism cannot select " + "experts and every rank would collect all of them", + ) + + +class Nvfp4Test(BlockScaledFormatTesterMixin, unittest.TestCase): + spec = NVFP4 + + +class Nvfp4TinyCheckpointTest(TinyCheckpointTesterMixin, TestCasePlus): + spec = NVFP4 + + +class Nvfp4LoadPathEquivalenceTest(LoadPathEquivalenceTesterMixin, TestCasePlus): + spec = NVFP4 diff --git a/tests/tensor_parallel/test_tensor_parallel.py b/tests/tensor_parallel/test_tensor_parallel.py index 906c1200c197..924973aeab92 100644 --- a/tests/tensor_parallel/test_tensor_parallel.py +++ b/tests/tensor_parallel/test_tensor_parallel.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import warnings +from types import SimpleNamespace from unittest.mock import patch import torch @@ -225,13 +226,12 @@ def _get_local_shape(self, global_shape, placement, world_size, rank): return tuple(local_shape) def _make_dtensor_shard_op(self, mesh, placement, param_shape, local_shape): - op = object.__new__(DtensorShardOperation) - op.device_mesh = mesh - op.placements = (placement,) - op.param_ndim = len(param_shape) - op._axis0_offset = 0 - op._axis0_local_size = local_shape[0] - return op + param = SimpleNamespace(device_mesh=mesh, placements=(placement,), ndim=len(param_shape), shape=param_shape) + with patch( + "transformers.distributed.sharding_utils.compute_local_shape_and_global_offset", + return_value=(local_shape, (0,) * len(param_shape)), + ): + return DtensorShardOperation(param) def test_colwise_gather_output_rejects_indivisible_out_features(self): model = torch.nn.Module() diff --git a/tests/utils/test_core_model_loading.py b/tests/utils/test_core_model_loading.py index 09ccc3b8c5c3..c490b05395c5 100644 --- a/tests/utils/test_core_model_loading.py +++ b/tests/utils/test_core_model_loading.py @@ -13,7 +13,6 @@ # limitations under the License. import copy import unittest -from types import SimpleNamespace import torch import torch.nn as nn @@ -33,7 +32,6 @@ GroupWeightRename, LinearToConv3d, MergeModulelist, - PermuteForRope, PrefixChange, Split, VisionFuseAndPermuteForRope, @@ -47,13 +45,19 @@ spawn_materialize, ) from transformers.modeling_utils import LoadStateDictConfig -from transformers.utils.import_utils import is_triton_available from ..test_modeling_common import compare_state_dicts from .test_distributed_sharding_utils import FakeMesh, _make_dtensor_shard_op class TestWeightGlobMatching(unittest.TestCase): + def test_many_to_many_requires_an_operation_that_supports_it(self): + with self.assertRaisesRegex(ValueError, "one to many"): + WeightConverter(["a", "b"], ["c", "d"], operations=[Concatenate(dim=0)]) + for operation in (ErnieFuseAndSplitTextVisionExperts(), ErnieFuseAndSplitTextVisionExperts().reverse_op): + with self.subTest(operation=operation): + WeightConverter(["a", "b"], ["c", "d"], operations=[operation]) + def setUp(self): self.weight_globs_digits = [ "model.layers.*.mlp.gate_up_proj.weight", @@ -518,173 +522,6 @@ def test_moe_and_qkv_conversion_reversed(self): # Make sure both saved state_dict are identical self.assertTrue(compare_state_dicts(reversed_state_dict, state_dict)) - def test_qkv_chunk_rope_permute_with_fp8_quantization(self): - if is_triton_available(): - from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize - else: - self.skipTest("Fine-grained FP8 integration tests require Triton to be installed.") - n_heads = 2 - head_dim = 4 - in_dim = 4 - out_dim = n_heads * head_dim - block_size = (4, 4) - - class RopeProjector(nn.Module): - def __init__(self, *, with_scale: bool = False): - super().__init__() - self.weight = nn.Parameter(torch.zeros(out_dim, in_dim)) - if with_scale: - scale_shape = (out_dim // block_size[0], in_dim // block_size[1]) - self.weight_scale_inv = nn.Parameter(torch.ones(scale_shape)) - - class RopeSelfAttn(nn.Module): - def __init__(self): - super().__init__() - self.q_proj = RopeProjector(with_scale=True) - self.k_proj = RopeProjector() - self.v_proj = RopeProjector() - - class RopeLayer(nn.Module): - def __init__(self): - super().__init__() - self.self_attn = RopeSelfAttn() - - class RopeModel(PreTrainedModel): - base_model_prefix = "model" - - def __init__(self, config): - super().__init__(config) - self.layers = nn.ModuleList([RopeLayer()]) - self.post_init() - - config = PreTrainedConfig() - config.num_attention_heads = n_heads - model = RopeModel(config) - - raw_q = torch.tensor( - [ - [1.0, -1.0, 1.0, -1.0], - [0.5, -0.5, 0.5, -0.5], - [-1.0, 1.0, -1.0, 1.0], - [-0.5, 0.5, -0.5, 0.5], - [1.0, 1.0, -1.0, -1.0], - [0.5, 0.5, -0.5, -0.5], - [-1.0, -1.0, 1.0, 1.0], - [-0.5, -0.5, 0.5, 0.5], - ], - dtype=torch.float32, - ) - raw_k = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) - raw_v = torch.arange(out_dim * in_dim, dtype=torch.float32).reshape(out_dim, in_dim) + 100.0 - raw_qkv = torch.cat([raw_q, raw_k, raw_v], dim=0) - state_dict = {"model.layers.0.self_attn.qkv_proj.weight": raw_qkv.clone()} - - quantizer_cls = type( - "FineGrainedFP8HfQuantizer", - (), - { - "__init__": lambda self, bs=block_size: setattr( - self, "quantization_config", SimpleNamespace(weight_block_size=bs) - ), - "param_needs_quantization": lambda self, _model, param_name: param_name.endswith("q_proj.weight"), - "get_quantize_ops": lambda self: Fp8Quantize(self), - "pre_quantized": False, - }, - ) - quantizer = quantizer_cls() - - weight_mapping = [ - WeightConverter( - "self_attn.qkv_proj.weight", - [ - "self_attn.q_proj.weight", - "self_attn.k_proj.weight", - "self_attn.v_proj.weight", - ], - operations=[Chunk(dim=0), PermuteForRope()], - ) - ] - load_config = LoadStateDictConfig(weight_mapping=weight_mapping, hf_quantizer=quantizer) - loading_info, _ = convert_and_load_state_dict_in_model(model, state_dict, load_config) - - self.assertEqual(loading_info.missing_keys, set()) - self.assertEqual(loading_info.unexpected_keys, set()) - self.assertEqual(loading_info.mismatched_keys, set()) - self.assertEqual(loading_info.conversion_errors, {}) - - permute_op = PermuteForRope() - permute_op.config = model.config - expected_q = permute_op._apply(raw_q) - expected_k = permute_op._apply(raw_k) - expected_v = permute_op._apply(raw_v) - - model_state = model.state_dict() - self.assertFalse(torch.allclose(raw_k, expected_k)) - torch.testing.assert_close(model_state["model.layers.0.self_attn.k_proj.weight"], expected_k) - torch.testing.assert_close(model_state["model.layers.0.self_attn.v_proj.weight"], expected_v) - - q_weight_key = "model.layers.0.self_attn.q_proj.weight" - scale_key = "model.layers.0.self_attn.q_proj.weight_scale_inv" - self.assertIn(scale_key, model_state) - expected_dtype = torch.float8_e4m3fn if hasattr(torch, "float8_e4m3fn") else torch.int8 - self.assertEqual(model_state[q_weight_key].dtype, expected_dtype) - self.assertEqual(model_state[q_weight_key].shape, torch.Size((out_dim, in_dim))) - self.assertEqual(model_state[scale_key].dtype, torch.float32) - self.assertEqual( - model_state[scale_key].shape, - torch.Size((out_dim // block_size[0], in_dim // block_size[1])), - ) - - dequant = Fp8Dequantize(block_size=block_size) - dequantized_q = dequant.convert( - [model_state[q_weight_key], model_state[scale_key]], - context={"quantization_config": quantizer.quantization_config}, - ) - torch.testing.assert_close(dequantized_q, expected_q, rtol=1e-2, atol=1e-2) - - def test_fp8_ue8m0_quantize_dequantize_round_trip(self): - # ``scale_fmt="ue8m0"`` rounds weight_scale_inv to a power of two; the weight has to be - # quantized with that same rounded scale, or ``weight * weight_scale_inv`` at dequant - # disagrees with the scale the weight was divided by (a silent per-block error up to an octave). - from transformers.integrations.finegrained_fp8 import Fp8Dequantize, Fp8Quantize - - if not hasattr(torch, "float8_e8m0fnu"): - self.skipTest("ue8m0 storage requires torch.float8_e8m0fnu") - - quantizer = SimpleNamespace( - quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="ue8m0") - ) - torch.manual_seed(0) - weight = torch.randn(128, 128, dtype=torch.float32) - - quantized = Fp8Quantize(quantizer)._quantize_one("layer.weight", weight) - recovered = Fp8Dequantize(quantizer)._dequantize_one( - quantized["layer.weight"], quantized["layer.weight_scale_inv"], output_dtype=torch.float32 - ) - rel_err = ((recovered - weight).abs().sum() / weight.abs().sum()).item() - self.assertLess(rel_err, 5e-2) # fp8 round-trip is ~2e-2; a scale mismatch inflates it past 0.2 - - def test_fp8_float_scale_fmt_quantization_unchanged(self): - # The ue8m0 reordering must leave the default scale_fmt="float" path bit-identical to the - # original single-division formula (scales = _FP8_MAX / max_abs). Any divergence only shows - # on rare fp32-ULP block maxima, so scan many random 128x128 blocks rather than one. - from transformers.integrations.finegrained_fp8 import _FP8_DTYPE, _FP8_MAX, _FP8_MIN, Fp8Quantize - - quantizer = Fp8Quantize( - SimpleNamespace(quantization_config=SimpleNamespace(weight_block_size=(128, 128), scale_fmt="float")) - ) - for seed in range(100): - torch.manual_seed(seed) - weight = torch.randn(128, 128, dtype=torch.float32) - out = quantizer._quantize_one("layer.weight", weight) - scales = _FP8_MAX / weight.abs().amax() - ref_q = torch.clamp(weight * scales, min=_FP8_MIN, max=_FP8_MAX).to(_FP8_DTYPE) - ref_inv = (1.0 / scales).to(torch.float32).reshape(1, 1) - self.assertTrue(torch.equal(out["layer.weight"], ref_q), f"float weight diverged at seed {seed}") - self.assertTrue( - torch.equal(out["layer.weight_scale_inv"], ref_inv), f"float scale diverged at seed {seed}" - ) - def test_scoped_renaming_does_not_leak_to_sibling_or_parent(self): """scope_prefix gates a WeightRenaming to keys under one submodel only — neither the sibling submodel nor the parent's own keys must be affected. diff --git a/tests/utils/test_distributed_sharding_utils.py b/tests/utils/test_distributed_sharding_utils.py index 9be7b2ebd753..decbad6b90b5 100644 --- a/tests/utils/test_distributed_sharding_utils.py +++ b/tests/utils/test_distributed_sharding_utils.py @@ -12,6 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. import unittest +from types import SimpleNamespace +from unittest import mock import torch @@ -65,25 +67,25 @@ def __getitem__(self, name): ) -def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape): +def _make_dtensor_shard_op(mesh, placements, param_shape, local_shape, source_axis_order=None): """Build a DtensorShardOperation without requiring a real DTensor / distributed init. The axis-0 ownership cache is computed by mimicking ``compute_local_shape_and_global_offset`` for the leading dim only: locate the mesh dim that shards param dim 0 (if any) and use its local rank. """ - op = object.__new__(DtensorShardOperation) - op.device_mesh = mesh - op.placements = tuple(placements) - op.param_ndim = len(param_shape) - op._axis0_offset = 0 - op._axis0_local_size = local_shape[0] + offsets = [0] * len(param_shape) for mesh_dim, p in enumerate(placements): if hasattr(p, "dim") and (p.dim % len(param_shape)) == 0: sub = mesh[mesh.mesh_dim_names[mesh_dim]] if mesh.ndim > 1 else mesh - op._axis0_offset = sub.get_local_rank() * local_shape[0] + offsets[0] = sub.get_local_rank() * local_shape[0] break - return op + param = SimpleNamespace(device_mesh=mesh, placements=tuple(placements), ndim=len(param_shape), shape=param_shape) + with mock.patch( + "transformers.distributed.sharding_utils.compute_local_shape_and_global_offset", + return_value=(local_shape, offsets), + ): + return DtensorShardOperation(param, source_axis_order=source_axis_order) class TestDtensorShardOperation(unittest.TestCase): @@ -131,6 +133,21 @@ def test_no_shard_placements_returns_full_copy(self): op = _make_dtensor_shard_op(mesh, [Replicate()], param_shape=(4, 4), local_shape=(4, 4)) torch.testing.assert_close(op.shard_tensor(tensor), expected[rank], msg=f"rank {rank}") + def test_transposed_checkpoint_reads_the_source_sharding_axis(self): + source = torch.arange(4 * 8 * 12).reshape(4, 8, 12) + target = source.transpose(1, 2) + for placement in (Shard(0), Shard(2), _StridedShard(1, split_factor=2)): + for rank in range(2): + with self.subTest(placement=placement, rank=rank): + mesh = FakeMesh(shape=2, rank=rank) + local_shape = list(target.shape) + local_shape[placement.dim] //= 2 + expected = _make_dtensor_shard_op(mesh, [placement], target.shape, local_shape).shard_tensor( + target + ) + reader = _make_dtensor_shard_op(mesh, [placement], target.shape, local_shape, (0, 2, 1)) + torch.testing.assert_close(reader.shard_tensor(source).transpose(1, 2), expected) + def test_1D_shard(self): tensor = torch.arange(16).reshape(4, 4).float() expected = { @@ -294,6 +311,13 @@ def test_moe_2D_shard_with_negative_dim_indices(self): else: torch.testing.assert_close(shard, expected[rank], msg=f"rank {rank}") + def test_0dim_per_expert_follows_expert_ownership(self): + """Per-expert 0-dim scales are dropped on the ranks that do not own the expert.""" + mesh = FakeMesh(2, rank=1) # owns experts [2, 4) + op = _make_dtensor_shard_op(mesh, [Shard(0)], param_shape=(4, 8), local_shape=(2, 8)) + self.assertEqual(op.shard_tensor(torch.tensor(3.5), tensor_idx=3).item(), 3.5) + self.assertIsNone(op.shard_tensor(torch.tensor(3.5), tensor_idx=1)) + def test_compute_strided_slice(self): # Direct tests for _compute_strided_slice(intervals, rank, world_size, split_factor). # Keys: (input_interval, rank, world_size, split_factor) -> expected output list. diff --git a/tests/utils/test_import_utils.py b/tests/utils/test_import_utils.py index 531ffab5a5c2..53d0508a81ff 100644 --- a/tests/utils/test_import_utils.py +++ b/tests/utils/test_import_utils.py @@ -392,9 +392,11 @@ def _compile_constant_helpers(): @parameterized.expand(_compile_constant_helpers()) def test_availability_helpers_are_compile_safe(helper_name: str, args: tuple): """ - These helpers get called from inside `torch.compile`d regions — e.g. `is_dtensor`, which every MoE - kernel integration reaches through `to_local`. Each carries `@_make_compile_constant`, so dynamo evaluates - it once at trace time and never enters the body; this checks the marker actually takes effect. + These helpers get called from inside `torch.compile`d regions — e.g. `is_deepgemm_loadable`, which + `prefers_deepgemm_linear` reaches from the traced `fp8_linear` dispatch gate. Each carries + `@_make_compile_constant`, so dynamo evaluates it once at trace time and never enters the body + (whose probes reach `importlib.metadata`, which dynamo cannot follow); this checks the marker + actually takes effect. Folding rather than keeping the bodies traceable is deliberate. Most bottom out in `_is_package_available`, whose `importlib.metadata` lookup dynamo cannot follow — and follows