Question: How best to allocate pipeline stages just for pre_process and post_process steps?

In training a pipelined GPT model, I find that the first and last stage of the pipeline end up with one transformer layer as well as the input/output embeddings like so:

stage=0 layers=4
     0: _to_float16
     1: EmbeddingPipe
     2: <lambda>
     3: ParallelTransformerLayerPipe
stage=1 layers=1
     4: ParallelTransformerLayerPipe
stage=2 layers=5
     4: ParallelTransformerLayerPipe
     5: <lambda>
     6: MixedFusedLayerNorm
     7: EmbeddingPipe
     8: float16_to_fp32
  loss: CrossEntropy

As a result, the memory footprint of the first/last stages is about 25-30% higher than all of the intermediate stages. For example, I see sizes like the following:

[Rank first] (after 10 iterations) memory (MB) | allocated: 6824.3671875 | max allocated: 10824.42236328125 | reserved: 13292.0 | max reserved: 13292.0

.. lots of ranks like below ...

[Rank middle] (after 10 iterations) memory (MB) | allocated: 4660.1181640625 | max allocated: 8041.30029296875 | reserved: 10044.0 | max reserved: 10044.0

.. lots of ranks like above ...

[Rank last] (after 10 iterations) memory (MB) | allocated: 7126.21044921875 | max allocated: 10980.66064453125 | reserved: 12992.0 | max reserved: 12992.0

For a long pipeline, I would like to dedicate a pipeline stage just for the pre_process step and another just for the post_process step so that each transformer layer is in its own stage. For example, I would like to run the above 3-layer transformer using 5 stages where the transformer layers are in the middle 3 stages:

stage=0 layers=3
     0: _to_float16
     1: EmbeddingPipe
     2: <lambda>
stage=1 layers=1
     3: ParallelTransformerLayerPipe
stage=2 layers=1
     4: ParallelTransformerLayerPipe
stage=3 layers=1
     4: ParallelTransformerLayerPipe
stage=4 layers=4
     5: <lambda>
     6: MixedFusedLayerNorm
     7: EmbeddingPipe
     8: float16_to_fp32
  loss: CrossEntropy

I tried to hack the _partition_layers() function here:

https://github.com/microsoft/DeepSpeed/blob/45a498d09886593bd4cbfc8c0f5f9463880e569e/deepspeed/runtime/pipe/module.py#L373-L380

to have the following:

        elif method.startswith('type:'):
            layertype = method.split(':')[1]
            binary_weights = [1] * len(self._layer_specs)
            for idx in self._find_layer_type(layertype):
                binary_weights[idx] = 2
            else:
                self.parts = ds_utils.partition_balanced(weights=binary_weights,
                                                         num_parts=num_stages)

That produces a layout close to what I want. For example, I get this for a 1-layer transformer using 3 stages:

stage=0 layers=3
     0: _to_float16
     1: EmbeddingPipe
     2: <lambda>
stage=1 layers=2
     3: ParallelTransformerLayerPipe
     4: <lambda>
stage=2 layers=3
     5: MixedFusedLayerNorm
     6: EmbeddingPipe
     7: float16_to_fp32
  loss: CrossEntropy

However, it seems to hang during training.

My goal in all of this is to allocate extra stages for the first and last stage so that I can maximize the memory used in all of the intermediate transformer stages. By moving the embedding layers and others to their own stage, I can increase the size of each transformer layer by 25-30%.

Is there a recommended way to force that layout?

Thanks.

Start in deepspeed/runtime/pipe/module.py, especially _partition_layers() and the type: partitioning path. Reproduce the shown 1-layer, 3-stage layout and investigate why training hangs after moving preprocessing and postprocessing into separate stages. Done means documenting or implementing a supported way to allocate those stages without breaking pipeline execution.