MPS support for doggettx-optimizations

[Any-Winter-4079]

Okay, so I've seen @lstein has added
x = x.contiguous() if x.device.type == 'mps' else x
to ldm/modules/attention.py in the doggettx-optimizations branch
but there's another error happening how
KeyError: 'active_bytes.all.current'
and this has to do with this function in attention.py

def forward(self, x, context=None, mask=None):
        h = self.heads

        q_in = self.to_q(x)
        context = default(context, x)
        k_in = self.to_k(context)
        v_in = self.to_v(context)
        del context, x

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q_in, k_in, v_in))
        del q_in, k_in, v_in

        r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device)

        stats = torch.cuda.memory_stats(q.device)
        mem_active = stats['active_bytes.all.current']
        mem_reserved = stats['reserved_bytes.all.current']
        mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
        mem_free_torch = mem_reserved - mem_active
        mem_free_total = mem_free_cuda + mem_free_torch

        gb = 1024 ** 3
        tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * 4
        mem_required = tensor_size * 2.5
        steps = 1

        if mem_required > mem_free_total:
            steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
            # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
            #       f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")

        if steps > 64:
            max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
            raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
                               f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')

        slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
        for i in range(0, q.shape[1], slice_size):
            end = i + slice_size
            s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * self.scale

            s2 = s1.softmax(dim=-1)
            del s1

            r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
            del s2

        del q, k, v

        r2 = rearrange(r1, '(b h) n d -> b n (h d)', h=h)
        del r1

        return self.to_out(r2)

Which is basically the code that detects your free memory, and then splits the softmax operation in steps, to allow to generate larger images.

Now, because we are on Mac, I'm not sure @lstein can help us much (unless he has one around), but I open this issue for anyone that wants to collaborate in porting this functionality to M1

[lstein]

If someone knows how to get free VRAM memory on MPS devices, we just need to replace the torch.cuda calls.

[lstein]

I Googled around, and there doesn't seem to be an equivalent set of memory interrogation calls for CPU.

I'm not sure how the M1 works, but if it is sharing main memory (i.e. RAM) you might be able to get the needed metrics using psutil

[Vargol]

I just hacked it all out in my fork and set slice_size to 1 :-) that gets me doing 1024x1024 (very slowly) on a 8Gb M1 mini.
Be interesting to see the results of that on a larger GPU.

[Any-Winter-4079]

I just hacked it all out in my fork and set slice_size to 1 :-) that gets me doing 1024x1024 (very slowly) on a 8Gb M1 mini.
Be interesting to see the results of that on a larger GPU.

It definitely works. I'll add the results below.
In this v1,
I've changed ldm/modules/diffusionmodules/model.py and ldm/modules/attention.py from Doggettx-optimization branch.
attention.py.zip
model.py.zip

In model.py I've commented out

# stats = torch.cuda.memory_stats(q.device)
# mem_active = stats['active_bytes.all.current']
# mem_reserved = stats['reserved_bytes.all.current']
# mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
# mem_free_torch = mem_reserved - mem_active
# mem_free_total = mem_free_cuda + mem_free_torch

and left steps at 1

tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * 4
mem_required = tensor_size * 2.5
steps = 1

And commented out

# if mem_required > mem_free_total:
#     steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))

so there's probably an improvement to be made using psutil here.

Where I did use psutil is in attention.py
Again, convention this out

# stats = torch.cuda.memory_stats(q.device)
# mem_active = stats['active_bytes.all.current']
# mem_reserved = stats['reserved_bytes.all.current']
# mem_free_cuda, _ = torch.cuda.mem_get_info(torch.cuda.current_device())
# mem_free_torch = mem_reserved - mem_active
# mem_free_total = mem_free_cuda + mem_free_torch

but importing psutil
import psutil
and using
mem_free_total = psutil.virtual_memory().available
So we can use it to calculate steps (the same way they do), instead of leaving it at steps=1

gb = 1024 ** 3
tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * 4
mem_required = tensor_size * 2.5
steps = 1

if mem_required > mem_free_total:
        steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))

[Any-Winter-4079]

Definitely a step in the right direction

[lstein]

This is looking pretty encouraging. When you are satisfied with the performance on MPS, could you make your changes conditional on the device type so that CUDA systems will work as well? Then make a PR against the doggettx-optimizations branch.

Think this might be done by tonight? I'm planning a development freeze, some testing, and then pulling into main over the weekend.

[Any-Winter-4079]

That's the plan, yes, to make the changes conditional based on device type.
I'm not sure about the PR tonight (b/c I've never done a pull request and it's already 00:00h here, so I might be a bit tired to look into how to make one tonight -fork, pull...), but I'll leave my code here in #431 in 1-1h30' max.

[Any-Winter-4079]

Okay, changes are done. I'm doing the testing.

[Any-Winter-4079]

attention 3.py.zip
model 2.py.zip

[Any-Winter-4079]

@lstein Above are the files.
ldm/modules/diffusionmodules/model.py and ldm/modules/attention.py

Performance seems comparable to CompViz vanilla with some M1 workarounds dbc8fc79008795875eb22ebf0c57927061af86bc (lstein fork) which is the best performance I've seen on M1.

Regarding memory, I have to do more digging because while this afternoon I could generate 896x896 and 1024x768 (results I couldn't generate before), now at night I'm back to memory errors.

In any case, this change should benefit CUDA users while allowing MPS devices to (apparently/presumably/hopefully) function at least as well as we currently do on the development branch

[Any-Winter-4079]

In the end, still awake :)
And I've found something very interesting for version 2 of these changes (post merge into development)
https://pullanswer.com/questions/mps-mpsndarray-error-product-of-dimension-sizes-2-31

The common error I think all M1 users get of Error: product of dimension sizes > 2**31' is referenced here

which made me think about tinkering with values, and setting steps = 64 (max value), it generates a 1024x1024 on my M1!
However, if we do it like this (currently the way in Doggettx-optimizations branch) steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2))) it fails

It takes a long time with steps=64, but testing around, it also works with steps=32, and even steps=4 (taking much less time).

Pretty nice, and calls for some testing tomorrow

PS: I'd just merge the 2 files above and leave this "finding" for a future PR.

[heurihermilab]

I can confirm that it's working and an improvement. Speedup was 2x over plain development branch, and now I'm testing larger image sizes...

Environment: Development branch with the two files above swapped in. Machine: MBP 14", M1 Pro, 16GB, latest OS, running miniforge with a base of Python 3.10.6. Browser: Firefox 104.0.2.

[Vargol]

which made me think about tinkering with values, and setting steps = 64 (max value), it generates a 1024x1024 on my M1! However, if we do it like this (currently the way in Doggettx-optimizations branch) steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2))) it fails

That's odd I've managed a very slow 1024x1024 from doggettx's optimaztions on my 8Gb M1
Here's the code I'm using, as I'm said before it was a cut of lstein's paint but with the Doggettx
code added and hard coded to 1 step.

https://github.com/Vargol/stable-diffusion_m1_8gb

Have you got an lot of other stuff running at the same time eating up Memory ?