Suggestions to train on low-end GPU

[ravioliasb]

I am somewhat new to PyTorch.
Could someone help me to run IJEPA on a low-end Nvidia GPU? It does not matter if it takes months to train. I just would like it to stop running out of memory so quickly.
I have reduced the batch size and number of workers to 2.

[spencermyoung513]

How much GPU VRAM do you have available? If you run nvidia-smi in your terminal, you should get a view showing the size of each GPU available to your system.

You should probably look into accumulating gradients. That will take some slight tinkering, but essentially what it does is wait to perform the backward pass until multiple forward passes have been executed. This gives you a larger effective batch size while still only requiring a few data samples to be loaded to your GPU at a time. Here's a helpful example (for an LLM): https://lightning.ai/blog/gradient-accumulation/.

I'd also recommend using an optimizer like SGD, since it has a lower memory overhead than Adam / AdamW. That will save a significant portion of memory.

[drywitte]

I have a couple 4090s and was trying to get FSDP to work, but ran into a lot of issues. In part it was because the momentum calculations of the target_encoder, which happen outside of the model's modules and thus don't get sharded correctly.

[FalsoMoralista]

I am somewhat new to PyTorch. Could someone help me to run IJEPA on a low-end Nvidia GPU? It does not matter if it takes months to train. I just would like it to stop running out of memory so quickly. I have reduced the batch size and number of workers to 2.

What is the role of number of workers, shouldn't it correspond to the number of cpu cores?

[lazarosgogos]

No, number of workers is the number of processes that simultaneously put data into RAM (as per source provided below). A CPU can handle many processes, independent of CPU cores (as any normal system nowadays can).
More info here

[FalsoMoralista]

No, number of workers is the number of processes that simultaneously put data into RAM (as per source provided below). A CPU can handle many processes, independent of CPU cores (as any normal system nowadays can). More info here

Excellent, thanks