apple / apple/coremltools

Slow conversion on transformer like model

Open
#1,900 6 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug
I found conversion time grows rapidly on transformer like model when layers/attention blocks increase.
But after adding an intermediate ops between blocks, time complexity is fixed to `O(blocks count)`.

```python
# workaround to speed up conversion
# (1, 1500, 768) -> (1, 1501, 768) -> (1, 1500, 768)
x = torch.cat([x, torch.zeros(1, 1, self.n_state)], dim=1).split(1500, dim=1)[0]
```

| attention blocks count | conversion time | conversion time with workaround|
| ------------- | ------------- | ------------- |
| n_layer = 1 | 4s | 2.7s |
| n_layer = 2 | 15s | 5.6s |
| n_layer = 3 | 36s | 8.0s |
| n_layer = 4 | 61s | 10.7s |

## To Reproduce

```Python
# minimal example, simplified AudioEncoder from openai/whisper
import torch
from torch import Tensor, nn

import numpy as np
import coremltools as ct
from timeit import default_timer as timer

class MultiHeadAttention(nn.Module):
def __init__(self, n_state: int):
super().__init__()
self.query = nn.Linear(n_state, n_state)
self.key = nn.Linear(n_state, n_state)
self.value = nn.Linear(n_state, n_state)
self.out = nn.Linear(n_state, n_state)

def forward(self, x: Tensor):
q = self.query(x)
k = self.key(x)
v = self.value(x)

q = q.view(1, 1500, 12, 64).permute(0, 2, 1, 3)
k = k.view(1, 1500, 12, 64).permute(0, 2, 3, 1)
v = v.view(1, 1500, 12, 64).permute(0, 2, 1, 3)

wv = (q @ k @ v).permute(0, 2, 1, 3).flatten(start_dim=2)
return self.out(wv)

class ResidualAttentionBlock(nn.Module):
def __init__(self, n_state: int):
super().__init__()

self.attn = MultiHeadAttention(n_state)
self.attn_ln = nn.LayerNorm(n_state)
self.n_state = n_state

def forward(self, x: Tensor):
# workaround to speed up converting
# (1, 1500, 768) -> (1, 1501, 768) -> (1, 1500, 768)
# x = torch.cat([x, torch.zeros(1, 1, self.n_state)], dim=1).split(1500, dim=1)[0]
return x + self.attn(self.attn_ln(x))

class AudioEncoder(nn.Module):
def __init__(self, n_state: int, n_layer: int):
super().__init__()
self.blocks = nn.ModuleList(
[ResidualAttentionBlock(n_state) for _ in range(n_layer)]
)

def forward(self, x: Tensor):
for block in self.blocks:
x = block(x)
return x

for n_layer in range(1, 5):
encoder = AudioEncoder(n_state=768, n_layer=n_layer)
encoder.eval()

x = torch.ones((1, 1500, 768))
traced_encoder = torch.jit.trace(encoder, x)

startT = timer()
mlmodel = ct.convert(
traced_encoder,
convert_to="mlprogram",
inputs=[ct.TensorType(name="x", shape=x.shape)],
outputs=[ct.TensorType(name="output")],
compute_units=ct.ComputeUnit.ALL,
)
print("----")
print(f"n_layer={n_layer}, conversion took {timer()-startT:.2f}s")
print("----\n")
```

## System environment:
- coremltools 7.0b1
- PyTorch 2.0.1
- macOS Ventura 13.4.1

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.