Skip to content

Commit 2b41021

Browse files
YufengShi-duduYufengShi-dudu
authored
Arm backend: Fix FP8 conv and matmul TOSA lowering (pytorch#20219)
Lower FP8 MATMUL with FP16 accumulation for TOSA 1.0 and keep the cast back to the exported graph dtype when the TOSA op widens the output. Create and validate FP16 bias tensors for FP8 conv-family TOSA ops, matching the TOSA output domain. This fixes invalid FP8 conv and matmul TOSA lowering. Change-Id: Ib2a12454e6df97535173eb8131c32c142a73544e cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>
1 parent e7c5415 commit 2b41021

12 files changed

Lines changed: 75 additions & 52 deletions

‎backends/arm/_passes/rewrite_conv_pass.py‎

Lines changed: 38 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ class RewriteConvPass(ArmPass):
4848
(CONV2D/DEPTHWISE/TRANSPOSE/CONV3D).
4949
"""
5050

51+
_FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
52+
5153
def __init__(self, exported_program: torch.export.ExportedProgram, *args, **kwargs):
5254
super().__init__(*args, **kwargs)
5355
self.exported_program = exported_program
@@ -146,11 +148,15 @@ def _add_bias(
146148
graph_module: torch.fx.GraphModule,
147149
node: torch.fx.Node,
148150
weight_node: torch.fx.Node,
151+
input_fake_tensor: torch.Tensor,
149152
) -> torch.fx.Node:
150153
output_channels = get_first_fake_tensor(node).shape[1]
151-
# add a node containing zeros if quantized, use int32, otherwise use float32
154+
# Add a zero bias with the dtype TOSA expects: int32 for
155+
# quantized conv, fp16 for FP8 conv, and the output dtype otherwise.
152156
if self._is_quantized_conv(node):
153157
bias_data = torch.zeros(size=(output_channels,), dtype=torch.int32)
158+
elif input_fake_tensor.dtype in self._FP8_DTYPES:
159+
bias_data = torch.zeros(size=(output_channels,), dtype=torch.float16)
154160
else:
155161
output_dtype = node.meta["val"].dtype
156162
bias_data = torch.zeros(size=(output_channels,), dtype=output_dtype)
@@ -174,6 +180,32 @@ def _add_bias(
174180
node.update_arg(2, bias_node)
175181
return bias_node
176182

183+
def _rewrite_fp8_bias(
184+
self,
185+
graph_module: torch.fx.GraphModule,
186+
node: torch.fx.Node,
187+
bias_node: torch.fx.Node,
188+
) -> torch.fx.Node:
189+
bias_tensor = get_param_tensor( # type: ignore[arg-type]
190+
self.exported_program, bias_node
191+
)
192+
if bias_tensor is None:
193+
raise RuntimeError(
194+
f"Bias node {bias_node.name} is not a parameter or buffer"
195+
)
196+
197+
kind = get_constant_placeholder_kind(self.exported_program, bias_node)
198+
persistent_buffer = is_persistent_buffer(self.exported_program, bias_node)
199+
with graph_module.graph.inserting_after(bias_node):
200+
return create_constant_placeholder(
201+
self.exported_program,
202+
graph=graph_module.graph,
203+
name=f"{node.name}_bias_fp16",
204+
kind=kind,
205+
data=bias_tensor.to(torch.float16),
206+
persistent_buffer=persistent_buffer,
207+
)
208+
177209
def _rewrite_weight(
178210
self,
179211
graph_module: torch.fx.GraphModule,
@@ -449,9 +481,12 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901
449481

450482
has_bias = bias is not None
451483
if not has_bias:
452-
bias = self._add_bias(graph_module, node, weight)
484+
bias = self._add_bias(graph_module, node, weight, input_fake_tensor)
453485
elif isinstance(bias, torch.fx.Node):
454-
self._mark_bias_as_int48_if_needed(node, bias)
486+
if input_fake_tensor.dtype in self._FP8_DTYPES:
487+
bias = self._rewrite_fp8_bias(graph_module, node, bias)
488+
else:
489+
self._mark_bias_as_int48_if_needed(node, bias)
455490

456491
conv_args: tuple[Any, ...]
457492
input_tensor_for_tosa_fake: torch.Tensor = input_fake_tensor

‎backends/arm/_passes/rewrite_matmul.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -105,11 +105,10 @@ def call(self, graph_module):
105105
elif (
106106
x1_fake_tensor.dtype in self._WIDENING_INPUT_DTYPES
107107
and x2_fake_tensor.dtype in self._WIDENING_INPUT_DTYPES
108-
and output_fake_tensor.dtype not in self._WIDENING_INPUT_DTYPES
108+
and output_fake_tensor.dtype != node_output_fake_tensor.dtype
109109
):
110-
# TOSA BF16/FP16/FP8 MATMUL outputs FP32, while the original
111-
# exported node outputs BF16/FP16/FP8. Cast back to preserve
112-
# the exported graph dtype.
110+
# TOSA BF16/FP16/FP8 MATMUL widens the output. Cast back to
111+
# preserve the exported graph dtype.
113112
with graph_module.graph.inserting_after(tosa_matmul_node):
114113
cast_node = create_node(
115114
graph_module.graph,

‎backends/arm/operators/op_tosa_conv2d.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,8 +69,22 @@ def define_node(
6969
valid_input_dtypes.append(ts.DType.BF16)
7070
if self.tosa_spec.support_extension("fp8e4m3"):
7171
valid_input_dtypes.append(ts.DType.FP8E4M3)
72+
if inputs[0].dtype == ts.DType.FP8E4M3:
73+
validate_valid_dtype(
74+
self.target, [inputs[1]], [ts.DType.FP8E4M3], self.tosa_spec
75+
)
76+
validate_valid_dtype(
77+
self.target, [inputs[2]], [ts.DType.FP16], self.tosa_spec
78+
)
7279
if self.tosa_spec.support_extension("fp8e5m2"):
7380
valid_input_dtypes.append(ts.DType.FP8E5M2)
81+
if inputs[0].dtype == ts.DType.FP8E5M2:
82+
validate_valid_dtype(
83+
self.target, [inputs[1]], [ts.DType.FP8E5M2], self.tosa_spec
84+
)
85+
validate_valid_dtype(
86+
self.target, [inputs[2]], [ts.DType.FP16], self.tosa_spec
87+
)
7488

7589
validate_valid_dtype(
7690
self.target,

‎backends/arm/operators/op_tosa_matmul.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ def define_node(
6262
validate_valid_dtype(
6363
self.target,
6464
[output],
65-
[ts.DType.INT32, ts.DType.INT48, ts.DType.FP32],
65+
[ts.DType.INT32, ts.DType.INT48, ts.DType.FP16, ts.DType.FP32],
6666
self.tosa_spec,
6767
)
6868

‎backends/arm/operators/op_tosa_transpose_conv2d.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ def define_node(
8080
self.target, [inputs[1]], [ts.DType.FP8E4M3], self.tosa_spec
8181
)
8282
validate_valid_dtype(
83-
self.target, [inputs[2]], [ts.DType.FP8E4M3], self.tosa_spec
83+
self.target, [inputs[2]], [ts.DType.FP16], self.tosa_spec
8484
)
8585
if self.tosa_spec.support_extension("fp8e5m2"):
8686
valid_input_dtypes.append(ts.DType.FP8E5M2)
@@ -89,7 +89,7 @@ def define_node(
8989
self.target, [inputs[1]], [ts.DType.FP8E5M2], self.tosa_spec
9090
)
9191
validate_valid_dtype(
92-
self.target, [inputs[2]], [ts.DType.FP8E5M2], self.tosa_spec
92+
self.target, [inputs[2]], [ts.DType.FP16], self.tosa_spec
9393
)
9494

9595
validate_valid_dtype(

‎backends/arm/test/ops/test_conv2d.py‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -553,10 +553,6 @@ def conv2d_fp16_1x1():
553553
"fp8e5m2",
554554
),
555555
}
556-
_fp8_conv2d_tosa_ref_model_xfails = {
557-
name: "MLETORCH-2238: Fix invalid FP8 CONV TOSA graphs" for name in test_data_FP_fp8
558-
}
559-
560556
# Generate a new test set paired with per_channel_quant=True/False.
561557
test_data_INT = {
562558
f"{k},per_channel_quant={q}": (lambda v=v, q=q: (v(), q))
@@ -611,9 +607,7 @@ def test_convolution_2d_tosa_FP(test_data):
611607
pipeline.run()
612608

613609

614-
@common.parametrize(
615-
"test_data", test_data_FP_fp8, xfails=_fp8_conv2d_tosa_ref_model_xfails
616-
)
610+
@common.parametrize("test_data", test_data_FP_fp8)
617611
def test_convolution_2d_tosa_FP_fp8(test_data):
618612
model, tosa_extension = test_data()
619613
pipeline = TosaPipelineFP[input_t](

‎backends/arm/test/ops/test_conv3d.py‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -515,10 +515,6 @@ def forward(self, x):
515515
"fp8e5m2",
516516
),
517517
}
518-
_fp8_conv3d_tosa_ref_model_xfails = {
519-
name: "MLETORCH-2238: Fix invalid FP8 CONV TOSA graphs" for name in test_data_FP_fp8
520-
}
521-
522518
test_data_FP_bf16 = {
523519
"bf16_3x3": lambda: Conv3d(
524520
height=10,
@@ -611,9 +607,7 @@ def test_convolution_3d_tosa_FP(test_data):
611607
pipeline.run()
612608

613609

614-
@common.parametrize(
615-
"test_data", test_data_FP_fp8, xfails=_fp8_conv3d_tosa_ref_model_xfails
616-
)
610+
@common.parametrize("test_data", test_data_FP_fp8)
617611
def test_convolution_3d_tosa_FP_fp8(test_data):
618612
model, tosa_extension = test_data()
619613
pipeline = TosaPipelineFP[input_t](

‎backends/arm/test/ops/test_depthwise_conv.py‎

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -227,11 +227,6 @@
227227
"fp8e5m2",
228228
),
229229
}
230-
_fp8_depthwise_conv_tosa_ref_model_xfails = {
231-
name: "MLETORCH-2238: Fix invalid FP8 CONV TOSA graphs"
232-
for name in test_data_conv2d_FP_fp8
233-
}
234-
235230
# Generate a new test set paired with per_channel_quant=True/False.
236231
test_data_conv2d_INT = {
237232
f"{k},per_channel_quant={q}": (lambda v=v, q=q: (v(), q))
@@ -293,11 +288,7 @@ def test_convolution_2d_tosa_FP_depthwise(test_data: torch.nn.Module):
293288
pipeline.run()
294289

295290

296-
@common.parametrize(
297-
"test_data",
298-
test_data_conv2d_FP_fp8,
299-
xfails=_fp8_depthwise_conv_tosa_ref_model_xfails,
300-
)
291+
@common.parametrize("test_data", test_data_conv2d_FP_fp8)
301292
def test_convolution_2d_tosa_FP_fp8_depthwise(test_data):
302293
model, tosa_extension = test_data()
303294
pipeline = TosaPipelineFP[input_t](

‎backends/arm/test/ops/test_matmul.py‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -375,10 +375,6 @@ def forward(self, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor):
375375
(exir_op_mm_2d, exir_op_mm_2d),
376376
),
377377
}
378-
fp8_xfails = {
379-
name: "MLETORCH-2239: Fix invalid FP8 MATMUL TOSA graphs" for name in test_suite_fp8
380-
}
381-
382378
xfails = {
383379
"double_input_randn_rand_1d_1d": "aten.dot.default is not supported",
384380
"double_input_randn_rand_2d_1d": "aten.mv.default is not supported",
@@ -401,7 +397,7 @@ def test_matmul_tosa_FP(test_case: test_case_t):
401397
pipeline.run()
402398

403399

404-
@common.parametrize("test_case", test_suite_fp8, xfails=fp8_xfails)
400+
@common.parametrize("test_case", test_suite_fp8)
405401
def test_matmul_tosa_FP_fp8(test_case: test_case_t):
406402
test_data = test_case()
407403
input_dtype = test_data.input_factory()[0].dtype

‎backends/arm/test/ops/test_transpose_conv2d.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -267,9 +267,6 @@ def _get_per_channel_observers(module: torch.nn.Module):
267267
"fp8e5m2",
268268
),
269269
}
270-
_fp8_transpose_conv_tosa_ref_model_xfails = {
271-
name: "MLETORCH-2238: Fix invalid FP8 CONV TOSA graphs" for name in test_data_FP8
272-
}
273270

274271

275272
@common.parametrize("test_data", test_data_FP | test_data_FP_fp16 | test_data_BF16)
@@ -287,9 +284,7 @@ def test_conv_transpose2d_tosa_FP(test_data):
287284
pipeline.run()
288285

289286

290-
@common.parametrize(
291-
"test_data", test_data_FP8, xfails=_fp8_transpose_conv_tosa_ref_model_xfails
292-
)
287+
@common.parametrize("test_data", test_data_FP8)
293288
def test_conv_transpose2d_tosa_FP_fp8(test_data):
294289
model, tosa_extension = test_data()
295290
pipeline = TosaPipelineFP[input_t](

0 commit comments

Comments
 (0)