I found this post which indicated that the problem could be that WSL doesn't allow you to install modules.
Hey, that's my answer! Keep in mind:
So ... don't trust it as being anywhere near reliable! Feel free to downvote it even! It's tough to say that about my own answer, but in retrospect it probably wasn't accurate. I'll add a link from it to this answer, however, so that it redirects to the latest info. Whether the info below is now correct or not, I'll need you to help me with!
I will be using PyTorch for a class and I want to use my GPU for compute.
Okay, that's the question I'm going to address, rather than how to build the kernel with ROCm support (since I no longer think that will help).
At this point, I still don't have an AMD GPU to be able to test this out on directly, but I have done some GPU compute work on my nVideo 2070 under WSL2 now, at least.
So here's my latest thinking, looking at that past question, your question, and some other information:
ROCm likely isn't needed under WSL2. WSL2 makes the GPU available to Linux through a passthrough system using the Windows (not Linux) driver. If you look under /usr/lib/wsl
, you'll find two directories which are mounted into your WSL2 instance -- ../drivers
and ../lib
. They are added to the library cache by the (also injected in), /etc/ld.so.conf/ld.wsl.conf
.
It's these libraries and drivers which give you access to the GPU.
When I looked at the other question last year, I didn't realize that ROCm was the Linux equivalent library for GPU compute. And if it does install as a module, then it's not likely to be useful under WSL2 anyway, as it won't have the direct access to the GPU that it needs. In retrospect, I should have been able to figure that out with what I knew then.
So what do you need to use Pytorch under WSL2 on your AMD GPU? Again, taking a guess here, but ...
This doc seems to indicate that you need Windows 11. While I thought that most of the GPU compute support was available in Windows 10, this may be one area that requires Windows 11. You don't mention your Windows version, so I'm not sure if this will be possible for you or not.
The latest AMD drivers for Windows. Most anything in the last year or so should be sufficient.
An appropriate Python environment. Microsoft recommends Conda. I ended up using Docker myself for TensorFlow, at least.
Dependencies are libblas3 libomp5 liblapack3
And then pip install pytorch-directml
Honestly, that should be about it. According to this doc, AMD is supported for PyTorch with DirectML.
When running TensorFlow on my nVidia under WSL2, I didn't require any additional native drivers/modules. I did have to provide the /usr/lib/wsl/lib
directory to Docker, but that's only because I needed to run a specific older Python version for TensorFlow, and the easiest way to do that was via Docker.