Skip to content

Commit abc5586

Browse files
authored
Support bf16 activations in ExecuTorch XNNPACK export (pytorch#21784)
Differential Revision: D115569455 Pull Request resolved: pytorch#21784
1 parent 7dd3335 commit abc5586

5 files changed

Lines changed: 43 additions & 9 deletions

File tree

‎backends/xnnpack/partition/xnnpack_partitioner.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,9 +144,15 @@ def generate_per_op_partitions(self, ep: ExportedProgram) -> List[Partition]:
144144

145145

146146
class XnnpackDynamicallyQuantizedPartitioner(XnnpackPartitioner):
147-
def __init__(self):
147+
def __init__(self, **kwargs):
148+
if "config_precisions" in kwargs:
149+
raise ValueError(
150+
"XnnpackDynamicallyQuantizedPartitioner pins config_precisions to "
151+
"DYNAMIC_QUANT and does not accept a config_precisions argument."
152+
)
148153
super().__init__(
149154
config_precisions=ConfigPrecisionType.DYNAMIC_QUANT,
155+
**kwargs,
150156
)
151157

152158

‎examples/models/llama/export_llama_lib.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -983,17 +983,24 @@ def _to_edge_and_lower_llama_xnnpack(
983983
generate_etrecord: bool = False,
984984
verbose: bool = False,
985985
gen_tag_fn: Optional[Callable[[torch.fx.Node], Optional[str]]] = None,
986+
enable_bf16: bool = False,
986987
) -> LLMEdgeManager: # noqa: C901
987988
partitioners = []
988989

989990
# Order matters here, dynamic quantization should be applied first when both xnnpack and xnnpack_extended_ops are enabled
990-
partitioners.append(get_xnnpack_partitioner(dynamic_quant_only_partitioner=True))
991+
partitioners.append(
992+
get_xnnpack_partitioner(
993+
dynamic_quant_only_partitioner=True, enable_bf16=enable_bf16
994+
)
995+
)
991996

992997
modelname = f"xnnpack_dq_{modelname}"
993998

994999
if xnnpack_extended_ops:
9951000
partitioners.append(
996-
get_xnnpack_partitioner(dynamic_quant_only_partitioner=False)
1001+
get_xnnpack_partitioner(
1002+
dynamic_quant_only_partitioner=False, enable_bf16=enable_bf16
1003+
)
9971004
)
9981005
modelname = f"xnnpack_{modelname}"
9991006

‎examples/models/llama/source_transformation/pre_quantization.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -279,13 +279,14 @@ def _replace_embedding_with_quantized_group_embedding_for_pre_quantization(
279279
dtype: torch.dtype,
280280
bit_width: int,
281281
group_size: Optional[int] = None,
282+
scales_precision: torch.dtype = torch.float32,
282283
):
283284
def filter_fn(child: torch.nn.Module, cur_fqn: str) -> bool:
284285
# Only replace embedding layers where the checkpoint contains explicit scales
285286
scales_key = f"{cur_fqn}.scales"
286287
if isinstance(child, nn.Embedding) and scales_key in checkpoint:
287288
assert checkpoint[f"{cur_fqn}.weight"].dtype == torch.int8
288-
assert checkpoint[scales_key].dtype == torch.float32
289+
assert checkpoint[scales_key].dtype == scales_precision
289290
return True
290291
return False
291292

@@ -297,6 +298,7 @@ def replacement_fn(child: torch.nn.Module) -> torch.nn.Module:
297298
group_size=group_size,
298299
dtype=dtype,
299300
packed=False, # TODO(lunwenh): support packed embedding for pre-quantized
301+
scales_precision=scales_precision,
300302
)
301303
return new_embedding
302304

@@ -309,10 +311,15 @@ def transform_embedding_for_pre_quantization(
309311
dtype: torch.dtype,
310312
bit_width: int,
311313
group_size: Optional[int] = None,
314+
scales_precision: Optional[torch.dtype] = None,
312315
) -> torch.nn.Module:
313316
"""
314317
Transform the model to be able to load pre-quantized checkpoints that
315318
are quantized with the given bit_width and group size for embedding.
319+
320+
``scales_precision`` is the dtype the checkpoint stores embedding scales in;
321+
it defaults to ``torch.float32`` to preserve the pre-existing contract
322+
unless explicitly overridden.
316323
"""
317324
if group_size is not None and group_size not in [0, 32, 64, 128, 256]:
318325
raise ValueError(
@@ -324,6 +331,7 @@ def transform_embedding_for_pre_quantization(
324331
dtype,
325332
bit_width,
326333
group_size,
334+
scales_precision if scales_precision is not None else torch.float32,
327335
)
328336
return module
329337

‎examples/models/llama/source_transformation/quantize.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -688,8 +688,11 @@ def __init__(
688688
dtype=torch.half,
689689
packed=False,
690690
bitwidth: int = 8,
691+
scales_precision: Optional[torch.dtype] = None,
691692
) -> None:
692693
super().__init__()
694+
if scales_precision is None:
695+
scales_precision = torch.float16
693696
if group_size is None or group_size == 0:
694697
group_size = embedding_dim
695698
self.group_size = group_size
@@ -728,12 +731,15 @@ def __init__(
728731
self.register_buffer(
729732
"scales",
730733
torch.ones(
731-
(vocab_size, groups_per_row), dtype=torch.float16, device=device
734+
(vocab_size, groups_per_row),
735+
dtype=scales_precision,
736+
device=device,
732737
),
733738
)
734739
else:
735740
self.register_buffer(
736-
"scales", torch.ones((vocab_size,), dtype=torch.float16, device=device)
741+
"scales",
742+
torch.ones((vocab_size,), dtype=scales_precision, device=device),
737743
)
738744

739745
@torch.no_grad()

‎extension/llm/export/partitioner_lib.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@
88
from typing import List, Optional
99

1010

11-
def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True):
11+
def get_xnnpack_partitioner(
12+
dynamic_quant_only_partitioner: bool = True,
13+
enable_bf16: bool = False,
14+
):
1215
"""
1316
Returns the XNNPACK partitioner.
1417
@@ -17,6 +20,10 @@ def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True):
1720
If dynamic_quant_only_partitioner is True, then only dynamically quantized
1821
linear layers will be partitioned.
1922
Else, anything which can be will be partitioned greedily.
23+
@arg enable_bf16:
24+
Opt in to delegating bf16 nodes. Required to keep quantized linears
25+
delegated when the model runs with bf16 activations; without it every
26+
bf16 node fails the partitioner dtype check and falls back to portable.
2027
"""
2128
from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
2229
XnnpackDynamicallyQuantizedPartitioner,
@@ -28,8 +35,8 @@ def get_xnnpack_partitioner(dynamic_quant_only_partitioner: bool = True):
2835
# 1. We need dynamically quantized partitioner for both pt2e_quantize options
2936
# as well as "qmode 8da4w" which is also dynamic quantizes linear layers.
3037
# 2. XNNPACK partitioner seems to result in seg fault for non dqlinear ops.
31-
return XnnpackDynamicallyQuantizedPartitioner()
32-
return XnnpackPartitioner()
38+
return XnnpackDynamicallyQuantizedPartitioner(enable_bf16=enable_bf16)
39+
return XnnpackPartitioner(enable_bf16=enable_bf16)
3340

3441

3542
def get_vulkan_partitioner(

0 commit comments

Comments
 (0)