HEAD requires a lot of memory causing OOM

1

Hi,

I have problems with OOM using ray and the dashboard shows that Head requires large amounts of memory. What does need Head so much memory for? object_store_memory seems to be unrelated. Instead of sharing my complex code, I share a minimum working example from ray (MNIST), that does not run out of memory, but still requires a lot of memory.

I am using ray 2.8.1 and the MNIST example from: Train a PyTorch model on Fashion MNIST — Ray 2.43.0

image
image1920×1058 94.5 KB

How can I reduce the memory requirements by Head? I get the problems on different systems (Linux, Mac) with and without GPUs.

Thank you very much in advance.