set_stage(stage_id: Optional[int] = None, placement=None)¶
Set stage id and placement of nn.Module/ModuleBlock 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.
# 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)