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.
Make the source op create the global tensor directly.
Make it legal for the “to(device)” API of the global tensor.
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