diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index 38d2ee193..272a045dd 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -1775,6 +1775,23 @@ def _parse_positional_args(context, node) -> Tuple[Var]: context.add(res) +def _stable_softplus_mil(x): + """Numerically stable softplus: max(x, 0) + log(1 + exp(-|x|)). + + Since -|x| <= 0, exp(-|x|) is always in (0, 1], so no overflow occurs + in any precision. This matches the formula used by coremltools' own + softplus MIL op ``value_inference``. + + See issue #2687 for the original fp16 overflow report. + """ + abs_x = mb.abs(x=x) + neg_abs_x = mb.mul(x=-1.0, y=abs_x) + exp_val = mb.exp(x=neg_abs_x) + log_val = mb.log(x=mb.add(x=1.0, y=exp_val)) + max_val = mb.maximum(x=x, y=0.0) + return mb.add(x=max_val, y=log_val) + + @register_torch_op def softplus(context, node): def _parse_positional_args(context, node) -> Tuple[Var]: @@ -1795,20 +1812,17 @@ def _parse_positional_args(context, node) -> Tuple[Var]: x, beta, threshold = _parse_positional_args(context, node) if beta == 1: - # this is the special case that Core ML softplus handles - res = mb.softplus(x=x, name=node.name) + sp = _stable_softplus_mil(x) + cond = mb.greater(x=x, y=threshold) else: - if x.rank == 4: - # can use Core ML softplus_parametric - C = x.shape[1] - alpha_br = np.repeat(1.0 / beta, C).astype("float32") - beta_br = np.repeat(beta, C).astype("float32") - res = mb.softplus_parametric(x=x, alpha=alpha_br, beta=beta_br, name=node.name) - else: - # have to generally decompose - beta_mul_x = mb.mul(x=beta, y=x) - softplus = mb.softplus(x=beta_mul_x) - res = mb.real_div(x=softplus, y=beta, name=node.name) + # For non-unit beta: softplus(x) = (1/beta) * softplus(beta * x) + beta_x = mb.mul(x=beta, y=x) + sp = mb.real_div(x=_stable_softplus_mil(beta_x), y=beta) + cond = mb.greater(x=beta_x, y=threshold) + + # Apply PyTorch's threshold: for beta * x > threshold, softplus(x) ≈ x, + # so return x directly. This matches PyTorch's exact behavior. + res = mb.select(cond=cond, a=x, b=sp, name=node.name) context.add(res) @@ -1817,8 +1831,8 @@ def mish(context, node): inputs = _get_inputs(context, node, expected=1) x = inputs[0] - softplus = mb.softplus(x=x) - tanh = mb.tanh(x=softplus) + # Mish(x) = x * tanh(softplus(x)) + tanh = mb.tanh(x=_stable_softplus_mil(x)) res = mb.mul(x=x, y=tanh, name=node.name) context.add(res) diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index dedb793f4..afc32e35b 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -6721,17 +6721,9 @@ def test_softplus( # executorch decomposes softplus to very basic log and exp target_op = "exp" else: - if beta is None or beta == 1: - # this is the special case that Core ML softplus handles - target_op = "softplus" - else: - if rank == 4: - # can use Core ML softplus_parametric - target_op = "softplus_parametric" - else: - # have to generally decompose to - # `x -> beta * x -> softplus(beta * x) -> softplus(beta * x) / beta` - target_op = "softplus" + # The converter now wraps softplus in a select for fp16 threshold safety, + # so we skip the target_op check for non-executorch frontends. + target_op = None self.run_compare_torch( input_shape, @@ -6743,6 +6735,48 @@ def test_softplus( target_op=target_op, ) + @pytest.mark.parametrize("frontend", frontends) + def test_softplus_fp16_stable_decomposition(self, frontend): + """Verify softplus is decomposed into overflow-safe ops for fp16 (#2687).""" + if frontend == TorchFrontend.EXECUTORCH: + pytest.skip("ExecuTorch decomposes softplus independently") + + # Demonstrate that naive fp16 softplus (log(1+exp(x))) overflows for + # large x, while the stable form (max(x,0)+log(1+exp(-|x|))) does not. + x_val = np.float16(15.0) + naive = np.float16(np.log(np.float16(1.0) + np.exp(x_val))) + assert not np.isfinite(naive), ( + f"Naive fp16 softplus should overflow at x=15, got {naive}" + ) + stable = np.float16( + np.maximum(x_val, np.float16(0)) + + np.log(np.float16(1.0) + np.exp(-np.abs(x_val))) + ) + assert np.isfinite(stable) and abs(float(stable) - 15.0) < 0.5, ( + f"Stable fp16 softplus should produce ~15.0, got {stable}" + ) + + # Verify the converter emits the stable decomposition, not the native op. + torch.manual_seed(0) + model = nn.Softplus().eval() + x = torch.randn(1, 10) + + torch_model = export_torch_model_to_frontend(model, (x,), frontend) + mlmodel = ct.convert( + torch_model, + inputs=[ct.TensorType(shape=x.shape)], + convert_to="mlprogram", + compute_precision=ct.precision.FLOAT16, + ) + + prog = mlmodel._mil_program + assert len(prog.find_ops(op_type="softplus")) == 0, ( + "Native softplus op should be replaced by stable decomposition" + ) + assert len(prog.find_ops(op_type="exp")) >= 1, ( + "Stable decomposition should contain at least one exp op" + ) + @pytest.mark.parametrize( "compute_unit, backend, frontend, shape", itertools.product(compute_units, backends, frontends, COMMON_SHAPES_ALL),