Skip to content

Multi device bugfix#206

Draft
asglover wants to merge 11 commits into
mainfrom
multi-device-bugfix
Draft

Multi device bugfix#206
asglover wants to merge 11 commits into
mainfrom
multi-device-bugfix

Conversation

@asglover

Copy link
Copy Markdown
Collaborator

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 asglover self-assigned this Jul 13, 2026
@asglover

Copy link
Copy Markdown
Collaborator Author

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant