r"""GraphModule is the graph representation of a nn.Module in a nn.Graph. When an nn.Module is added into an nn.Graph, it is wrapped into a ProxyModule. The ProxyModule has a GraphModule inside it. You can get and set the GraphModule to enable graph optimization on the nn.Module.
| 85 | |
| 86 | |
| 87 | class GraphModule(GraphBlock): |
| 88 | r"""GraphModule is the graph representation of a nn.Module in a nn.Graph. |
| 89 | |
| 90 | When an nn.Module is added into an nn.Graph, it is wrapped into a ProxyModule. The ProxyModule has a GraphModule inside it. |
| 91 | You can get and set the GraphModule to enable graph optimization on the nn.Module. |
| 92 | """ |
| 93 | |
| 94 | def __init__( |
| 95 | self, |
| 96 | prefix: str = "", |
| 97 | name: str = "", |
| 98 | belonged_graph: weakref.ProxyTypes = None, |
| 99 | belonged_proxy: weakref.ProxyTypes = None, |
| 100 | ): |
| 101 | super().__init__( |
| 102 | prefix, name, belonged_graph, belonged_proxy, GraphBlockType.MODULE |
| 103 | ) |
| 104 | self._is_null = True |
| 105 | self._stage_id = None |
| 106 | self._stage_placement = None |
| 107 | self._activation_checkpointing = None |
| 108 | |
| 109 | self._debug = False |
| 110 | self._debug_min_s_level = 2 |
| 111 | self._debug_max_v_level = 0 |
| 112 | self._debug_max_py_stack_depth = 2 |
| 113 | self._debug_only_user_py_stack = True |
| 114 | self._debug_op_repr_with_py_stack = False |
| 115 | self._is_executing_forward = False |
| 116 | self._args_repr = [] |
| 117 | self._outs_repr = [] |
| 118 | |
| 119 | def set_stage(self, stage_id: int = None, placement=None): |
| 120 | r"""Set stage id and placement of nn.Module in pipeline parallelism. |
| 121 | |
| 122 | Args: |
| 123 | stage_id (int): stage id of this module. |
| 124 | placement (flow.placement): the placement of all tensor in this module. |
| 125 | |
| 126 | Note: |
| 127 | There will be automatically do tensor.to_global(placement) for all input tensor of |
| 128 | this module. So there is no need to write to_global() in the module forward when using |
| 129 | Pipeline Parallelism which is not recommended. |
| 130 | |
| 131 | For example: |
| 132 | |
| 133 | .. code-block:: python |
| 134 | |
| 135 | # module0 and module1 are two nn.Module in a nn.Graph. |
| 136 | # When a nn.Module is added into a nn.Graph, it is wrapped into a ProxyModule. |
| 137 | # We can set Stage ID and Placement by using ProxyModule.to(GraphModule).set_stage() |
| 138 | # The Stage ID is numbered starting from 0 and increasing by 1. |
| 139 | # The Placement is all tensors placement of this module. |
| 140 | import oneflow as flow |
| 141 | from oneflow.nn.graph import GraphModule |
| 142 | P_0 = flow.placement(type = "cuda", ranks = [0, 1]) |
| 143 | P_1 = flow.placement(type = "cuda", ranks = [2, 3]) |
| 144 | self.module0.to(GraphModule).set_stage(stage_id = 0, placement = P0) |