diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index be74338d754..7f471832a30 100644 --- a/examples/models/phonon/README.md +++ b/examples/models/phonon/README.md @@ -57,7 +57,9 @@ cmake-out/examples/models/phonon/phonon_runner --model_path=... --text="..." --s ``` XNNPACK delegates the linear and convolution ops; the flow LM's attention runs -as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. +as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. The export also +rewrites the Mimi decoder's transposed convolutions as convolutions, which +XNNPACK runs faster. `--quantize 4bit` or `--quantize 8bit` quantizes the linear layers of the flow LM's transformer and flow net with the scheme of the backend; the codec stays in diff --git a/examples/models/phonon/export_phonon.py b/examples/models/phonon/export_phonon.py index 7784fc65d33..2bae25d8f47 100644 --- a/examples/models/phonon/export_phonon.py +++ b/examples/models/phonon/export_phonon.py @@ -82,6 +82,7 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool): raise ValueError("--webgpu-compat-ops only applies to --backend webgpu") return [st.use_webgpu_compat_ops] return { + "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d], "mlx": [st.decompose_elu], "vulkan": [ st.decompose_elu, diff --git a/examples/models/phonon/source_transformations.py b/examples/models/phonon/source_transformations.py index 69a2c228a54..4c6a1dcdb9b 100644 --- a/examples/models/phonon/source_transformations.py +++ b/examples/models/phonon/source_transformations.py @@ -4,7 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Function-preserving module swaps for backends that lack an op. +"""Function-preserving module swaps for backends that lack an op or run it +slowly. The model in model.py uses plain PyTorch ops. Each swap below computes the same function from more widely supported ops; export_phonon.py selects them @@ -13,10 +14,15 @@ decompose_elu MLX and Vulkan have no aten.elu kernel; without it the Mimi decoder splits at every ELU. conv_transpose1d_as_2d Vulkan runs transposed convolution on 4-D inputs only. + conv_transpose1d_as_conv1d + XNNPACK's deconvolution is slow; this runs the Mimi + decoder's transposed convolutions (kernel = + 2 * stride) as a convolution and an interleave. depthwise_conv_transpose_as_mul Vulkan computes the Mimi upsampler (depthwise, kernel = 2 * stride) incorrectly; this runs it as a - multiply and an overlap-add. + multiply and an overlap-add. XNNPACK also runs it + faster this way. use_webgpu_compat_ops WebGPU runtimes built before these runtime fixes (off by default; --webgpu-compat-ops): * native_layer_norm resize hook (the flow LM's @@ -93,6 +99,36 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return y.squeeze(2) +class ConvTranspose1dAsConv1d(nn.Module): + """ConvTranspose1d with kernel_size == 2 * stride as a kernel-2 Conv1d + with stride * out_channels outputs (one per output phase), interleaved + into the output steps.""" + + def __init__(self, convtr: nn.ConvTranspose1d): + super().__init__() + (self.stride,) = convtr.stride + s = self.stride + w = convtr.weight.detach() # [cin, cout, 2s] + cin, cout, _ = w.shape + # Phase j of output step p reads input p - 1 through tap j + s and + # input p through tap j. + taps = torch.stack([w[..., s:], w[..., :s]], -1) # [cin, cout, s, 2] + self.weight = nn.Parameter( + taps.permute(1, 2, 0, 3).reshape(cout * s, cin, 2).contiguous(), + requires_grad=False, + ) + self.bias = convtr.bias + self.out_channels = cout + + def forward(self, x: torch.Tensor) -> torch.Tensor: + z = F.conv1d(x, self.weight, padding=1) # [b, cout * s, n + 1] + b, _, n = z.shape + s = self.stride + y = z.view(b, self.out_channels, s, n).transpose(2, 3) + y = y.reshape(b, self.out_channels, n * s) + return y if self.bias is None else y + self.bias.view(1, -1, 1) + + class ComputeIn(nn.Module): """Runs `module` with its weights in `dtype` on a float32 input and returns float32. Other attributes (bias, kernel_size, ...) are read from @@ -171,6 +207,23 @@ def conv_transpose1d_as_2d(programs: dict[str, nn.Module]) -> int: return count +def conv_transpose1d_as_conv1d(programs: dict[str, nn.Module]) -> int: + count = 0 + for program in programs.values(): + for module in program.modules(): + if not isinstance(module, StreamingConvTranspose1d): + continue + convtr = module.convtr + if ( + isinstance(convtr, nn.ConvTranspose1d) + and convtr.groups == 1 + and convtr.kernel_size[0] == 2 * convtr.stride[0] + ): + module.convtr = ConvTranspose1dAsConv1d(convtr) + count += 1 + return count + + def depthwise_conv_transpose_as_mul(programs: dict[str, nn.Module]) -> int: count = 0 for program in programs.values(): diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index ce952bfec70..63b79dee939 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/examples/models/phonon/tests/test_phonon.py @@ -43,6 +43,7 @@ ) from executorch.examples.models.phonon.source_transformations import ( compute_weights_in, + ConvTranspose1dAsConv1d, use_webgpu_compat_ops, ) from executorch.examples.models.phonon.tests.reference import ( @@ -236,19 +237,43 @@ def test_webgpu_compat_ops_preserve_outputs(self): for a, b in zip(outputs, programs[name](*inputs[name])): self.assertLess(max_diff(a, b), 1e-4) - def test_vulkan_swaps_preserve_outputs(self): + def _check_swaps(self, backend: str, counts: list[int]): inputs = example_inputs(TINY) programs = build_programs(self.model) with torch.no_grad(): before = programs["mimi_decode"](*inputs["mimi_decode"]) - # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. - counts = [t(programs) for t in source_transforms_for("vulkan", False)] - self.assertEqual(counts, [7, 1, 2]) + self.assertEqual( + [t(programs) for t in source_transforms_for(backend, False)], counts + ) with torch.no_grad(): after = programs["mimi_decode"](*inputs["mimi_decode"]) for a, b in zip(before, after): self.assertLess(max_diff(a, b), 1e-4) + def test_vulkan_swaps_preserve_outputs(self): + # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("vulkan", [7, 1, 2]) + + def test_xnnpack_swaps_preserve_outputs(self): + # The depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("xnnpack", [1, 2]) + + def test_conv_transpose_as_conv_matches(self): + torch.manual_seed(0) + for stride, steps, bias in ( + (1, 1, True), + (4, 16, False), + (5, 1, True), + (6, 16, True), + ): + convtr = torch.nn.ConvTranspose1d( + 6, 3, 2 * stride, stride=stride, bias=bias + ) + x = torch.randn(1, 6, steps) + with torch.no_grad(): + got = ConvTranspose1dAsConv1d(convtr)(x) + self.assertLess(max_diff(got, convtr(x)), 1e-5, stride) + def test_compute_weights_in_keeps_float32_io(self): inputs = example_inputs(TINY) programs = build_programs(self.model)