Add allreduce runtime support for nvshmem reduction on-stream api #21973
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
Tixxx
force-pushed
the
tixxx/nvshmem_ar
branch
2 times, most recently
from
11f965b to
74f118c
Compare
| is_nvshmem_collective = backend_config.backend() == CollectiveBackendConfig::NVSHMEM; | ||
| } | ||
| return (alias->instruction()->opcode() == HloOpcode::kCustomCall && | ||
| alias->instruction()->custom_call_target() == "mosaic_gpu") || is_nvshmem_collective; |
There was a problem hiding this comment.
why are mosaic_gpu custom calls in the same category as nvshmem ops?
There was a problem hiding this comment.
This is from the allocator branch. We use nvshmem allocator for mosaic gpu calls or native nvshmem collectives. I think instead of using mosaic_gpu custom call, the pallas kernel will use a more specific custom call name so this will change after it's changed in the other pr and rebase on that.
Tixxx
force-pushed
the
tixxx/nvshmem_ar
branch
4 times, most recently
from
36e5e7d to
52701d0
Compare
Tixxx
force-pushed
the
tixxx/nvshmem_ar
branch
from
52701d0 to
e4b7e73
Compare
Added more fixes on top of base branch, needs one last rebase
This pr adds runtime support for nvshmem reduction on-stream api.
It adds a new backend config to instruction the emitter to lower to nvshmem with corresponding collective thunks that call nvshmem reduction api.