[REQUEST] Dict input and output of Pipeline Parallelism

Hi I found it confusing when handling multiple inputs of a layer in pipeline parallelism. The lack of flexbility make it difficult to use tuple input. Is there any plans to use dict input?

e.g.

class Layer1(nn.Module):
    ...
    def forward(self, input_ids, attention_mask):
        ...
        return {"hidden_states": h, "position_ids": pos_ids}

class Layer2(nn.Module):
    ...
    def forward(self, hidden_states, position_ids):
        ...
        return {"hidden_states": h}

Also use can set some parameters with default values in forward to use them alternatively.