[REQUEST] implementing interleaved 1F1B in pipeline parallelism

Deepspeed doesn't have interleaved 1F1B implementation, and I haven't found anyone elsewhere using Deepspeed to implement this. So I wonder whether it is possible to implement interleaved 1F1B without too many modifications.

It seems difficult to implement interleaved 1F1B in deepspeed because of some reasons:

• each device needs to store two different copies of stage parameters (if num_stages=num_gpus*2)
• it has to be specified which copy of the parameters is executed when performing forward propagation or back propagation on each device

I see Megatron using model_chunk to specify stage parameters when executing forward propagation. Here is code.

However, deepspeed doesn't have this design code

Can anyone help or give any suggestion?

Start by comparing DeepSpeed's stage execution in deepspeed/runtime/engine.py around line 1709 with Megatron-LM's interleaved schedule in schedules.py around line 547. Determine how multiple model copies and their forward/backward selection would fit the existing pipeline design; done would require a reviewed implementation of interleaved 1F1B, with validation of the affected schedule.