diff --git a/backends/xnnpack/operators/node_visitor.py b/backends/xnnpack/operators/node_visitor.py index a0f03205ed5..160ee03c765 100644 --- a/backends/xnnpack/operators/node_visitor.py +++ b/backends/xnnpack/operators/node_visitor.py @@ -265,12 +265,15 @@ def get_per_channel_dtype( ) else: node_dtype = get_node_dtype(node) - if node_dtype is not None and node_dtype == torch.float16: - dtype = ( - XNNDatatype.xnn_datatype_fp32 - if force_fp32 - else XNNDatatype.xnn_datatype_fp16 - ) + # fp16/bf16 tensors keep their datatype unless we've been asked to + # force fp32 (e.g. biases for dynamic-quant or bf16 fully-connected), + # in which case they fall back to the default fp32. + float_dtype_map = { + torch.float16: XNNDatatype.xnn_datatype_fp16, + torch.bfloat16: XNNDatatype.xnn_datatype_bf16, + } + if not force_fp32: + dtype = float_dtype_map.get(node_dtype, dtype) return dtype @@ -591,7 +594,7 @@ def get_serialized_buffer_index( # Quantize buffer if static data is indeed quantized if quant_params is not None and not quant_params.is_dynamic: const_val = quant_params.quantize_tensor(const_val).contiguous() - elif const_val.dtype != torch.float16 or force_fp32: + elif const_val.dtype not in (torch.float16, torch.bfloat16) or force_fp32: # ensure that the const is fp32 const_val = const_val.to(dtype=torch.float32).contiguous() @@ -712,12 +715,19 @@ def define_nodes_tensor_inputs_outputs( bias_quant_params = QuantParams.from_bias( bias_node, weight_quant_params, input_quant_params ) + # XNNPACK's bf16 fully-connected (bf16_bf16_f32) takes a bf16 + # activation/weight but an fp32 bias, so force the bias to fp32. + weight_val = weight_node.meta.get("val", None) + bias_force_fp32 = ( + weight_val is not None and weight_val.dtype == torch.bfloat16 + ) self.define_tensor( bias_node, xnn_graph, vals_to_ids, quant_params=bias_quant_params, convert_to_nhwc=False, # Bias is generally 1d and can not be in NHWC + force_fp32=bias_force_fp32, ) def define_node( diff --git a/backends/xnnpack/operators/op_linear.py b/backends/xnnpack/operators/op_linear.py index dda1d3e53ef..31eb0d7950f 100644 --- a/backends/xnnpack/operators/op_linear.py +++ b/backends/xnnpack/operators/op_linear.py @@ -73,6 +73,11 @@ def define_node( force_fp32 = False if input_quant_params is not None and input_quant_params.is_dynamic: force_fp32 = True + # XNNPACK's bf16 fully-connected (bf16_bf16_f32) takes a bf16 + # activation/weight but an fp32 bias, so force the bias to fp32. + weight_val = weight_node.meta.get("val", None) + if weight_val is not None and weight_val.dtype == torch.bfloat16: + force_fp32 = True self.define_tensor( get_input_node(node, 2), diff --git a/backends/xnnpack/partition/config/xnnpack_config.py b/backends/xnnpack/partition/config/xnnpack_config.py index 817f9d1cf50..e0a575e000a 100644 --- a/backends/xnnpack/partition/config/xnnpack_config.py +++ b/backends/xnnpack/partition/config/xnnpack_config.py @@ -50,6 +50,11 @@ def __init__(self, **kwargs): self.force_non_static_weights_for_f32_linear = kwargs.get( "force_non_static_weights_for_f32_linear", False ) + # Opt-in flag for bf16 delegation (e.g. bf16 fully-connected). XNNPACK + # only supports bf16 fully-connected on new enough revisions, so this + # stays off by default and bf16 nodes fall back to the portable op + # unless the caller explicitly enables it. + self.enable_bf16 = kwargs.get("enable_bf16", False) def get_partition( self, node: torch.fx.Node, ep: ExportedProgram @@ -229,6 +234,8 @@ def _check_node_has_valid_dtype(self, node): torch.float32, torch.float16, } + if self.enable_bf16: + valid_dtypes.add(torch.bfloat16) # Only allow int8 and quant dtypes for quant operations if is_quant(node) or is_dequant(node) or is_qparam(node): valid_dtypes.update( diff --git a/backends/xnnpack/runtime/XNNCompiler.cpp b/backends/xnnpack/runtime/XNNCompiler.cpp index 3f2fb5a1aaa..96f3316999c 100644 --- a/backends/xnnpack/runtime/XNNCompiler.cpp +++ b/backends/xnnpack/runtime/XNNCompiler.cpp @@ -791,6 +791,32 @@ Error defineConvertNode( return Error::Ok; }; +/* +Look up a serialized tensor value (plain or quantized wrapper) by its +output id. Returns nullptr if not found. +*/ +const fb_xnnpack::XNNTensorValue* getSerializedTensorValue( + const fb_xnnpack::XNNGraph* graph, + uint32_t id) noexcept { + if (graph == nullptr || graph->xvalues() == nullptr) { + return nullptr; + } + for (auto value : *graph->xvalues()) { + const fb_xnnpack::XNNTensorValue* tv = nullptr; + if (value->xvalue_union_type() == fb_xnnpack::XValueUnion::XNNTensorValue) { + tv = value->xvalue_union_as_XNNTensorValue(); + } else if ( + value->xvalue_union_type() == + fb_xnnpack::XValueUnion::XNNQuantizedTensorValue) { + tv = value->xvalue_union_as_XNNQuantizedTensorValue()->tensor_value(); + } + if (tv != nullptr && tv->id_out() == id) { + return tv; + } + } + return nullptr; +} + /* Define serialized linear(fully-connected) node into the subgraph using the remapped ids to map the serialized ids, to the new ids generated @@ -810,6 +836,44 @@ Error defineFullyConnectedNode( REMAP_ID(remapped_ids, graph_node->bias_id(), fc_bias); REMAP_ID(remapped_ids, graph_node->output_id(), fc_output); + // XNNPACK only provides a bf16 fully-connected of type bf16_bf16_f32: + // bf16 activation x bf16 weight -> fp32 output. When the serialized graph + // asks for a bf16 output (e.g. a fully bf16 model), define the FC with an + // fp32 intermediate output and append a convert (fp32 -> bf16) so the + // delegate boundary stays bf16. + const auto* in_tv = getSerializedTensorValue(graph, graph_node->input1_id()); + const auto* filt_tv = + getSerializedTensorValue(graph, graph_node->filter_id()); + const auto* out_tv = getSerializedTensorValue(graph, graph_node->output_id()); + const bool needs_bf16_output_convert = in_tv != nullptr && + filt_tv != nullptr && out_tv != nullptr && + in_tv->datatype() == DataType::xnn_datatype_bf16 && + filt_tv->datatype() == DataType::xnn_datatype_bf16 && + out_tv->datatype() == DataType::xnn_datatype_bf16; + + uint32_t fc_compute_output = fc_output; + if (needs_bf16_output_convert) { + std::vector out_dims = + flatbufferDimsToVector(out_tv->dims()); + uint32_t intermediate_id = XNN_INVALID_VALUE_ID; + xnn_status ts = xnn_define_tensor_value( + subgraph_ptr, + xnn_datatype_fp32, + out_dims.size(), + out_dims.data(), + /*data=*/nullptr, + /*external_id=*/XNN_INVALID_VALUE_ID, + /*flags=*/0, + &intermediate_id); + ET_CHECK_OR_RETURN_ERROR( + ts == xnn_status_success, + Internal, + "Failed to define fp32 intermediate for bf16 linear node %i: %s", + node->debug_handle(), + xnn_status_to_string(ts)); + fc_compute_output = intermediate_id; + } + xnn_status status = xnn_define_fully_connected( subgraph_ptr, min_max.first, @@ -817,7 +881,7 @@ Error defineFullyConnectedNode( fc_input1, fc_filter, fc_bias, - fc_output, + fc_compute_output, graph_node->flags()); ET_CHECK_OR_RETURN_ERROR( status == xnn_status_success, @@ -826,6 +890,17 @@ Error defineFullyConnectedNode( node->debug_handle(), xnn_status_to_string(status)); + if (needs_bf16_output_convert) { + xnn_status cs = xnn_define_convert( + subgraph_ptr, fc_compute_output, fc_output, /*flags=*/0); + ET_CHECK_OR_RETURN_ERROR( + cs == xnn_status_success, + Internal, + "Failed to define bf16 output convert for linear node %i: %s", + node->debug_handle(), + xnn_status_to_string(cs)); + } + return Error::Ok; }; diff --git a/backends/xnnpack/test/ops/test_linear.py b/backends/xnnpack/test/ops/test_linear.py index dc92a9542a9..6e0fde705c0 100644 --- a/backends/xnnpack/test/ops/test_linear.py +++ b/backends/xnnpack/test/ops/test_linear.py @@ -46,6 +46,12 @@ torchao_installed = False +def is_fbcode() -> bool: + # torch.version.git_version is only set in OSS PyTorch; the internal + # Buck-built torch omits it. + return not hasattr(torch.version, "git_version") + + # Pytorch Modules Used for Testing class BaseLinear(torch.nn.Module): def __init__( @@ -64,7 +70,11 @@ def __init__( self.ic = input_channels self.oc = output_channels - assert dtype in [torch.float, torch.half], "Unsupported op dtype" + assert dtype in [ + torch.float, + torch.half, + torch.bfloat16, + ], "Unsupported op dtype" self.op_dtype = dtype self.in_size = in_size @@ -388,6 +398,7 @@ def _test_groupwise_dq_linear( num_linears: int = 1, atol: float = 5e-3, # TODO(T212995726): Investigate right atol for rand[n] inputs rtol: float = 5e-3, # TODO(T212995726): Investigate right rtol for rand[n] inputs + enable_bf16: bool = False, ): """ Helper function to test groupwise dynamic quantized linear op with different configurations. @@ -404,6 +415,7 @@ def _test_groupwise_dq_linear( DynamicallyQuantizedPartitioner = XnnpackPartitioner( config_precisions=ConfigPrecisionType.DYNAMIC_QUANT, per_op_mode=True, + enable_bf16=enable_bf16, ) tester = ( Tester(mod, inputs) @@ -706,11 +718,22 @@ def _test_qd8_per_token_weight_per_channel_group_int4( # Mean: 0.2373046875, 0.237060546875 # Max: 1.0078125, 1.0078125 # Min: -0.08465576171875, -0.08441162109375 - atol = ( - 1e-2 if dtype == torch.half else 5e-3 - ) # TODO(T212995726): Investigate right atol for rand[n] inputs + # bf16 has ~8x coarser mantissa than fp16, so it needs a + # looser atol. + # TODO(T212995726): Investigate right atol for rand[n] inputs + if dtype == torch.bfloat16: + atol = 8e-2 + elif dtype == torch.half: + atol = 1e-2 + else: + atol = 5e-3 self._test_groupwise_dq_linear( - lin_mod, inputs, group_size=bl, use_bias=use_bias, atol=atol + lin_mod, + inputs, + group_size=bl, + use_bias=use_bias, + atol=atol, + enable_bf16=dtype == torch.bfloat16, ) def test_fp16_linear(self): @@ -839,6 +862,17 @@ def test_linear_qd8_f16_per_token_weight_per_channel_group_int4(self): def test_linear_qd8_f32_per_token_weight_per_channel_group_int4(self): self._test_qd8_per_token_weight_per_channel_group_int4(dtype=torch.float) + # Tests for q[dp]8-bf16-qb4w + @unittest.skipIf( + not torchao_installed, "Per Channel Group Quantization Required TorchAO" + ) + @unittest.skipIf( + is_fbcode(), + "wait for XNNPACK pin bump to enable", + ) + def test_linear_qd8_bf16_per_token_weight_per_channel_group_int4(self): + self._test_qd8_per_token_weight_per_channel_group_int4(dtype=torch.bfloat16) + @unittest.skipIf( not torchao_installed, "Per Channel Group Quantization Required TorchAO" )