fairseq2.device

This module provides abstractions for managing PyTorch devices and handling CUDA contexts.

ABCs

class fairseq2.device.CudaContext

Bases: ABC

Represents an interface for interacting with CUDA runtime and device information.

abstract is_available() → bool

abstract device_count() → int

abstract get_device_properties(device: device) → Any

abstract memory_stats(device: device) → dict[str, Any]

abstract reset_peak_memory_stats() → None

class fairseq2.device.DeviceContext

Bases: ABC

Provides methods to get and set the current device of the calling thread.

This interface can be used as an alternative to the corresponding standalone functions in object-oriented code.

abstract get_current_device() → device

See get_current_device().

abstract set_device(device: device) → AbstractContextManager[None]

See set_device().

class fairseq2.device.SupportsDeviceTransfer

Bases: ABC

Represents an object that can be transferred between devices.

abstract to(device: device, *, non_blocking: bool = False) → None

Transfers this object to the specified device.

If non_blocking is True, the transfer will be performed asynchronously if possible.

Functions

fairseq2.device.detect_default_device() → device

Detects the default device of this process from environment variables.

The default device is determined by the following precedence:

1. If FAIRSEQ2_DEVICE environment variable is set, the specified device will be used.

2. If CUDA is enabled and CUDA_VISIBLE_DEVICES environment variable contains a single device, the specified device will be used.

3. If CUDA is enabled and LOCAL_RANK environment variable is set, the CUDA device at the specified index will be used.

4. CPU will be used.

Raises:

• EnvironmentVariableError – FAIRSEQ2_DEVICE environment variable does not represent a device.

• EnvironmentVariableError – LOCAL_RANK environment variable is not a positive integer.

• LocalRankOutOfRangeError – LOCAL_RANK environment variable exceeds the number of available devices.

fairseq2.device.get_current_device() → device

Returns the current device of the calling thread.

Warning

This function might impose a slight performance cost. Avoid calling it in hot code paths.

fairseq2.device.set_device(device: device) → AbstractContextManager[None]

Changes the device of the calling thread to the specified device.

This function acts as a context manager, ensuring that within its scope, any operation that constructs tensors uses the specified device - unless an explicit device argument is provided.

Note

This function is equivalent to using a torch.device as a context manager.

import torch

from fairseq2.device import set_device

cuda0_device = torch.device("cuda:0")

with set_device(cuda0_device):
    t = torch.ones((4,4))

    assert t.device == cuda0_device

    cuda1_device = torch.device("cuda:1")

    with set_device(cuda1_device):
        t = torch.ones((4, 4))

        assert t.device == cuda1_device

t = torch.ones((4, 4))

assert t.device == torch.device("cpu")

Exceptions

class fairseq2.device.LocalRankOutOfRangeError(rank: int, num_devices: int, device_type: str)

Bases: Exception