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.
# 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 ProxyModule.to(GraphModule).set_stage() # 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]) self.module0.to(GraphModule).set_stage(stage_id = 0, placement = P0) self.module1.to(GraphModule).set_stage(stage_id = 1, placement = P1)