Question about support for sequence parallel

[zigzagcai]

Hi,

I recently learnt about this selective SSM architecture, and it was awesome!
But I have some questions. We know that the Transformer architecture supports sequence parallelism, so does Mamba (the potential alternative of Transformer) support sequence parallelism?

[tridao]

In general, yes.
Which flavor of sequence parallelism are you referring to? The one in Megatron-LM?

[zigzagcai]

In general, yes. Which flavor of sequence parallelism are you referring to? The one in Megatron-LM?

Thanks for your timely response!
Sure. I am referring to the one in Megatron-LM. I am wondering does Mamba has built-in support for this kind of sequence parallel, or we need to implement it manually?

[tridao]

Nothing is built-in, but it'll be implemented in the future.

[zigzagcai]

Got it. Thanks!

[josiahbjorgaard]

I'm looking into implementing this, has there been any work done yet?

[zigzagcai]

I'm looking into implementing this, has there been any work done yet?

Hello, I and my colleagues have already implemented this feature.

You can just try my PR or try my colleagues' PR

[josiahbjorgaard]

Hi @zigzagcai thanks for your reply. The PR looks very interesting, but I think you may be referring to a different kind of sequence parallelism than referenced above in Megatron-LM. Rather than in your implementation where a batch of sequences is aggregated into a single sequence for the purpose of computation, I'm referring to splitting a massive sequence onto multiple GPUs. Please correct me if I'm wrong, though.

Specifically I'm working on implementing the kind of context parallelism referenced in the Mamba 2 paper in #664. https://arxiv.org/pdf/2405.21060