[EP] add initial support for NVSHMEM-based all-to-all #1569
Conversation
meta-cla
bot
added
the
CLA Signed
| # TODO: why do we need this clone? | ||
| return out.clone() |
There was a problem hiding this comment.
Can you try removing this clone after we added out_buffer.detach() ?
There was a problem hiding this comment.
still erroring out if removing this clone
RuntimeError: Output 0 of AllToAllVDev2dBackward is a view and its base or another view of its base has been modified inplace. This view was created inside a custom Function (or because an input was returned as-is) and the autograd logic to handle view+inplace would override the custom backward associated with the custom Function, leading to incorrect gradients. This behavior is forbidden. You can fix this by cloning the output of the custom Function.
| self.output_splits = None | ||
|
|
||
| # performing all-to-all dispatch on the input | ||
| def _token_dispatch(self, mod, inputs, device_mesh): |
There was a problem hiding this comment.
i think this new implementation will get rid of the need of torch._dynamo.config.capture_scalar_outputs, avoiding the need to handle unbacked symints
As titled. This PR also does some refactoring around grouped_mm calling, as NVSHMEM-based all-to-all takes
num_tokens_per_expertand preparesoffsets.What works
num_local_experts == 1What doesn't work and needs debugging
num_local_experts > 1other TODOs
ExpertTensorParallelsupport (currently only supports ETP=1)