oneflow.placement

OneFlow placement Class

class oneflow.placement

A oneflow.palcement is an object representing the device group on which a oneflow.Tensor is or will be allocated. The oneflow.palcement contains a device type (‘cpu’ or ‘cuda’) and corresponding device sequence.

A oneflow.Tensor’s palcement can be accessed via the Tensor.palcement property.

A oneflow.palcement can be constructed in several ways:

>>> import oneflow as flow

>>> p = flow.placement("cuda", ranks=[0, 1, 2, 3])
>>> p
oneflow.placement(type="cuda", ranks=[0, 1, 2, 3])
>>> p = flow.placement("cuda", ranks=[[0, 1], [2, 3]])
>>> p
oneflow.placement(type="cuda", ranks=[[0, 1], [2, 3]])

oneflow.env.all_device_placement(device_type: str) → oneflow._oneflow_internal.placement

Return a placement containing all devices of all machines under env.

Parameters

device_type (str) – cuda or cpu

For examples:

# world_size = 4, node_size = 1
import oneflow as flow

p = flow.env.all_device_placement("cuda") # oneflow.placement(device_type="cuda", ranks=[0, 1, 2, 3])
p = flow.env.all_device_placement("cpu") # oneflow.placement(device_type="cpu", ranks=[0, 1, 2, 3])