Pipeline parallelism for models in transformers

[KaiLv69]

Hi, there. I'm working on using pipeline parallelism in DeepSpeed with models in transformers (like llama or mistral). I found that the pipeline engine in DeepSpeed only supports send/receive tensor between layers. While the output of transformers model's layer is a tuple of tensor.

How can I use pipeline parallelism with huggingface models? Is there any elegant way to apply patches?

Thanks for any suggestions.

[HackGiter]

You just need to wrap the transformer layer with custom nn.Module and set the engine has_attention_mask.