[REQUEST] Custom partition option for PipelineModule

Is your feature request related to a problem? Please describe.
The currently pipeline parallel module gives limited control on layer partition -- we pretty much only have an option of partitioning by "parameters" or "uniform" (number of layers). Both these methods are often suboptimal for typical language transformers as ideally we want to balance GPU memory and utilization uniformly. In typical language models, both the memory and compute is much larger on the "head" layer -- they have the highest number of activations proportional to vocab size). At the same time, the overhead of the first "embedding" layer is minimal even though it has the same number of parameters as the "head". Furthermore, the "uniform" partition code uses a "floor" function to round, which ends up prioritizing the embedding layer more, in turn leading to worse gpu utilization.

Describe the solution you'd like
Allow an option to provide custom partition list, say 'parts', which can be directly applied to self.parts here --- only minimal checks are needed on the list to make sure len(parts)=num_parts+1 and parts[-1] = num_items.

Start in deepspeed/runtime/pipe/module.py at the PipelineModule partition setup, then review the uniform partition logic in deepspeed/runtime/utils.py. Trace how self.parts is built and validate the requested parts length and final value. Done means a custom partition list can be supplied and is applied for pipeline partitioning.