Support bf16 activations in Llama4 ExecuTorch XNNPACK export (#21784)

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,
         )
 
 

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}"
 

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 ``dtype`` (mirroring the linear path, where computation and
+    scale precision are both driven by the model dtype).
     """
     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 dtype,
     )
     return module
 

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()

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(