From 5c4898361efbab5d58b7035d265d580df66b42bd Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Thu, 13 Aug 2026 06:59:49 -0700 Subject: [PATCH] Support bf16 activations in ExecuTorch XNNPACK export (#21784) Summary: fb: Reviewed By: digantdesai, billmguo, telgamal-1 Differential Revision: D115569455 --- backends/xnnpack/partition/xnnpack_partitioner.py | 8 +++++++- examples/models/llama/export_llama_lib.py | 11 +++++++++-- .../llama/source_transformation/pre_quantization.py | 10 +++++++++- .../models/llama/source_transformation/quantize.py | 10 ++++++++-- extension/llm/export/partitioner_lib.py | 13 ++++++++++--- 5 files changed, 43 insertions(+), 9 deletions(-) diff --git a/backends/xnnpack/partition/xnnpack_partitioner.py b/backends/xnnpack/partition/xnnpack_partitioner.py index 44207e2247a..5327121b158 100644 --- a/backends/xnnpack/partition/xnnpack_partitioner.py +++ b/backends/xnnpack/partition/xnnpack_partitioner.py @@ -144,9 +144,15 @@ def generate_per_op_partitions(self, ep: ExportedProgram) -> List[Partition]: class XnnpackDynamicallyQuantizedPartitioner(XnnpackPartitioner): - def __init__(self): + def __init__(self, **kwargs): + if "config_precisions" in kwargs: + raise ValueError( + "XnnpackDynamicallyQuantizedPartitioner pins config_precisions to " + "DYNAMIC_QUANT and does not accept a config_precisions argument." + ) super().__init__( config_precisions=ConfigPrecisionType.DYNAMIC_QUANT, + **kwargs, ) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index a15268e9751..61bf1961166 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -983,17 +983,24 @@ def _to_edge_and_lower_llama_xnnpack( generate_etrecord: bool = False, verbose: bool = False, gen_tag_fn: Optional[Callable[[torch.fx.Node], Optional[str]]] = None, + enable_bf16: bool = False, ) -> LLMEdgeManager: # noqa: C901 partitioners = [] # Order matters here, dynamic quantization should be applied first when both xnnpack and xnnpack_extended_ops are enabled - partitioners.append(get_xnnpack_partitioner(dynamic_quant_only_partitioner=True)) + partitioners.append( + get_xnnpack_partitioner( + dynamic_quant_only_partitioner=True, enable_bf16=enable_bf16 + ) + ) modelname = f"xnnpack_dq_{modelname}" if xnnpack_extended_ops: partitioners.append( - get_xnnpack_partitioner(dynamic_quant_only_partitioner=False) + get_xnnpack_partitioner( + dynamic_quant_only_partitioner=False, enable_bf16=enable_bf16 + ) ) modelname = f"xnnpack_{modelname}" diff --git a/examples/models/llama/source_transformation/pre_quantization.py b/examples/models/llama/source_transformation/pre_quantization.py index 8f6bcab8046..26450340071 100644 --- a/examples/models/llama/source_transformation/pre_quantization.py +++ b/examples/models/llama/source_transformation/pre_quantization.py @@ -279,13 +279,14 @@ def _replace_embedding_with_quantized_group_embedding_for_pre_quantization( dtype: torch.dtype, bit_width: int, group_size: Optional[int] = None, + scales_precision: torch.dtype = torch.float32, ): def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool: # Only replace embedding layers where the checkpoint contains explicit scales scales_key = f"{cur_fqn}.scales" if isinstance(child, nn.Embedding) and scales_key in checkpoint: assert checkpoint[f"{cur_fqn}.weight"].dtype == torch.int8 - assert checkpoint[scales_key].dtype == torch.float32 + assert checkpoint[scales_key].dtype == scales_precision return True return False @@ -297,6 +298,7 @@ def replacement_fn(child: torch.nn.Module) -> torch.nn.Module: group_size=group_size, dtype=dtype, packed=False, # TODO(lunwenh): support packed embedding for pre-quantized + scales_precision=scales_precision, ) return new_embedding @@ -309,10 +311,15 @@ def transform_embedding_for_pre_quantization( dtype: torch.dtype, bit_width: int, group_size: Optional[int] = None, + scales_precision: Optional[torch.dtype] = None, ) -> torch.nn.Module: """ Transform the model to be able to load pre-quantized checkpoints that are quantized with the given bit_width and group size for embedding. + + ``scales_precision`` is the dtype the checkpoint stores embedding scales in; + it defaults to ``torch.float32`` to preserve the pre-existing contract + unless explicitly overridden. """ if group_size is not None and group_size not in [0, 32, 64, 128, 256]: raise ValueError( @@ -324,6 +331,7 @@ def transform_embedding_for_pre_quantization( dtype, bit_width, group_size, + scales_precision if scales_precision is not None else torch.float32, ) return module diff --git a/examples/models/llama/source_transformation/quantize.py b/examples/models/llama/source_transformation/quantize.py index 70afc7dcbcc..d820428c2bc 100644 --- a/examples/models/llama/source_transformation/quantize.py +++ b/examples/models/llama/source_transformation/quantize.py @@ -688,8 +688,11 @@ def __init__( dtype=torch.half, packed=False, bitwidth: int = 8, + scales_precision: Optional[torch.dtype] = None, ) -> None: super().__init__() + if scales_precision is None: + scales_precision = torch.float16 if group_size is None or group_size == 0: group_size = embedding_dim self.group_size = group_size @@ -728,12 +731,15 @@ def __init__( self.register_buffer( "scales", torch.ones( - (vocab_size, groups_per_row), dtype=torch.float16, device=device + (vocab_size, groups_per_row), + dtype=scales_precision, + device=device, ), ) else: self.register_buffer( - "scales", torch.ones((vocab_size,), dtype=torch.float16, device=device) + "scales", + torch.ones((vocab_size,), dtype=scales_precision, device=device), ) @torch.no_grad() diff --git a/extension/llm/export/partitioner_lib.py b/extension/llm/export/partitioner_lib.py index 19c0b7fdcfb..f319dc8dc19 100644 --- a/extension/llm/export/partitioner_lib.py +++ b/extension/llm/export/partitioner_lib.py @@ -8,7 +8,10 @@ from typing import List, Optional -def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True): +def get_xnnpack_partitioner( + dynamic_quant_only_partitioner: bool = True, + enable_bf16: bool = False, +): """ Returns the XNNPACK partitioner. @@ -17,6 +20,10 @@ def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True): If dynamic_quant_only_partitioner is True, then only dynamically quantized linear layers will be partitioned. Else, anything which can be will be partitioned greedily. + @arg enable_bf16: + Opt in to delegating bf16 nodes. Required to keep quantized linears + delegated when the model runs with bf16 activations; without it every + bf16 node fails the partitioner dtype check and falls back to portable. """ from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( XnnpackDynamicallyQuantizedPartitioner, @@ -28,8 +35,8 @@ def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True): # 1. We need dynamically quantized partitioner for both pt2e_quantize options # as well as "qmode 8da4w" which is also dynamic quantizes linear layers. # 2. XNNPACK partitioner seems to result in seg fault for non dqlinear ops. - return XnnpackDynamicallyQuantizedPartitioner() - return XnnpackPartitioner() + return XnnpackDynamicallyQuantizedPartitioner(enable_bf16=enable_bf16) + return XnnpackPartitioner(enable_bf16=enable_bf16) def get_vulkan_partitioner(