How to split the num_layers unevenly when using pipeline parallelism?

[aoyulong]

For now, Megatron-LM requires the num_layers must be divisible by pipeline_model_parallel_size:
assert args.num_layers % args.transformer_pipeline_model_parallel_size == 0, \ 'num_layers must be divisible by transformer_pipeline_model_parallel_size'
How to split the num_layers unevenly? For example, given num_layers=7 and pipeline_model_parallel_size=2, the results will be 3 num_layers for stage 0 and 4 num_layers for stage 1.

[jon-barker]

As you note, this isn't currently supported. Are you looking for guidance on how it should be implemented?

[aoyulong]

@jon-barker I will try to implement it and some guidances will be very helpful. I have found the direct place that needs to be modified. But I am not sure whether there are some other places that need to be modified.

[jon-barker]

There is a proposed approach to this here: #274

This is still in review and may not be the solution we adopt though.

[github-actions]

Marking as stale. No activity in 60 days.

[edgeinfinity1]

This has been implemented by 77732c3 , in case anyone wants it.