oneflow.utils.global_view.global_mode

class oneflow.utils.global_view.global_mode(enabled, placement=None, sbp=None)

Create a scope to provide global information for the computation process within it.

It provides convinience for converting from local execution to global execution, especially for converting to ddp global execution.

  1. Make the source op create the global tensor directly.

  2. Make it legal for the “to(device)” API of the global tensor.

  3. Make it legal to use “.device” to get the device type of the global tensor.

Note

Both placement and sbp are required if the global mode is enabled.

Parameters
  • enabled (bool) – whether the global mode is enbaled.

  • placement (oneflow.placement, optional) – the desired placement of the input. Default: None

  • sbp (oneflow.sbp.sbp, list/tuple of oneflow.sbp.sbp, optional) – the desired sbp of the input or self-defined functions in order to specify SBP. Default: None

For example:

class LinearEvalGraphWithDDP(flow.nn.Graph):
    def __init__(self):
        super().__init__()
        self.linear_dp = linear_dp

    def build(self, x):
        with global_mode(True, placement=P, sbp=B):
            device = self.linear_dp.weight.device

            x = x.to(device)

            out = self.linear_dp(x)

            # The local tensor will be converted to global
            sample = flow.randn(out.shape, device="cpu").to(device)
            out = out + sample * 100
            out = out - sample * 100

        return out
with global_mode(False):
    # The tensor will be keeped as local.
    sample = flow.randn(out.shape, device="cpu").to(device)
    out = out + sample * 100
    out = out - sample * 100