Model parallel (splitting layers vertically)

[Androsimus]

Hi everyone!

Tell me please, is it possible to train model on many GPUs with some layers, that should be splitted vertically?
Are there best practices or advices for such tasks in Pytorch Lightning?

[akihironitta]

Hi @Androsimus, Lightning supports multiple strategies for model parallel training. There's an independent documentation page at https://pytorch-lightning.readthedocs.io/en/latest/advanced/model_parallel.html#model-parallel :)

[Androsimus]

@akihironitta Nevertheless, I don't see anything about vertical splitting. Besides I need further customization of processing different shards of splitted layers.

[ananthsub]

@Androsimus does this issue describe your current scenario? if so, could you see if the proposed solution there would satisfy what you're looking for? #11922

[Androsimus]

@ananthsub well, I haven't dived deep in your suggested code, but yes, your proposed ideas are really close to what I need!