tenstorrent/tt-buda

Huge memory requirement compiling a rather small UNet

urbste opened this issue · 4 comments

This issue is related to a discussion on discord with Milan Kordic.

I have been trying to compile a Unet with different input shapes: 256x256, 512x512 and 1024x1024.
When compiling the 512x512 model, the compiler uses > 64GB (I even had to increase the system SWAP to 16GB).
For the 1024x1024 model he wants to allocate 1024GB of Ram (which is obiously not feasible).

Milan mentioned to open an issue so it can be tracked.

If you need my code or my model please don't hestiate to ask :)

Can you please try to set environment flag PYBUDA_CONVERT_PARAMS_TO_TVM=0 and see if this helps?

I tested the UNet with this call and input shape 1024x1024:
PYBUDA_CONVERT_PARAMS_TO_TVM=0 python test_unet.py
And I verified in my python script that the env variable ist set:
print("PYBUDA_CONVERT_PARAMS_TO_TVM: ",os.environ["PYBUDA_CONVERT_PARAMS_TO_TVM"])

The resource utilization is still huge:
image

And then at some point it crashes:
image

Thanks for the follow-up let me have the team take another look and see.

We spoke as a team about this, for our UNET testing we used an input size of 256x256 and never tested input shapes this large.

High Res images is something we want to work towards but we haven't begun this development yet. As we continue down our roadmap we will keep this in mind.

To give you some additional details our memory consumption on 256x256 is around 6.7 gb. If we assume for simplicity sake we have linear scaling we would end up with 6.7 x 16 giving us a memory requirement of around 108gb of usage.