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 ?

[Any-Winter-4079]

This is my memory usage right after booting the computer. Only the model loaded + 1 VS tab with the code

So, I introduced 2 prints

print('mem_required', mem_required / 10**9)
print('mem_free_total', mem_free_total / 10**9)

and I get

mem_required 25.17630976
mem_free_total 51.890225152

The mem_free_total is calculated with mem_free_total = psutil.virtual_memory().available and it makes sense with the picture above from Activity Monitor. (64 - 13) GB = 51 GB
However, it crashes with Error: product of dimension sizes > 2**31 even though in theory, it only needs 25GB. Which doesn't make sense, does it?

In this discussion https://pullanswer.com/questions/mps-mpsndarray-error-product-of-dimension-sizes-2-31 they were saying that the problem was with Metal and that depending on the size/number of dimensions of the operation (e.g. einsum), a different algorithm might get selected.

So maybe you give it a smaller array and it fails but feed it a bigger array and chooses a different algorithm that doesn't have the Error: product of dimension sizes > 2**31 bug and it works. That's my understanding.

[Any-Winter-4079]

For example, setting the steps as you do (with fixed value instead of calculating it), and "banana sushi" -s50 -C7.5 -n3 -W896 -H896 here are my results:

With steps = 1

mem_required 25.17630976
mem_free_total 53.156200448

Result: Error: product of dimension sizes > 2**31

With steps = 2

mem_required 25.17630976
mem_free_total 52.520927232

It works.

Why, it's really not apparent to me.

---

The problem with hard-setting the steps, though, is that, as the code progresses, mem_free_total reduces (note: leak or expected behavior ?).

mem_required 25.17630976
mem_free_total 32.847593472

So, maybe there could be a point where it failed mid-execution, because we hard-set steps = 2 and it can't do it anymore.

The solution I'm thinking is a mix of both techniques. Setting the steps dynamically (so it doesn't run out of memory), but also setting steps = max (2, steps) -not letting it be steps = 1, where it throws the error- for images larger than 512x512 or something like that.

[Vargol]

The 2**31 seems to be einsum trying to use a Tensor with more than 2,147,483,648 values as part of its calculation
it not a memory thing, just a bug or limitation in the einsum implementation somewhere.

I remember have a similar issue when I simply set steps to 1 but allowed the slice_size calculation to go ahead

slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]

which wasn't in an older cut of the code..

my code doing this , is that what you tried ?

        steps=1

        for i in range(0, q.shape[0], steps):
            end = i + steps 

And yes I appreciate that if people try even bigger images they may run out of memory but for me more steps just means slower renders and 1024x1024 is already 50 S / IT, as in n_sample=50, n_iter=1 takes 40 odd Minutes to generate an image .

[lstein]

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

Looks like I was monitoring the wrong thread! I'll fold in these changes this morning and freeze development for testing. Thanks so much for this.

[Any-Winter-4079]

@Vargol I can take slice_size up to ~10k. More than that, it seems to fail.

Have you tried bigger slice_size? You are basically using 1 as slice_size, which seems very low (even for 8GB)

I tried something very similar to your code, simply with larger slice_size, in my case using this formula they have slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1] which depends on steps.

For example, for "banana sushi" -s5 -C7.5 -n3 -W896 -H896 then q.shape = torch.Size([16, 12544, 40]) so q.shape[1] is 12544 and given that my mem_free_total >= mem_required, it maintains steps = 1, so with the formula above, keeps slice_size = 12544 and fails.

We should be able to find a sweet spot, shouldn't we?

[Any-Winter-4079]

For example,
slice_size = 8192 (or below) allows me to run 896x896
slice_size = 6000 (or below) allows me to run 1024x1024

Hopefully there's some formula we can come up with for all M1 machines (8GB to 128GB)

---

Update:
So for 1024x1024, doggettx-optimizations branch suggests I use slice_size = 8192 (which fails)
But manually hard-setting slice_size = 8185 works.
So their calculation is not far off, but not completely precise for M1

[Vargol]

So with my fixed value steps = slice_size, running
dream> "banana sushi" -s1 -C7.5 -n1 -W832 -H768

1 - 6 steps work, 6 steps are over 5x slower than 1 step

step = slice_size = 1
 1/1 [00:20<00:00, 20.45s/it

step = slice_size = 6
1/1 [01:50<00:00, 110.50s/it]

7 - 10 steps blow memory while sampling,

step = slice_size = 7 - 10
	The Metal Performance Shaders operations encoded on it may not have completed.
	Error: 
	(null)
	Insufficient Memory (00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory)
	<AGXG13GFamilyCommandBuffer: 0x16d1049d0>

steps >= 11 fail with a oversized? buffer before sampling shows up in dreams.py

step = slice_size = 11
RuntimeError: Invalid buffer size: 5.94 GB

[Any-Winter-4079]

@Vargol Hmm I'll study your case too.
To me, it happened the weirdest thing. Doggettx branch suggested 8192 slice size. Guess what? It failed. But 8191 works for 1024x1024

[Any-Winter-4079]

Oh, this is interesting. So my computer can take slice_size = 8191 for 1024x1024. Found by trial and error.

Okay, so what slice_size could It take for 896x896? Well, I did (1024x1024 / 896x896) * 8191 = 10698.4489796, which rounding up is 10699. I tried that value and... it works!

But, I tried 10700 (one more) and it fails!

I'm sure there is a formula to be found (including RAM), but at least we seem to be able to hack the max slice_size for our own devices, which is awesome!

---

Update: So, I picked a random size. I wanted a 3200x1600 image. I used the formula and slice_size = 1677.5168. This time, I could not round up, but rounding down to 1677, it works again! (I only did it for 1 step but hey, it completed successfully)
And 1678 fails.

[Any-Winter-4079]

I'll study it a bit more, but the problem with Doggettx (besides 8192 vs 8191) is that sometimes it suggests an even larger slice_size (I guess when it computes steps = 1 instead of steps = 2 based on memory), and then it breaks. If it weren't for that, maybe we could have used Doggettx's slice_size - 1

[Any-Winter-4079]

@i3oc9i can you try your max slice_size for say 1024x1024 in your Mac with 128GB? We might be able to work out a formula including the RAM. Or someone else with a Mac different than 64 GB (which I have)

[i3oc9i]

@Any-Winter-4079
sorry, this last two days I was busy at my work, I just checkout the development branch 75f633c

and fail, may be I'm missing somethibg ? whre is the new code to test ?

dream> "a don on the moon" -s50 -W1024 -H1024 -C7.5 -Ak_lms
>> This input is larger than your defaults. If you run out of memory, please use a smaller image.
/Users/ivano/Code/Ai/dream-dev/ldm/modules/embedding_manager.py:152: UserWarning: The operator 'aten::nonzero' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1659484612588/work/aten/src/ATen/mps/MPSFallback.mm:11.)
  placeholder_idx = torch.where(
Generating:   0%|                                                                                                                                                                                                                                                                                                                          | 0/1 [00:00<?, ?it/s/AppleInternal/Library/BuildRoots/20d6c351-ee94-11ec-bcaf-7247572f23b4/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:705: failed assertion `[MPSNDArray initWithDevice:descriptor:] Error: product of dimension sizes > 2**31'                                                                  | 0/50 [00:00<?, ?it/s]
zsh: abort      python scripts/dream.py --full_precision --outdir ../@Stuffs/images/samples
/Users/ivano/.miniconda/envs/dream-dev/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

dream> "test" -s50 -W832 -H832 -C7.5 -Ak_lms -S12345678 (run but I get noise)

[ryudrigo]

Could someone with a Mac please run these lines?

import torch
print (torch.cuda.get_device_name(0))

That's the best way I know to detect if it's a Mac GPU, but I couldn't find what to check it against. Thanks!

[Any-Winter-4079]

@ryudrigo

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/eduardoarinopelegrin/opt/anaconda3/envs/do_not_touch-osx-arm64-stable-diffusion/lib/python3.9/site-packages/torch/cuda/__init__.py", line 329, in get_device_name
    return get_device_properties(device).name
  File "/Users/eduardoarinopelegrin/opt/anaconda3/envs/do_not_touch-osx-arm64-stable-diffusion/lib/python3.9/site-packages/torch/cuda/__init__.py", line 359, in get_device_properties
    _lazy_init()  # will define _get_device_properties
  File "/Users/eduardoarinopelegrin/opt/anaconda3/envs/do_not_touch-osx-arm64-stable-diffusion/lib/python3.9/site-packages/torch/cuda/__init__.py", line 211, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

Basically the command torch.cuda is going to fail because we don't have cuda.
You can detect it like this: device_type = 'mps' if x.device.type == 'mps' else 'cuda' Here you have 2 files that check for Mac GPU #431 (comment)

[Doggettx]

I'll study it a bit more, but the problem with Doggettx (besides 8192 vs 8191) is that sometimes it suggests an even larger slice_size (I guess when it computes steps = 1 instead of steps = 2 based on memory), and then it breaks. If it weren't for that, maybe we could have used Doggettx's slice_size - 1

I wouldn't adjust the slice_size because then it starts running incomplete parts of the whole array. It's best to increase the multiplier, which is probably too low then. So this part:

mem_required = tensor_size * 2.5

Probably needs more than .5 extra, could try 2.6 or if you want to be safe just put it at 3, it'll just scale up the steps a bit earlier than needed. Which scales down the slice_size

On a side note, it doesn't really have to step up in powers of 2, I just found that that was faster on average. You could change this part:

    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

To something like

    slice_size = q.shape[1] // steps
    for i in range(0, q.shape[1], slice_size):
        end = min(q.shape[1], i + slice_size)

then it can run at any step or slice_size (even higher than 64, but you'll crash later then anyhow due to other parts running out of memory)

[netsvetaev]

@netsvetaev do you have other stuff running ?

I would have expected a chunk more available, 8 - 9 Gb based on what by 8Gb has available at that point.

Sorry, I had a safari window. Now 6.99gb and 1.30s/it. But 768 is very slow/doesn't work.

[Any-Winter-4079]

As a new threshold, I propose maybe 3GB?
Vargol had a peak of 2.08GB
The other option is to use psutil.virtual_memory().total instead of psutil.virtual_memory().available

[Any-Winter-4079]

@netsvetaev can you run 1024x1024? Does it take 10 minutes? 10 minutes is what the other person with Mac and 16GB reported.
Also, this should produce the same results but I have hopefully fixed the 512x512 issue using mem_total so you get around 1.3s/it
attention 12.py.zip

[netsvetaev]

@netsvetaev can you run 1024x1024? Does it take 10 minutes?

It takes 30 at 35-37s/it. Hm.
Noticed that after I close terminal and open it again, it shows more ram available, +200-500mb.

768 gives an error, then goes for 150s/it.

rror: command buffer exited with error status. | 0/50 [00:00<?, ?it/s] The Metal Performance Shaders operations encoded on it may not have completed. Error: (null) Insufficient Memory (00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory) <AGXG13XFamilyCommandBuffer: 0x288c234a0> label = <none> device = <AGXG13XDevice: 0x13f646400> name = Apple M1 Pro commandQueue = <AGXG13XFamilyCommandQueue: 0x12a34fe00> label = <none> device = <AGXG13XDevice: 0x13f646400> name = Apple M1 Pro retainedReferences = 1

[Any-Winter-4079]

Some advice is, every time you run a dream command and completes, exit with q because at least a few days ago, there were some memory leaks. And then you can re-enter again and have a bit more RAM

[Any-Winter-4079]

In any case, have you been able to run 1024x1024 with better results with some other code?

[netsvetaev]

In any case, have you been able to run 1024x1024 with better results with some other code?

No, my best was around 35-40, no differences with your code. Maybe 1-2s.

[Any-Winter-4079]

You know, it might be because you have a bit less RAM available than the other person with 16GB. I can't find any other reason. Architecture? Pytorch version?

[Any-Winter-4079]

As a final test, do you get better results for large images with slice_size = 2 instead of slice_size = 1 It might break though.
Here

[netsvetaev]

You know, it might be because you have a bit less RAM available than the other person with 16GB. I can't find any other reason. Architecture? Pytorch version?

MacOS 13 beta, I think.

[Any-Winter-4079]

You know, it might be because you have a bit less RAM available than the other person with 16GB. I can't find any other reason. Architecture? Pytorch version?

MacOS 13 beta, I think.

I'm running 12.5.1. No idea if there is any performance improvement/loss with the beta. I'd assume these results are more RAM-dependant than OS version-dependant, but who knows :)

[Any-Winter-4079]

This is the version I'm planning on doing a PR with https://github.com/lstein/stable-diffusion/discussions/457#discussioncomment-3635644 If someone experiences a downgrade in performance vs. before, let me know.

[netsvetaev]

attention 12.py.zip

Getting better:
512 — 1.6 (from 1.8 to 1.2)
768 — 9.8, was 13
1024 — 35.

[Any-Winter-4079]

Oh, so happy to read that!

[i3oc9i]

attention 12.py.zip

Sorry bu I dont see speed difference with this.
I get the same execution time for 512x512 and 896x576

[Any-Winter-4079]

Not for 64-128GB. The update is to give more speed for 16 -32GB Mac.
For us it should be the same speed.
I've done the PR #540

[netsvetaev]

Oh, so happy to read that!

Is it ok that I have better results on a main branch? Seems like it less ram hungry and also faster. 512 at 1.48 (yours is 1.35), 1024 at 29.7 with up to 4.2gb swap.

[netsvetaev]

I'm happy to add that on the latest macos 13 beta 1.14 main is got faster as 1.15s/it on 512px (58s total, always was 1:15-1:25), 8.4it/s on 768 (7:18, was 10-12 mins), and still 35s/it on 1024.

UPD. After a fresh install I've got 1.38s/it 512, 6.20s/it 768 and 20.5s/it 1024. So it were mine problems.

[lstein]

On my HPC node, I also see a remarkable variation in VRAM usage in the doggettx branch as I make small adjustments to image size. Fortunately, it is pretty stable at lower image sizes where most people will be working.

[Any-Winter-4079]

@lstein There was a later update to attention.py in #582
These are the two major PR's I believe.