From e1275afe205c78f67fa28bdb7a38ef4c503b9947 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 09:56:18 +0900 Subject: [PATCH 01/27] first cleanup --- .../integrations/finegrained_fp8.py | 11 ----- .../quantizers/quantizer_finegrained_fp8.py | 49 +++++-------------- 2 files changed, 12 insertions(+), 48 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 084b17733b16..f00d024b36c8 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -607,18 +607,7 @@ def fp8_grouped_mm_experts_forward( class FP8Experts(nn.Module): - # 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 - - # Per-`_experts_implementation` rewrite of parallel-layer kinds in the TP/EP plan. - # The plan dicts store `{module-path-pattern: parallel-layer-kind}`; this maps an - # old kind to a new kind, and the quantizer rewrites every plan VALUE that matches. - # The default `MoeTensorParalellExperts` kind is impl-agnostic; some impls need a - # distinct TP layer (e.g. megamoe needs no gradient-sync hooks and an EP - # `process_group` injection). Declared here so the quantizer doesn't have to know - # about impl-specific TP needs — extend this dict when adding new impls. _impl_tp_layer_overrides: dict[str, dict[str, str]] = { "deepgemm_megamoe": { "moe_tp_experts": "megamoe_experts", diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index b0b4d24c74b3..1f9adf120436 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -165,47 +165,22 @@ def _process_model_after_weight_loading(self, model, **kwargs): _disable_deepgemm_on_multi_device(model) return model - def update_tp_plan(self, config): - if "Qwen3" in config.__class__.__name__: - text_plan = { - "layers.*.self_attn.q_proj.weight": "colwise", - "layers.*.self_attn.q_proj.weight_scale_inv": "colwise", - "layers.*.self_attn.k_proj.weight": "colwise", - "layers.*.self_attn.k_proj.weight_scale_inv": "colwise", - "layers.*.self_attn.v_proj.weight": "colwise", - "layers.*.self_attn.v_proj.weight_scale_inv": "colwise", - "layers.*.self_attn.o_proj.weight": "rowwise", - "layers.*.self_attn.o_proj.weight_scale_inv": "rowwise", - "layers.*.mlp.gate_proj.weight": "colwise", - "layers.*.mlp.gate_proj.weight_scale_inv": "colwise", - "layers.*.mlp.up_proj.weight": "colwise", - "layers.*.mlp.up_proj.weight_scale_inv": "colwise", - "layers.*.mlp.down_proj.weight": "rowwise", - "layers.*.mlp.down_proj.weight_scale_inv": "rowwise", - } - - config.base_model_tp_plan = text_plan + def update_plan(self, config, plan_name, overwrite): + plan = getattr(config, plan_name) + for key, style in plan.items().copy(): + plan[key] = overwrite[style] + # TODO: we need to define a global "style": "new_key" for all quantization plans + if style == "grouped_gemm": + plan.setdefault(f"{key}_scale_inv", style) + setattr(config, plan_name, plan) + def update_tp_plan(self, config): from ..integrations.finegrained_fp8 import FP8Experts impl = getattr(config, "_experts_implementation", None) - layer_overrides = FP8Experts._impl_tp_layer_overrides.get(impl, {}) - for plan_attr in ("base_model_tp_plan", "base_model_ep_plan"): - base_plan = getattr(config, plan_attr, None) or {} - # Per-impl rewrite of the experts parallel-layer kind. Applied LAST so it composes - # on top of any plan written above (e.g. the Qwen3 dense plan). Models carry the - # experts mapping under `base_model_tp_plan` and/or `base_model_ep_plan` — rewrite - # both. See `FP8Experts._impl_tp_layer_overrides`. - updated_plan = {k: layer_overrides.get(v, v) for k, v in base_plan.items()} - - # Expert scales must be sharded along with their corresponding weights. - for key, style in list(updated_plan.items()): - if style == "grouped_gemm": - updated_plan.setdefault(f"{key}_scale_inv", style) - - if updated_plan != base_plan: - setattr(config, plan_attr, updated_plan) - + overwrite = FP8Experts._impl_tp_layer_overrides.get(impl, {}) + self.update_plan(config, "tp_plan", overwrite) + self.update_plan(config, "ep_plan", overwrite) return config def is_serializable(self): From 82d660008d01f6da4e34641eb0d42f7340ff2445 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 11:36:38 +0900 Subject: [PATCH 02/27] more cleanup The various agent PRs most probably did not understand the core logic / idea behind the weight converter, this left a lot of method patching here and ther. 1. Deserialize: some method just need a cast, other need more work. The ConversionOps added as the "first" operation, updating the weights conversion (each one of them) are here for that. 2. TP: bloated code to just update a plan. 3. Patch: I have no idea how, but we accepted plans that had explicit overwrite. This is legacy --- src/transformers/conversion_mapping.py | 10 +- src/transformers/core_model_loading.py | 20 +- src/transformers/integrations/bitnet.py | 53 ++---- src/transformers/integrations/mxfp4.py | 92 ++++----- src/transformers/quantizers/base.py | 11 ++ .../quantizer_compressed_tensors.py | 16 -- .../quantizers/quantizer_fbgemm_fp8.py | 44 ----- .../quantizers/quantizer_finegrained_fp8.py | 123 +++--------- .../quantizers/quantizer_higgs.py | 36 +--- src/transformers/quantizers/quantizer_hqq.py | 180 ------------------ .../quantizers/quantizer_mxfp4.py | 81 +------- tests/quantization/hqq/test_hqq.py | 10 - tests/quantization/mxfp4/test_mxfp4.py | 9 +- 13 files changed, 98 insertions(+), 587 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 5e3cf0be76b1..01ba673a7dad 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -1001,11 +1001,7 @@ def _build_checkpoint_conversion_mapping(): source_patterns=r"\.moe\.router_bias$", target_patterns=".mlp.gate.e_score_correction_bias", ), - # A plain `WeightRenaming`, so FP8's `update_weight_conversions` won't auto-extend it - # with a `weight_scale_inv` sibling the way it does for the `gate_up_proj` - # `WeightConverter` below — the explicit rename right after this one covers that instead. WeightRenaming(source_patterns=r"\.moe\.down_proj\.weight$", target_patterns=".mlp.experts.down_proj"), - # `down_proj`'s FP8 dequant scale; not auto-added since the rename above isn't a `WeightConverter`. WeightRenaming( source_patterns=r"\.moe\.down_proj\.weight_scale_inv$", target_patterns=".mlp.experts.down_proj_scale_inv", @@ -1995,7 +1991,7 @@ def _build_checkpoint_conversion_mapping(): mapping["qwen4_exp_text"] += [ WeightConverter( source_patterns="ngram_embedding.shard_*.weight", - target_patterns="ngram_embedding.weight", + target_patterns="ngram_embedding.weight$", operations=[Concatenate(dim=0, num_shards_attribute="split_ngram_parts")], # The size of the embedding is ~95 GiB, so we cannot afford to perform the Cat on device, as it will need # a temporary memory buffer of the same size @@ -2218,10 +2214,6 @@ def get_model_conversion_mapping( if add_legacy: weight_conversions.extend(get_checkpoint_conversion_mapping("legacy")) - # Let the quantizer rewrite / augment the conversion pipeline. This is where the - # FP8 dequantizer (when `dequantize=True`) prepends a `Fp8Dequantize` op to - # every existing converter so that per-block scales are applied *before* any - # expert-merge / concat ops flatten the per-expert structure away. if hf_quantizer is not None: weight_conversions = hf_quantizer.update_weight_conversions(weight_conversions) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index cef38f7fca26..13d3203b8ef0 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -1712,28 +1712,14 @@ def convert_and_load_state_dict_in_model( mapping.quantization_operation = hf_quantizer.get_quantize_ops() _dtype = dtype - if ( - hf_quantizer - and hf_quantizer.pre_quantized - and ( - original_key != renamed_key - or not ( - tensor.get_dtype().startswith(("F", "BF")) - if hasattr(tensor, "get_dtype") - else tensor.is_floating_point() - ) - ) - ): - # if the key was renamed as it is not available in the state dict otherwise, it means that we are deserializing it, - # so we need to make sure to load the tensor with the same dtype from the checkpoint - # TODO: make the condition more srict for native fp8 model such as qwen2moe fp8 - _dtype = None - elif dtype_plan != {} and dtype_policy_alt.search(renamed_key): + if dtype_plan != {} and dtype_policy_alt.search(renamed_key): 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 + if hf_quantizer and hf_quantizer.pre_quantized: + _dtype = hf_quantizer.get_new_dtype(original_key, renamed_key, tensor, _dtype) # Per-expert sharding (EP) needs `tensor_idx` = the expert index so the # distributed op selects whole experts. The signal is a `MergeModulelist` diff --git a/src/transformers/integrations/bitnet.py b/src/transformers/integrations/bitnet.py index 0465c54b5377..ee2bd5e150ee 100644 --- a/src/transformers/integrations/bitnet.py +++ b/src/transformers/integrations/bitnet.py @@ -1,4 +1,5 @@ -from ..quantizers.quantizers_utils import should_convert_module +from ..core_model_loading import ConversionOps +from ..quantizers.quantizers_utils import get_module_from_name, should_convert_module from ..utils import is_torch_available, logging @@ -254,7 +255,7 @@ def __init__( use_rms_norm: bool = False, rms_norm_eps: float = 1e-6, ): - super().__init__(in_features, out_features, bias) + super().__init__(in_features, out_features, bias, device=device, dtype=dtype) self.online_quant = online_quant # Optional RMSNorm self.rms_norm = None @@ -264,18 +265,6 @@ def __init__( self.rms_norm = LlamaRMSNorm(in_features, eps=rms_norm_eps) if not online_quant: self.weight_scale = nn.Buffer(torch.ones((1), dtype=dtype, device=device)) - self._register_load_state_dict_pre_hook(self.load_hook) - - def load_hook( - self, - state_dict, - prefix, - *args, - **kwargs, - ): - if (prefix + "weight") in state_dict and state_dict[prefix + "weight"].dtype != self.weight.dtype: - state_dict[prefix + "weight"] = unpack_weights(state_dict[prefix + "weight"], dtype=self.weight.dtype) - return state_dict def forward(self, input): # Optional RMSNorm on activations prior to quantization. @@ -351,7 +340,7 @@ def replace_with_bitnet_linear(model, modules_to_not_convert: list[str] | None = return model -class BitNetDeserialize: +class BitNetDeserialize(ConversionOps): def __init__(self, hf_quantizer): self.hf_quantizer = hf_quantizer @@ -362,30 +351,10 @@ def convert( full_layer_name: str | None = None, **kwargs, ) -> dict[str, torch.Tensor]: - for key, value in input_dict.items(): - if isinstance(value, list): - input_dict[key] = value[0] - key_weight = "weight" - weight = input_dict.pop(key_weight) - from ..quantizers.quantizers_utils import get_module_from_name - - needs_unpacking = False - target_dtype = weight.dtype - if model is not None and full_layer_name is not None: - module, _ = get_module_from_name(model, full_layer_name) - if hasattr(module, "out_features") and hasattr(module, "in_features"): - # Packed: shape[0] * VALUES_PER_ITEM == out_features - # Unpacked: shape[0] == out_features - expected_out = module.out_features - actual_out = weight.shape[0] - if actual_out * VALUES_PER_ITEM == expected_out: - needs_unpacking = True - # Unpack into the module's compute dtype, not the packed uint8 dtype, - # otherwise the ternary weights stay uint8 and F.linear fails with a - # dtype mismatch (e.g. BFloat16 != unsigned char). - if hasattr(module, "weight_scale"): - target_dtype = module.weight_scale.dtype - if needs_unpacking: - weight_uint8 = weight.to(torch.uint8) - weight = unpack_weights(weight_uint8, dtype=target_dtype) - return {key_weight: weight} + weight = input_dict["weight"] + weight = weight[0] if isinstance(weight, list) else weight + module, _ = get_module_from_name(model, full_layer_name) + # packed checkpoints hold VALUES_PER_ITEM ternary rows per stored row + if isinstance(module, AutoBitLinear) and weight.shape[0] * VALUES_PER_ITEM == module.out_features: + weight = unpack_weights(weight.to(torch.uint8), dtype=module.weight.dtype) + return {"weight": weight} diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 02b17035fe91..8d9e04441dfc 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -22,6 +22,7 @@ from ..core_model_loading import ConversionOps, _IdentityOp from ..distributed.utils import _is_torch_distributed_initialized from ..quantizers.quantizers_utils import get_module_from_name, on_device, should_convert_module +from .hub_kernels import use_kernel_forward_from_hub logger = logging.get_logger(__name__) @@ -188,43 +189,13 @@ def convert( missing_keys: list[str] | None = None, **kwargs, ) -> dict[str, torch.Tensor]: - num_local_experts = getattr(model.config, "num_local_experts", 32) - hidden_size = getattr(model.config, "hidden_size", 2880) - - proj = "gate_up_proj" if "gate_up_proj" in full_layer_name else "down_proj" - - name = full_layer_name.rsplit("_", 1)[0] module, _ = get_module_from_name(model, full_layer_name) - state_dict = {} - if isinstance(module, Mxfp4GptOssExperts): - if "bias" in full_layer_name: - name = full_layer_name.replace("_blocks", "") - state_dict[name] = getattr(module, proj + "_bias") - return state_dict - if "gate_up_proj" in full_layer_name: - state_dict[f"{name}_blocks"] = ( - module.gate_up_proj.storage.layout.unswizzle_data(module.gate_up_proj.storage.data) - .transpose(-1, -2) - .reshape(num_local_experts, -1, 90, 16) - ) - state_dict[f"{name}_scales"] = ( - module.gate_up_proj_precision_config.weight_scale.storage.layout.unswizzle_data( - module.gate_up_proj_precision_config.weight_scale.storage.data - ).transpose(-1, -2) - ) - else: - state_dict[f"{name}_blocks"] = ( - module.down_proj.storage.layout.unswizzle_data(module.down_proj.storage.data) - .transpose(-1, -2) - .reshape(num_local_experts, hidden_size, 90, -1) - ) - state_dict[f"{name}_scales"] = ( - module.down_proj_precision_config.weight_scale.storage.layout.unswizzle_data( - module.down_proj_precision_config.weight_scale.storage.data - ).transpose(-1, -2) - ) - - return state_dict + name = full_layer_name.removesuffix("_blocks") + proj = name.rsplit(".", 1)[-1] + weight, scales = getattr(module, proj), getattr(module, f"{proj}_precision_config").weight_scale + scales = scales.storage.layout.unswizzle_data(scales.storage.data).transpose(-1, -2) + blocks = weight.storage.layout.unswizzle_data(weight.storage.data).transpose(-1, -2) + return {f"{name}_blocks": blocks.reshape(*scales.shape, -1), f"{name}_scales": scales} # Copied from GPT_OSS repo and vllm @@ -473,22 +444,29 @@ def topk(vals, k): return RoutingData(gate_scal, hist, n_local_experts, hit_experts, expt_data), gather_indx, scatter_indx -def mlp_forward(self, hidden_states): - if _is_torch_distributed_initialized() and hasattr(self, "_is_hooked"): - routing = routing_torch_dist - else: - routing = triton_kernels_hub.routing.routing +@use_kernel_forward_from_hub("MegaBlocksMoeMLP") +class Mxfp4GptOssMLP(nn.Module): + def __init__(self, router, experts): + super().__init__() + self.router = router + self.experts = experts + + def forward(self, hidden_states): + if _is_torch_distributed_initialized() and hasattr(self, "_is_hooked"): + routing = routing_torch_dist + else: + routing = triton_kernels_hub.routing.routing - batch_size = hidden_states.shape[0] - hidden_states = hidden_states.reshape(-1, self.router.hidden_dim) - router_logits = nn.functional.linear(hidden_states, self.router.weight, self.router.bias) + batch_size = hidden_states.shape[0] + hidden_states = hidden_states.reshape(-1, self.router.hidden_dim) + router_logits = nn.functional.linear(hidden_states, self.router.weight, self.router.bias) - with on_device(router_logits.device): - routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k) + with on_device(router_logits.device): + routing_data, gather_idx, scatter_idx = routing(router_logits, self.router.top_k) - routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx=scatter_idx) - routed_out = routed_out.reshape(batch_size, -1, self.router.hidden_dim) - return routed_out, router_logits + routed_out = self.experts(hidden_states, routing_data, gather_idx, scatter_idx=scatter_idx) + routed_out = routed_out.reshape(batch_size, -1, self.router.hidden_dim) + return routed_out, router_logits def dequantize_convertops(blocks, scales): @@ -566,17 +544,13 @@ def replace_with_mxfp4_linear(model, quantization_config=None, modules_to_not_co triton_kernels_hub = get_kernel("kernels-community/gpt-oss-triton-kernels", version=1) has_been_replaced = False - for module_name, module in model.named_modules(): - if not should_convert_module(module_name, modules_to_not_convert): - continue - if module.__class__.__name__ == "GptOssExperts" and not quantization_config.dequantize: + for module_name, module in list(model.named_modules()): + if module.__class__.__name__ == "GptOssMLP" and should_convert_module( + f"{module_name}.experts", modules_to_not_convert + ): with torch.device("meta"): - model.set_submodule(module_name, Mxfp4GptOssExperts(model.config)) - has_been_replaced = True - if module.__class__.__name__ == "GptOssMLP" and not quantization_config.dequantize: - from types import MethodType - - module.forward = MethodType(mlp_forward, module) + model.set_submodule(module_name, Mxfp4GptOssMLP(module.router, Mxfp4GptOssExperts(model.config))) + has_been_replaced = True if not has_been_replaced: logger.warning( diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index be9353bfbba5..52452207e0fd 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -123,6 +123,17 @@ def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: return param.element_size() + def get_new_dtype( + self, original_key: str, renamed_key: str, param, dtype: "torch.dtype | None" + ) -> "torch.dtype | None": + """dtype to materialize the pre-quantized checkpoint tensor `param` in, `None` keeps the checkpoint's. + `dtype` is the one the model declares for `renamed_key`. By default, quantized (non-float) and renamed + tensors are deserialized by conversion ops, so they keep the checkpoint dtype.""" + is_float = ( + param.get_dtype().startswith(("F", "BF")) if hasattr(param, "get_dtype") else param.is_floating_point() + ) + return dtype if is_float and original_key == renamed_key else None + def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization""" return max_memory diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 055fbb2f244c..dff536ecd3a4 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -63,7 +63,6 @@ def __init__(self, quantization_config: CompressedTensorsConfig, **kwargs): from compressed_tensors.compressors import ModelCompressor self.compressor = ModelCompressor.from_compression_config(quantization_config) - self.quantization_config = quantization_config def validate_environment(self, *args, **kwargs): if not is_compressed_tensors_available(): @@ -131,21 +130,6 @@ def _process_model_after_weight_loading(self, model, **kwargs): if self.quantization_config.dequantize and self.quantization_config.is_quantization_compressed: self.compressor.decompress_model(model=model) - # NOTE: TP plan override for compressed tensors removed - unsupported styles were used. - # TODO: Implement proper TP support for compressed tensors quantization - def update_tp_plan(self, config): - additional_plan = { - "layers.*.feed_forward.experts.*.gate_proj.weight": "colwise", - "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "colwise", - "layers.*.feed_forward.experts.*.up_proj.weight": "colwise", - "layers.*.feed_forward.experts.*.up_proj.weight_scale": "colwise", - "layers.*.feed_forward.experts.*.down_proj.weight": "rowwise", - } - if config.get_text_config() is not None and config.get_text_config().base_model_tp_plan is not None: - config.get_text_config().base_model_tp_plan.update(additional_plan) - - return config - @property def is_trainable(self): # The FP8 kernel path is inference-only; load with `dequantize=True` to fine-tune. diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 03c655fc8189..01721c1dd34f 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -149,50 +149,6 @@ def _process_model_after_weight_loading(self, model, **kwargs): m.input_scale_ub.fill_(self.quantization_config.activation_scale_ub) return model - def update_tp_plan(self, config): - if "Llama4" in config.__class__.__name__: - text_plan = { - # We are using a different tp plan with local_colwise and local_rowwise for the attention because fbgemm operations cannot be parallelized - # With local_colwise and local_rowwise, all the operations are done locally, and we add a gather operation to gather the results instead of - # using dtensors - "layers.*.self_attn.q_proj.weight": "colwise", - "layers.*.self_attn.q_proj.weight_scale": "colwise", - "layers.*.self_attn.k_proj.weight": "colwise", - "layers.*.self_attn.k_proj.weight_scale": "colwise", - "layers.*.self_attn.v_proj.weight": "colwise", - "layers.*.self_attn.v_proj.weight_scale": "colwise", - "layers.*.self_attn.o_proj.weight": "rowwise", - # We keep the same sequence_parallel plan for layernorms - "layers.*.input_layernorm.weight": "sequence_parallel", - "layers.*.post_attention_layernorm.weight": "sequence_parallel", - "norm.weight": "sequence_parallel", - # We keep the same local_colwise and local_rowwise plan for the feed forward shared expert - # We also add scales for the shared expert, for local_colwise the scale is also local_colwise - # For local_rowwise the scale is replicated, so we don't need to add it - "layers.*.feed_forward.shared_expert.gate_proj.weight": "colwise", - "layers.*.feed_forward.shared_expert.gate_proj.weight_scale": "colwise", - "layers.*.feed_forward.shared_expert.up_proj.weight": "colwise", - "layers.*.feed_forward.shared_expert.up_proj.weight_scale": "colwise", - "layers.*.feed_forward.shared_expert.down_proj.weight": "rowwise", - "layers.*.feed_forward.experts.*.gate_proj.weight": "colwise", - "layers.*.feed_forward.experts.*.gate_proj.weight_scale": "colwise", - "layers.*.feed_forward.experts.*.up_proj.weight": "colwise", - "layers.*.feed_forward.experts.*.up_proj.weight_scale": "colwise", - "layers.*.feed_forward.experts.*.down_proj.weight": "rowwise", - # For Fused implementation we use local_packed_rowwise for the gate_up_proj, and the same for the packed scales - # We use local_colwise for the down_proj, and the scales are replicated so we don't add them - "layers.*.feed_forward.experts.gate_up_proj": "packed_rowwise", - "layers.*.feed_forward.experts.gate_up_proj_scale": "packed_rowwise", - "layers.*.feed_forward.experts.down_proj": "colwise", - } - if config.get_text_config() is not None: - config.get_text_config().base_model_tp_plan = text_plan - else: - config.base_model_tp_plan = text_plan - return config - - return config - def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 1f9adf120436..d35c36741b2e 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -24,9 +24,6 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FineGrainedFP8Config" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") @@ -87,13 +84,6 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** return True return False - def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: - "Return the element size (in bytes) for `param_name`." - if self.param_needs_quantization(model, param_name): - # 8 bit, this is needed as when `pre_quantized`` is False, we don't set the dtype of the FP8Linear in order to correctly load the weights - return 1 - return super().param_element_size(model, param_name, param) - def _normalize_modules_to_not_convert(self, model: "PreTrainedModel"): """Rewrite the skip-list to the model's own module tree. For models that were already released, if they have a list of modules to not quantize @@ -139,25 +129,6 @@ def _process_model_before_weight_loading( ) def _process_model_after_weight_loading(self, model, **kwargs): - # dsv4-flash-base stores its (power-of-two) ue8m0 scales in a float32 container under - # `.scale`; those renamed keys keep the on-disk float32 dtype, so cast them to the UE8M0 - # dtype the kernels expect (exact, since the values are powers of two). Checkpoints that - # already ship the native float8 E8M0 dtype (e.g. dsv4-flash) are left untouched. - if self.quantization_config.scale_fmt == "ue8m0": - from ..integrations.finegrained_fp8 import _get_ue8m0_dtype - - ue8m0 = _get_ue8m0_dtype() - float32_scales = [ - name - for name, param in model.named_parameters() - if name.endswith("_scale_inv") and param.dtype == torch.float32 - ] - for name in float32_scales: - module_name, _, attr = name.rpartition(".") - module = model.get_submodule(module_name) - scale = getattr(module, attr) - setattr(module, attr, torch.nn.Parameter(scale.data.to(ue8m0), requires_grad=False)) - # Single-process multi-device is unsafe for DeepGEMM (its kernels are bound to one CUDA # context); route those models through Triton/grouped_mm instead. from ..integrations.finegrained_fp8 import _disable_deepgemm_on_multi_device @@ -165,23 +136,24 @@ def _process_model_after_weight_loading(self, model, **kwargs): _disable_deepgemm_on_multi_device(model) return model - def update_plan(self, config, plan_name, overwrite): - plan = getattr(config, plan_name) - for key, style in plan.items().copy(): - plan[key] = overwrite[style] - # TODO: we need to define a global "style": "new_key" for all quantization plans + def _update_plan(self, config, plan_name): + from ..integrations.finegrained_fp8 import FP8Experts + + overwrite = FP8Experts._impl_tp_layer_overrides.get(getattr(config, "_experts_implementation", None), {}) + plan = {key: overwrite.get(style, style) for key, style in (getattr(config, plan_name, None) or {}).items()} + # TODO: we need to define a global "style": "new_key" for all quantization plans + for key, style in list(plan.items()): if style == "grouped_gemm": plan.setdefault(f"{key}_scale_inv", style) - setattr(config, plan_name, plan) + if plan: + setattr(config, plan_name, plan) + return config def update_tp_plan(self, config): - from ..integrations.finegrained_fp8 import FP8Experts + return self._update_plan(config, "base_model_tp_plan") - impl = getattr(config, "_experts_implementation", None) - overwrite = FP8Experts._impl_tp_layer_overrides.get(impl, {}) - self.update_plan(config, "tp_plan", overwrite) - self.update_plan(config, "ep_plan", overwrite) - return config + def update_ep_plan(self, config): + return self._update_plan(config, "base_model_ep_plan") def is_serializable(self): return True @@ -199,6 +171,12 @@ def get_quantize_ops(self): return Fp8Quantize(self) + def get_new_dtype(self, original_key, renamed_key, param, dtype): + if self.quantization_config.dequantize: + return super().get_new_dtype(original_key, renamed_key, param, dtype) + # `FP8Linear` / `FP8Experts` declare the dtype of every param, e.g. ue8m0 scales stored as float32 + return dtype + def get_weight_conversions(self): from ..core_model_loading import WeightConverter from ..integrations.finegrained_fp8 import Fp8Dequantize @@ -216,66 +194,25 @@ def get_weight_conversions(self): return [] def update_weight_conversions(self, weight_conversions): - """When loading with ``dequantize=True``, attach an :class:`Fp8Dequantize` op to - every existing :class:`WeightConverter` so that per-block scales are folded into - the weight *before* any later merge/concat ops collapse the per-expert structure. - - For each model-supplied converter that has a ``.weight`` source, we: - 1. anchor the existing weight patterns with ``$`` so they don't accidentally - also match the ``.weight_scale_inv`` keys (the regex is searched, so the - unanchored prefix would match both, sending scales to the wrong bucket); - 2. add anchored ``*.weight_scale_inv`` sources next to each weight pattern so - the loader collects scale tensors alongside the weight tensors into the - *same* converter bucket (both keys rewrite to the same target); - 3. prepend a fresh :class:`Fp8Dequantize` op so dequant runs first, before - any merge/concat collapses the per-expert structure. - - The generic ``weight$ + weight_scale_inv → weight`` converter from - :meth:`get_weight_conversions` is still appended at the end as a fallback for - plain ``nn.Linear`` weights with no model-specific converter. - """ + # A key goes through at most one `WeightConverter`, so per-expert weights/scales merged by a model converter + # can only be dequantized by an op prepended to that converter. from ..core_model_loading import WeightConverter, WeightRenaming from ..integrations.finegrained_fp8 import Fp8Dequantize - # `*.scale` → `*.weight_scale_inv`. Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) - # ship per-block scales under `.scale`; the model expects `.weight_scale_inv`. - # Lives here (not in each model's `conversion_mapping`) so non-FP8 round-trips - # don't see a stray rule. Needed in both dequantize modes — `dequantize=False` - # loads scales as parameters, `dequantize=True` feeds them into `Fp8Dequantize`. + # Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) ship per-block scales under `.scale` scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") - weight_conversions = [scale_rename] + list(weight_conversions) - - # Some checkpoints shard weights (e.g. Qwen4-Exp's `ngram_embedding`). Since WeightConverter targets become source - # patterns when saving, anchor them to avoid matching the corresponding scale parameters. - for conv in weight_conversions: - if isinstance(conv, WeightConverter): - conv._original_target_patterns = [ - f"{p}$" if p.endswith(".weight") else p for p in conv._original_target_patterns - ] - - if not (self.pre_quantized and self.quantization_config.dequantize): - return weight_conversions + self.get_weight_conversions() - - updated: list = [] + updated = [scale_rename] + dequantize = self.pre_quantized and self.quantization_config.dequantize for conv in weight_conversions: - # Only WeightConverter has ``.operations`` to extend with the dequant op; - # WeightRenaming (e.g. the ``scale_rename`` we prepended) just passes through. - if not isinstance(conv, WeightConverter): - updated.append(conv) - continue weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] - if weight_sources: - anchored_weight = [p + "$" for p in weight_sources] - scale_sources = [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] + if isinstance(conv, WeightConverter) and dequantize and weight_sources: other = [p for p in conv.source_patterns if not p.endswith(".weight")] - new_sources = anchored_weight + scale_sources + other - new_ops = [Fp8Dequantize(self)] + list(conv.operations) conv = WeightConverter( - source_patterns=new_sources, + source_patterns=[p + "$" for p in weight_sources] + + [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] + + other, target_patterns=conv._original_target_patterns, - operations=new_ops, + operations=[Fp8Dequantize(self)] + conv.operations, ) updated.append(conv) - # Generic fallback for plain ``nn.Linear`` weights with no model-specific converter. - updated.extend(self.get_weight_conversions()) - return updated + return updated + self.get_weight_conversions() diff --git a/src/transformers/quantizers/quantizer_higgs.py b/src/transformers/quantizers/quantizer_higgs.py index 78d168aab59c..4667c9304437 100644 --- a/src/transformers/quantizers/quantizer_higgs.py +++ b/src/transformers/quantizers/quantizer_higgs.py @@ -78,40 +78,6 @@ def update_dtype(self, dtype: "torch.dtype") -> "torch.dtype": return dtype - # TODO: to remove - # Kept here in case we see some interest in adding support for it - # def create_quantized_param( - # self, - # model: "PreTrainedModel", - # param_value: "torch.Tensor", - # param_name: str, - # target_device: "torch.device", - # **kwargs, - # ): - # from ..integrations import quantize_with_higgs - - # flute_dict = quantize_with_higgs( - # param_value.to(target_device), - # self.quantization_config.bits, - # self.quantization_config.p, - # self.quantization_config.group_size, - # self.quantization_config.hadamard_size, - # ) - # del param_value - - # module, _ = get_module_from_name(model, param_name) - # module_name = ".".join(param_name.split(".")[:-1]) - # for key, value in flute_dict.items(): - # if key in module._parameters: - # module._parameters[key] = torch.nn.Parameter(value, requires_grad=False) - # elif key in module._buffers: - # module._buffers[key] = torch.nn.Buffer(value) - # elif key == "tune_metadata": - # module.tune_metadata = value - # self.quantization_config.tune_metadata[module_name] = value.to_dict() - # else: - # raise ValueError(f"Unexpected key {key} in module {module}") - def _process_model_before_weight_loading( self, model: "PreTrainedModel", @@ -171,7 +137,7 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** else: return False - def _dequantize(self, model): + def _dequantize(self, model, dtype=None): from ..integrations import dequantize_higgs model = dequantize_higgs(model) diff --git a/src/transformers/quantizers/quantizer_hqq.py b/src/transformers/quantizers/quantizer_hqq.py index 05dce3d996a0..bc3a86208fc2 100755 --- a/src/transformers/quantizers/quantizer_hqq.py +++ b/src/transformers/quantizers/quantizer_hqq.py @@ -28,18 +28,6 @@ if is_torch_available(): import torch -if is_hqq_available(): - from hqq.core.quantize import HQQLinear - - # This is a compatibility hack. HQQ-quantized linear layers do not have a `weight` attribute, - # but some models attempt to access `weight.dtype` during the forward pass. To prevent runtime errors, - # we patch HQQLinear with a dummy `weight` property that returns an empty tensor with the correct dtype and device. - @property - def weight(self): - return torch.empty(0, dtype=self.compute_dtype, device=self.device) - - HQQLinear.weight = weight - logger = logging.get_logger(__name__) @@ -58,19 +46,8 @@ def __init__(self, quantization_config, **kwargs): "A valid HQQ version (>=0.2.1) is not available. Please follow the instructions to install it: `https://github.com/mobiusml/hqq/`." ) super().__init__(quantization_config, **kwargs) - self.dtype = None - self.using_multi_gpu = False - # Keys that are serialized specifically by hqq - self.hqq_keys = HQQLinear(None, None).state_dict_keys() - {"bias"} def validate_environment(self, *args, **kwargs): - if self.dtype is None: - if "dtype" in kwargs: - self.dtype = kwargs["dtype"] - else: - self.dtype = torch.float32 - logger.info("Setting dtype to torch.float32 as the default value since it was not specified.") - device_map = kwargs.get("device_map") if isinstance(device_map, dict): if "cpu" in device_map.values() or "disk" in device_map.values(): @@ -78,70 +55,6 @@ def validate_environment(self, *args, **kwargs): "You are attempting to use an HQQ model with a device_map that contains a CPU or disk device." " This is not supported. Please remove the CPU or disk device from the device_map." ) - else: - self.using_multi_gpu = len(set(device_map.values())) > 1 - - # TODO: to remove - # Kept here in case we see some interest in adding support for it - # # Adds missing keys for HQQLinear modules that are loaded but the model with initialized with torch.nn.Linear - # def update_expected_keys( - # self, model: "PreTrainedModel", expected_keys: list[str], loaded_keys: list[str] - # ) -> list[str]: - # if not self.pre_quantized: - # return expected_keys - - # # Collects all quantizable (linear) layers - # def _find_hqq_quantizable_layers(model, layers): - # for name, module in model.named_children(): - # if isinstance(module, (torch.nn.Linear)): - # layers.add(module.name) - # _find_hqq_quantizable_layers(module, layers) - - # new_keys = set(expected_keys) - - # # Name modules - # for name, module in model.named_modules(): - # module.name = name - - # # valid modules are Linear layers that have HQQLinear state_dict. We ignore skip_modules and any layers with Linear state_dict() params - # _valid_modules = set() - # _find_hqq_quantizable_layers(model, _valid_modules) - - # # Remove skipped modules - # _skipped_modules = set() - # for _module in _valid_modules: - # for _skip_module in model.config.quantization_config["skip_modules"]: - # if _skip_module in _module: - # _skipped_modules.add(_module) - # _valid_modules -= _skipped_modules - - # # Append new expected layers based on _ref_keys - # _ref_keys = HQQLinear( - # linear_layer=None, - # quant_config=None, - # compute_dtype=torch.float16, - # device="cpu", - # del_orig=False, - # ).state_dict_keys() - {"bias"} - - # # Clean-up - # _rm_keys = set() - # for key in new_keys: - # if any(_module in key for _module in _valid_modules): - # _rm_keys.add(key) - # new_keys -= _rm_keys - # # At this point, new_keys contains all the keys of the layers that are NOT HQQLinear or torch.nn.Linear - - # # Re-populate Linear/HQQLinear - # for _module in _valid_modules: - # if _module + ".weight" in loaded_keys: - # new_keys.add(_module + ".weight") - # else: - # new_keys.update({_module + "." + _ref_key for _ref_key in _ref_keys}) - # if _module + ".bias" in loaded_keys: - # new_keys.add(_module + ".bias") - - # return list(new_keys) def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: module, _ = get_module_from_name(model, param_name) @@ -149,99 +62,6 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** # `create_quantized_param`, even when `self.is_quantized == True` return isinstance(module, torch.nn.Linear) - # TODO: to remove - # def create_quantized_param( - # self, - # model: "PreTrainedModel", - # param_value: "torch.Tensor", - # param_name: str, - # target_device: "torch.device", - # **kwargs, - # ): - # module, tensor_name = get_module_from_name(model, param_name) - # module_name = param_name.rsplit(".", 1)[0] - # parent_module, node = get_module_from_name(model, module_name) - - # quant_config = model.config.quantization_config["quant_config"] - # skip_modules = model.config.quantization_config["skip_modules"] - - # # In this case we do not quantize this layer (it's explicitly skipped) -> simply load param - # if any(skip_module in module.name for skip_module in skip_modules): - # module.load_state_dict( - # {tensor_name: param_value.to(device=target_device, dtype=self.dtype)}, strict=False, assign=True - # ) - # return - - # # We need this hack as the model is not pre-prepared as an empty skeleton on meta device - # if self.pre_quantized: - # # Save them for later - # if not hasattr(self, "hqq_params"): - # self.hqq_params = defaultdict(dict) - # self.hqq_params[module_name].update({tensor_name: param_value}) - # hqq_params = self.hqq_params[module_name] - - # # If they are all present and saved, make it a HQQLinear layer! (we cannot do it param after param because - # # hqq does not support it...) - # if all(k in hqq_params for k in self.hqq_keys) and ("bias" in hqq_params or module.bias is None): - # hqq_layer = HQQLinear( - # linear_layer=None, - # quant_config=None, - # compute_dtype=self.dtype, - # device=target_device, - # del_orig=False, - # ) - # hqq_layer.load_state_dict(hqq_params) - - # if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor): - # hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias) - # if self.using_multi_gpu: - # hqq_layer = self._patch_layer_for_multigpu(hqq_layer) - - # setattr(parent_module, node, hqq_layer) - # del self.hqq_params[module_name], module - # return - - # # Load param in the module (without caring about device or dtype, it will be changed later) - # module.load_state_dict({tensor_name: param_value}, strict=False, assign=True) - - # # If both the weight and bias have already been loaded, time to quantize! - # module_is_ready = module.weight.device.type != "meta" and ( - # module.bias is None or module.bias.device.type != "meta" - # ) - - # if module_is_ready: - # module_tag = ".".join(module.name.split(".")[-2:]) - # if "weight_quant_params" in quant_config: - # module_quant_config = quant_config - # elif module_tag in quant_config: - # module_quant_config = quant_config[module_tag] - - # hqq_layer = HQQLinear( - # module, - # quant_config=module_quant_config, - # compute_dtype=self.dtype, - # device=target_device, - # del_orig=True, - # ) - - # if hqq_layer.bias is not None and isinstance(hqq_layer.bias, torch.Tensor): - # hqq_layer.bias = torch.nn.Parameter(hqq_layer.bias) - - # if self.using_multi_gpu: - # hqq_layer = self._patch_layer_for_multigpu(hqq_layer) - - # setattr(parent_module, node, hqq_layer) - - def _patch_layer_for_multigpu(self, hqq_layer): - def forward_with_device(self, x): - out = torch.matmul(x.to(self.device), self.dequantize().t()) - if self.bias is not None: - out += self.bias - return out - - hqq_layer.forward = lambda x: forward_with_device(hqq_layer, x) - return hqq_layer - def _process_model_before_weight_loading( self, model: "PreTrainedModel", diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index 84dd0bfeee20..ed3c03997f0e 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -38,7 +38,6 @@ from ..core_model_loading import WeightConverter logger = logging.get_logger(__name__) -triton_kernels_hub = None class Mxfp4HfQuantizer(HfQuantizer): @@ -49,21 +48,6 @@ class Mxfp4HfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "Mxfp4Config" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - self.triton_kernels_hub = None - - def _lazy_import_kernels(self): - """Lazy import and initialize kernels only when needed""" - if self.triton_kernels_hub is None: - try: - from ..integrations.hub_kernels import get_kernel - - self.triton_kernels_hub = get_kernel("kernels-community/gpt-oss-triton-kernels", version=1) - except ImportError: - raise ImportError("kernels package is required for MXFP4 quantization") - return self.triton_kernels_hub - def validate_environment(self, *args, **kwargs): if not is_torch_available(): raise ImportError( @@ -149,9 +133,6 @@ def validate_environment(self, *args, **kwargs): elif not kernels_installed: raise ValueError("MXFP4 quantization requires the `kernels` package: `pip install kernels>=0.12.0`") - if not self.pre_quantized: - self._lazy_import_kernels() - 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(): @@ -211,66 +192,16 @@ def _process_model_before_weight_loading( model, modules_to_not_convert=self.modules_to_not_convert, quantization_config=self.quantization_config ) - def update_tp_plan(self, config): - if "GptOssConfig" in config.__class__.__name__: - if getattr(config, "base_model_tp_plan", None) is not None: - config.base_model_tp_plan.update( - { - "layers.*.mlp.experts.gate_up_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.gate_up_proj_scales": "grouped_gemm", - "layers.*.mlp.experts.down_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.down_proj_scales": "grouped_gemm", - } - ) - return config - - def update_ep_plan(self, config): - if "GptOssConfig" in config.__class__.__name__: - if getattr(config, "base_model_ep_plan", None) is not None: - config.base_model_ep_plan.update( - { - "layers.*.mlp.experts.gate_up_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.gate_up_proj_scales": "grouped_gemm", - "layers.*.mlp.experts.down_proj_blocks": "grouped_gemm", - "layers.*.mlp.experts.down_proj_scales": "grouped_gemm", - } - ) - return config - def get_state_dict_and_metadata(self, model): - from ..integrations import Mxfp4GptOssExperts + from ..integrations.mxfp4 import Mxfp4GptOssExperts, Mxfp4ReverseDeserialize state_dict = model.state_dict() - num_local_experts = getattr(model.config, "num_local_experts", 32) - hidden_size = getattr(model.config, "hidden_size", 2880) - + reverse_op = Mxfp4ReverseDeserialize(self) for name, module in model.named_modules(): - if not ( - isinstance(module, Mxfp4GptOssExperts) - and hasattr(module, "gate_up_proj") - and hasattr(module, "down_proj") - ): - continue - - for proj in ("gate_up_proj", "down_proj"): - triton_tensor = getattr(module, proj) - precision_config = getattr(module, f"{proj}_precision_config") - - blocks = triton_tensor.storage.layout.unswizzle_data(triton_tensor.storage.data).transpose(-1, -2) - if proj == "gate_up_proj": - blocks = blocks.reshape(num_local_experts, -1, 90, 16) - else: - blocks = blocks.reshape(num_local_experts, hidden_size, 90, -1) - - scales = precision_config.weight_scale.storage.layout.unswizzle_data( - precision_config.weight_scale.storage.data - ).transpose(-1, -2) - - state_dict[f"{name}.{proj}_blocks"] = blocks - state_dict[f"{name}.{proj}_scales"] = scales - - metadata = {} - return state_dict, metadata + 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 diff --git a/tests/quantization/hqq/test_hqq.py b/tests/quantization/hqq/test_hqq.py index 913bf6bf9e75..3e6473004f88 100755 --- a/tests/quantization/hqq/test_hqq.py +++ b/tests/quantization/hqq/test_hqq.py @@ -148,16 +148,6 @@ def test_quantized_model_to_new_device_and_new_dtype(self): check_hqqlayer(self, hqq_runner.model.model.layers[0].self_attn.v_proj) check_forward(self, hqq_runner.model) - def test_quantized_model_fake_weight_dtype(self): - quant_config = HqqConfig(nbits=8, group_size=64) - - hqq_runner = HQQLLMRunner( - model_id=MODEL_ID, quant_config=quant_config, compute_dtype=torch.float16, device=torch_device - ) - - # We use a hack to inject a fake weight to HQQLinear. Check that it works - self.assertEqual(hqq_runner.model.model.layers[0].self_attn.v_proj.weight.dtype, torch.float16) - @slow @require_torch_accelerator diff --git a/tests/quantization/mxfp4/test_mxfp4.py b/tests/quantization/mxfp4/test_mxfp4.py index e9dbfbf2e202..f7ba17da40ea 100644 --- a/tests/quantization/mxfp4/test_mxfp4.py +++ b/tests/quantization/mxfp4/test_mxfp4.py @@ -133,7 +133,6 @@ def test_quantizer_validation_no_accelerator(self): _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), - patch("transformers.quantizers.quantizer_mxfp4.Mxfp4HfQuantizer._lazy_import_kernels"), ): from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer @@ -192,7 +191,6 @@ def test_quantizer_validation_order_dequantize_before_accelerator_check(self): _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), - patch("transformers.quantizers.quantizer_mxfp4.Mxfp4HfQuantizer._lazy_import_kernels"), ): from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer @@ -356,17 +354,14 @@ def test_convert_moe_packed_tensors(self): @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 - from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer - - config = Mxfp4Config() - quantizer = Mxfp4HfQuantizer(config) # 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, quantizer._lazy_import_kernels()) + quantized_w, w_scale = quantize_to_mxfp4(w, get_kernel("kernels-community/gpt-oss-triton-kernels", version=1)) # Check that shapes are reasonable self.assertEqual(quantized_w.dtype, torch.uint8) From 55a30d4973351b9fa344fe701d47c0d379d0d2ec Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 11:48:28 +0900 Subject: [PATCH 03/27] update --- src/transformers/distributed/tensor_parallel.py | 4 +++- .../quantizers/quantizer_fbgemm_fp8.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 00b44df39b63..66add4f5f64d 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -204,6 +204,8 @@ def shard_param(self, module, param, mesh): if meta is None: return placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2) + if meta.ndim == 0: # e.g. per-tensor quantization scales + placement = Replicate() module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, @@ -278,7 +280,7 @@ def shard_param(self, module, param, mesh): placement = Shard(0) else: # bias is replicated (added after the row-reduce); weight shards on input dim (-1) - placement = Replicate() if param == "bias" else Shard(-1) + placement = Replicate() if param == "bias" or meta.ndim == 0 else Shard(-1) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 01721c1dd34f..665a0d011764 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -149,6 +149,22 @@ def _process_model_after_weight_loading(self, model, **kwargs): m.input_scale_ub.fill_(self.quantization_config.activation_scale_ub) return model + def _update_plan(self, config, plan_name): + text_config = config.get_text_config() + plan = dict(getattr(text_config, plan_name, None) or {}) + for key, style in list(plan.items()): + if style == "grouped_gemm": + plan.setdefault(f"{key}_scale", style) + if plan: + setattr(text_config, plan_name, plan) + return config + + def update_tp_plan(self, config): + return self._update_plan(config, "base_model_tp_plan") + + def update_ep_plan(self, config): + return self._update_plan(config, "base_model_ep_plan") + def is_serializable(self): return True From 09c4a310bae406a403659be84c750bbde4eb1778 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 11:53:18 +0900 Subject: [PATCH 04/27] sub config fixes Saw another PR recently: #48759, had the same bug. this adresses it --- src/transformers/modeling_utils.py | 2 +- src/transformers/quantizers/auto.py | 2 -- src/transformers/quantizers/base.py | 12 ++++++------ .../quantizers/quantizer_fbgemm_fp8.py | 15 +++++---------- .../quantizers/quantizer_finegrained_fp8.py | 16 ++++++---------- src/transformers/quantizers/quantizer_nvfp4.py | 5 ++--- tests/quantization/nvfp4/test_nvfp4.py | 6 +++--- 7 files changed, 23 insertions(+), 35 deletions(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index 4d5afe1df9a2..a70470322d03 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -2086,7 +2086,7 @@ def set_experts_implementation(self, experts_implementation: str | dict): else experts_implementation.get("", self.config._experts_implementation) ) - # MegaMoE is locked at load time: its TP plan is baked into `base_model_tp_plan` by `update_tp_plan` (and isn't + # MegaMoE is locked at load time: its TP plan is baked into `model.tp_plan` by `update_tp_plan` (and isn't # re-evaluated) and `setup_megamoe_weights` mutates the expert weights into UTCCP layout on first forward. # Either side of a switch would silently produce garbage, so reject it with a clear pointer. current = self.config._experts_implementation diff --git a/src/transformers/quantizers/auto.py b/src/transformers/quantizers/auto.py index 892518e65e41..f5dc3fe5d5d8 100644 --- a/src/transformers/quantizers/auto.py +++ b/src/transformers/quantizers/auto.py @@ -380,8 +380,6 @@ def get_hf_quantizer(config, quantization_config, device_map, weights_only, user weights_only=weights_only, ) device_map = hf_quantizer.update_device_map(device_map) - config = hf_quantizer.update_tp_plan(config) - config = hf_quantizer.update_ep_plan(config) config = hf_quantizer.update_attn_implementation(config) # In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry` diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 52452207e0fd..17cd229dfe23 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -152,13 +152,11 @@ def validate_environment(self, *args, **kwargs): """ return - def update_tp_plan(self, config): - "updates the tp plan for the scales" - return config + def update_tp_plan(self, model): + "Updates `model.tp_plan` in place: the full plan, merged from the config and all sub-configs." - def update_ep_plan(self, config): - "updates the tp plan for the scales" - return config + def update_ep_plan(self, model): + "Updates `model.ep_plan` in place: the full plan, merged from the config and all sub-configs." def update_attn_implementation(self, config): """Sees the requested `attn_implementation` before the model is built, so a quantizer can default it.""" @@ -184,6 +182,8 @@ def preprocess_model(self, model: "PreTrainedModel", dtype=None, **kwargs): if self.pre_quantized: self._convert_model_for_quantization(model) self._process_model_before_weight_loading(model, **kwargs) + self.update_tp_plan(model) + self.update_ep_plan(model) def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): return model diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 665a0d011764..91e929b3bb8d 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -149,21 +149,16 @@ def _process_model_after_weight_loading(self, model, **kwargs): m.input_scale_ub.fill_(self.quantization_config.activation_scale_ub) return model - def _update_plan(self, config, plan_name): - text_config = config.get_text_config() - plan = dict(getattr(text_config, plan_name, None) or {}) + def _update_plan(self, plan): for key, style in list(plan.items()): if style == "grouped_gemm": plan.setdefault(f"{key}_scale", style) - if plan: - setattr(text_config, plan_name, plan) - return config - def update_tp_plan(self, config): - return self._update_plan(config, "base_model_tp_plan") + def update_tp_plan(self, model): + self._update_plan(model.tp_plan) - def update_ep_plan(self, config): - return self._update_plan(config, "base_model_ep_plan") + def update_ep_plan(self, model): + self._update_plan(model.ep_plan) def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index d35c36741b2e..4079926075ba 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -136,24 +136,20 @@ def _process_model_after_weight_loading(self, model, **kwargs): _disable_deepgemm_on_multi_device(model) return model - def _update_plan(self, config, plan_name): + def _update_plan(self, plan, config): from ..integrations.finegrained_fp8 import FP8Experts overwrite = FP8Experts._impl_tp_layer_overrides.get(getattr(config, "_experts_implementation", None), {}) - plan = {key: overwrite.get(style, style) for key, style in (getattr(config, plan_name, None) or {}).items()} - # TODO: we need to define a global "style": "new_key" for all quantization plans for key, style in list(plan.items()): + plan[key] = overwrite.get(style, style) if style == "grouped_gemm": plan.setdefault(f"{key}_scale_inv", style) - if plan: - setattr(config, plan_name, plan) - return config - def update_tp_plan(self, config): - return self._update_plan(config, "base_model_tp_plan") + def update_tp_plan(self, model): + self._update_plan(model.tp_plan, model.config) - def update_ep_plan(self, config): - return self._update_plan(config, "base_model_ep_plan") + def update_ep_plan(self, model): + self._update_plan(model.ep_plan, model.config) def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index cb58ca213578..b2ac84cd3565 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -74,15 +74,14 @@ def validate_environment(self, device_map, **kwargs): if major < 10: raise RuntimeError(f"NVFP4 requires Blackwell (sm100+) block-scaled tensor cores; found sm{major}{minor}.") - def update_tp_plan(self, config): - distributed_config = getattr(config, "distributed_config", None) + def update_tp_plan(self, model): + 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." ) - return config def param_needs_quantization(self, model: PreTrainedModel, param_name: str, **kwargs) -> bool: from ..integrations.nvfp4 import NVFP4Linear diff --git a/tests/quantization/nvfp4/test_nvfp4.py b/tests/quantization/nvfp4/test_nvfp4.py index 5847e352d7c2..323ae2de0547 100644 --- a/tests/quantization/nvfp4/test_nvfp4.py +++ b/tests/quantization/nvfp4/test_nvfp4.py @@ -145,10 +145,10 @@ def test_rejects_multi_device_map_until_tp_metadata_is_supported(self, _, __): 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): - config = SimpleNamespace(distributed_config=SimpleNamespace(tp_size=2)) + 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_tp_plan(config) + self.quantizer.update_tp_plan(model) def test_update_tp_plan_keeps_non_tp_configs(self): configs = ( @@ -158,7 +158,7 @@ def test_update_tp_plan_keeps_non_tp_configs(self): for config in configs: with self.subTest(config=config): - self.assertIs(self.quantizer.update_tp_plan(config), config) + self.quantizer.update_tp_plan(SimpleNamespace(config=config)) @patch("transformers.quantizers.quantizer_nvfp4.is_kernels_available", return_value=True) @patch("torch.cuda.is_available", return_value=True) From b0f1a47a40590ccaf507b1d84328044e1e335cf1 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 12:18:27 +0900 Subject: [PATCH 05/27] more simplification --- src/transformers/quantizers/base.py | 48 +++++++++++++---- .../quantizer_compressed_tensors.py | 51 ++++--------------- .../quantizers/quantizer_fbgemm_fp8.py | 12 +---- .../quantizers/quantizer_finegrained_fp8.py | 44 +++++----------- 4 files changed, 61 insertions(+), 94 deletions(-) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 17cd229dfe23..5e5010b639d5 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -11,6 +11,7 @@ # 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 re from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any @@ -154,9 +155,19 @@ def validate_environment(self, *args, **kwargs): def update_tp_plan(self, model): "Updates `model.tp_plan` in place: the full plan, merged from the config and all sub-configs." + self._update_plan(model, model.tp_plan) def update_ep_plan(self, model): "Updates `model.ep_plan` in place: the full plan, merged from the config and all sub-configs." + self._update_plan(model, model.ep_plan) + + def _update_plan(self, model, plan): + "Serialized params (e.g. `gate_up_proj_scale_inv`) are sharded like their param (`gate_up_proj`)." + from ..distributed.tensor_parallel import replace_layer_number_by_wildcard + + params = {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()} + extra = {key + n.removeprefix("weight"): style for key, style in plan.items() for n in self.serialized_names} + plan.update({key: style for key, style in extra.items() if key in params and key not in plan}) def update_attn_implementation(self, config): """Sees the requested `attn_implementation` before the model is built, so a quantizer can default it.""" @@ -307,20 +318,35 @@ def get_quantize_ops(self): ) def get_weight_conversions(self): + "General converters, for the serialized tensors that no model converter claims." return [] + def get_load_op(self, converter): + """Op turning the serialized tensors of a pre-quantized `converter` into its params: a deserialization, or a + dequantization when `dequantize=True`. It is prepended to the converter's ops, i.e. before any merge.""" + return None + + # How a `weight` is saved in pre-quantized checkpoints, e.g. `["weight", "weight_scale_inv"]` + serialized_names: list[str] = [] + def update_weight_conversions(self, weight_conversions): - """Give the quantizer a chance to rewrite the weight conversion pipeline. - - Loading runs ``renamings → converters → (dequant → merge → concat)``. Dequant - has to happen *before* any merge/concat op because those operations aren't - aware of per-block scales, so the per-expert (weight, scale) pairs need to be - collapsed into full-precision tensors first. Subclasses (e.g. the FP8 - quantizer in ``dequantize=True`` mode) override this to inject a dequantize - op at the start of each model-provided :class:`WeightConverter` and attach the - matching scale source patterns. Default: no-op. - """ - return weight_conversions + self.get_weight_conversions() + """Prepends `get_load_op` to the model converters, with the serialized `weight` sources, then appends the + general `get_weight_conversions`.""" + from ..core_model_loading import WeightConverter + + updated = [] + for conv in weight_conversions: + op = self.get_load_op(conv) if self.pre_quantized and isinstance(conv, WeightConverter) else None + sources = conv._original_source_patterns + new_sources = list( + dict.fromkeys(re.sub(r"weight$", f"{n}$", p) for n in self.serialized_names for p in sources) + ) + if op is not None and new_sources != sources: + new = WeightConverter(new_sources, conv._original_target_patterns, [op] + conv.operations) + new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix + conv = new + updated.append(conv) + return updated + self.get_weight_conversions() class SequentialLlama4TextExperts(ModuleList): diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index dff536ecd3a4..27683ee707ad 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -167,44 +167,13 @@ def get_weight_conversions(self): ) ] - def update_weight_conversions(self, weight_conversions): - """Attach the quantization sources (scales, packed weights) of MoE expert converters - to their bucket and prepend a :class:`DecompressExperts` op, so the per-expert - (weight, scale) pairs are dequantized *before* the merge / concat ops collapse the - per-expert structure. FP8 checkpoints keep the plain ``weight`` name; packed formats - use ``weight_packed`` / ``weight_shape``. - - The generic converters from :meth:`get_weight_conversions` are appended at the end: - converters are matched in order, so the expert converters keep their scale keys. - """ - updated: list = [] - for conv in weight_conversions: - # Only WeightConverter for experts have ``.operations`` to extend with the dequant op - if not isinstance(conv, WeightConverter) or any("experts" not in p for p in conv.source_patterns): - updated.append(conv) - continue - weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] - if weight_sources: - scheme = get_experts_scheme(self.quantization_config.quantization_config) - scale_sources = [p + "_scale$" for p in weight_sources] - other = [p for p in conv.source_patterns if not p.endswith(".weight")] - if _is_fp8_scheme(scheme): - # Merged experts cannot stay FP8 (they are not nn.Linear): they are - # dequantized to the model dtype before the merge. The weight patterns - # must be anchored with `$`: patterns are regex-searched, so unanchored - # `.weight` would also match the `.weight_scale` keys. - new_sources = [p + "$" for p in weight_sources] + scale_sources + other - else: - packed_weight = [p + "_packed$" for p in weight_sources] - shape_sources = [p + "_shape$" for p in weight_sources] - new_sources = packed_weight + scale_sources + shape_sources + other - new_ops = [DecompressExperts(self, scheme=scheme)] + list(conv.operations) - conv = WeightConverter( - source_patterns=new_sources, - target_patterns=conv._original_target_patterns, - operations=new_ops, - ) - updated.append(conv) - - updated.extend(self.get_weight_conversions()) - return updated + @property + def serialized_names(self): + if _is_fp8_scheme(get_experts_scheme(self.quantization_config.quantization_config)): + return ["weight", "weight_scale"] + return ["weight_packed", "weight_scale", "weight_shape"] + + def get_load_op(self, converter): + # Merged experts are not `nn.Linear`, so they can't stay compressed: decompress them before the merge + if all("experts" in p for p in converter.source_patterns): + return DecompressExperts(self, scheme=get_experts_scheme(self.quantization_config.quantization_config)) diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 91e929b3bb8d..7bb725d28cfd 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -44,6 +44,7 @@ class FbgemmFp8HfQuantizer(HfQuantizer): """ requires_calibration = False + serialized_names = ["weight", "weight_scale"] quantization_config: "FbgemmFp8Config" def __init__(self, quantization_config, **kwargs): @@ -149,17 +150,6 @@ def _process_model_after_weight_loading(self, model, **kwargs): m.input_scale_ub.fill_(self.quantization_config.activation_scale_ub) return model - def _update_plan(self, plan): - for key, style in list(plan.items()): - if style == "grouped_gemm": - plan.setdefault(f"{key}_scale", style) - - def update_tp_plan(self, model): - self._update_plan(model.tp_plan) - - def update_ep_plan(self, model): - self._update_plan(model.ep_plan) - def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 4079926075ba..3c655e6dc9a2 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -136,20 +136,12 @@ def _process_model_after_weight_loading(self, model, **kwargs): _disable_deepgemm_on_multi_device(model) return model - def _update_plan(self, plan, config): + def _update_plan(self, model, plan): from ..integrations.finegrained_fp8 import FP8Experts - overwrite = FP8Experts._impl_tp_layer_overrides.get(getattr(config, "_experts_implementation", None), {}) - for key, style in list(plan.items()): - plan[key] = overwrite.get(style, style) - if style == "grouped_gemm": - plan.setdefault(f"{key}_scale_inv", style) - - def update_tp_plan(self, model): - self._update_plan(model.tp_plan, model.config) - - def update_ep_plan(self, model): - self._update_plan(model.ep_plan, model.config) + super()._update_plan(model, plan) + overwrite = FP8Experts._impl_tp_layer_overrides.get(getattr(model.config, "_experts_implementation", None), {}) + plan.update({key: overwrite.get(style, style) for key, style in plan.items()}) def is_serializable(self): return True @@ -189,26 +181,16 @@ def get_weight_conversions(self): ] return [] - def update_weight_conversions(self, weight_conversions): - # A key goes through at most one `WeightConverter`, so per-expert weights/scales merged by a model converter - # can only be dequantized by an op prepended to that converter. - from ..core_model_loading import WeightConverter, WeightRenaming + serialized_names = ["weight", "weight_scale_inv"] + + def get_load_op(self, converter): from ..integrations.finegrained_fp8 import Fp8Dequantize + return Fp8Dequantize(self) if self.quantization_config.dequantize else None + + def update_weight_conversions(self, weight_conversions): + from ..core_model_loading import WeightRenaming + # Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) ship per-block scales under `.scale` scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") - updated = [scale_rename] - dequantize = self.pre_quantized and self.quantization_config.dequantize - for conv in weight_conversions: - weight_sources = [p for p in conv.source_patterns if p.endswith(".weight")] - if isinstance(conv, WeightConverter) and dequantize and weight_sources: - other = [p for p in conv.source_patterns if not p.endswith(".weight")] - conv = WeightConverter( - source_patterns=[p + "$" for p in weight_sources] - + [p[: -len(".weight")] + ".weight_scale_inv$" for p in weight_sources] - + other, - target_patterns=conv._original_target_patterns, - operations=[Fp8Dequantize(self)] + conv.operations, - ) - updated.append(conv) - return updated + self.get_weight_conversions() + return [scale_rename] + super().update_weight_conversions(weight_conversions) From d59c65f8ddccf7051542197a8649850db8f01453 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 12:23:26 +0900 Subject: [PATCH 06/27] debloat more shit, standardize --- src/transformers/quantizers/base.py | 47 ++++++++++++------- .../quantizers/quantizer_bitnet.py | 20 +++----- .../quantizers/quantizer_bnb_4bit.py | 34 ++++++-------- .../quantizers/quantizer_bnb_8bit.py | 16 ++----- .../quantizer_compressed_tensors.py | 8 ++-- .../quantizers/quantizer_fbgemm_fp8.py | 2 +- .../quantizers/quantizer_finegrained_fp8.py | 23 ++------- .../quantizers/quantizer_fp_quant.py | 28 ++++------- .../quantizers/quantizer_metal.py | 16 ++----- .../quantizers/quantizer_mxfp4.py | 36 ++++---------- .../quantizers/quantizer_quark.py | 26 +++------- src/transformers/quantizers/quantizer_sinq.py | 32 +++---------- .../quantizers/quantizer_torchao.py | 33 ++++--------- 13 files changed, 105 insertions(+), 216 deletions(-) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 5e5010b639d5..88fae536ecaf 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -11,7 +11,6 @@ # 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 re from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any @@ -166,7 +165,10 @@ def _update_plan(self, model, plan): from ..distributed.tensor_parallel import replace_layer_number_by_wildcard params = {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()} - extra = {key + n.removeprefix("weight"): style for key, style in plan.items() for n in self.serialized_names} + suffixes = { + n.rstrip("$").removeprefix(p.rstrip("$")) for p, names in self.serialized_names.items() for n in names + } + extra = {key + suffix: style for key, style in plan.items() for suffix in suffixes} plan.update({key: style for key, style in extra.items() if key in params and key not in plan}) def update_attn_implementation(self, config): @@ -317,32 +319,41 @@ def get_quantize_ops(self): f"{self.quantization_config.quant_method} is not available yet and will be supported soon." ) + # How params are saved in pre-quantized checkpoints: `{param: [serialized source patterns]}`, + # e.g. `{"weight": ["weight$", "weight_scale_inv"]}` + serialized_names: dict[str, list[str]] = {} + + def get_load_op(self, converter=None): + """Op loading the serialized tensors: a deserialization, or a dequantization when `dequantize=True`. + Without `converter`, the op of the general converters. With a model `converter` (e.g. a per-expert merge), + the op to prepend to it, if the quantizer's op supports it.""" + return None + def get_weight_conversions(self): "General converters, for the serialized tensors that no model converter claims." - return [] + from ..core_model_loading import WeightConverter - def get_load_op(self, converter): - """Op turning the serialized tensors of a pre-quantized `converter` into its params: a deserialization, or a - dequantization when `dequantize=True`. It is prepended to the converter's ops, i.e. before any merge.""" - return None + op = self.get_load_op() + return [WeightConverter(names, param, [op]) for param, names in self.serialized_names.items()] if op else [] - # How a `weight` is saved in pre-quantized checkpoints, e.g. `["weight", "weight_scale_inv"]` - serialized_names: list[str] = [] + def _serialized_sources(self, pattern: str) -> list[str]: + "`experts.*.gate_proj.weight` -> `[experts.*.gate_proj.weight$, experts.*.gate_proj.weight_scale_inv$]`" + for param, names in self.serialized_names.items(): + if pattern.endswith(param := param.rstrip("$")): + return [pattern.removesuffix(param) + name.rstrip("$") + "$" for name in names] + return [pattern] def update_weight_conversions(self, weight_conversions): - """Prepends `get_load_op` to the model converters, with the serialized `weight` sources, then appends the - general `get_weight_conversions`.""" + """Prepends `get_load_op` to the model converters, with their serialized sources, then appends the general + `get_weight_conversions`.""" from ..core_model_loading import WeightConverter updated = [] for conv in weight_conversions: - op = self.get_load_op(conv) if self.pre_quantized and isinstance(conv, WeightConverter) else None - sources = conv._original_source_patterns - new_sources = list( - dict.fromkeys(re.sub(r"weight$", f"{n}$", p) for n in self.serialized_names for p in sources) - ) - if op is not None and new_sources != sources: - new = WeightConverter(new_sources, conv._original_target_patterns, [op] + conv.operations) + op = self.get_load_op(conv) if isinstance(conv, WeightConverter) else None + sources = [s for p in conv._original_source_patterns for s in self._serialized_sources(p)] + if op is not None and sources != conv._original_source_patterns: + new = WeightConverter(sources, conv._original_target_patterns, [op] + conv.operations) new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix conv = new updated.append(conv) diff --git a/src/transformers/quantizers/quantizer_bitnet.py b/src/transformers/quantizers/quantizer_bitnet.py index ec22976939ee..2432afa96e21 100644 --- a/src/transformers/quantizers/quantizer_bitnet.py +++ b/src/transformers/quantizers/quantizer_bitnet.py @@ -106,19 +106,11 @@ def is_qat_trainable(self) -> bool: and self.quantization_config.quantization_mode == "online" ) - def get_weight_conversions(self): - from ..core_model_loading import WeightConverter + serialized_names = {"weight": ["weight"]} + + def get_load_op(self, converter=None): from ..integrations.bitnet import BitNetDeserialize - if ( - self.quantization_config.linear_class == "autobitlinear" - and self.quantization_config.quantization_mode == "offline" - ): - return [ - WeightConverter( - source_patterns=["weight"], - target_patterns=["weight"], - operations=[BitNetDeserialize(self)], - ) - ] - return [] + config = self.quantization_config + if converter is None and config.linear_class == "autobitlinear" and config.quantization_mode == "offline": + return BitNetDeserialize(self) diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py index 5fd040bed10e..e2e2d7886db0 100644 --- a/src/transformers/quantizers/quantizer_bnb_4bit.py +++ b/src/transformers/quantizers/quantizer_bnb_4bit.py @@ -37,7 +37,6 @@ if is_torch_available(): import torch - from ..core_model_loading import WeightConverter logger = logging.get_logger(__name__) @@ -165,23 +164,20 @@ def get_quantize_ops(self): return Bnb4bitQuantize(self) - def get_weight_conversions(self): + serialized_names = { + "weight": [ + "weight.nested_absmax", + "weight.nested_quant_map", + "weight.quant_map", + "weight.absmax", + "weight.quant_state.bitsandbytes__nf4", + "weight.quant_state.bitsandbytes__fp4", + "weight", + ] + } + + def get_load_op(self, converter=None): from ..integrations.bitsandbytes import Bnb4bitDeserialize - if self.pre_quantized: - return [ - WeightConverter( - source_patterns=[ - "weight.nested_absmax", - "weight.nested_quant_map", - "weight.quant_map", - "weight.absmax", - "weight.quant_state.bitsandbytes__nf4", - "weight.quant_state.bitsandbytes__fp4", - "weight", - ], - target_patterns="weight", - operations=[Bnb4bitDeserialize(self)], - ) - ] - return [] + if converter is None and self.pre_quantized: + return Bnb4bitDeserialize(self) diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py index cf3192bcd72a..57ee3377ab35 100644 --- a/src/transformers/quantizers/quantizer_bnb_8bit.py +++ b/src/transformers/quantizers/quantizer_bnb_8bit.py @@ -37,7 +37,6 @@ if is_torch_available(): import torch - from ..core_model_loading import WeightConverter logger = logging.get_logger(__name__) @@ -164,15 +163,10 @@ def get_quantize_ops(self): return Bnb8bitQuantize(self) - def get_weight_conversions(self): + serialized_names = {"weight": ["SCB", "weight_format", "weight"]} + + def get_load_op(self, converter=None): from ..integrations.bitsandbytes import Bnb8bitDeserialize - if self.pre_quantized: - return [ - WeightConverter( - source_patterns=["SCB", "weight_format", "weight"], - target_patterns="weight", - operations=[Bnb8bitDeserialize(self)], - ) - ] - return [] + if converter is None and self.pre_quantized: + return Bnb8bitDeserialize(self) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 27683ee707ad..afc623cef4e7 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -170,10 +170,10 @@ def get_weight_conversions(self): @property def serialized_names(self): if _is_fp8_scheme(get_experts_scheme(self.quantization_config.quantization_config)): - return ["weight", "weight_scale"] - return ["weight_packed", "weight_scale", "weight_shape"] + return {"weight": ["weight", "weight_scale"]} + return {"weight": ["weight_packed", "weight_scale", "weight_shape"]} - def get_load_op(self, converter): + def get_load_op(self, converter=None): # Merged experts are not `nn.Linear`, so they can't stay compressed: decompress them before the merge - if all("experts" in p for p in converter.source_patterns): + if converter is not None and all("experts" in p for p in converter.source_patterns): return DecompressExperts(self, scheme=get_experts_scheme(self.quantization_config.quantization_config)) diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 7bb725d28cfd..29a1a8cc8e2c 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -44,7 +44,7 @@ class FbgemmFp8HfQuantizer(HfQuantizer): """ requires_calibration = False - serialized_names = ["weight", "weight_scale"] + serialized_names = {"weight": ["weight", "weight_scale"]} quantization_config: "FbgemmFp8Config" def __init__(self, quantization_config, **kwargs): diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 3c655e6dc9a2..bd95cf1114eb 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -165,28 +165,13 @@ def get_new_dtype(self, original_key, renamed_key, param, dtype): # `FP8Linear` / `FP8Experts` declare the dtype of every param, e.g. ue8m0 scales stored as float32 return dtype - def get_weight_conversions(self): - from ..core_model_loading import WeightConverter - from ..integrations.finegrained_fp8 import Fp8Dequantize - - if self.pre_quantized and self.quantization_config.dequantize: - return [ - # either use the dollar sign, or permute the source patterns to start matching against the scales first - # We also collect the activation scales, they will not be used - WeightConverter( - source_patterns=["weight$", "weight_scale_inv", "activation_scale"], - target_patterns="weight", - operations=[Fp8Dequantize(self)], - ) - ] - return [] - - serialized_names = ["weight", "weight_scale_inv"] + serialized_names = {"weight": ["weight$", "weight_scale_inv", "activation_scale"]} - def get_load_op(self, converter): + def get_load_op(self, converter=None): from ..integrations.finegrained_fp8 import Fp8Dequantize - return Fp8Dequantize(self) if self.quantization_config.dequantize else None + if self.pre_quantized and self.quantization_config.dequantize: + return Fp8Dequantize(self) def update_weight_conversions(self, weight_conversions): from ..core_model_loading import WeightRenaming diff --git a/src/transformers/quantizers/quantizer_fp_quant.py b/src/transformers/quantizers/quantizer_fp_quant.py index c367b957d206..e3c36b2f58cf 100644 --- a/src/transformers/quantizers/quantizer_fp_quant.py +++ b/src/transformers/quantizers/quantizer_fp_quant.py @@ -140,25 +140,13 @@ def get_quantize_ops(self): return FpQuantQuantize(self) - def get_weight_conversions(self): - from ..core_model_loading import WeightConverter + @property + def serialized_names(self): + name = ".dqweight" if self.quantization_config.pseudoquantization else ".qweight" + return {name: [name]} + + def get_load_op(self, converter=None): from ..integrations.fp_quant import FpQuantDeserialize - if self.pre_quantized: - if self.quantization_config.pseudoquantization: - return [ - WeightConverter( - source_patterns=[".dqweight"], - target_patterns=".dqweight", - operations=[FpQuantDeserialize(self)], - ), - ] - else: - return [ - WeightConverter( - source_patterns=[".qweight"], - target_patterns=".qweight", - operations=[FpQuantDeserialize(self)], - ), - ] - return [] + if converter is None and self.pre_quantized: + return FpQuantDeserialize(self) diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py index a3706ed7e2b4..6520118049cf 100644 --- a/src/transformers/quantizers/quantizer_metal.py +++ b/src/transformers/quantizers/quantizer_metal.py @@ -115,16 +115,10 @@ def get_quantize_ops(self): return MetalQuantize(self) - def get_weight_conversions(self): - from ..core_model_loading import WeightConverter + serialized_names = {"weight": ["weight$", "scales", "qbiases"]} + + def get_load_op(self, converter=None): from ..integrations.metal_quantization import MetalDequantize - if self.pre_quantized and self.quantization_config.dequantize: - return [ - WeightConverter( - source_patterns=["weight$", "scales", "qbiases"], - target_patterns="weight", - operations=[MetalDequantize(self)], - ) - ] - return [] + if converter is None and self.pre_quantized and self.quantization_config.dequantize: + return MetalDequantize(self) diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index ed3c03997f0e..d6c93030227f 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -35,7 +35,6 @@ if is_torch_available(): import torch - from ..core_model_loading import WeightConverter logger = logging.get_logger(__name__) @@ -218,32 +217,13 @@ def get_quantize_ops(self): return Mxfp4Quantize(self) - def get_weight_conversions(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_op(self, converter=None): from ..integrations.mxfp4 import Mxfp4Dequantize, Mxfp4Deserialize - if self.pre_quantized and self.quantization_config.dequantize: - return [ - WeightConverter( - source_patterns=["down_proj_blocks", "down_proj_scales"], - target_patterns=r"down_proj$", - operations=[Mxfp4Dequantize(self)], - ), - WeightConverter( - source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"], - target_patterns=["gate_up_proj$"], - operations=[Mxfp4Dequantize(self)], - ), - ] - - return [ - WeightConverter( - source_patterns=["gate_up_proj_blocks", "gate_up_proj_scales"], - target_patterns=r"gate_up_proj$", - operations=[Mxfp4Deserialize(self)], - ), - WeightConverter( - source_patterns=["down_proj_blocks", "down_proj_scales"], - target_patterns=r"down_proj$", - operations=[Mxfp4Deserialize(self)], - ), - ] + if converter is None and self.pre_quantized: + return Mxfp4Dequantize(self) if self.quantization_config.dequantize else Mxfp4Deserialize(self) diff --git a/src/transformers/quantizers/quantizer_quark.py b/src/transformers/quantizers/quantizer_quark.py index 590159cce77c..4c4df00c0ae2 100644 --- a/src/transformers/quantizers/quantizer_quark.py +++ b/src/transformers/quantizers/quantizer_quark.py @@ -80,25 +80,11 @@ def is_serializable(self): def is_trainable(self): return False - def get_weight_conversions(self): - from ..core_model_loading import WeightConverter + # The checkpoint stores `weight_scale`, ... while `QParamsLinear` holds them in its quantizers (`weight_quantizer.scale`) + serialized_names = {target: [source] for source, target in CHECKPOINT_KEYS.items()} + + def get_load_op(self, converter=None): from ..integrations.quark import QuarkDeserialize - # In Quark, quantization is managed through a QParamsLinear module, which holds - # separate quantizers for the weights, inputs, and biases (e.g. weight_quantizer - # input_quantizer, bias_quantizer, etc.). - # - # The checkpoint stores keys like `weight_scale`, `input_scale`, etc. - # but the model's state_dict() exposes `weight_quantizer.scale`, `input_quantizer.scale`, etc. - # We rename from checkpoint format to model format, and the QuarkDeserialize operation - # handles assigning values into the corresponding quantizer attributes. - converters = [] - for source_key, target_key in CHECKPOINT_KEYS.items(): - converters.append( - WeightConverter( - source_patterns=[source_key], - target_patterns=target_key, - operations=[QuarkDeserialize(self)], - ) - ) - return converters + if converter is None: + return QuarkDeserialize(self) diff --git a/src/transformers/quantizers/quantizer_sinq.py b/src/transformers/quantizers/quantizer_sinq.py index d42d4cc8296d..544ec93fec5c 100644 --- a/src/transformers/quantizers/quantizer_sinq.py +++ b/src/transformers/quantizers/quantizer_sinq.py @@ -168,34 +168,14 @@ def get_quantize_ops(self): return SinqQuantize(self) - def get_weight_conversions(self): - """ - If `pre_quantized=True`, interpret a checkpoint produced by SINQLinear.state_dict: - - .W_q - .bias - .meta + # A checkpoint produced by `SINQLinear.state_dict`, rebuilt into a `SINQLinear` + serialized_names = {".weight": [".W_q", ".meta", ".bias"]} - via a WeightConverter + SinqDeserialize so that we reconstruct a SINQLinear - module instead of a plain nn.Linear. - """ - from ..core_model_loading import WeightConverter + def get_load_op(self, converter=None): + from ..integrations.sinq import SinqDeserialize - if self.pre_quantized: - from ..integrations.sinq import SinqDeserialize - - return [ - WeightConverter( - source_patterns=[ - ".W_q", - ".meta", - ".bias", - ], - target_patterns=[".weight"], - operations=[SinqDeserialize(self)], - ) - ] - return [] + if converter is None and self.pre_quantized: + return SinqDeserialize(self) def _process_model_before_weight_loading( self, diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py index 864757afca59..5b3780bb1004 100644 --- a/src/transformers/quantizers/quantizer_torchao.py +++ b/src/transformers/quantizers/quantizer_torchao.py @@ -31,10 +31,6 @@ MIN_TORCH_VERSION = "2.6.0" -if is_torch_available(): - from ..core_model_loading import WeightConverter - - if is_torch_available(): import torch @@ -203,26 +199,13 @@ def _discover_quantized_param_names(self) -> set[str]: return param_names or {"weight"} - def get_weight_conversions(self): + @property + def serialized_names(self): + suffixes = ["qdata", "scale_and_zero", "per_tensor_scale", "scale", "zero_point", "act_pre_scale"] + return {name: [f"_{name}_{s}" for s in suffixes] for name in sorted(self._discover_quantized_param_names())} + + def get_load_op(self, converter=None): from ..integrations.torchao import TorchAoDeserialize - if self.pre_quantized: - param_names = self._discover_quantized_param_names() - converters = [] - for param_name in sorted(param_names): - converters.append( - WeightConverter( - source_patterns=[ - f"_{param_name}_qdata", - f"_{param_name}_scale_and_zero", - f"_{param_name}_per_tensor_scale", - f"_{param_name}_scale", - f"_{param_name}_zero_point", - f"_{param_name}_act_pre_scale", - ], - target_patterns=param_name, - operations=[TorchAoDeserialize(self)], - ), - ) - return converters - return [] + if converter is None and self.pre_quantized: + return TorchAoDeserialize(self) From d8276525c92d72d9ab78384f99142b29e2194bee Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 13:31:04 +0900 Subject: [PATCH 07/27] nits --- src/transformers/core_model_loading.py | 4 +-- .../integrations/compressed_tensors.py | 2 +- src/transformers/quantizers/base.py | 36 +++++++++---------- .../quantizers/quantizer_bitnet.py | 2 +- .../quantizers/quantizer_bnb_4bit.py | 2 +- .../quantizers/quantizer_bnb_8bit.py | 2 +- .../quantizer_compressed_tensors.py | 10 +++--- .../quantizers/quantizer_finegrained_fp8.py | 8 +---- .../quantizers/quantizer_fp_quant.py | 2 +- src/transformers/quantizers/quantizer_gguf.py | 4 +++ .../quantizers/quantizer_metal.py | 2 +- .../quantizers/quantizer_mxfp4.py | 2 +- .../quantizers/quantizer_quark.py | 2 +- src/transformers/quantizers/quantizer_sinq.py | 2 +- .../quantizers/quantizer_torchao.py | 2 +- 15 files changed, 39 insertions(+), 43 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 13d3203b8ef0..a831e00cb64a 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -278,7 +278,7 @@ def convert( # (such as the MoEs' gate_proj/up_proj merging), we are wasting quite some memory tensors = input_dict.pop(source_pattern) target_pattern = self.get_target_pattern(input_size, source_pattern, target_patterns) - # DecompressExperts pre-allocates a stacked tensor to avoid holding N individual + # CompressedTensorsDequantize pre-allocates a stacked tensor to avoid holding N individual # decompressed tensors simultaneously. Pass it through to skip the redundant copy # that torch.stack would otherwise make. if isinstance(tensors, torch.Tensor): @@ -1719,7 +1719,7 @@ def convert_and_load_state_dict_in_model( elif empty_param is not None and empty_param.dtype != _dtype: _dtype = empty_param.dtype # usually correct when initializing if hf_quantizer and hf_quantizer.pre_quantized: - _dtype = hf_quantizer.get_new_dtype(original_key, renamed_key, tensor, _dtype) + _dtype = hf_quantizer.get_new_dtype(mapping, tensor, _dtype) # Per-expert sharding (EP) needs `tensor_idx` = the expert index so the # distributed op selects whole experts. The signal is a `MergeModulelist` diff --git a/src/transformers/integrations/compressed_tensors.py b/src/transformers/integrations/compressed_tensors.py index 9bfe989442fd..86ba9a20dad9 100644 --- a/src/transformers/integrations/compressed_tensors.py +++ b/src/transformers/integrations/compressed_tensors.py @@ -42,7 +42,7 @@ def get_experts_scheme(quantization_config): return groups[0] -class DecompressExperts(ConversionOps): +class CompressedTensorsDequantize(ConversionOps): """ Dequantize MoE layers when they are in new layout, because they aren't `nn.Module` anymore! diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 88fae536ecaf..a25920715011 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -123,16 +123,15 @@ def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: return param.element_size() - def get_new_dtype( - self, original_key: str, renamed_key: str, param, dtype: "torch.dtype | None" - ) -> "torch.dtype | None": - """dtype to materialize the pre-quantized checkpoint tensor `param` in, `None` keeps the checkpoint's. - `dtype` is the one the model declares for `renamed_key`. By default, quantized (non-float) and renamed - tensors are deserialized by conversion ops, so they keep the checkpoint dtype.""" + def get_new_dtype(self, converter, param, dtype: "torch.dtype | None") -> "torch.dtype | None": + """dtype to materialize the pre-quantized checkpoint tensor `param` in, `None` keeps the checkpoint's: the one + its module declares (`dtype`), unless `param` is quantized (non-float) or an input of a load operation.""" is_float = ( param.get_dtype().startswith(("F", "BF")) if hasattr(param, "get_dtype") else param.is_floating_point() ) - return dtype if is_float and original_key == renamed_key else None + # The quantizer's ops (load operations) are the ones holding it + is_load_input = any(getattr(op, "hf_quantizer", None) for op in getattr(converter, "operations", ())) + return dtype if is_float and not is_load_input else None def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: """adjust max_memory argument for infer_auto_device_map() if extra memory is needed for quantization""" @@ -323,7 +322,7 @@ def get_quantize_ops(self): # e.g. `{"weight": ["weight$", "weight_scale_inv"]}` serialized_names: dict[str, list[str]] = {} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): """Op loading the serialized tensors: a deserialization, or a dequantization when `dequantize=True`. Without `converter`, the op of the general converters. With a model `converter` (e.g. a per-expert merge), the op to prepend to it, if the quantizer's op supports it.""" @@ -333,25 +332,22 @@ def get_weight_conversions(self): "General converters, for the serialized tensors that no model converter claims." from ..core_model_loading import WeightConverter - op = self.get_load_op() + op = self.get_load_operation() return [WeightConverter(names, param, [op]) for param, names in self.serialized_names.items()] if op else [] - def _serialized_sources(self, pattern: str) -> list[str]: - "`experts.*.gate_proj.weight` -> `[experts.*.gate_proj.weight$, experts.*.gate_proj.weight_scale_inv$]`" - for param, names in self.serialized_names.items(): - if pattern.endswith(param := param.rstrip("$")): - return [pattern.removesuffix(param) + name.rstrip("$") + "$" for name in names] - return [pattern] - def update_weight_conversions(self, weight_conversions): - """Prepends `get_load_op` to the model converters, with their serialized sources, then appends the general - `get_weight_conversions`.""" + """Prepends `get_load_operation` to the model converters, with their serialized sources, then appends the + general `get_weight_conversions`.""" from ..core_model_loading import WeightConverter updated = [] for conv in weight_conversions: - op = self.get_load_op(conv) if isinstance(conv, WeightConverter) else None - sources = [s for p in conv._original_source_patterns for s in self._serialized_sources(p)] + op = self.get_load_operation(conv) if isinstance(conv, WeightConverter) else None + # `experts.*.gate_proj.weight` -> `experts.*.gate_proj.weight$`, `experts.*.gate_proj.weight_scale_inv$` + sources = [] + for pattern in conv._original_source_patterns: + prefix, _, param = pattern.rpartition(".") + sources += [f"{prefix}.{n.rstrip('$')}$" for n in self.serialized_names.get(param, [])] or [pattern] if op is not None and sources != conv._original_source_patterns: new = WeightConverter(sources, conv._original_target_patterns, [op] + conv.operations) new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix diff --git a/src/transformers/quantizers/quantizer_bitnet.py b/src/transformers/quantizers/quantizer_bitnet.py index 2432afa96e21..be4bb1e99d38 100644 --- a/src/transformers/quantizers/quantizer_bitnet.py +++ b/src/transformers/quantizers/quantizer_bitnet.py @@ -108,7 +108,7 @@ def is_qat_trainable(self) -> bool: serialized_names = {"weight": ["weight"]} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.bitnet import BitNetDeserialize config = self.quantization_config diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py index e2e2d7886db0..63725f876204 100644 --- a/src/transformers/quantizers/quantizer_bnb_4bit.py +++ b/src/transformers/quantizers/quantizer_bnb_4bit.py @@ -176,7 +176,7 @@ def get_quantize_ops(self): ] } - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.bitsandbytes import Bnb4bitDeserialize if converter is None and self.pre_quantized: diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py index 57ee3377ab35..d8b6caffae22 100644 --- a/src/transformers/quantizers/quantizer_bnb_8bit.py +++ b/src/transformers/quantizers/quantizer_bnb_8bit.py @@ -165,7 +165,7 @@ def get_quantize_ops(self): serialized_names = {"weight": ["SCB", "weight_format", "weight"]} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.bitsandbytes import Bnb8bitDeserialize if converter is None and self.pre_quantized: diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index afc623cef4e7..6a5295832576 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -23,7 +23,7 @@ import torch from ..core_model_loading import WeightConverter - from ..integrations.compressed_tensors import DecompressExperts, get_experts_scheme + from ..integrations.compressed_tensors import CompressedTensorsDequantize, get_experts_scheme logger = logging.get_logger(__name__) @@ -173,7 +173,9 @@ def serialized_names(self): return {"weight": ["weight", "weight_scale"]} return {"weight": ["weight_packed", "weight_scale", "weight_shape"]} - def get_load_op(self, converter=None): - # Merged experts are not `nn.Linear`, so they can't stay compressed: decompress them before the merge + def get_load_operation(self, converter=None): + # Merged experts are not quantized modules: dequantize their per-expert weights before the merge if converter is not None and all("experts" in p for p in converter.source_patterns): - return DecompressExperts(self, scheme=get_experts_scheme(self.quantization_config.quantization_config)) + return CompressedTensorsDequantize( + self, scheme=get_experts_scheme(self.quantization_config.quantization_config) + ) diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index bd95cf1114eb..e44dfab5f522 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -159,15 +159,9 @@ def get_quantize_ops(self): return Fp8Quantize(self) - def get_new_dtype(self, original_key, renamed_key, param, dtype): - if self.quantization_config.dequantize: - return super().get_new_dtype(original_key, renamed_key, param, dtype) - # `FP8Linear` / `FP8Experts` declare the dtype of every param, e.g. ue8m0 scales stored as float32 - return dtype - serialized_names = {"weight": ["weight$", "weight_scale_inv", "activation_scale"]} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.finegrained_fp8 import Fp8Dequantize if self.pre_quantized and self.quantization_config.dequantize: diff --git a/src/transformers/quantizers/quantizer_fp_quant.py b/src/transformers/quantizers/quantizer_fp_quant.py index e3c36b2f58cf..d0f88b536a1c 100644 --- a/src/transformers/quantizers/quantizer_fp_quant.py +++ b/src/transformers/quantizers/quantizer_fp_quant.py @@ -145,7 +145,7 @@ def serialized_names(self): name = ".dqweight" if self.quantization_config.pseudoquantization else ".qweight" return {name: [name]} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.fp_quant import FpQuantDeserialize if converter is None and self.pre_quantized: diff --git a/src/transformers/quantizers/quantizer_gguf.py b/src/transformers/quantizers/quantizer_gguf.py index 362ab1c69f2e..33cd5c26c970 100644 --- a/src/transformers/quantizers/quantizer_gguf.py +++ b/src/transformers/quantizers/quantizer_gguf.py @@ -126,6 +126,10 @@ def update_dtype(self, dtype): self.dtype = dtype return dtype + def get_new_dtype(self, converter, param, dtype): + # The transforms keep the file's precision, `Cast` (last op of every chain) lands in `self.dtype` + return None if self.supported else super().get_new_dtype(converter, param, dtype) + def get_state_dict(self, checkpoint_file: str, model): """The file's tensors, quantized ones kept as raw blocks.""" if self.supported: diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py index 6520118049cf..ff8bd9ccf96a 100644 --- a/src/transformers/quantizers/quantizer_metal.py +++ b/src/transformers/quantizers/quantizer_metal.py @@ -117,7 +117,7 @@ def get_quantize_ops(self): serialized_names = {"weight": ["weight$", "scales", "qbiases"]} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.metal_quantization import MetalDequantize if converter is None and self.pre_quantized and self.quantization_config.dequantize: diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index d6c93030227f..d5f8dd831bd6 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -222,7 +222,7 @@ def get_quantize_ops(self): "down_proj$": ["down_proj_blocks", "down_proj_scales"], } - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.mxfp4 import Mxfp4Dequantize, Mxfp4Deserialize if converter is None and self.pre_quantized: diff --git a/src/transformers/quantizers/quantizer_quark.py b/src/transformers/quantizers/quantizer_quark.py index 4c4df00c0ae2..c140f464fdd6 100644 --- a/src/transformers/quantizers/quantizer_quark.py +++ b/src/transformers/quantizers/quantizer_quark.py @@ -83,7 +83,7 @@ def is_trainable(self): # The checkpoint stores `weight_scale`, ... while `QParamsLinear` holds them in its quantizers (`weight_quantizer.scale`) serialized_names = {target: [source] for source, target in CHECKPOINT_KEYS.items()} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.quark import QuarkDeserialize if converter is None: diff --git a/src/transformers/quantizers/quantizer_sinq.py b/src/transformers/quantizers/quantizer_sinq.py index 544ec93fec5c..a6fa15205dfa 100644 --- a/src/transformers/quantizers/quantizer_sinq.py +++ b/src/transformers/quantizers/quantizer_sinq.py @@ -171,7 +171,7 @@ def get_quantize_ops(self): # A checkpoint produced by `SINQLinear.state_dict`, rebuilt into a `SINQLinear` serialized_names = {".weight": [".W_q", ".meta", ".bias"]} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.sinq import SinqDeserialize if converter is None and self.pre_quantized: diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py index 5b3780bb1004..9b71586092bb 100644 --- a/src/transformers/quantizers/quantizer_torchao.py +++ b/src/transformers/quantizers/quantizer_torchao.py @@ -204,7 +204,7 @@ def serialized_names(self): suffixes = ["qdata", "scale_and_zero", "per_tensor_scale", "scale", "zero_point", "act_pre_scale"] return {name: [f"_{name}_{s}" for s in suffixes] for name in sorted(self._discover_quantized_param_names())} - def get_load_op(self, converter=None): + def get_load_operation(self, converter=None): from ..integrations.torchao import TorchAoDeserialize if converter is None and self.pre_quantized: From 22306f02f0305ca6a1e7da3be10d1edf696b57eb Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 13:40:35 +0900 Subject: [PATCH 08/27] fucking bloatland --- src/transformers/core_model_loading.py | 4 ++++ .../integrations/finegrained_fp8.py | 24 ++++--------------- .../integrations/metal_quantization.py | 4 ++-- src/transformers/quantizers/base.py | 10 ++++---- .../quantizers/quantizer_bnb_4bit.py | 2 +- .../quantizers/quantizer_bnb_8bit.py | 2 +- .../quantizers/quantizer_finegrained_fp8.py | 2 +- .../quantizers/quantizer_metal.py | 2 +- .../quantizers/quantizer_mxfp4.py | 4 ++-- 9 files changed, 21 insertions(+), 33 deletions(-) diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index a831e00cb64a..838a01440640 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -880,6 +880,10 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list for i, source_pattern in enumerate(self.source_patterns): group_name = f"g{i}" pattern = source_pattern.replace("*.", r".*\.") + # A pattern ends on a name boundary: `weight` matches `q_proj.weight`, not `q_proj.weight_scale_inv`. + # One ending with `_` is the start of a name, e.g. `hc_` for `hc_fn` + if re.search(r"[a-zA-Z0-9]$", pattern): + pattern += r"(?!\w)" branches.append(f"(?P<{group_name}>{pattern})") self.compiled_sources = re.compile("|".join(branches)) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index f00d024b36c8..525c53fb4d73 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -1000,18 +1000,6 @@ class Fp8Dequantize(ConversionOps): def __init__(self, hf_quantizer): self.hf_quantizer = hf_quantizer - def _scale_pattern_for(self, weight_pattern: str) -> str: - # Strip the optional ``$`` regex anchor so we can match the underlying name. - anchored = weight_pattern.endswith("$") - base = weight_pattern[:-1] if anchored else weight_pattern - if base.endswith(".weight"): - scale = base[: -len(".weight")] + ".weight_scale_inv" - elif base == "weight": - scale = "weight_scale_inv" - else: - scale = base + "_scale_inv" - return scale + "$" if anchored else scale - # E2M1 (FP4) value table — checkpoints sometimes ship MoE experts as packed FP4 # (two e2m1 nibbles per int8 byte), so the "weight" dtype lands as ``int8`` / # ``float4_e2m1fn_x2`` and we have to unpack before applying the scale grid. @@ -1086,16 +1074,14 @@ def convert( **kwargs, ) -> dict[str, list[torch.Tensor] | torch.Tensor]: output_dtype = self._get_target_dtype(model, full_layer_name) - # Backward-compatible single-tensor path (the legacy fallback converter declares - # ``["weight$", "weight_scale_inv", "activation_scale"]`` and produces a single - # ``weight`` target). Also handles the no-scale case (e.g. RMSNorm weights that - # match ``weight$`` but ship no ``weight_scale_inv`` alongside). - if "weight$" in input_dict: + # Single-tensor path (the general converter declares ``["weight", "weight_scale_inv", "activation_scale"]`` + # and produces a single ``weight`` target). Also handles the no-scale case (e.g. RMSNorm weights). + if "weight" in input_dict: # The downstream renamer in `core_model_loading._convert_one_module` uses the # output dict's *key*, not its content, to derive prefix/suffix; if `full_layer_name` # is unset (direct invocation / tests) fall back to the legacy converter's target. target_key = full_layer_name if full_layer_name is not None else "weight" - quantized = input_dict["weight$"] + quantized = input_dict["weight"] quantized = quantized[0] if isinstance(quantized, list) else quantized if "weight_scale_inv" in input_dict: scales = input_dict["weight_scale_inv"] @@ -1108,7 +1094,7 @@ def convert( for key, value in input_dict.items(): if "activation_scale" in key or "weight_scale_inv" in key: continue # consumed by the dequant; drop from the chain - scale_key = self._scale_pattern_for(key) + scale_key = f"{key}_scale_inv" if scale_key not in input_dict: # No scale to apply (e.g. unrelated entry) — pass through untouched. result[key] = value diff --git a/src/transformers/integrations/metal_quantization.py b/src/transformers/integrations/metal_quantization.py index c51797e40a2c..6bccc212aa53 100644 --- a/src/transformers/integrations/metal_quantization.py +++ b/src/transformers/integrations/metal_quantization.py @@ -286,9 +286,9 @@ def convert(self, input_dict: dict, full_layer_name: str | None = None, **kwargs group_size = self.hf_quantizer.quantization_config.group_size if len(input_dict) < 2: - return {full_layer_name: input_dict["weight$"]} + return {full_layer_name: input_dict["weight"]} - quantized = input_dict["weight$"][0] + quantized = input_dict["weight"][0] scales = input_dict["scales"][0] qbiases = input_dict["qbiases"][0] diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index a25920715011..47e753b64091 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -164,9 +164,7 @@ def _update_plan(self, model, plan): from ..distributed.tensor_parallel import replace_layer_number_by_wildcard params = {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()} - suffixes = { - n.rstrip("$").removeprefix(p.rstrip("$")) for p, names in self.serialized_names.items() for n in names - } + suffixes = {n.removeprefix(p) for p, names in self.serialized_names.items() for n in names} extra = {key + suffix: style for key, style in plan.items() for suffix in suffixes} plan.update({key: style for key, style in extra.items() if key in params and key not in plan}) @@ -319,7 +317,7 @@ def get_quantize_ops(self): ) # How params are saved in pre-quantized checkpoints: `{param: [serialized source patterns]}`, - # e.g. `{"weight": ["weight$", "weight_scale_inv"]}` + # e.g. `{"weight": ["weight", "weight_scale_inv"]}` serialized_names: dict[str, list[str]] = {} def get_load_operation(self, converter=None): @@ -343,11 +341,11 @@ def update_weight_conversions(self, weight_conversions): updated = [] for conv in weight_conversions: op = self.get_load_operation(conv) if isinstance(conv, WeightConverter) else None - # `experts.*.gate_proj.weight` -> `experts.*.gate_proj.weight$`, `experts.*.gate_proj.weight_scale_inv$` + # `experts.*.gate_proj.weight` -> `experts.*.gate_proj.weight`, `experts.*.gate_proj.weight_scale_inv` sources = [] for pattern in conv._original_source_patterns: prefix, _, param = pattern.rpartition(".") - sources += [f"{prefix}.{n.rstrip('$')}$" for n in self.serialized_names.get(param, [])] or [pattern] + sources += [f"{prefix}.{n}" for n in self.serialized_names.get(param, [])] or [pattern] if op is not None and sources != conv._original_source_patterns: new = WeightConverter(sources, conv._original_target_patterns, [op] + conv.operations) new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py index 63725f876204..791ef9f4e9ef 100644 --- a/src/transformers/quantizers/quantizer_bnb_4bit.py +++ b/src/transformers/quantizers/quantizer_bnb_4bit.py @@ -166,13 +166,13 @@ def get_quantize_ops(self): serialized_names = { "weight": [ + "weight", "weight.nested_absmax", "weight.nested_quant_map", "weight.quant_map", "weight.absmax", "weight.quant_state.bitsandbytes__nf4", "weight.quant_state.bitsandbytes__fp4", - "weight", ] } diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py index d8b6caffae22..098c723defdd 100644 --- a/src/transformers/quantizers/quantizer_bnb_8bit.py +++ b/src/transformers/quantizers/quantizer_bnb_8bit.py @@ -163,7 +163,7 @@ def get_quantize_ops(self): return Bnb8bitQuantize(self) - serialized_names = {"weight": ["SCB", "weight_format", "weight"]} + serialized_names = {"weight": ["weight", "SCB", "weight_format"]} def get_load_operation(self, converter=None): from ..integrations.bitsandbytes import Bnb8bitDeserialize diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index e44dfab5f522..ba49f3baacbf 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -159,7 +159,7 @@ def get_quantize_ops(self): return Fp8Quantize(self) - serialized_names = {"weight": ["weight$", "weight_scale_inv", "activation_scale"]} + serialized_names = {"weight": ["weight", "weight_scale_inv", "activation_scale"]} def get_load_operation(self, converter=None): from ..integrations.finegrained_fp8 import Fp8Dequantize diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py index ff8bd9ccf96a..7d03e0b6267c 100644 --- a/src/transformers/quantizers/quantizer_metal.py +++ b/src/transformers/quantizers/quantizer_metal.py @@ -115,7 +115,7 @@ def get_quantize_ops(self): return MetalQuantize(self) - serialized_names = {"weight": ["weight$", "scales", "qbiases"]} + serialized_names = {"weight": ["weight", "scales", "qbiases"]} def get_load_operation(self, converter=None): from ..integrations.metal_quantization import MetalDequantize diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index d5f8dd831bd6..541f6ceaec4b 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -218,8 +218,8 @@ def get_quantize_ops(self): return Mxfp4Quantize(self) serialized_names = { - "gate_up_proj$": ["gate_up_proj_blocks", "gate_up_proj_scales"], - "down_proj$": ["down_proj_blocks", "down_proj_scales"], + "gate_up_proj": ["gate_up_proj_blocks", "gate_up_proj_scales"], + "down_proj": ["down_proj_blocks", "down_proj_scales"], } def get_load_operation(self, converter=None): From 1986e7713974bfc736afef939ab8f62825806cfd Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 13:41:39 +0900 Subject: [PATCH 09/27] Fix PR slow CI suggestion install: non-editable install --- .github/workflows/pr_slow_ci_suggestion.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pr_slow_ci_suggestion.yml b/.github/workflows/pr_slow_ci_suggestion.yml index 050822430208..3f951b9d123a 100644 --- a/.github/workflows/pr_slow_ci_suggestion.yml +++ b/.github/workflows/pr_slow_ci_suggestion.yml @@ -97,7 +97,7 @@ jobs: # The script imports `transformers`. This is the base branch checkout, not the PR's code. - name: Install transformers - run: pip install -e . + run: pip install . - name: Run script to get jobs to run id: get_jobs From 37a685f284cceaa73def067dc50c2d20d6229db7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 13:44:29 +0900 Subject: [PATCH 10/27] update now that the $ bloat is gone --- src/transformers/conversion_mapping.py | 66 +++++++++---------- src/transformers/core_model_loading.py | 14 ++-- src/transformers/fusion_mapping.py | 8 +-- .../gguf/gguf_conversion_mapping.py | 4 +- .../quantizers/quantizer_finegrained_fp8.py | 2 +- 5 files changed, 45 insertions(+), 49 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 01ba673a7dad..9e51f58d07e2 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -151,7 +151,7 @@ def _build_checkpoint_conversion_mapping(): # General HC prefix which is dropped WeightRenaming(r"\.hc_pre\.hc_", ".hc_"), # Attn gating + sinks - WeightRenaming(r"\.learnable_sink_param$", ".sinks"), + WeightRenaming(r"\.learnable_sink_param", ".sinks"), WeightRenaming(r"\.linear_gate", ".gate_proj"), # Follow DSv4 HC standards WeightRenaming(r"\.hc_attn_layer\.hc_fn", ".attn_hc.fn"), @@ -246,8 +246,8 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(source_patterns=r"\.attn\.k_sconv", target_patterns=r".self_attn.k_sconv.conv1d"), WeightRenaming(source_patterns=r"\.attn\.v_sconv", target_patterns=r".self_attn.v_sconv.conv1d"), WeightRenaming(source_patterns=r"\.attn\.rel_logits_proj", target_patterns=r".self_attn.rel_logits_proj"), - WeightRenaming(source_patterns=r"attn_sconv\.weight$", target_patterns=r"attn_sconv.conv1d.weight"), - WeightRenaming(source_patterns=r"mlp_sconv\.weight$", target_patterns=r"mlp_sconv.conv1d.weight"), + WeightRenaming(source_patterns=r"attn_sconv\.weight", target_patterns=r"attn_sconv.conv1d.weight"), + WeightRenaming(source_patterns=r"mlp_sconv\.weight", target_patterns=r"mlp_sconv.conv1d.weight"), WeightRenaming(source_patterns=r"mlp_norm", target_patterns=r"post_attention_layernorm"), WeightRenaming(source_patterns=r"attn_norm", target_patterns=r"input_layernorm"), ], @@ -469,8 +469,8 @@ def _build_checkpoint_conversion_mapping(): ], "SegformerForSemanticSegmentation": [WeightRenaming("decode_head.linear_c", "decode_head.linear_projections")], "videomae": [ - WeightRenaming(r"attention\.attention\.q_bias$", "attention.attention.query.bias"), - WeightRenaming(r"attention\.attention\.v_bias$", "attention.attention.value.bias"), + WeightRenaming(r"attention\.attention\.q_bias", "attention.attention.query.bias"), + WeightRenaming(r"attention\.attention\.v_bias", "attention.attention.value.bias"), ], "swin": [ WeightRenaming("attention.self.query", "attention.q_proj"), @@ -550,27 +550,27 @@ def _build_checkpoint_conversion_mapping(): 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"), - 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"^embed\.weight", target_patterns="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"\.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"), - WeightRenaming(source_patterns=r"\.hc_attn_base$", target_patterns=r".attn_hc.base"), - WeightRenaming(source_patterns=r"\.hc_attn_scale$", target_patterns=r".attn_hc.scale"), - WeightRenaming(source_patterns=r"\.hc_ffn_fn$", target_patterns=r".ffn_hc.fn"), - WeightRenaming(source_patterns=r"\.hc_ffn_base$", target_patterns=r".ffn_hc.base"), - WeightRenaming(source_patterns=r"\.hc_ffn_scale$", target_patterns=r".ffn_hc.scale"), - WeightRenaming(source_patterns=r"\.attn_sink$", target_patterns=r".sinks"), + WeightRenaming(source_patterns=r"\.hc_attn_fn", target_patterns=r".attn_hc.fn"), + WeightRenaming(source_patterns=r"\.hc_attn_base", target_patterns=r".attn_hc.base"), + WeightRenaming(source_patterns=r"\.hc_attn_scale", target_patterns=r".attn_hc.scale"), + WeightRenaming(source_patterns=r"\.hc_ffn_fn", target_patterns=r".ffn_hc.fn"), + WeightRenaming(source_patterns=r"\.hc_ffn_base", target_patterns=r".ffn_hc.base"), + WeightRenaming(source_patterns=r"\.hc_ffn_scale", target_patterns=r".ffn_hc.scale"), + WeightRenaming(source_patterns=r"\.attn_sink", target_patterns=r".sinks"), WeightRenaming( source_patterns=r"\.indexer\.weights_proj\.", 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\."), - WeightRenaming(source_patterns=r"\.ape$", target_patterns=r"\.position_bias"), + 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."), @@ -578,7 +578,7 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(source_patterns=r"\.wo_a\.", target_patterns=r".o_a_proj."), WeightRenaming(source_patterns=r"\.wo_b\.", target_patterns=r".o_b_proj."), WeightRenaming(source_patterns=r"\.q_norm\.", target_patterns=r".q_a_norm."), - WeightRenaming(source_patterns=r"\.gate\.bias$", target_patterns=r".gate.e_score_correction_bias"), + WeightRenaming(source_patterns=r"\.gate\.bias", target_patterns=r".gate.e_score_correction_bias"), WeightRenaming(source_patterns=r"\.shared_experts\.w1\.", target_patterns=r".shared_experts.gate_proj."), WeightRenaming(source_patterns=r"\.shared_experts\.w2\.", target_patterns=r".shared_experts.down_proj."), WeightRenaming(source_patterns=r"\.shared_experts\.w3\.", target_patterns=r".shared_experts.up_proj."), @@ -996,14 +996,14 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(source_patterns=r"^model\.layers\.", target_patterns="model.language_model.layers."), WeightRenaming(source_patterns=r"^model\.norm\.", target_patterns="model.language_model.norm."), # ---- Pass 2: MoE renames ---- - WeightRenaming(source_patterns=r"\.moe\.gate\.weight$", target_patterns=".mlp.gate.weight"), + WeightRenaming(source_patterns=r"\.moe\.gate\.weight", target_patterns=".mlp.gate.weight"), WeightRenaming( - source_patterns=r"\.moe\.router_bias$", + source_patterns=r"\.moe\.router_bias", target_patterns=".mlp.gate.e_score_correction_bias", ), - WeightRenaming(source_patterns=r"\.moe\.down_proj\.weight$", target_patterns=".mlp.experts.down_proj"), + WeightRenaming(source_patterns=r"\.moe\.down_proj\.weight", target_patterns=".mlp.experts.down_proj"), WeightRenaming( - source_patterns=r"\.moe\.down_proj\.weight_scale_inv$", + source_patterns=r"\.moe\.down_proj\.weight_scale_inv", target_patterns=".mlp.experts.down_proj_scale_inv", ), WeightRenaming(source_patterns=r"\.share_expert\.", target_patterns=".mlp.shared_experts."), @@ -1018,17 +1018,17 @@ def _build_checkpoint_conversion_mapping(): # Scoped separately from "step3p7" above: bare substring patterns like `\.attn\.` -> `.self_attn.` # would otherwise also match the text decoder's `language_model.layers.*.self_attn.*` keys. "step3p5_vision": [ - WeightRenaming(source_patterns=r"^conv1\.weight$", target_patterns="embeddings.patch_embedding.weight"), + WeightRenaming(source_patterns=r"^conv1\.weight", target_patterns="embeddings.patch_embedding.weight"), WeightRenaming( - source_patterns=r"^positional_embedding$", + source_patterns=r"^positional_embedding", target_patterns="embeddings.position_embedding.weight", ), WeightRenaming(source_patterns=r"^transformer\.resblocks\.", target_patterns="layers."), WeightRenaming(source_patterns=r"^ln_pre\.", target_patterns="pre_layernorm."), WeightRenaming(source_patterns=r"^vit_downsampler1\.", target_patterns="downsampler1."), WeightRenaming(source_patterns=r"^vit_downsampler2\.", target_patterns="downsampler2."), - WeightRenaming(source_patterns=r"\.ls_1\.gamma$", target_patterns=".lambda_1"), - WeightRenaming(source_patterns=r"\.ls_2\.gamma$", target_patterns=".lambda_2"), + WeightRenaming(source_patterns=r"\.ls_1\.gamma", target_patterns=".lambda_1"), + WeightRenaming(source_patterns=r"\.ls_2\.gamma", target_patterns=".lambda_2"), WeightRenaming(source_patterns=r"\.mlp\.c_fc\.", target_patterns=".mlp.fc1."), WeightRenaming(source_patterns=r"\.mlp\.c_proj\.", target_patterns=".mlp.fc2."), WeightRenaming(source_patterns=r"\.attn\.", target_patterns=".self_attn."), @@ -1586,8 +1586,8 @@ def _build_checkpoint_conversion_mapping(): ), ], "sapiens2": [ - WeightRenaming(r"^cls_token$", r"embeddings.cls_token"), - WeightRenaming(r"^storage_tokens$", r"embeddings.register_tokens"), + WeightRenaming(r"^cls_token", r"embeddings.cls_token"), + WeightRenaming(r"^storage_tokens", r"embeddings.register_tokens"), WeightRenaming(r"^patch_embed\.projection\.", r"embeddings.patch_embeddings."), WeightRenaming(r"blocks\.(\d+)\.", r"model.layer.\1."), WeightRenaming(r"attn\.proj\.", r"attention.o_proj."), @@ -1596,7 +1596,7 @@ def _build_checkpoint_conversion_mapping(): WeightRenaming(r"attn\.wv\.", r"attention.v_proj."), WeightRenaming(r"attn\.q_norm\.", r"attention.q_norm."), WeightRenaming(r"attn\.k_norm\.", r"attention.k_norm."), - WeightRenaming(r"attn\.gamma\.weight$", r"layer_scale1.lambda1"), + WeightRenaming(r"attn\.gamma\.weight", r"layer_scale1.lambda1"), WeightRenaming(r"ffn\.w3\.", r"mlp.down_proj."), WeightRenaming(r"\.ln1\.", r".norm1."), WeightRenaming(r"\.ln2\.", r".norm2."), @@ -1899,11 +1899,11 @@ def _build_checkpoint_conversion_mapping(): mapping["legacy"] += [ WeightRenaming( - source_patterns=".weight_g$", + source_patterns=".weight_g", target_patterns=".parametrizations.weight.original0", ), WeightRenaming( - source_patterns=".weight_v$", + source_patterns=".weight_v", target_patterns=".parametrizations.weight.original1", ), ] @@ -1991,7 +1991,7 @@ def _build_checkpoint_conversion_mapping(): mapping["qwen4_exp_text"] += [ WeightConverter( source_patterns="ngram_embedding.shard_*.weight", - target_patterns="ngram_embedding.weight$", + target_patterns="ngram_embedding.weight", operations=[Concatenate(dim=0, num_shards_attribute="split_ngram_parts")], # The size of the embedding is ~95 GiB, so we cannot afford to perform the Cat on device, as it will need # a temporary memory buffer of the same size diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 838a01440640..8873f27e1e6c 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -754,7 +754,7 @@ def process_target_pattern(pattern: str) -> tuple[str, str | None]: Process a target pattern for reverse mapping (when targets become sources). This handles several edge cases in checkpoint conversion mappings: - - Removes `^` prefix and `$` suffix (start/end of string anchors) + - Removes the `^` prefix (start of string anchor) - Removes negative lookahead/lookbehind assertions - Detects capturing groups and replaces them with `\\1` backreference @@ -767,8 +767,6 @@ def process_target_pattern(pattern: str) -> tuple[str, str | None]: """ # Some mapping contains `^` to notify start of string when matching -> remove it during reverse mapping pattern = pattern.removeprefix("^") - # Some mapping contains `$` to notify end of string when matching -> remove it during reverse mapping - pattern = pattern.removesuffix("$") # Remove negative lookahead/behind if any. This is ugly but needed for reverse mapping of # Qwen2.5, Sam3, Ernie4.5 VL MoE! It needs to be non greedy in case there are several pattern = re.sub(r"\(\?.+?\)?\)", "", pattern) @@ -786,15 +784,13 @@ def process_target_pattern(pattern: str) -> tuple[str, str | None]: def process_source_pattern(source_pattern: str, target_pattern: str) -> str: """ Process a source pattern for reverse mapping (when sources become targets). - This is useful because usually if the original source (so now the target in reverse mode) had a `^` or `$` - to restrict to start/end of string, we should do the same in reverse mode. This is why this method is conditioned + This is useful because usually if the original source (so now the target in reverse mode) had a `^` + to restrict to the start of string, we should do the same in reverse mode. This is why this method is conditioned on the target pattern, we want to do it only for pairs (source, target) when the original source (so the current target in reverse mode) had it. """ if target_pattern.startswith("^"): source_pattern = f"^{source_pattern}" if not source_pattern.startswith("^") else source_pattern - if target_pattern.endswith("$"): - source_pattern = f"{source_pattern}$" if not source_pattern.endswith("$") else source_pattern return source_pattern @@ -1158,11 +1154,11 @@ def __init__( if prefix_to_add is not None: super().__init__( # We use a lookbehind to avoid adding the prefix if we detect that it's already present - source_patterns=rf"^{prefix}(?:(?!{prefix_to_add}\.))(.+)$", + source_patterns=rf"^{prefix}(?:(?!{prefix_to_add}\.))(.+)", target_patterns=rf"{prefix}{prefix_to_add}\.\1", ) else: - super().__init__(source_patterns=rf"^{prefix}{prefix_to_remove}\.(.+)$", target_patterns=rf"{prefix}\1") + super().__init__(source_patterns=rf"^{prefix}{prefix_to_remove}\.(.+)", target_patterns=rf"{prefix}\1") def reverse_transform(self) -> WeightTransform: """Reverse the current `WeightTransform` instance, to be able to save with the opposite weight transformations.""" diff --git a/src/transformers/fusion_mapping.py b/src/transformers/fusion_mapping.py index c3e0157024a3..edb425cbc7fe 100644 --- a/src/transformers/fusion_mapping.py +++ b/src/transformers/fusion_mapping.py @@ -123,8 +123,8 @@ def make_transforms(self, config: "PretrainedConfig") -> list[WeightTransform]: return [ WeightConverter( - source_patterns=r"patch_embed\.proj\.weight$", - target_patterns=r"patch_embed\.linear_proj\.weight$", + source_patterns=r"patch_embed\.proj\.weight", + target_patterns=r"patch_embed\.linear_proj\.weight", operations=[ Conv3dToLinear( in_channels=in_channels, @@ -133,8 +133,8 @@ def make_transforms(self, config: "PretrainedConfig") -> list[WeightTransform]: ], ), WeightRenaming( - source_patterns=r"patch_embed\.proj\.bias$", - target_patterns=r"patch_embed\.linear_proj\.bias$", + source_patterns=r"patch_embed\.proj\.bias", + target_patterns=r"patch_embed\.linear_proj\.bias", ), ] diff --git a/src/transformers/integrations/gguf/gguf_conversion_mapping.py b/src/transformers/integrations/gguf/gguf_conversion_mapping.py index ce038a9d9d3e..42732cfc15c9 100644 --- a/src/transformers/integrations/gguf/gguf_conversion_mapping.py +++ b/src/transformers/integrations/gguf/gguf_conversion_mapping.py @@ -90,8 +90,8 @@ def _qwen35(config) -> list[WeightTransform]: WeightRenaming(r"\.ssm_conv1d\.", ".linear_attn.conv1d."), WeightRenaming(r"\.ssm_norm\.", ".linear_attn.norm."), # the one norm with no offset WeightRenaming(r"\.ssm_out\.", ".linear_attn.out_proj."), - WeightRenaming(r"\.ssm_a$", ".linear_attn.A_log"), - WeightRenaming(r"\.ssm_dt\.bias$", ".linear_attn.dt_bias"), + WeightRenaming(r"\.ssm_a", ".linear_attn.A_log"), + WeightRenaming(r"\.ssm_dt\.bias", ".linear_attn.dt_bias"), ] ) diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index ba49f3baacbf..38caa07581bd 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -171,5 +171,5 @@ def update_weight_conversions(self, weight_conversions): from ..core_model_loading import WeightRenaming # Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) ship per-block scales under `.scale` - scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale$", target_patterns=r"\1.weight_scale_inv") + scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale", target_patterns=r"\1.weight_scale_inv") return [scale_rename] + super().update_weight_conversions(weight_conversions) From 1b89c27fa27bc7211b9097829970212d9c395030 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 13:50:32 +0900 Subject: [PATCH 11/27] Fix PR slow CI suggestion: install transformers in a venv (system site-packages not writable) --- .github/workflows/pr_slow_ci_suggestion.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pr_slow_ci_suggestion.yml b/.github/workflows/pr_slow_ci_suggestion.yml index 3f951b9d123a..07a5d0eeb8c3 100644 --- a/.github/workflows/pr_slow_ci_suggestion.yml +++ b/.github/workflows/pr_slow_ci_suggestion.yml @@ -97,13 +97,13 @@ jobs: # The script imports `transformers`. This is the base branch checkout, not the PR's code. - name: Install transformers - run: pip install . + run: python -m venv ~/venv && ~/venv/bin/pip install . - name: Run script to get jobs to run id: get_jobs run: | set -o pipefail - python utils/get_pr_run_slow_jobs.py | tee output.txt + ~/venv/bin/python utils/get_pr_run_slow_jobs.py | tee output.txt echo "jobs_to_run: $(tail -n 1 output.txt)" echo "jobs_to_run=$(tail -n 1 output.txt)" >> $GITHUB_OUTPUT From 54b90378b3532715a92e5965015644da189ca83e Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 13:51:41 +0900 Subject: [PATCH 12/27] metal tests: unanchored weight key --- tests/quantization/metal/test_metal.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/quantization/metal/test_metal.py b/tests/quantization/metal/test_metal.py index 5bafb2a45d84..f876194afc93 100644 --- a/tests/quantization/metal/test_metal.py +++ b/tests/quantization/metal/test_metal.py @@ -460,7 +460,7 @@ def test_metal_dequantize_returns_target_dtype(self): dq_result = dq_op.convert( { - "weight$": [q_result["layer.weight"]], + "weight": [q_result["layer.weight"]], "scales": [q_result["layer.scales"]], "qbiases": [q_result["layer.qbiases"]], }, @@ -482,7 +482,7 @@ def test_quantize_then_dequantize_roundtrip(self): dq_op = MetalDequantize(dq_quantizer) dq_result = dq_op.convert( { - "weight$": [q_result["layer.weight"]], + "weight": [q_result["layer.weight"]], "scales": [q_result["layer.scales"]], "qbiases": [q_result["layer.qbiases"]], }, From 6ec139fd4852902da72c0580e8958d96dcd93a33 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 14:02:36 +0900 Subject: [PATCH 13/27] nits this time --- .../integrations/compressed_tensors.py | 7 ++++- src/transformers/integrations/fbgemm_fp8.py | 4 ++- .../integrations/finegrained_fp8.py | 2 -- src/transformers/integrations/mxfp4.py | 5 ++-- src/transformers/modeling_utils.py | 2 +- src/transformers/quantizers/base.py | 22 +++++++--------- src/transformers/quantizers/quantizer_aqlm.py | 4 --- .../quantizers/quantizer_auto_round.py | 4 --- src/transformers/quantizers/quantizer_awq.py | 3 --- .../quantizers/quantizer_bitnet.py | 3 --- .../quantizers/quantizer_bnb_4bit.py | 13 +--------- .../quantizers/quantizer_bnb_8bit.py | 12 +-------- src/transformers/quantizers/quantizer_eetq.py | 3 --- .../quantizers/quantizer_fbgemm_fp8.py | 11 +------- .../quantizers/quantizer_finegrained_fp8.py | 4 +-- .../quantizers/quantizer_fouroversix.py | 26 ++++--------------- .../quantizers/quantizer_fp_quant.py | 4 --- .../quantizers/quantizer_higgs.py | 4 --- src/transformers/quantizers/quantizer_hqq.py | 5 ---- .../quantizers/quantizer_metal.py | 3 --- .../quantizers/quantizer_nvfp4.py | 9 ++----- src/transformers/quantizers/quantizer_spqr.py | 4 --- .../quantizers/quantizer_torchao.py | 9 +------ src/transformers/quantizers/quantizer_vptq.py | 4 --- tests/quantization/nvfp4/test_nvfp4.py | 6 ++--- 25 files changed, 38 insertions(+), 135 deletions(-) diff --git a/src/transformers/integrations/compressed_tensors.py b/src/transformers/integrations/compressed_tensors.py index 86ba9a20dad9..ef64a2ca3350 100644 --- a/src/transformers/integrations/compressed_tensors.py +++ b/src/transformers/integrations/compressed_tensors.py @@ -71,6 +71,7 @@ def convert( from compressed_tensors.compressors.format import infer_module_format ct_quantization_config = self.hf_quantizer.compressor.quantization_config + dtype = kwargs["model"].get_parameter(full_layer_name).dtype quantization_scheme = self.scheme or get_experts_scheme(ct_quantization_config) @@ -92,7 +93,7 @@ def convert( scales = scales if isinstance(scales, list) else [scales] output = None for i, (quant, scale) in enumerate(zip(quantized, scales)): - dense = (quant.to(torch.float32) * scale.to(torch.float32)).to(torch.bfloat16) + dense = (quant.to(torch.float32) * scale.to(torch.float32)).to(dtype) if output is None: # Pre-allocate the stacked buffer (see the packed path below). output = torch.empty((len(quantized), *dense.shape), dtype=dense.dtype, device=dense.device) @@ -169,6 +170,10 @@ def reverse_op(self) -> "ConversionOps": return _IdentityOp() +# Backward compatibility, the op used to be named after compressed-tensors' "decompress" +DecompressExperts = CompressedTensorsDequantize + + _FP8_DTYPE = torch.float8_e4m3fn _FP8_MIN = torch.finfo(_FP8_DTYPE).min _FP8_MAX = torch.finfo(_FP8_DTYPE).max diff --git a/src/transformers/integrations/fbgemm_fp8.py b/src/transformers/integrations/fbgemm_fp8.py index 21bb5d576549..159de5eff84a 100644 --- a/src/transformers/integrations/fbgemm_fp8.py +++ b/src/transformers/integrations/fbgemm_fp8.py @@ -283,6 +283,8 @@ def replace_with_fbgemm_fp8_linear( pre_quantized (`bool`, defaults to `False`): Whether the model is pre-quantized or not """ + from ..models.llama4.modeling_llama4 import Llama4TextExperts + global quantize_fp8_per_row quantize_fp8_per_row = get_quantize_fp8_per_row() @@ -295,7 +297,7 @@ def replace_with_fbgemm_fp8_linear( new_module = None with init_empty_weights(include_buffers=True): - if module.__class__.__name__ == "Llama4TextExperts": + if isinstance(module, Llama4TextExperts): # TODO: make sure tp works later text_config = getattr(model.config, "text_config", model.config) new_module = FbgemmFp8Llama4TextExperts(text_config or model.config) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index 525c53fb4d73..adc1ffc7711e 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -1074,8 +1074,6 @@ def convert( **kwargs, ) -> dict[str, list[torch.Tensor] | torch.Tensor]: output_dtype = self._get_target_dtype(model, full_layer_name) - # Single-tensor path (the general converter declares ``["weight", "weight_scale_inv", "activation_scale"]`` - # and produces a single ``weight`` target). Also handles the no-scale case (e.g. RMSNorm weights). if "weight" in input_dict: # The downstream renamer in `core_model_loading._convert_one_module` uses the # output dict's *key*, not its content, to derive prefix/suffix; if `full_layer_name` diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 8d9e04441dfc..6bb3849b85a8 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -538,6 +538,7 @@ def replace_with_mxfp4_linear(model, quantization_config=None, modules_to_not_co if quantization_config.dequantize: return model + from ..models.gpt_oss.modeling_gpt_oss import GptOssMLP from .hub_kernels import get_kernel global triton_kernels_hub @@ -545,9 +546,7 @@ def replace_with_mxfp4_linear(model, quantization_config=None, modules_to_not_co has_been_replaced = False for module_name, module in list(model.named_modules()): - if module.__class__.__name__ == "GptOssMLP" and should_convert_module( - f"{module_name}.experts", modules_to_not_convert - ): + if isinstance(module, GptOssMLP) and should_convert_module(f"{module_name}.experts", modules_to_not_convert): with torch.device("meta"): model.set_submodule(module_name, Mxfp4GptOssMLP(module.router, Mxfp4GptOssExperts(model.config))) has_been_replaced = True diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index a70470322d03..81895d70eeea 100644 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -2086,7 +2086,7 @@ def set_experts_implementation(self, experts_implementation: str | dict): else experts_implementation.get("", self.config._experts_implementation) ) - # MegaMoE is locked at load time: its TP plan is baked into `model.tp_plan` by `update_tp_plan` (and isn't + # MegaMoE is locked at load time: its TP plan is baked into `model.tp_plan` by `update_parallel_plan` (and isn't # re-evaluated) and `setup_megamoe_weights` mutates the expert weights into UTCCP layout on first forward. # Either side of a switch would silently produce garbage, so reject it with a clear pointer. current = self.config._experts_implementation diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 47e753b64091..ecfd6f91a7f1 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -120,7 +120,12 @@ def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] """ return device_map + # Bytes per element of the params quantized on the fly: their module still declares the full precision dtype + quantized_element_size: float | None = None + def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: + if self.quantized_element_size is not None and self.param_needs_quantization(model, param_name): + return self.quantized_element_size return param.element_size() def get_new_dtype(self, converter, param, dtype: "torch.dtype | None") -> "torch.dtype | None": @@ -151,16 +156,9 @@ def validate_environment(self, *args, **kwargs): """ return - def update_tp_plan(self, model): - "Updates `model.tp_plan` in place: the full plan, merged from the config and all sub-configs." - self._update_plan(model, model.tp_plan) - - def update_ep_plan(self, model): - "Updates `model.ep_plan` in place: the full plan, merged from the config and all sub-configs." - self._update_plan(model, model.ep_plan) - - def _update_plan(self, model, plan): - "Serialized params (e.g. `gate_up_proj_scale_inv`) are sharded like their param (`gate_up_proj`)." + def update_parallel_plan(self, model, plan: dict[str, str]): + """Updates in place `model.tp_plan` or `model.ep_plan` (the full plans, merged from the config and all + sub-configs): serialized params (e.g. `gate_up_proj_scale_inv`) are sharded like their param (`gate_up_proj`).""" from ..distributed.tensor_parallel import replace_layer_number_by_wildcard params = {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()} @@ -192,8 +190,8 @@ def preprocess_model(self, model: "PreTrainedModel", dtype=None, **kwargs): if self.pre_quantized: self._convert_model_for_quantization(model) self._process_model_before_weight_loading(model, **kwargs) - self.update_tp_plan(model) - self.update_ep_plan(model) + self.update_parallel_plan(model, model.tp_plan) + self.update_parallel_plan(model, model.ep_plan) def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): return model diff --git a/src/transformers/quantizers/quantizer_aqlm.py b/src/transformers/quantizers/quantizer_aqlm.py index 2d0fd7939d88..6b51491cbc92 100644 --- a/src/transformers/quantizers/quantizer_aqlm.py +++ b/src/transformers/quantizers/quantizer_aqlm.py @@ -25,7 +25,6 @@ from ..integrations import replace_with_aqlm_linear from ..utils import is_accelerate_available, is_aqlm_available, logging -from ..utils.quantization_config import QuantizationConfigMixin logger = logging.get_logger(__name__) @@ -39,9 +38,6 @@ class AqlmHfQuantizer(HfQuantizer): requires_calibration = True quantization_config: "AqlmConfig" - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Using `aqlm` quantization requires Accelerate: `pip install accelerate`") diff --git a/src/transformers/quantizers/quantizer_auto_round.py b/src/transformers/quantizers/quantizer_auto_round.py index 3a5231401306..2d7f7b4acf23 100644 --- a/src/transformers/quantizers/quantizer_auto_round.py +++ b/src/transformers/quantizers/quantizer_auto_round.py @@ -20,7 +20,6 @@ from ..modeling_utils import PreTrainedModel from ..utils import is_auto_round_available, logging -from ..utils.quantization_config import QuantizationConfigMixin logger = logging.get_logger(__name__) @@ -34,9 +33,6 @@ class AutoRoundQuantizer(HfQuantizer): # AutoRound requires data calibration - we support only inference requires_calibration = True - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): self.device_map = kwargs.get("device_map") if not is_auto_round_available(): diff --git a/src/transformers/quantizers/quantizer_awq.py b/src/transformers/quantizers/quantizer_awq.py index 3a574ef13f8f..9188daa508a6 100644 --- a/src/transformers/quantizers/quantizer_awq.py +++ b/src/transformers/quantizers/quantizer_awq.py @@ -42,9 +42,6 @@ class AwqQuantizer(HfQuantizer): requires_calibration = True quantization_config: "AwqConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, **kwargs): if not is_gptqmodel_available(): raise ImportError( diff --git a/src/transformers/quantizers/quantizer_bitnet.py b/src/transformers/quantizers/quantizer_bitnet.py index be4bb1e99d38..815ddc628b66 100644 --- a/src/transformers/quantizers/quantizer_bitnet.py +++ b/src/transformers/quantizers/quantizer_bitnet.py @@ -41,9 +41,6 @@ class BitNetHfQuantizer(HfQuantizer): requires_calibration = True quantization_config: "BitNetQuantConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Loading a BitNet quantized model requires accelerate (`pip install accelerate`)") diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py index 791ef9f4e9ef..655c8b0f98a3 100644 --- a/src/transformers/quantizers/quantizer_bnb_4bit.py +++ b/src/transformers/quantizers/quantizer_bnb_4bit.py @@ -47,11 +47,9 @@ class Bnb4BitHfQuantizer(HfQuantizer): """ requires_calibration = False + quantized_element_size = 0.5 quantization_config: "BitsAndBytesConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError( @@ -79,14 +77,6 @@ def validate_environment(self, *args, **kwargs): "for more details. " ) - def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: - "Return the element size (in bytes) for `param_name`." - if self.param_needs_quantization(model, param_name): - # 4 bit - return 0.5 - - return super().param_element_size(model, param_name, param) - def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: import bitsandbytes as bnb @@ -143,7 +133,6 @@ def _process_model_before_weight_loading( def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): setattr(model, "is_loaded_in_4bit", True) - setattr(model, "is_4bit_serializable", self.is_serializable()) return model def is_serializable(self): diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py index 098c723defdd..3851fb32862d 100644 --- a/src/transformers/quantizers/quantizer_bnb_8bit.py +++ b/src/transformers/quantizers/quantizer_bnb_8bit.py @@ -47,11 +47,9 @@ class Bnb8BitHfQuantizer(HfQuantizer): """ requires_calibration = False + quantized_element_size = 1 quantization_config: "BitsAndBytesConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError( @@ -103,13 +101,6 @@ def update_device_map(self, device_map): ) return device_map - def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: - "Return the element size (in bytes) for `param_name`." - if self.param_needs_quantization(model, param_name): - # 8-bit - return 1 - return super().param_element_size(model, param_name, param) - def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: import bitsandbytes as bnb @@ -118,7 +109,6 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): setattr(model, "is_loaded_in_8bit", True) - model.is_8bit_serializable = self.is_serializable() return model def _process_model_before_weight_loading( diff --git a/src/transformers/quantizers/quantizer_eetq.py b/src/transformers/quantizers/quantizer_eetq.py index ce1d350f3d96..720c8fa75370 100644 --- a/src/transformers/quantizers/quantizer_eetq.py +++ b/src/transformers/quantizers/quantizer_eetq.py @@ -39,9 +39,6 @@ class EetqHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "EetqConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_kernels_available(): raise ImportError("Loading an EETQ quantized model requires kernels (`pip install kernels`)") diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index 29a1a8cc8e2c..b51f568fed33 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -44,12 +44,10 @@ class FbgemmFp8HfQuantizer(HfQuantizer): """ requires_calibration = False + quantized_element_size = 1 serialized_names = {"weight": ["weight", "weight_scale"]} quantization_config: "FbgemmFp8Config" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_torch_cuda_available() and not is_torch_xpu_available(): raise ImportError("Using fbgemm fp8 quantization requires a GPU or XPU") @@ -111,13 +109,6 @@ def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, ** return True return False - def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: - "Return the element size (in bytes) for `param_name`." - if self.param_needs_quantization(model, param_name): - # 8 bit, this is needed as when `pre_quantized`` is False, we don't set the dtype of the FP8Linear in order to correctly load the weights - return 1 - return super().param_element_size(model, param_name, param) - def _process_model_before_weight_loading( self, model: "PreTrainedModel", diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 38caa07581bd..19bcb876250f 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -136,10 +136,10 @@ def _process_model_after_weight_loading(self, model, **kwargs): _disable_deepgemm_on_multi_device(model) return model - def _update_plan(self, model, plan): + def update_parallel_plan(self, model, plan): from ..integrations.finegrained_fp8 import FP8Experts - super()._update_plan(model, plan) + super().update_parallel_plan(model, plan) overwrite = FP8Experts._impl_tp_layer_overrides.get(getattr(model.config, "_experts_implementation", None), {}) plan.update({key: overwrite.get(style, style) for key, style in plan.items()}) diff --git a/src/transformers/quantizers/quantizer_fouroversix.py b/src/transformers/quantizers/quantizer_fouroversix.py index 79acb8cbb216..87b4d5c8989b 100644 --- a/src/transformers/quantizers/quantizer_fouroversix.py +++ b/src/transformers/quantizers/quantizer_fouroversix.py @@ -26,9 +26,6 @@ class FourOverSixHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "FourOverSixConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_fouroversix_available(): raise ImportError( @@ -101,23 +98,10 @@ def get_quantize_ops(self): return FourOverSixQuantize(self) def get_weight_conversions(self): - """ - Return weight conversions for loading pre-quantized checkpoints of - other pre-quantized models (not fouroversix models). After first use, - the pre_quantized_model_config_type attribute is set to None to ensure - subsequent calls (e.g., during save_pretrained) return an empty list - since, by then, the model will be saved with our framework's format - so weight conversions are no longer needed. - """ + # Pre-quantized checkpoints of other formats are converted by fouroversix itself + model_config_type = getattr(self.quantization_config, "pre_quantized_model_config_type", None) + if model_config_type is None: + return [] from fouroversix import WeightConversions - # pre_quantized_model_config_type is only set if we are loading a - # pre-quantized model so it is not guaranteed to exist. - if hasattr(self.quantization_config, "pre_quantized_model_config_type"): - model_config_type = self.quantization_config.pre_quantized_model_config_type - weight_conversions = WeightConversions.get_weight_conversions( - model_config_type, - ) - return weight_conversions - - return [] + return WeightConversions.get_weight_conversions(model_config_type) diff --git a/src/transformers/quantizers/quantizer_fp_quant.py b/src/transformers/quantizers/quantizer_fp_quant.py index d0f88b536a1c..4d4b9bcda88a 100644 --- a/src/transformers/quantizers/quantizer_fp_quant.py +++ b/src/transformers/quantizers/quantizer_fp_quant.py @@ -22,7 +22,6 @@ from ..utils.quantization_config import FPQuantConfig from ..utils import is_fp_quant_available, is_qutlass_available, is_torch_available, is_torch_xpu_available, logging -from ..utils.quantization_config import QuantizationConfigMixin if is_torch_available(): @@ -40,9 +39,6 @@ class FPQuantHfQuantizer(HfQuantizer): is_qat_trainable = True quantization_config: "FPQuantConfig" - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, device_map, **kwargs): if not torch.cuda.is_available() and not is_torch_xpu_available(): raise NotImplementedError( diff --git a/src/transformers/quantizers/quantizer_higgs.py b/src/transformers/quantizers/quantizer_higgs.py index 4667c9304437..d053d23164a2 100644 --- a/src/transformers/quantizers/quantizer_higgs.py +++ b/src/transformers/quantizers/quantizer_higgs.py @@ -23,7 +23,6 @@ from ..utils.quantization_config import HiggsConfig from ..utils import is_accelerate_available, is_flute_available, is_hadamard_available, is_torch_available, logging -from ..utils.quantization_config import QuantizationConfigMixin if is_torch_available(): @@ -40,9 +39,6 @@ class HiggsHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "HiggsConfig" - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, device_map, **kwargs): if not torch.cuda.is_available(): raise NotImplementedError("HIGGS quantization is only supported on GPU. Please use a different quantizer.") diff --git a/src/transformers/quantizers/quantizer_hqq.py b/src/transformers/quantizers/quantizer_hqq.py index bc3a86208fc2..7457a032c2bd 100755 --- a/src/transformers/quantizers/quantizer_hqq.py +++ b/src/transformers/quantizers/quantizer_hqq.py @@ -71,11 +71,6 @@ def _process_model_before_weight_loading( # prepare_for_hqq_linear() also sets the right quantization config inside the model (model.config.quantization_config) and the layers (hqq_layer.quant_config) model = prepare_for_hqq_linear(model, quantization_config=self.quantization_config) - def _process_model_after_weight_loading(self, model: "PreTrainedModel", **kwargs): - setattr(model, "is_hqq_quantized", True) - setattr(model, "is_hqq_serializable", self.is_serializable()) - return model - def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py index 7d03e0b6267c..0444dee127da 100644 --- a/src/transformers/quantizers/quantizer_metal.py +++ b/src/transformers/quantizers/quantizer_metal.py @@ -40,9 +40,6 @@ class MetalHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "MetalConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if self.quantization_config.dequantize: return diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index b2ac84cd3565..dcc1afadf9ae 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -43,6 +43,7 @@ class NVFP4HfQuantizer(HfQuantizer): """Quantize eligible linear weights to NVFP4 while loading a full-precision checkpoint.""" requires_calibration = False + quantized_element_size = 0.5 quantization_config: NVFP4Config def validate_environment(self, device_map, **kwargs): @@ -74,7 +75,7 @@ def validate_environment(self, device_map, **kwargs): if major < 10: raise RuntimeError(f"NVFP4 requires Blackwell (sm100+) block-scaled tensor cores; found sm{major}{minor}.") - def update_tp_plan(self, model): + 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: @@ -89,12 +90,6 @@ def param_needs_quantization(self, model: PreTrainedModel, param_name: str, **kw module, tensor_name = get_module_from_name(model, param_name) return isinstance(module, NVFP4Linear) and tensor_name == "weight" - def param_element_size(self, model: PreTrainedModel, param_name: str, param: torch.Tensor) -> float: - if self.param_needs_quantization(model, param_name): - # NVFP4 packs two 4-bit weights per byte; this estimate excludes the smaller scale metadata. - return 0.5 - return super().param_element_size(model, param_name, param) - def _process_model_before_weight_loading(self, model: PreTrainedModel, **kwargs): from ..integrations.nvfp4 import replace_with_nvfp4_linear diff --git a/src/transformers/quantizers/quantizer_spqr.py b/src/transformers/quantizers/quantizer_spqr.py index 918b7c6a01b5..b9457520fa25 100644 --- a/src/transformers/quantizers/quantizer_spqr.py +++ b/src/transformers/quantizers/quantizer_spqr.py @@ -22,7 +22,6 @@ from ..integrations import replace_with_spqr_linear from ..utils import is_accelerate_available, is_spqr_available, is_torch_available, logging -from ..utils.quantization_config import QuantizationConfigMixin if is_torch_available(): @@ -39,9 +38,6 @@ class SpQRHfQuantizer(HfQuantizer): requires_calibration = True quantization_config: "SpQRConfig" - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not torch.cuda.is_available(): raise RuntimeError("GPU is required to run SpQR quantized model.") diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py index 9b71586092bb..1cb0366a2a35 100644 --- a/src/transformers/quantizers/quantizer_torchao.py +++ b/src/transformers/quantizers/quantizer_torchao.py @@ -64,7 +64,7 @@ def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) size_digit = _fuzzy_match_size(type(self.quantization_config.quant_type).__name__) - self.quantized_param_size = 0.5 if size_digit == "4" else 1 + self.quantized_element_size = 0.5 if size_digit == "4" else 1 def validate_environment(self, *args, **kwargs): if not is_torchao_available(): @@ -87,13 +87,6 @@ def get_state_dict_and_metadata(self, model): """ return flatten_tensor_state_dict(model.state_dict()) - def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: - "Return the element size (in bytes) for `param_name`." - if self.param_needs_quantization(model, param_name) and self.quantized_param_size is not None: - return self.quantized_param_size - - return super().param_element_size(model, param_name, param) - def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: # need more space for the quantization parameters (e.g. scale). Tested with int4 wo and group size = 128 max_memory = {key: val * 0.9 for key, val in max_memory.items()} diff --git a/src/transformers/quantizers/quantizer_vptq.py b/src/transformers/quantizers/quantizer_vptq.py index 8d60559b09d3..993c4feea322 100644 --- a/src/transformers/quantizers/quantizer_vptq.py +++ b/src/transformers/quantizers/quantizer_vptq.py @@ -21,7 +21,6 @@ from ..utils.quantization_config import VptqConfig from ..utils import is_accelerate_available, is_torch_available, is_vptq_available, logging -from ..utils.quantization_config import QuantizationConfigMixin if is_torch_available(): @@ -38,9 +37,6 @@ class VptqHfQuantizer(HfQuantizer): requires_calibration = True quantization_config: "VptqConfig" - def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): - super().__init__(quantization_config, **kwargs) - def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Using `vptq` quantization requires Accelerate: `pip install accelerate`") diff --git a/tests/quantization/nvfp4/test_nvfp4.py b/tests/quantization/nvfp4/test_nvfp4.py index 323ae2de0547..b903a68257dc 100644 --- a/tests/quantization/nvfp4/test_nvfp4.py +++ b/tests/quantization/nvfp4/test_nvfp4.py @@ -148,9 +148,9 @@ 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_tp_plan(model) + self.quantizer.update_parallel_plan(model, {}) - def test_update_tp_plan_keeps_non_tp_configs(self): + def test_update_parallel_plan_keeps_non_tp_configs(self): configs = ( SimpleNamespace(), SimpleNamespace(distributed_config=SimpleNamespace(tp_size=1)), @@ -158,7 +158,7 @@ def test_update_tp_plan_keeps_non_tp_configs(self): for config in configs: with self.subTest(config=config): - self.quantizer.update_tp_plan(SimpleNamespace(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) From 1989c65830c57aff87e51ddc1094f49ca6d190d6 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 14:03:09 +0900 Subject: [PATCH 14/27] up --- src/transformers/quantizers/quantizer_quanto.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/transformers/quantizers/quantizer_quanto.py b/src/transformers/quantizers/quantizer_quanto.py index 2ba2217cb912..7201a0b5494e 100644 --- a/src/transformers/quantizers/quantizer_quanto.py +++ b/src/transformers/quantizers/quantizer_quanto.py @@ -23,15 +23,11 @@ from ..utils import ( is_accelerate_available, is_optimum_quanto_available, - is_torch_available, logging, ) from ..utils.quantization_config import QuantoConfig -if is_torch_available(): - import torch - logger = logging.get_logger(__name__) @@ -51,7 +47,7 @@ def __init__(self, quantization_config: QuantoConfig, **kwargs): "int4": 0.5, "int2": 0.25, } - self.quantized_param_size = map_to_param_size.get(self.quantization_config.weights) + self.quantized_element_size = map_to_param_size.get(self.quantization_config.weights) def validate_environment(self, *args, **kwargs): if not is_optimum_quanto_available(): @@ -91,13 +87,6 @@ def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | max_memory = {key: val * 0.90 for key, val in max_memory.items()} return max_memory - def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: - "Return the element size (in bytes) for `param_name`." - if self.param_needs_quantization(model, param_name) and self.quantized_param_size is not None: - return self.quantized_param_size - - return super().param_element_size(model, param_name, param) - def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwargs): from ..integrations import replace_with_quanto_layers From 2a78fb434c6936bd6931108c6ba61c406eed6f9b Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 14:16:45 +0900 Subject: [PATCH 15/27] remove last pieces of dead code --- .../custom_quantization_int8_example.py | 1 - src/transformers/integrations/quark.py | 36 ------------------- src/transformers/quantizers/base.py | 6 ---- src/transformers/quantizers/quantizer_gguf.py | 1 - .../quantizers/quantizer_quark.py | 14 +++----- src/transformers/quantizers/quantizer_sinq.py | 1 - 6 files changed, 4 insertions(+), 55 deletions(-) delete mode 100644 src/transformers/integrations/quark.py diff --git a/examples/quantization/custom_quantization_int8_example.py b/examples/quantization/custom_quantization_int8_example.py index 02244fde4947..f2a027907451 100644 --- a/examples/quantization/custom_quantization_int8_example.py +++ b/examples/quantization/custom_quantization_int8_example.py @@ -140,7 +140,6 @@ class Int8SymmetricQuantizer(HfQuantizer): """ requires_calibration = False - requires_parameters_quantization = True def __init__(self, quantization_config: QuantizationConfigMixin, **kwargs): super().__init__(quantization_config, **kwargs) diff --git a/src/transformers/integrations/quark.py b/src/transformers/integrations/quark.py deleted file mode 100644 index 05c40065aea2..000000000000 --- a/src/transformers/integrations/quark.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright 2025 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. - - -from ..core_model_loading import ConversionOps -from ..utils import is_torch_available - - -if is_torch_available(): - import torch - - -class QuarkDeserialize(ConversionOps): - def __init__(self, hf_quantizer): - self.hf_quantizer = hf_quantizer - - def convert( - self, - input_dict: torch.Tensor, - full_layer_name: str | None = None, - **kwargs, - ) -> dict[str, torch.Tensor]: - value = list(input_dict.values())[0] - value = value[0] if isinstance(value, list) else value - return {full_layer_name: value} diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index ecfd6f91a7f1..3ba0c376cd75 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -247,12 +247,6 @@ def _dequantize(self, model, dtype=None): f"{self.quantization_config.quant_method} has no implementation of `dequantize`, please raise an issue on GitHub." ) - def get_param_name(self, param_name: str) -> str: - """ - Override this method if you want to adjust the `param_name`. - """ - return param_name - @staticmethod def get_modules_to_not_convert( model: "PreTrainedModel", diff --git a/src/transformers/quantizers/quantizer_gguf.py b/src/transformers/quantizers/quantizer_gguf.py index 33cd5c26c970..265c0138b66c 100644 --- a/src/transformers/quantizers/quantizer_gguf.py +++ b/src/transformers/quantizers/quantizer_gguf.py @@ -42,7 +42,6 @@ class GgufHfQuantizer(HfQuantizer): quantization_config: "GgufConfig" header: "GgufHeader" # set by `read_header`, before any of the loading hooks run requires_calibration = False - requires_parameters_quantization = False def __init__(self, quantization_config, **kwargs): super().__init__(quantization_config, **kwargs) diff --git a/src/transformers/quantizers/quantizer_quark.py b/src/transformers/quantizers/quantizer_quark.py index c140f464fdd6..e55e88845dcb 100644 --- a/src/transformers/quantizers/quantizer_quark.py +++ b/src/transformers/quantizers/quantizer_quark.py @@ -70,9 +70,6 @@ def _process_model_before_weight_loading(self, model: "PreTrainedModel", **kwarg return model - def param_needs_quantization(self, model: "PreTrainedModel", param_name: str, **kwargs) -> bool: - return True - def is_serializable(self): return False @@ -80,11 +77,8 @@ def is_serializable(self): def is_trainable(self): return False - # The checkpoint stores `weight_scale`, ... while `QParamsLinear` holds them in its quantizers (`weight_quantizer.scale`) - serialized_names = {target: [source] for source, target in CHECKPOINT_KEYS.items()} - - def get_load_operation(self, converter=None): - from ..integrations.quark import QuarkDeserialize + def get_weight_conversions(self): + from ..core_model_loading import WeightRenaming - if converter is None: - return QuarkDeserialize(self) + # The checkpoint stores `weight_scale`, ... while `QParamsLinear` holds them in its quantizers (`weight_quantizer.scale`) + return [WeightRenaming(source, target) for source, target in CHECKPOINT_KEYS.items()] diff --git a/src/transformers/quantizers/quantizer_sinq.py b/src/transformers/quantizers/quantizer_sinq.py index a6fa15205dfa..d794441e7ada 100644 --- a/src/transformers/quantizers/quantizer_sinq.py +++ b/src/transformers/quantizers/quantizer_sinq.py @@ -44,7 +44,6 @@ class SinqHfQuantizer(HfQuantizer): * A-SINQ (activation-aware) SINQ quantization """ - requires_parameters_quantization: bool = True quantization_config: SinqConfig def __init__(self, quantization_config: SinqConfig, **kwargs): From 52af127e9c0d0d057719bcf299c1c10541fe1465 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 14:21:29 +0900 Subject: [PATCH 16/27] fix failing test --- src/transformers/quantizers/quantizer_compressed_tensors.py | 2 ++ src/transformers/quantizers/quantizer_torchao.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 6a5295832576..152733638c8b 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -169,6 +169,8 @@ def get_weight_conversions(self): @property def serialized_names(self): + if self.quantization_config.quantization_config is None: # e.g. sparsity only + return {} if _is_fp8_scheme(get_experts_scheme(self.quantization_config.quantization_config)): return {"weight": ["weight", "weight_scale"]} return {"weight": ["weight_packed", "weight_scale", "weight_shape"]} diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py index 1cb0366a2a35..cbf11d4b5659 100644 --- a/src/transformers/quantizers/quantizer_torchao.py +++ b/src/transformers/quantizers/quantizer_torchao.py @@ -194,6 +194,9 @@ def _discover_quantized_param_names(self) -> set[str]: @property def serialized_names(self): + # Read from the checkpoint's torchao metadata, so only pre-quantized checkpoints have them + if not self.pre_quantized: + return {} suffixes = ["qdata", "scale_and_zero", "per_tensor_scale", "scale", "zero_point", "act_pre_scale"] return {name: [f"_{name}_{s}" for s in suffixes] for name in sorted(self._discover_quantized_param_names())} From c656891311db2fbbd674d96694c26180b33a27a2 Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 13 Oct 2026 14:39:16 +0900 Subject: [PATCH 17/27] fix test --- tests/utils/test_fusion_mapping.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/utils/test_fusion_mapping.py b/tests/utils/test_fusion_mapping.py index f83cf041998a..656a1bf8cb89 100644 --- a/tests/utils/test_fusion_mapping.py +++ b/tests/utils/test_fusion_mapping.py @@ -159,8 +159,8 @@ def test_register_fusion_patches_raises_on_transform_conflicts(self): model_type, [ WeightConverter( - source_patterns=r"patch_embed\.proj\.weight$", - target_patterns=r"patch_embed\.other_linear_proj\.weight$", + source_patterns=r"patch_embed\.proj\.weight", + target_patterns=r"patch_embed\.other_linear_proj\.weight", operations=[Conv3dToLinear(in_channels=3, kernel_size=(2, 2, 2))], ) ], From 84ae4998f7c7e8077d9deef7813419dac7000129 Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 14 Oct 2026 17:04:48 +0900 Subject: [PATCH 18/27] pr comments --- .github/workflows/pr_slow_ci_suggestion.yml | 4 ++-- src/transformers/integrations/compressed_tensors.py | 2 +- src/transformers/integrations/mxfp4.py | 1 - 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pr_slow_ci_suggestion.yml b/.github/workflows/pr_slow_ci_suggestion.yml index b0b07108e93c..e3f0047a9249 100644 --- a/.github/workflows/pr_slow_ci_suggestion.yml +++ b/.github/workflows/pr_slow_ci_suggestion.yml @@ -102,13 +102,13 @@ jobs: # The script imports `transformers`. This is the base branch checkout, not the PR's code. - name: Install transformers - run: python -m venv ~/venv && ~/venv/bin/pip install . + run: pip install -e . - name: Run script to get jobs to run id: get_jobs run: | set -o pipefail - ~/venv/bin/python utils/get_pr_run_slow_jobs.py | tee output.txt + python utils/get_pr_run_slow_jobs.py | tee output.txt echo "jobs_to_run: $(tail -n 1 output.txt)" echo "jobs_to_run=$(tail -n 1 output.txt)" >> $GITHUB_OUTPUT diff --git a/src/transformers/integrations/compressed_tensors.py b/src/transformers/integrations/compressed_tensors.py index ef64a2ca3350..45c86f3a6ddf 100644 --- a/src/transformers/integrations/compressed_tensors.py +++ b/src/transformers/integrations/compressed_tensors.py @@ -170,7 +170,7 @@ def reverse_op(self) -> "ConversionOps": return _IdentityOp() -# Backward compatibility, the op used to be named after compressed-tensors' "decompress" +# BC, the name was not aligned with other quantizers. DecompressExperts = CompressedTensorsDequantize diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index 6bb3849b85a8..f207d358540c 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -444,7 +444,6 @@ def topk(vals, k): return RoutingData(gate_scal, hist, n_local_experts, hit_experts, expt_data), gather_indx, scatter_indx -@use_kernel_forward_from_hub("MegaBlocksMoeMLP") class Mxfp4GptOssMLP(nn.Module): def __init__(self, router, experts): super().__init__() From 3c6324932fbbaab23aedfeb0cd64f6487c57484f Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 14 Oct 2026 17:22:08 +0900 Subject: [PATCH 19/27] human coments --- src/transformers/quantizers/base.py | 42 ++++++++++++------- src/transformers/quantizers/quantizer_gguf.py | 2 +- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 3ba0c376cd75..b5e76c54f540 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -120,7 +120,6 @@ def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] """ return device_map - # Bytes per element of the params quantized on the fly: their module still declares the full precision dtype quantized_element_size: float | None = None def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: @@ -129,12 +128,15 @@ def param_element_size(self, model: "PreTrainedModel", param_name: str, param: " return param.element_size() def get_new_dtype(self, converter, param, dtype: "torch.dtype | None") -> "torch.dtype | None": - """dtype to materialize the pre-quantized checkpoint tensor `param` in, `None` keeps the checkpoint's: the one - its module declares (`dtype`), unless `param` is quantized (non-float) or an input of a load operation.""" + """ + `dtype` to cast the pre-quantized checkpoint tensor `param` in. + If that param is not quantized (float type) we use the input `dtype` that comes from the initialized `nn.Module`. + + If the quantization method has a `Deserializer` or `Dequantizer` (load function) we return None so it's handled by them. + """ is_float = ( param.get_dtype().startswith(("F", "BF")) if hasattr(param, "get_dtype") else param.is_floating_point() ) - # The quantizer's ops (load operations) are the ones holding it is_load_input = any(getattr(op, "hf_quantizer", None) for op in getattr(converter, "operations", ())) return dtype if is_float and not is_load_input else None @@ -157,8 +159,9 @@ def validate_environment(self, *args, **kwargs): return def update_parallel_plan(self, model, plan: dict[str, str]): - """Updates in place `model.tp_plan` or `model.ep_plan` (the full plans, merged from the config and all - sub-configs): serialized params (e.g. `gate_up_proj_scale_inv`) are sharded like their param (`gate_up_proj`).""" + """ + Updates in place `model.tp_plan` or `model.ep_plan` (the full plans, merged from the config and all + sub-configs) to take into account the quantizer's custom keys like 'inv_scale'""" from ..distributed.tensor_parallel import replace_layer_number_by_wildcard params = {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()} @@ -308,32 +311,41 @@ def get_quantize_ops(self): f"{self.quantization_config.quant_method} is not available yet and will be supported soon." ) - # How params are saved in pre-quantized checkpoints: `{param: [serialized source patterns]}`, - # e.g. `{"weight": ["weight", "weight_scale_inv"]}` + # Maps a non quantized parameter name to its corresponding quantized names + # e.g. `{"weight": ["weight", "weight_scale_inv"]}` means a weight becomes weight + scale inv. serialized_names: dict[str, list[str]] = {} def get_load_operation(self, converter=None): - """Op loading the serialized tensors: a deserialization, or a dequantization when `dequantize=True`. + """ + Operation loading the serialized tensors: a deserialization, or a dequantization when `dequantize=True`. Without `converter`, the op of the general converters. With a model `converter` (e.g. a per-expert merge), the op to prepend to it, if the quantizer's op supports it.""" return None def get_weight_conversions(self): - "General converters, for the serialized tensors that no model converter claims." + "Return the loading to perform on any parameter that is not matched by the model's converters" from ..core_model_loading import WeightConverter - op = self.get_load_operation() - return [WeightConverter(names, param, [op]) for param, names in self.serialized_names.items()] if op else [] + operation = self.get_load_operation() + return ( + [WeightConverter(names, param, [operation]) for param, names in self.serialized_names.items()] + if operation + else [] + ) def update_weight_conversions(self, weight_conversions): - """Prepends `get_load_operation` to the model converters, with their serialized sources, then appends the - general `get_weight_conversions`.""" + """ + Prepends either Deserialization or Dequantization to the model's specific converters (from + the conversion mapping). This ensures operations are applied to the quantizer's custom quantization parameters. + + Say you have a `Concatenation` operation on "mlp.weight", this + updates it to also apply on "mlp.weight_scale_inv". + """ from ..core_model_loading import WeightConverter updated = [] for conv in weight_conversions: op = self.get_load_operation(conv) if isinstance(conv, WeightConverter) else None - # `experts.*.gate_proj.weight` -> `experts.*.gate_proj.weight`, `experts.*.gate_proj.weight_scale_inv` sources = [] for pattern in conv._original_source_patterns: prefix, _, param = pattern.rpartition(".") diff --git a/src/transformers/quantizers/quantizer_gguf.py b/src/transformers/quantizers/quantizer_gguf.py index 265c0138b66c..817346e1f3d7 100644 --- a/src/transformers/quantizers/quantizer_gguf.py +++ b/src/transformers/quantizers/quantizer_gguf.py @@ -126,7 +126,7 @@ def update_dtype(self, dtype): return dtype def get_new_dtype(self, converter, param, dtype): - # The transforms keep the file's precision, `Cast` (last op of every chain) lands in `self.dtype` + # supperted architectures will cast themselves. return None if self.supported else super().get_new_dtype(converter, param, dtype) def get_state_dict(self, checkpoint_file: str, model): From 5b6eab19c7c8c2863f1d9696dae4da3e7311673d Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 14 Oct 2026 18:05:55 +0900 Subject: [PATCH 20/27] Treat scales and etc individually? --- src/transformers/quantizers/base.py | 41 ++++++++++++++++++----------- 1 file changed, 26 insertions(+), 15 deletions(-) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index b5e76c54f540..be8d0a956eaa 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -11,6 +11,7 @@ # 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 re from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any @@ -319,7 +320,8 @@ def get_load_operation(self, converter=None): """ Operation loading the serialized tensors: a deserialization, or a dequantization when `dequantize=True`. Without `converter`, the op of the general converters. With a model `converter` (e.g. a per-expert merge), - the op to prepend to it, if the quantizer's op supports it.""" + the op to prepend when the serialized tensors are combined before its ops (a dequantization). With None, + its ops run on each serialized name separately (see `update_weight_conversions`).""" return None def get_weight_conversions(self): @@ -335,26 +337,35 @@ def get_weight_conversions(self): def update_weight_conversions(self, weight_conversions): """ - Prepends either Deserialization or Dequantization to the model's specific converters (from - the conversion mapping). This ensures operations are applied to the quantizer's custom quantization parameters. - - Say you have a `Concatenation` operation on "mlp.weight", this - updates it to also apply on "mlp.weight_scale_inv". + Makes the model's converters (from the conversion mapping) load the serialized names, e.g. a merge of + `experts.*.gate_proj.weight` into `experts.gate_up_proj`: + - dequantize: the converter also collects `experts.*.gate_proj.weight_scale_inv`, its op runs first + - otherwise: a copy merges `experts.*.gate_proj.weight_scale_inv` into `experts.gate_up_proj_scale_inv` """ from ..core_model_loading import WeightConverter updated = [] for conv in weight_conversions: - op = self.get_load_operation(conv) if isinstance(conv, WeightConverter) else None - sources = [] - for pattern in conv._original_source_patterns: - prefix, _, param = pattern.rpartition(".") - sources += [f"{prefix}.{n}" for n in self.serialized_names.get(param, [])] or [pattern] - if op is not None and sources != conv._original_source_patterns: - new = WeightConverter(sources, conv._original_target_patterns, [op] + conv.operations) - new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix - conv = new updated.append(conv) + if not isinstance(conv, WeightConverter): + continue + sources, targets = conv._original_source_patterns, conv._original_target_patterns + param = sources[0].rpartition(".")[2] + if not all(source.endswith(f".{param}") for source in sources): + continue + names = self.serialized_names.get(param, []) + if names and (op := self.get_load_operation(conv)) is not None: + sources = [re.sub(rf"{param}$", name, source) for source in sources for name in names] + updated[-1] = WeightConverter(sources, targets, [op] + conv.operations) + updated[-1].scope_prefix, updated[-1].base_model_prefix = conv.scope_prefix, conv.base_model_prefix + continue + # `weight_scale_inv` -> `_scale_inv`, e.g. `gate_proj.weight_scale_inv` merged into `gate_up_proj_scale_inv` + for suffix in [name.removeprefix(param) for name in names if name.startswith(param) and name != param]: + new = WeightConverter( + [s + suffix for s in sources], [t + suffix for t in targets], list(conv.operations) + ) + new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix + updated.append(new) return updated + self.get_weight_conversions() From bfb63e280f9ec80e200ae10665ae9f350279c20b Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 14 Oct 2026 18:12:07 +0900 Subject: [PATCH 21/27] use a map --- .../integrations/metal_quantization.py | 26 ++++++++++++------- src/transformers/quantizers/base.py | 18 +++++++++---- .../quantizers/quantizer_bitnet.py | 7 ++--- .../quantizers/quantizer_bnb_4bit.py | 5 ++-- .../quantizers/quantizer_bnb_8bit.py | 5 ++-- .../quantizers/quantizer_finegrained_fp8.py | 5 ++-- .../quantizers/quantizer_fp_quant.py | 5 ++-- .../quantizers/quantizer_metal.py | 5 ++-- .../quantizers/quantizer_mxfp4.py | 5 ++-- src/transformers/quantizers/quantizer_sinq.py | 5 ++-- .../quantizers/quantizer_torchao.py | 5 ++-- 11 files changed, 50 insertions(+), 41 deletions(-) diff --git a/src/transformers/integrations/metal_quantization.py b/src/transformers/integrations/metal_quantization.py index 6bccc212aa53..6db0f9a65c4b 100644 --- a/src/transformers/integrations/metal_quantization.py +++ b/src/transformers/integrations/metal_quantization.py @@ -285,15 +285,23 @@ def convert(self, input_dict: dict, full_layer_name: str | None = None, **kwargs bits = self.hf_quantizer.quantization_config.bits group_size = self.hf_quantizer.quantization_config.group_size - if len(input_dict) < 2: - return {full_layer_name: input_dict["weight"]} - - quantized = input_dict["weight"][0] - scales = input_dict["scales"][0] - qbiases = input_dict["qbiases"][0] - - w_deq = _affine_dequantize_tensor(quantized, scales, qbiases, group_size, bits) - return {full_layer_name: w_deq.to(scales.dtype)} + # A model converter collects several weights (e.g. per expert): `*.weight` comes with `*.scales`, `*.qbiases` + result = {} + for key, weights in input_dict.items(): + prefix = key.removesuffix("weight") + if key.endswith(("scales", "qbiases")): + continue + if prefix + "scales" not in input_dict: + result[key] = weights + continue + scales, qbiases = input_dict[prefix + "scales"], input_dict[prefix + "qbiases"] + result[key] = [ + _affine_dequantize_tensor(w, s, b, group_size, bits).to(s.dtype) + for w, s, b in zip(weights, scales, qbiases) + ] + if "weight" in result: # the general converter: one module + return {full_layer_name: result["weight"][0]} + return result @property def reverse_op(self) -> "ConversionOps": diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index be8d0a956eaa..34b9665afdc1 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -316,13 +316,21 @@ def get_quantize_ops(self): # e.g. `{"weight": ["weight", "weight_scale_inv"]}` means a weight becomes weight + scale inv. serialized_names: dict[str, list[str]] = {} + def get_load_operations(self) -> dict: + """Op classes loading the serialized tensors, e.g. `{"deserialize": Bnb4bitDeserialize}`""" + return {} + def get_load_operation(self, converter=None): """ - Operation loading the serialized tensors: a deserialization, or a dequantization when `dequantize=True`. - Without `converter`, the op of the general converters. With a model `converter` (e.g. a per-expert merge), - the op to prepend when the serialized tensors are combined before its ops (a dequantization). With None, - its ops run on each serialized name separately (see `update_weight_conversions`).""" - return None + Op loading the serialized tensors of `converter`, or of the general converters without one. A deserialization + builds the params of one module, so it only runs on the general converters. A dequantization also runs before + the ops of the model's converters, e.g. before a per-expert merge. + """ + mode = "dequantize" if getattr(self.quantization_config, "dequantize", False) else "deserialize" + op = self.get_load_operations().get(mode) if self.pre_quantized else None + if op is None or (converter is not None and mode == "deserialize"): + return None + return op(self) def get_weight_conversions(self): "Return the loading to perform on any parameter that is not matched by the model's converters" diff --git a/src/transformers/quantizers/quantizer_bitnet.py b/src/transformers/quantizers/quantizer_bitnet.py index 815ddc628b66..060e881bf0d1 100644 --- a/src/transformers/quantizers/quantizer_bitnet.py +++ b/src/transformers/quantizers/quantizer_bitnet.py @@ -105,9 +105,10 @@ def is_qat_trainable(self) -> bool: serialized_names = {"weight": ["weight"]} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.bitnet import BitNetDeserialize config = self.quantization_config - if converter is None and config.linear_class == "autobitlinear" and config.quantization_mode == "offline": - return BitNetDeserialize(self) + if config.linear_class == "autobitlinear" and config.quantization_mode == "offline": + return {"deserialize": BitNetDeserialize} + return {} diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py index 655c8b0f98a3..c5e31f419b4f 100644 --- a/src/transformers/quantizers/quantizer_bnb_4bit.py +++ b/src/transformers/quantizers/quantizer_bnb_4bit.py @@ -165,8 +165,7 @@ def get_quantize_ops(self): ] } - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.bitsandbytes import Bnb4bitDeserialize - if converter is None and self.pre_quantized: - return Bnb4bitDeserialize(self) + return {"deserialize": Bnb4bitDeserialize} diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py index 3851fb32862d..e7fcb6c267cc 100644 --- a/src/transformers/quantizers/quantizer_bnb_8bit.py +++ b/src/transformers/quantizers/quantizer_bnb_8bit.py @@ -155,8 +155,7 @@ def get_quantize_ops(self): serialized_names = {"weight": ["weight", "SCB", "weight_format"]} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.bitsandbytes import Bnb8bitDeserialize - if converter is None and self.pre_quantized: - return Bnb8bitDeserialize(self) + return {"deserialize": Bnb8bitDeserialize} diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 19bcb876250f..5ed2a0ddb655 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -161,11 +161,10 @@ def get_quantize_ops(self): serialized_names = {"weight": ["weight", "weight_scale_inv", "activation_scale"]} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.finegrained_fp8 import Fp8Dequantize - if self.pre_quantized and self.quantization_config.dequantize: - return Fp8Dequantize(self) + return {"dequantize": Fp8Dequantize} def update_weight_conversions(self, weight_conversions): from ..core_model_loading import WeightRenaming diff --git a/src/transformers/quantizers/quantizer_fp_quant.py b/src/transformers/quantizers/quantizer_fp_quant.py index 4d4b9bcda88a..038869d885af 100644 --- a/src/transformers/quantizers/quantizer_fp_quant.py +++ b/src/transformers/quantizers/quantizer_fp_quant.py @@ -141,8 +141,7 @@ def serialized_names(self): name = ".dqweight" if self.quantization_config.pseudoquantization else ".qweight" return {name: [name]} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.fp_quant import FpQuantDeserialize - if converter is None and self.pre_quantized: - return FpQuantDeserialize(self) + return {"deserialize": FpQuantDeserialize} diff --git a/src/transformers/quantizers/quantizer_metal.py b/src/transformers/quantizers/quantizer_metal.py index 0444dee127da..99ef8a9b6d9c 100644 --- a/src/transformers/quantizers/quantizer_metal.py +++ b/src/transformers/quantizers/quantizer_metal.py @@ -114,8 +114,7 @@ def get_quantize_ops(self): serialized_names = {"weight": ["weight", "scales", "qbiases"]} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.metal_quantization import MetalDequantize - if converter is None and self.pre_quantized and self.quantization_config.dequantize: - return MetalDequantize(self) + return {"dequantize": MetalDequantize} diff --git a/src/transformers/quantizers/quantizer_mxfp4.py b/src/transformers/quantizers/quantizer_mxfp4.py index 541f6ceaec4b..0714aaf198cd 100644 --- a/src/transformers/quantizers/quantizer_mxfp4.py +++ b/src/transformers/quantizers/quantizer_mxfp4.py @@ -222,8 +222,7 @@ def get_quantize_ops(self): "down_proj": ["down_proj_blocks", "down_proj_scales"], } - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.mxfp4 import Mxfp4Dequantize, Mxfp4Deserialize - if converter is None and self.pre_quantized: - return Mxfp4Dequantize(self) if self.quantization_config.dequantize else Mxfp4Deserialize(self) + return {"deserialize": Mxfp4Deserialize, "dequantize": Mxfp4Dequantize} diff --git a/src/transformers/quantizers/quantizer_sinq.py b/src/transformers/quantizers/quantizer_sinq.py index d794441e7ada..79ecd4304691 100644 --- a/src/transformers/quantizers/quantizer_sinq.py +++ b/src/transformers/quantizers/quantizer_sinq.py @@ -170,11 +170,10 @@ def get_quantize_ops(self): # A checkpoint produced by `SINQLinear.state_dict`, rebuilt into a `SINQLinear` serialized_names = {".weight": [".W_q", ".meta", ".bias"]} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.sinq import SinqDeserialize - if converter is None and self.pre_quantized: - return SinqDeserialize(self) + return {"deserialize": SinqDeserialize} def _process_model_before_weight_loading( self, diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py index cbf11d4b5659..7dfd155e1aa9 100644 --- a/src/transformers/quantizers/quantizer_torchao.py +++ b/src/transformers/quantizers/quantizer_torchao.py @@ -200,8 +200,7 @@ def serialized_names(self): suffixes = ["qdata", "scale_and_zero", "per_tensor_scale", "scale", "zero_point", "act_pre_scale"] return {name: [f"_{name}_{s}" for s in suffixes] for name in sorted(self._discover_quantized_param_names())} - def get_load_operation(self, converter=None): + def get_load_operations(self): from ..integrations.torchao import TorchAoDeserialize - if converter is None and self.pre_quantized: - return TorchAoDeserialize(self) + return {"deserialize": TorchAoDeserialize} From f850c82359eb4b25f26c1ef6fe195b315e15503b Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 14 Oct 2026 18:24:09 +0900 Subject: [PATCH 22/27] replicate explicitely --- src/transformers/distributed/tensor_parallel.py | 6 +++--- src/transformers/quantizers/base.py | 1 + src/transformers/quantizers/quantizer_finegrained_fp8.py | 4 ++++ 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/transformers/distributed/tensor_parallel.py b/src/transformers/distributed/tensor_parallel.py index 66add4f5f64d..c6c11c041ab2 100644 --- a/src/transformers/distributed/tensor_parallel.py +++ b/src/transformers/distributed/tensor_parallel.py @@ -204,8 +204,6 @@ def shard_param(self, module, param, mesh): if meta is None: return placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2) - if meta.ndim == 0: # e.g. per-tensor quantization scales - placement = Replicate() module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, @@ -280,7 +278,7 @@ def shard_param(self, module, param, mesh): placement = Shard(0) else: # bias is replicated (added after the row-reduce); weight shards on input dim (-1) - placement = Replicate() if param == "bias" or meta.ndim == 0 else Shard(-1) + placement = Replicate() if param == "bias" else Shard(-1) module._parameters[param] = torch.nn.Parameter( distribute_tensor(meta, mesh, [placement], src_data_rank=None), requires_grad=meta.requires_grad, @@ -949,6 +947,8 @@ class ParallelInterface(GeneralInterface): "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(), "mla_kv_a_proj": MlaKvAProjParallel(), "all_reduce": AllReduceParallel(), + # Every rank keeps the full param, e.g. a per-tensor quantization scale + "replicate": TensorParallelLayer(), } if is_torch_distributed_available() else {} diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 34b9665afdc1..fc2d775dc595 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -130,6 +130,7 @@ def param_element_size(self, model: "PreTrainedModel", param_name: str, param: " def get_new_dtype(self, converter, param, dtype: "torch.dtype | None") -> "torch.dtype | None": """ + This happens at materialization time, when the tensor is sent from disk to device. `dtype` to cast the pre-quantized checkpoint tensor `param` in. If that param is not quantized (float type) we use the input `dtype` that comes from the initialized `nn.Module`. diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 5ed2a0ddb655..658cb693e40f 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -137,9 +137,13 @@ def _process_model_after_weight_loading(self, model, **kwargs): return model def update_parallel_plan(self, model, plan): + from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard from ..integrations.finegrained_fp8 import FP8Experts super().update_parallel_plan(model, plan) + # Per-tensor scales (`activation_scale`, `weight_scale_inv` without blocks) are scalars, not sharded + scalars = {replace_layer_number_by_wildcard(n) for n, p in model.named_parameters() if p.ndim == 0} + plan.update({name: "replicate" for name in scalars if _get_parameter_plan(name, plan) is not None}) overwrite = FP8Experts._impl_tp_layer_overrides.get(getattr(model.config, "_experts_implementation", None), {}) plan.update({key: overwrite.get(style, style) for key, style in plan.items()}) From 599614fdffd39b2f949c8250ad50f59c9b9b521f Mon Sep 17 00:00:00 2001 From: Arthur Date: Wed, 14 Oct 2026 22:23:05 +0900 Subject: [PATCH 23/27] impl copy --- src/transformers/conversion_mapping.py | 2 +- src/transformers/core_model_loading.py | 16 ++-- src/transformers/quantizers/base.py | 77 +++++++++---------- .../quantizer_compressed_tensors.py | 10 +-- .../quantizers/quantizer_finegrained_fp8.py | 4 +- .../quantizers/quantizer_gemma.py | 11 +++ src/transformers/quantizers/quantizer_gguf.py | 2 +- 7 files changed, 65 insertions(+), 57 deletions(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 9e51f58d07e2..2fd6ec543435 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -2215,6 +2215,6 @@ def get_model_conversion_mapping( weight_conversions.extend(get_checkpoint_conversion_mapping("legacy")) if hf_quantizer is not None: - weight_conversions = hf_quantizer.update_weight_conversions(weight_conversions) + weight_conversions = hf_quantizer.update_weight_conversions(weight_conversions, model) return weight_conversions diff --git a/src/transformers/core_model_loading.py b/src/transformers/core_model_loading.py index 8873f27e1e6c..9ff820a656e4 100644 --- a/src/transformers/core_model_loading.py +++ b/src/transformers/core_model_loading.py @@ -873,7 +873,8 @@ def __init__(self, source_patterns: str | list[str], target_patterns: str | list # Construct the regex we will use to rename keys from the sources to the targets branches = [] - for i, source_pattern in enumerate(self.source_patterns): + # Longest first: a key is matched by the first branch, so `weight.absmax` must come before `weight` + for i, source_pattern in sorted(enumerate(self.source_patterns), key=lambda item: -len(item[1])): group_name = f"g{i}" pattern = source_pattern.replace("*.", r".*\.") # A pattern ends on a name boundary: `weight` matches `q_proj.weight`, not `q_proj.weight_scale_inv`. @@ -976,12 +977,13 @@ def reverse_transform(self) -> WeightTransform: # All reverse ops, in reverse order kwargs["operations"] = [op.reverse_op for op in self.operations[::-1]] - reverse_transform = self.__class__( - source_patterns=self._original_target_patterns, target_patterns=self._original_source_patterns, **kwargs - ) - reverse_transform.scope_prefix = self.scope_prefix - reverse_transform.base_model_prefix = self.base_model_prefix - return reverse_transform + return self.copy(self._original_target_patterns, self._original_source_patterns, **kwargs) + + 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`.""" + 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 def materialize_tensors(self) -> dict[str, list[torch.Tensor]]: """ diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index fc2d775dc595..51bd4646f4e9 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -318,63 +318,60 @@ def get_quantize_ops(self): serialized_names: dict[str, list[str]] = {} def get_load_operations(self) -> dict: - """Op classes loading the serialized tensors, e.g. `{"deserialize": Bnb4bitDeserialize}`""" - return {} - - def get_load_operation(self, converter=None): """ - Op loading the serialized tensors of `converter`, or of the general converters without one. A deserialization - builds the params of one module, so it only runs on the general converters. A dequantization also runs before - the ops of the model's converters, e.g. before a per-expert merge. + Op classes loading the serialized tensors, e.g. `{"deserialize": Mxfp4Deserialize, "dequantize": Mxfp4Dequantize}`. + A deserialization builds the params of one module, so it only runs on the general converters. A dequantization + also runs before the ops of the model's converters, e.g. before a per-expert merge. """ - mode = "dequantize" if getattr(self.quantization_config, "dequantize", False) else "deserialize" - op = self.get_load_operations().get(mode) if self.pre_quantized else None - if op is None or (converter is not None and mode == "deserialize"): - return None - return op(self) + return {} def get_weight_conversions(self): "Return the loading to perform on any parameter that is not matched by the model's converters" from ..core_model_loading import WeightConverter - operation = self.get_load_operation() - return ( - [WeightConverter(names, param, [operation]) for param, names in self.serialized_names.items()] - if operation - else [] - ) + mode = "dequantize" if getattr(self.quantization_config, "dequantize", False) else "deserialize" + operation = self.get_load_operations().get(mode) if self.pre_quantized else None + if operation is None: + return [] + return [WeightConverter(names, param, [operation(self)]) for param, names in self.serialized_names.items()] - def update_weight_conversions(self, weight_conversions): + def update_weight_conversions(self, weight_conversions, model): """ - Makes the model's converters (from the conversion mapping) load the serialized names, e.g. a merge of - `experts.*.gate_proj.weight` into `experts.gate_up_proj`: - - dequantize: the converter also collects `experts.*.gate_proj.weight_scale_inv`, its op runs first - - otherwise: a copy merges `experts.*.gate_proj.weight_scale_inv` into `experts.gate_up_proj_scale_inv` + Makes the model's converters (from the conversion mapping) load the serialized names. The model's params + decide, e.g. for a merge of `experts.*.gate_proj.weight` into `experts.gate_up_proj`: + - `experts.gate_up_proj_scale_inv` is a param: a copy merges `experts.*.gate_proj.weight_scale_inv` into it + - otherwise the converter also collects `experts.*.gate_proj.weight_scale_inv`, the dequantize op runs first """ from ..core_model_loading import WeightConverter + param_names = "\n".join(name for name, _ in model.named_parameters()) + dequantize = self.get_load_operations().get("dequantize") if self.pre_quantized else None updated = [] - for conv in weight_conversions: - updated.append(conv) - if not isinstance(conv, WeightConverter): + for converter in weight_conversions: + updated.append(converter) + if not isinstance(converter, WeightConverter): continue - sources, targets = conv._original_source_patterns, conv._original_target_patterns + sources, targets = converter._original_source_patterns, converter._original_target_patterns param = sources[0].rpartition(".")[2] - if not all(source.endswith(f".{param}") for source in sources): + serialized = self.serialized_names.get(param) + if not serialized or not all(source.endswith(f".{param}") for source in sources): continue - names = self.serialized_names.get(param, []) - if names and (op := self.get_load_operation(conv)) is not None: - sources = [re.sub(rf"{param}$", name, source) for source in sources for name in names] - updated[-1] = WeightConverter(sources, targets, [op] + conv.operations) - updated[-1].scope_prefix, updated[-1].base_model_prefix = conv.scope_prefix, conv.base_model_prefix - continue - # `weight_scale_inv` -> `_scale_inv`, e.g. `gate_proj.weight_scale_inv` merged into `gate_up_proj_scale_inv` - for suffix in [name.removeprefix(param) for name in names if name.startswith(param) and name != param]: - new = WeightConverter( - [s + suffix for s in sources], [t + suffix for t in targets], list(conv.operations) + # `weight_scale_inv` -> `_scale_inv` + suffixes = [name.removeprefix(param) for name in serialized if name.startswith(param) and name != param] + if suffixes and re.search(f"{targets[0]}{suffixes[0]}$", param_names, re.MULTILINE): + # The module keeps the companions as params: each one gets its own copy of the converter + for suffix in suffixes: + suffixed_sources = [source + suffix for source in sources] + suffixed_targets = [target + suffix for target in targets] + updated.append( + converter.copy(suffixed_sources, suffixed_targets, operations=list(converter.operations)) + ) + elif dequantize is not None: + # The module only has `param`: the converter also collects the companions and dequantizes first + all_sources = [re.sub(rf"{param}$", name, source) for source in sources for name in serialized] + updated[-1] = converter.copy( + all_sources, targets, operations=[dequantize(self)] + converter.operations ) - new.scope_prefix, new.base_model_prefix = conv.scope_prefix, conv.base_model_prefix - updated.append(new) return updated + self.get_weight_conversions() diff --git a/src/transformers/quantizers/quantizer_compressed_tensors.py b/src/transformers/quantizers/quantizer_compressed_tensors.py index 152733638c8b..8c6eb71a63fb 100644 --- a/src/transformers/quantizers/quantizer_compressed_tensors.py +++ b/src/transformers/quantizers/quantizer_compressed_tensors.py @@ -13,6 +13,7 @@ # limitations under the License. from copy import deepcopy +from functools import partial from ..utils import is_compressed_tensors_available, is_torch_available, is_torch_greater_or_equal, logging from ..utils.quantization_config import CompressedTensorsConfig @@ -175,9 +176,6 @@ def serialized_names(self): return {"weight": ["weight", "weight_scale"]} return {"weight": ["weight_packed", "weight_scale", "weight_shape"]} - def get_load_operation(self, converter=None): - # Merged experts are not quantized modules: dequantize their per-expert weights before the merge - if converter is not None and all("experts" in p for p in converter.source_patterns): - return CompressedTensorsDequantize( - self, scheme=get_experts_scheme(self.quantization_config.quantization_config) - ) + def get_load_operations(self): + scheme = get_experts_scheme(self.quantization_config.quantization_config) + return {"dequantize": partial(CompressedTensorsDequantize, scheme=scheme)} diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 658cb693e40f..2c2483c7eb5a 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -170,9 +170,9 @@ def get_load_operations(self): return {"dequantize": Fp8Dequantize} - def update_weight_conversions(self, weight_conversions): + def update_weight_conversions(self, weight_conversions, model): from ..core_model_loading import WeightRenaming # Some FP8 checkpoints (e.g. DeepSeek-V4-Flash) ship per-block scales under `.scale` scale_rename = WeightRenaming(source_patterns=r"^(.+)\.scale", target_patterns=r"\1.weight_scale_inv") - return [scale_rename] + super().update_weight_conversions(weight_conversions) + return [scale_rename] + super().update_weight_conversions(weight_conversions, model) diff --git a/src/transformers/quantizers/quantizer_gemma.py b/src/transformers/quantizers/quantizer_gemma.py index 5c26b35a35c6..6b05b08ab7b6 100644 --- a/src/transformers/quantizers/quantizer_gemma.py +++ b/src/transformers/quantizers/quantizer_gemma.py @@ -62,6 +62,17 @@ def _process_model_before_weight_loading(self, model, **kwargs): ignored.update([r".*\.k_cache_scale$", r".*\.v_cache_scale$"]) model._keys_to_ignore_on_load_unexpected = ignored # type: ignore[unresolved-attribute] + def update_parallel_plan(self, model, plan): + from ..distributed.tensor_parallel import _get_parameter_plan, replace_layer_number_by_wildcard + + super().update_parallel_plan(model, plan) + # Every rank keeps the full scalar activation scales, and the per-output-channel `weight_scale` of a rowwise + # layer (it shards the input dim) + for name in {replace_layer_number_by_wildcard(name) for name, _ in model.named_parameters()}: + style = _get_parameter_plan(name, plan) + if style and (name.endswith("activation_scale") or name.endswith(".weight_scale") and "rowwise" in style): + plan[name] = "replicate" + @property def is_serializable(self): return True diff --git a/src/transformers/quantizers/quantizer_gguf.py b/src/transformers/quantizers/quantizer_gguf.py index 817346e1f3d7..3fe20948bbf9 100644 --- a/src/transformers/quantizers/quantizer_gguf.py +++ b/src/transformers/quantizers/quantizer_gguf.py @@ -175,7 +175,7 @@ def _process_model_after_weight_loading(self, model, **kwargs): model._weight_conversions = None return model - def update_weight_conversions(self, weight_conversions): + def update_weight_conversions(self, weight_conversions, model): """Swap in this file's conversions, dropping the model's own. The model's describe a safetensors checkpoint, and match nothing in a GGUF. From 1dd7efde93f9e1e481cc03d9b465bea8c43982af Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 16 Oct 2026 15:56:54 +0900 Subject: [PATCH 24/27] Quantized element size as property, one-line FP8Experts comments, ue8m0 float32 scale loading test --- .../integrations/finegrained_fp8.py | 2 ++ src/transformers/integrations/mxfp4.py | 1 - .../quantizers/quantizer_quanto.py | 12 ++----- .../quantizers/quantizer_torchao.py | 8 ++--- .../quantization/finegrained_fp8/test_fp8.py | 35 +++++++++++++++++++ 5 files changed, 43 insertions(+), 15 deletions(-) diff --git a/src/transformers/integrations/finegrained_fp8.py b/src/transformers/integrations/finegrained_fp8.py index adc1ffc7711e..d8f0cdf6419b 100644 --- a/src/transformers/integrations/finegrained_fp8.py +++ b/src/transformers/integrations/finegrained_fp8.py @@ -607,7 +607,9 @@ def fp8_grouped_mm_experts_forward( class FP8Experts(nn.Module): + # Internal: set True by `_disable_deepgemm_on_multi_device` when the model spans >1 CUDA device. _deepgemm_disabled = False + # `_experts_implementation` -> {default TP/EP plan style: impl-specific style}, applied by the quantizer. _impl_tp_layer_overrides: dict[str, dict[str, str]] = { "deepgemm_megamoe": { "moe_tp_experts": "megamoe_experts", diff --git a/src/transformers/integrations/mxfp4.py b/src/transformers/integrations/mxfp4.py index f207d358540c..749e40117f78 100644 --- a/src/transformers/integrations/mxfp4.py +++ b/src/transformers/integrations/mxfp4.py @@ -22,7 +22,6 @@ from ..core_model_loading import ConversionOps, _IdentityOp from ..distributed.utils import _is_torch_distributed_initialized from ..quantizers.quantizers_utils import get_module_from_name, on_device, should_convert_module -from .hub_kernels import use_kernel_forward_from_hub logger = logging.get_logger(__name__) diff --git a/src/transformers/quantizers/quantizer_quanto.py b/src/transformers/quantizers/quantizer_quanto.py index 7201a0b5494e..f02675c68d3f 100644 --- a/src/transformers/quantizers/quantizer_quanto.py +++ b/src/transformers/quantizers/quantizer_quanto.py @@ -39,15 +39,9 @@ class QuantoHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "QuantoConfig" - def __init__(self, quantization_config: QuantoConfig, **kwargs): - super().__init__(quantization_config, **kwargs) - map_to_param_size = { - "int8": 1, - "float8": 1, - "int4": 0.5, - "int2": 0.25, - } - self.quantized_element_size = map_to_param_size.get(self.quantization_config.weights) + @property + def quantized_element_size(self) -> float | None: + return {"int8": 1, "float8": 1, "int4": 0.5, "int2": 0.25}.get(self.quantization_config.weights) def validate_environment(self, *args, **kwargs): if not is_optimum_quanto_available(): diff --git a/src/transformers/quantizers/quantizer_torchao.py b/src/transformers/quantizers/quantizer_torchao.py index 7dfd155e1aa9..d6b31c8f76a1 100644 --- a/src/transformers/quantizers/quantizer_torchao.py +++ b/src/transformers/quantizers/quantizer_torchao.py @@ -60,11 +60,9 @@ class TorchAoHfQuantizer(HfQuantizer): requires_calibration = False quantization_config: "TorchAoConfig" - def __init__(self, quantization_config, **kwargs): - super().__init__(quantization_config, **kwargs) - - size_digit = _fuzzy_match_size(type(self.quantization_config.quant_type).__name__) - self.quantized_element_size = 0.5 if size_digit == "4" else 1 + @property + def quantized_element_size(self) -> float: + return 0.5 if _fuzzy_match_size(type(self.quantization_config.quant_type).__name__) == "4" else 1 def validate_environment(self, *args, **kwargs): if not is_torchao_available(): diff --git a/tests/quantization/finegrained_fp8/test_fp8.py b/tests/quantization/finegrained_fp8/test_fp8.py index 4e0d03d0340e..0017cc967588 100644 --- a/tests/quantization/finegrained_fp8/test_fp8.py +++ b/tests/quantization/finegrained_fp8/test_fp8.py @@ -520,3 +520,38 @@ def test_single_device_keeps_deepgemm(self): _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 + from safetensors.torch import load_file, save_file + + from transformers import MixtralConfig, MixtralForCausalLM + + 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, + ) # 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" + ) + expected = {k: v.float() for k, v in quantized.state_dict().items() if k.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() + } + 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")} + self.assertEqual(scales.keys(), expected.keys()) + for name, scale in scales.items(): + self.assertEqual(scale.dtype, torch.float8_e8m0fnu, name) + torch.testing.assert_close(scale.float(), expected[name], rtol=0, atol=0) From 602fd8bac37c91e7b9c93f42325beaec2e98fda7 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 16 Oct 2026 16:16:34 +0900 Subject: [PATCH 25/27] Drop the $ anchor from the bailing_hybrid expert_bias renaming, patterns end on a name boundary now --- src/transformers/conversion_mapping.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/conversion_mapping.py b/src/transformers/conversion_mapping.py index 4e3a05222cc6..f83c580fbe2b 100755 --- a/src/transformers/conversion_mapping.py +++ b/src/transformers/conversion_mapping.py @@ -2015,7 +2015,7 @@ def _build_checkpoint_conversion_mapping(): operations=[Concatenate(dim=0)], ), WeightRenaming(r"\.self_attn\.dense\.", ".self_attn.o_proj."), - WeightRenaming(r"\.mlp\.gate\.expert_bias$", ".mlp.gate.e_score_correction_bias"), + WeightRenaming(r"\.mlp\.gate\.expert_bias", ".mlp.gate.e_score_correction_bias"), ] mapping["ernie4_5_moe"] = mapping["qwen2_moe"].copy() From e185b4526afa96d1a4f892cc2a07f169750da5f8 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 16 Oct 2026 18:17:35 +0900 Subject: [PATCH 26/27] Restore the 1 byte quantized element size of FineGrainedFP8 --- src/transformers/quantizers/quantizer_finegrained_fp8.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 2c2483c7eb5a..8bbde0288a13 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -22,6 +22,7 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): """ requires_calibration = False + quantized_element_size = 1 quantization_config: "FineGrainedFP8Config" def validate_environment(self, *args, **kwargs): From d6624ae505e26e95788b97aebd09fa3efa9304c3 Mon Sep 17 00:00:00 2001 From: Arthur Date: Fri, 16 Oct 2026 18:24:40 +0900 Subject: [PATCH 27/27] Quantized element size as a property in every quantizer --- src/transformers/quantizers/base.py | 5 ++++- src/transformers/quantizers/quantizer_bnb_4bit.py | 5 ++++- src/transformers/quantizers/quantizer_bnb_8bit.py | 5 ++++- src/transformers/quantizers/quantizer_fbgemm_fp8.py | 5 ++++- src/transformers/quantizers/quantizer_finegrained_fp8.py | 5 ++++- src/transformers/quantizers/quantizer_nvfp4.py | 5 ++++- 6 files changed, 24 insertions(+), 6 deletions(-) diff --git a/src/transformers/quantizers/base.py b/src/transformers/quantizers/base.py index 51bd4646f4e9..72c14f792de1 100644 --- a/src/transformers/quantizers/base.py +++ b/src/transformers/quantizers/base.py @@ -121,7 +121,10 @@ def update_device_map(self, device_map: dict[str, Any] | None) -> dict[str, Any] """ return device_map - quantized_element_size: float | None = None + @property + def quantized_element_size(self) -> float | None: + "Bytes per element of a quantized weight, `None` to use the param's own element size." + return None def param_element_size(self, model: "PreTrainedModel", param_name: str, param: "torch.Tensor") -> float: if self.quantized_element_size is not None and self.param_needs_quantization(model, param_name): diff --git a/src/transformers/quantizers/quantizer_bnb_4bit.py b/src/transformers/quantizers/quantizer_bnb_4bit.py index c5e31f419b4f..0af8ba4b65ae 100644 --- a/src/transformers/quantizers/quantizer_bnb_4bit.py +++ b/src/transformers/quantizers/quantizer_bnb_4bit.py @@ -47,9 +47,12 @@ class Bnb4BitHfQuantizer(HfQuantizer): """ requires_calibration = False - quantized_element_size = 0.5 quantization_config: "BitsAndBytesConfig" + @property + def quantized_element_size(self) -> float: + return 0.5 + def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError( diff --git a/src/transformers/quantizers/quantizer_bnb_8bit.py b/src/transformers/quantizers/quantizer_bnb_8bit.py index e7fcb6c267cc..e075676d2a06 100644 --- a/src/transformers/quantizers/quantizer_bnb_8bit.py +++ b/src/transformers/quantizers/quantizer_bnb_8bit.py @@ -47,9 +47,12 @@ class Bnb8BitHfQuantizer(HfQuantizer): """ requires_calibration = False - quantized_element_size = 1 quantization_config: "BitsAndBytesConfig" + @property + def quantized_element_size(self) -> float: + return 1 + def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError( diff --git a/src/transformers/quantizers/quantizer_fbgemm_fp8.py b/src/transformers/quantizers/quantizer_fbgemm_fp8.py index b51f568fed33..2980696894c6 100644 --- a/src/transformers/quantizers/quantizer_fbgemm_fp8.py +++ b/src/transformers/quantizers/quantizer_fbgemm_fp8.py @@ -44,10 +44,13 @@ class FbgemmFp8HfQuantizer(HfQuantizer): """ requires_calibration = False - quantized_element_size = 1 serialized_names = {"weight": ["weight", "weight_scale"]} quantization_config: "FbgemmFp8Config" + @property + def quantized_element_size(self) -> float: + return 1 + def validate_environment(self, *args, **kwargs): if not is_torch_cuda_available() and not is_torch_xpu_available(): raise ImportError("Using fbgemm fp8 quantization requires a GPU or XPU") diff --git a/src/transformers/quantizers/quantizer_finegrained_fp8.py b/src/transformers/quantizers/quantizer_finegrained_fp8.py index 8bbde0288a13..823328a48a41 100644 --- a/src/transformers/quantizers/quantizer_finegrained_fp8.py +++ b/src/transformers/quantizers/quantizer_finegrained_fp8.py @@ -22,9 +22,12 @@ class FineGrainedFP8HfQuantizer(HfQuantizer): """ requires_calibration = False - quantized_element_size = 1 quantization_config: "FineGrainedFP8Config" + @property + def quantized_element_size(self) -> float: + return 1 + def validate_environment(self, *args, **kwargs): if not is_accelerate_available(): raise ImportError("Loading an FP8 quantized model requires accelerate (`pip install accelerate`)") diff --git a/src/transformers/quantizers/quantizer_nvfp4.py b/src/transformers/quantizers/quantizer_nvfp4.py index dcc1afadf9ae..0779c9d60ebf 100644 --- a/src/transformers/quantizers/quantizer_nvfp4.py +++ b/src/transformers/quantizers/quantizer_nvfp4.py @@ -43,9 +43,12 @@ class NVFP4HfQuantizer(HfQuantizer): """Quantize eligible linear weights to NVFP4 while loading a full-precision checkpoint.""" requires_calibration = False - quantized_element_size = 0.5 quantization_config: NVFP4Config + @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.")