The torch.distributed.pipeline.sync.Pipe use a torch.nn.Sequential object as its underlying model. However, it seems that torch.nn.Sequential can only have one input. For the Pipe doesn’t modify the input signature of the underlying module, does it means that the Pipe can only handle models with single input?
If so, is there any method to walk around this limitation and pipeline parallelize models with multiple inputs?
Many thanks!!
Hey @JiayiFeng
For the
Pipedoesn’t modify the input signature of the underlying module, does it means that thePipecan only handle models with single input?
I believe this is true, as of v1.13. However, for multi-input model, I assume you can create a wrapper that packs multiple input/output into one customized object and unpack before calling into the original model?
Hi @mrshenli, thanks for your replay! But I’m afraid it won’t work. Because Pipe requires that the inputs contain at least one tensor. If I wrap all inputs into one object, it means that there is only one input for Pipe and it is not a tensor, Pipe will reject it.