Slow conversion on transformer like model
- 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
Assessment
This issue has not been assessed yet.