Pipeline Parallel Engine expects outputs of length 2

It seems there are some assumptions in the pipeline parallel engine that the number of outputs of each stage is 2.

In our model @ https://github.com/EleutherAI/gpt-neox/ we occasionally need to send other tensors, such as positional embeddings, across stages. At the pipeline boundary, at some point the third output tensor gets dropped, and we run into errors later on in the model.

This will also be a problem for inference, presumably, if we're sending around layer pasts and presents.

There are even places where the name of the class is checked, to handle cases relating to the boolean mask, etc, and so our model would not work if it didn't also have the name 'GPT2ModelPipe'.

        # NCCL does not like to send torch.BoolTensor types, so cast the mask to half().
        # We could do char, but with half() we can eventually flatten with other fp16
        # messages (TODO)
        if self.module.__class__.__name__ == 'GPT2ModelPipe':
            outputs = list(outputs)
            outputs[-1] = outputs[-1].half()
            outputs = tuple(outputs)

Are there any plans to increase the generality of the pipeline engine? We're running into lots of problems relating to this. I could help put together a PR for it if someone would have any ideas as to the best approach to take.

Start by tracing the pipeline parallel engine's handling of stage outputs across pipeline boundaries, including the shown GPT2ModelPipe class-name conditional. Check how additional tensors are dropped and how boolean masks are converted before communication. Done means arbitrary output counts, including positional embeddings and layer pasts/presents, survive stage transfers for training and inference.