Multi device bugfix

[asglover]

This is a PR to harden multi device behavior.

• It adds testing to make sure all tensors are on the same device
• it adds usage of "device guards". This is how PyTorch signals to apis which do not explicitly take device, which device to use.
• It also adjust how BLAS handles work. BLAS handles must differ across devices. PyTorch keeps a pool of BLAS handles. This PR switches OEQ to using PyTorch's BLAS handles. Just a note that this current design does not set the BLAS precision mode which means it inherits from pytorch settings. I think that's a reasonable design because people can control precision with the idiomatic pytorch methods. If you think this calculation should never be in TF32, then I can add a manual setting command.
• It adds stream tests to that should fail before the changes from the PR come in on multi device. And should be fixed later.

[asglover]

This is a draft PR for now because I'm making these changes based on pytorch documentation and recommendations for multi device. I need to get access to multi device machines in order to test properly. The lack of AMD machine availability for testing is very annoying. I don't have any access through work (currently). And AMD didn't respond to my application for developer grants for this project.

[asglover]

@vbharadwaj-bk I tested this on a multi device host and verified the problems the branch is trying to address exist in main and verified that using the device guard works as intended. There are a few aspects to this change.

• having the streams come from the correct stream for the device.
  - This happens at the C++ level by using the torch APIs defined for the purpose.
• using the pytorch cuBLAS handles.
  • BLAS handles are one per device, and the easiest way to handle this is to reuse the handles that pytorch hands out to other pytorch ops for calling blas
  • from rocBLAS to hipBLAS because torch uses hipBLAS
• It makes the workspace be a "buffer" associated with a module, this will make the buffer follow the module around as you send it from device to device. It also selects "persist" = false, so that the the workspace is not serialized during export. Which is nice because they are just zeros.
• I added an all on the same device check to the C++ level. It's different than the existing check tensor, I can unify the mechanisms or leave it separate.
• Adds a check that CPU pointers are on the right device.

I have not tested this on a AMD machine. I cannot get access. I need help testing this aspect of the change.

[vbharadwaj-bk]

This rolls in a few different changes. Some comments:

1. Prefer the stream_ordering_test.py bio e in a separate PR (and would like a few more details about the failure mode you are trying to catch here).

2. The multi device changes are good.

3. [No action needed] Regarding the BLAS handle - there's a tricky problem with the existing codebase. cuBLAS is not ABI-stable across major versions, so if the user has a CUDAToolkit that's not the same as PyTorch's version (which could well be the case, as PyTorch installs just the NVIDIA libraries it needs from a Python distribution channel), then we have a mismatch and may not be able to find certain symbols at link time. So far, it has worked because we load torch first and the symbols it loads are consistent with the ones we have compiled against.

Switching to PyTorch's BLAS handle is a step in the right direction; there's no good way to solve this problem without distributing multiple package versions.

Can run some of the testing in a bit.

[asglover]

Thanks for the review.

1. I'll take a look at this and come back with some changes.

2. Thanks!

3. This is an interesting point. I'll think more about this. It would be nice to be able to somehow use the same cuBLAS library as PyTorch when compiling our package. If we precompile, we'd have to do it for multiple versions. I wonder if by luck, the subsection of the API that we use happens to be stable? hmmm, thanks for calling this out.

Whatever the solution, it would be nice to have a strategy that we can port for other accelerators.

[asglover]

I went back and thought about this more.

The core problem is that we get the cublas handle but don't know which version of cublas it's from.
I did do some tests and it looks like the cublas handle at least is stable enough where we don't actually need to know which version it's from. Just pretending it's the right version happened to work, because someone at Nvidia must be very careful. But this is incredibly sketchy, can break, and I haven't tested rocm.

I think a cleaner solution is there is a stable torch c shim that allows you to do batched matmul "inplace" aka with out specified. Which is all we were doing with the BLAS calls. This would let us make torch do the proper handling of the different torch versions for us. I think this is great. This is the only way we can really get away from needing to do the oeq-cuda-128 , oeq-cuda-130. Which would have been a real drag.

So on a separate branch, I'm going to try to move us off of the direct BLAS calls onto a torch stable abi c shim. That will separately resolve the most sketchy part of this PR. then this PR can just focus on device guards, using the right streams, ect.

[vbharadwaj-bk]

Cool; or if there is a solution using python-callable group_gemm or any higher level Python API (https://pytorch.org/blog/accelerating-moes-with-a-triton-persistent-cache-aware-grouped-gemm-kernel/), that would be preferable. The only reason this is C++ at all is to express something that Python could not, at the time.

[asglover]

The aten grouped gemm op is incomplete / not really real. We recreate the grouped gemm behavior by calling batched matmul many times. We could move this into the python level with a for loop and batched matmul calls. Or with a custom triton op.

I wanted to make a smaller change so it's easier to review, so the PR I'm making keeps the structure the same, a loop inside the op in c++: instead of calling blas, we call torch's out of place batched matmul op.

Long term I know how and want to get rid of the grouped gemm calls by making a specialized kernel for symmetric.

If we want something sooner, because we want to bring this to Jax as well, then yeah, I think we'd have to move the batched matmul calls up to the framework level.

I think freeing us from having to link BLAS is a good step on its own. Let me polish up the PR and then we can discuss it directly.

[asglover]

see #220