Is this a starting point to enable processing across multiple GPUs?

[lethargus]

If I understand the changes you've made by splitting up the work... How hard would it be to send segments to different GPUs at the same time, rather than serially feeding a single GPU? Thinking of those of us with used K80s or SLI-like configurations.

[Doggettx]

Yea technically that should be no problem. Unfortunately I don't have a SLI configuration to test with and I've only been using Python and torch/cuda for a few weeks so my knowledge isn't high enough yet to make such a change blindly without testing a bunch of stuff to see how it works.

But theoretically you should just be able to run the for loop in parallel since there's no write overlap between the work done

[lethargus]

You mean in the model.py AttnBlock forward def? I'll give it a shot when I get some free time. It looks like you pull the memory free from the current device to calculate the slice sizes so I'll probably just pull the minimum free from all available devices as a starting point. I have a feeling the hard part will be creating a generator for joblib that picks a unique GPU for each process.

Impressive work for only using Python for a few weeks.