Performance degradation in `mps` after `einsum` replacement

[pcuenca]

Before #445 was merged I was getting ~31s inference time in mps. After the change, time goes up to 42s. I verified again in main @ b2b3b1a, and time is again 31s.

I haven't checked other platforms yet.

Any ideas, @patil-suraj?

[abesmon]

I'm not sure if its related, but I found another performance degradation and it seems related to PyTorch version:

previously (1.13.0.dev20220604) inference step took ~ 2s
now (1.13.0.dev20220909) it takes ~7-8s

(found while working with #457)

[patil-suraj]

@pcuenca I tried this on v100, and didn't see any degradation. Also pinging @patrickvonplaten @anton-l

[pcuenca]

@patil-suraj There's no degradation in 3090 either, or in CPU (Linux). So this appears to only affect the mps device.

[Birch-san]

yes, I tried making the same change in CompVis attention.py, on MPS backend:
Birch-san/stable-diffusion@d2d533d
matmul does indeed seem to be slower.

inference speed dropped from ~1.25s/it to ~1.82s/it.
8 step sample regressed from ~10.7s to ~14.5s.

also, a different picture is generated. despite same seed each time.

einsum:

subsequent runs of einsum() produce the same image, so it is deterministic.

matmul:

subsequent runs of matmul() produce the same image, so it is deterministic.

Using a semi-recent pytorch nightly, 1.13.0.dev20220826.

[patil-suraj]

also, a different picture is generated. despite same seed each time.

how are you using the seeds ? diffusers pipelines uses the torch.Generator objects for seeds. To get reproducible results we need to reinit the torch.Generator with the same seed as using the same generator multiple times advances the rng state.

The correct way to check this would be running this same block multiple times.

with autocast("cuda"):
    images = pipe(prompt, generator=torch.Generator(device="cuda").manual_seed(1024)).images
images[0]

[patrickvonplaten]

Just to be sure, performance degradation is only noticeable for MPS here no?

[patrickvonplaten]

Can someone from PyTorch or Apple help us here a bit? I'm not sure we have the know-how to dive into core-operations of PyTorch to figure out what is faster. Are there some common gotcha's on how to use PyTorch for MPS?

[pcuenca]

Just to be sure, performance degradation is only noticeable for MPS here no?

Yes.

[Birch-san]

how are you using the seeds?

pretty much the same way as CompVis txt2img works out-of-the-box:

https://github.com/Birch-san/stable-diffusion/blob/18bb5f8fad9148b162b61a4b94d9cadba3cae721/scripts/txt2img_fork.py#L711

I run --n_iter 1 --n_samples 1 --seed 68673924. the program terminates. I run it again with the same parameters. it produces the same image each time.
the problem is that changing einsum() to matmul() produces a different image. the new image is reproduced consistently each time I restart the program.

[pcuenca]

I was trying to isolate a a reproducible test case to report in a PyTorch issue but couldn't, the following two functions behave similarly.

def test_einsum(dim, device):
    q = torch.rand((dim, 4096, 40)).to(device)
    k = torch.rand((dim, 4096, 40)).to(device)
    v = torch.rand((dim, 4096, 40)).to(device)

    x = torch.einsum('b i d, b j d -> b i j', q, k)
    x = x.softmax(dim=-1)
    x = torch.einsum("b i j, b j d -> b i d", x, v)

def test_matmul(dim, device):
    q = torch.rand((dim, 4096, 40)).to(device)
    k = torch.rand((dim, 4096, 40)).to(device)
    v = torch.rand((dim, 4096, 40)).to(device)

    x = torch.matmul(q, k.transpose(1, 2))
    x = x.softmax(dim=-1)
    x = torch.matmul(x, v)

I don't know if there's something else we are doing that is preventing the mps graphs from being cached.

[patrickvonplaten]

@pcuenca - let's ask the PyTorch team here

[Birch-san]

here's why my einsum() returned a different image than matmul():
pytorch/pytorch#85224

it's because the first invocation of einsum() returns an incorrect result, the first time you invoke it. on nightly builds of pytorch, but not on stable.
the matmul seems to be the more correct output. I expect subsequent samples to produce the same result as matmul().

if this seems incongruous with my assurance that for each strategy I produced multiple samples and each strategy was self-consistent on each run: it's because I rebooted Python between samples. so einsum() would always be afflicted by the "first sample is wrong" problem, but still produce the same (wrong) image on the next run.