Full DualPipe schedule support

[meganfrisella]

Problem

Piper's interface and scheduling backend do not currently support the full DualPipe schedule, which duplicates pipeline stages to enable a two-way pipeline.

Expected behavior

Scheduling interface

Piper's scheduling interface should support DualPipe's device placement and microbatch ordering patterns.

Specifying placement/replication/sharding with the scheduling interface involves passing the correct device lists for each PP stage. For example, the following place directives describe the DualPipe (top) and DualPipeV (bottom) placements in the image below (assuming 2-way DP).

DualPipe. Devices 0-7 represent DP group 0 and devices 8-15 represent DP group 1.

[
  {"op": "place", "filter": {"PP": 0}, "devices": [0, 7, 8, 15], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 1}, "devices": [1, 6, 9, 14], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 2}, "devices": [2, 5, 10, 13], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 3}, "devices": [3, 4, 11, 12], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 4}, "devices": [3, 4, 11, 12], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 5}, "devices": [2, 5, 10, 13], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 6}, "devices": [1, 6, 9, 14], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 7}, "devices": [0, 7, 8, 15], "stream": "pp_stream"},
]

DualPipeV. Devices 0-3 represent DP group 0 and devices 4-7 represent DP group 1.

[
  {"op": "place", "filter": {"PP": 0}, "devices": [0, 4], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 1}, "devices": [1, 5], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 2}, "devices": [2, 6], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 3}, "devices": [3, 7], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 4}, "devices": [3, 7], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 5}, "devices": [2, 6], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 6}, "devices": [1, 5], "stream": "pp_stream"},
  {"op": "place", "filter": {"PP": 7}, "devices": [0, 4], "stream": "pp_stream"},
]

Expressing different valid paths through the pipeline is not currently supported by the interface. Currently, the split directive assumes there is one path through the pipeline (i.e. one copy of each pipeline stage per DP group) and adds new microbatches that follow the identical pipeline path. It will require some modification or addition to the interface to introduce microbatches that flow through one of multiple valid pipeline paths.

Scheduling backend

Piper's backend should insert the correct synchronization. Similar to the synchronization for tied embeddings, some care must be taken to ensure synchronization happens after all devices with a copy of a stage have completed their final backward for that stage (this is in the case of gradient accumulation, which is probably the simplest case because it involves synchronizing all the replicas at once rather than just the DP vs PP replicas).

The actors need to construct the appropriate process groups for synchronization.