diff --git a/examples/models/phonon/README.md b/examples/models/phonon/README.md index 7f471832a30..35d4406422b 100644 --- a/examples/models/phonon/README.md +++ b/examples/models/phonon/README.md @@ -57,9 +57,10 @@ 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`. The export also -rewrites the Mimi decoder's transposed convolutions as convolutions, which -XNNPACK runs faster. +as the ExecuTorch custom op `llama::sdpa_with_kv_cache`. + +For XNNPACK, Vulkan, MLX and CUDA, the export rewrites the Mimi decoder's +transposed convolutions as convolutions, which those backends run 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 2bae25d8f47..e66c2849ce2 100644 --- a/examples/models/phonon/export_phonon.py +++ b/examples/models/phonon/export_phonon.py @@ -81,14 +81,12 @@ def source_transforms_for(backend: str, webgpu_compat_ops: bool): if backend != "webgpu": raise ValueError("--webgpu-compat-ops only applies to --backend webgpu") return [st.use_webgpu_compat_ops] + mimi_convs = [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d] return { - "xnnpack": [st.depthwise_conv_transpose_as_mul, st.conv_transpose1d_as_conv1d], - "mlx": [st.decompose_elu], - "vulkan": [ - st.decompose_elu, - st.depthwise_conv_transpose_as_mul, - st.conv_transpose1d_as_2d, - ], + "xnnpack": mimi_convs, + "cuda": mimi_convs, + "mlx": [st.decompose_elu, *mimi_convs], + "vulkan": [st.decompose_elu, *mimi_convs], }.get(backend, []) diff --git a/examples/models/phonon/source_transformations.py b/examples/models/phonon/source_transformations.py index 4c6a1dcdb9b..28e73b005c8 100644 --- a/examples/models/phonon/source_transformations.py +++ b/examples/models/phonon/source_transformations.py @@ -13,16 +13,19 @@ 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_2d Older WebGPU runtimes have no 1-D transposed + convolution (use_webgpu_compat_ops). 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. + Runs the Mimi decoder's transposed convolutions + (kernel = 2 * stride) as a convolution and an + interleave, which XNNPACK, Vulkan, MLX and CUDA run + faster than their transposed convolution. 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. XNNPACK also runs it - faster this way. + multiply and an overlap-add. XNNPACK, MLX and CUDA + use it for the upsampler too, which + conv_transpose1d_as_conv1d skips. 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 @@ -171,7 +174,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: y = x.unsqueeze(3) * self.weight # [b, c, n, 2s] head = y[..., :s].reshape(b, c, n * s) tail = y[..., s:].reshape(b, c, n * s) - zeros = torch.zeros(b, c, s, dtype=x.dtype) + zeros = x.new_zeros(b, c, s) return torch.cat([head, zeros], 2) + torch.cat([zeros, tail], 2) @@ -262,6 +265,8 @@ def decompose_dynamic_layer_norm(programs: dict[str, nn.Module]) -> int: nn.ConvTranspose1d, nn.Embedding, ConvTranspose1dAs2d, + ConvTranspose1dAsConv1d, + DepthwiseConvTranspose1dAsMul, ) diff --git a/examples/models/phonon/tests/test_phonon.py b/examples/models/phonon/tests/test_phonon.py index 63b79dee939..a551b871ef6 100644 --- a/examples/models/phonon/tests/test_phonon.py +++ b/examples/models/phonon/tests/test_phonon.py @@ -32,6 +32,7 @@ Phonon, PhononConfig, program_metadata, + StreamingConvTranspose1d, ) from executorch.examples.models.phonon.run_phonon import ( EagerPrograms, @@ -43,6 +44,7 @@ ) from executorch.examples.models.phonon.source_transformations import ( compute_weights_in, + ComputeIn, ConvTranspose1dAsConv1d, use_webgpu_compat_ops, ) @@ -213,6 +215,8 @@ def test_mimi_matches_reference(self): def test_weight_names_are_unique_across_programs(self): # The CUDA backend shares a weight between methods when the names match. programs = build_programs(self.model, "sdpa") + for transform in source_transforms_for("cuda", False): + transform(programs) compute_weights_in(programs, torch.bfloat16) owner = {} for name, program in programs.items(): @@ -258,6 +262,13 @@ def test_xnnpack_swaps_preserve_outputs(self): # The depthwise upsampler, 2 decoder transposed convs. self._check_swaps("xnnpack", [1, 2]) + def test_cuda_swaps_preserve_outputs(self): + self._check_swaps("cuda", [1, 2]) + + def test_mlx_swaps_preserve_outputs(self): + # 7 ELUs, the depthwise upsampler, 2 decoder transposed convs. + self._check_swaps("mlx", [7, 1, 2]) + def test_conv_transpose_as_conv_matches(self): torch.manual_seed(0) for stride, steps, bias in ( @@ -275,13 +286,23 @@ def test_conv_transpose_as_conv_matches(self): self.assertLess(max_diff(got, convtr(x)), 1e-5, stride) def test_compute_weights_in_keeps_float32_io(self): + # --dtype bf16 is for CUDA and MLX, which apply the Mimi swaps first. inputs = example_inputs(TINY) programs = build_programs(self.model) with torch.no_grad(): before = programs["mimi_decode"](*inputs["mimi_decode"]) + for transform in source_transforms_for("cuda", False): + transform(programs) self.assertGreater(compute_weights_in(programs, torch.bfloat16), 0) + mimi = programs["mimi_decode"] + convtrs = [ + m.convtr for m in mimi.modules() if isinstance(m, StreamingConvTranspose1d) + ] + self.assertEqual(len(convtrs), 3) + for convtr in convtrs: + self.assertIsInstance(convtr, ComputeIn) with torch.no_grad(): - after = programs["mimi_decode"](*inputs["mimi_decode"]) + after = mimi(*inputs["mimi_decode"]) for a, b in zip(before, after): self.assertEqual(b.dtype, torch.float32) self.assertLess(max_diff(a, b), 0.05 * (a.abs().max().item() + 1e-6))