GraphModule.set_stage(stage_id: Optional[int] = None, placement=None)

Set stage id and placement of nn.Module in pipeline parallelism.

  • stage_id (int) – stage id of this module.

  • placement (flow.placement) – the placement of all tensor in this module.


There will be automatically do tensor.to_global(placement) for all input tensor of this module. So there is no need to write to_global() in the module forward when using Pipeline Parallelism which is not recommended.

For example:

# module0 and module1 are two nn.Module in a nn.Graph.
# When a nn.Module is added into a nn.Graph, it is wrapped into a ProxyModule.
# We can set Stage ID and Placement by using
# The Stage ID is numbered starting from 0 and increasing by 1.
# The Placement is all tensors placement of this module.
import oneflow as flow
from oneflow.nn.graph import GraphModule
P_0 = flow.placement(type = "cuda", ranks = [0, 1])
P_1 = flow.placement(type = "cuda", ranks = [2, 3]) = 0, placement = P0) = 1, placement = P1)