oneflow.nn.graph.block_config.BlockConfig.set_stage

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

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

Parameters

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

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

Note

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:

# m_stage0 and m_stage1 are the two pipeline stages of the network, respectively.
# We can set Stage ID and Placement by using Module.config.set_stage()
# The Stage ID is numbered starting from 0 and increasing by 1.
# The Placement is all tensors placement of this module.
P_0 = flow.placement(type = "cuda", ranks = [0, 1])
P_1 = flow.placement(type = "cuda", ranks = [2, 3])
self.module_pipeline.m_stage0.config.set_stage(stage_id = 0, placement = P0)
self.module_pipeline.m_stage1.config.set_stage(stage_id = 1, placement = P1)