Tensor Attributes¶

Each local oneflow.Tensor has a oneflow.dtype, oneflow.device, and global oneflow.Tensor has a oneflow.dtype, oneflow.placement, oneflow.sbp.

oneflow.device¶

class oneflow.device¶

The documentation is referenced from: https://pytorch.org/docs/1.10/tensor_attributes.html#torch.torch.device.

A oneflow.device is an object representing the device on which a oneflow.Tensor is or will be allocated.

The oneflow.device contains a device type (‘cpu’ or ‘cuda’) and optional device ordinal for the device type. If the device ordinal is not present, this object will always represent the current device for the device type.

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

A oneflow.device can be constructed via a string or via a string and device ordinal

Via a string:

>>> import oneflow as flow
>>> flow.device('cuda:0')
device(type='cuda', index=0)

>>> flow.device('cpu')
device(type='cpu', index=0)

>>> flow.device('cuda')  # current cuda device
device(type='cuda', index=0)

Via a string and device ordinal:

>>> import oneflow as flow
>>> flow.device('cuda', 0)
device(type='cuda', index=0)

>>> flow.device('cpu', 0)
device(type='cpu', index=0)

Note

The oneflow.device argument in functions can generally be substituted with a string. This allows for fast prototyping of code.

>>> import oneflow as flow
>>> # Example of a function that takes in a oneflow.device
>>> cuda0 = flow.device('cuda:0')
>>> x = flow.randn(2,3, device=cuda0)

>>> # You can substitute the flow.device with a string
>>> x = flow.randn(2,3, device='cuda:0')

oneflow.placement¶

class oneflow.placement¶

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

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

A oneflow.placement can be constructed in several ways:

>>> import oneflow as flow

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

oneflow.env.all_device_placement¶

oneflow.env.all_device_placement(device_type) → oneflow.placement ¶

Returns a placement that contains all available devices.

Parameters: device_type (str) – cuda or cpu

For examples:

# Runs on 4 ranks
import oneflow as flow

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

oneflow.sbp.sbp¶

class oneflow.sbp.sbp¶

A oneflow.sbp is an object representing that how the data of the global tensor is distributed across the ranks of the Tensor placement.

oneflow.sbp includes three types:

oneflow.sbp.split(dim)

Indicates that the global tensor is evenly divided according to the dimension dim and distributed on each rank.

oneflow.sbp.broadcast()

Indicates that the global tensor is replicated on each rank.

oneflow.sbp.partial_sum()

Indicates that the value of the global tensor is element-wise sum of the local tensors distributed in each rank.

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

A oneflow.sbp can be constructed in several ways:

>>> import oneflow as flow

>>> s = flow.sbp.split(0)
>>> s
oneflow.sbp.split(dim=0)
>>> b = flow.sbp.broadcast()
>>> b
oneflow.sbp.broadcast
>>> p = flow.sbp.partial_sum()
>>> p
oneflow.sbp.partial_sum